Chains SDK Reference
Python SDK Reference for Chains
Chainlet classes
APIs for creating user-defined Chainlets.
class truss_chains.ChainletBase
Base class for all chainlets.
Inheriting from this class adds validations to make sure subclasses adhere to the chainlet pattern and facilitates remote chainlet deployment.
Refer to the docs and this example chainlet for more guidance on how to create subclasses.
truss_chains.depends
Sets a “symbolic marker” to indicate to the framework that a chainlet is a
dependency of another chainlet. The return value of depends
is intended to be
used as a default argument in a chainlet’s __init__
-method.
When deploying a chain remotely, a corresponding stub to the remote is injected in
its place. In run_local
mode an instance of a local chainlet is injected.
Refer to the docs and this example chainlet for more guidance on how make one chainlet depend on another chainlet.
Despite the type annotation, this does not immediately provide a
chainlet instance. Only when deploying remotely or using run_local
a
chainlet instance is provided.
Parameters:
Name | Type | Description |
---|---|---|
chainlet_cls | Type[ChainletBase] | The chainlet class of the dependency. |
retries | int | The number of times to retry the remote chainlet in case of failures (e.g. due to transient network issues). |
timeout_sec | int | Timeout for the HTTP request to this chainlet. |
use_binary | bool | Whether to send data in binary format. This can give a parsing speedup and message size reduction (~25%) for numpy arrays. Use NumpyArrayField as a field type on pydantic models for integration and set this option to True . For simple text data, there is no significant benefit. |
- Returns: A “symbolic marker” to be used as a default argument in a chainlet’s initializer.
truss_chains.depends_context
Sets a “symbolic marker” for injecting a context object at runtime.
Refer to the docs and this
example chainlet
for more guidance on the __init__
-signature of chainlets.
Despite the type annotation, this does not immediately provide a
context instance. Only when deploying remotely or using run_local
a
context instance is provided.
- Returns: A “symbolic marker” to be used as a default argument in a chainlet’s initializer.
class truss_chains.DeploymentContext
Bases: pydantic.BaseModel
Bundles config values and resources needed to instantiate Chainlets.
The context can optionally added as a trailing argument in a Chainlet’s
__init__
method and then used to set up the chainlet (e.g. using a secret as
an access token for downloading model weights).
Parameters:
Name | Type | Description |
---|---|---|
data_dir | Path|None | The directory where the chainlet can store and access data, e.g. for downloading model weights. |
chainlet_to_service | Mapping[str,[DeployedServiceDescriptor(#class-truss-chains-deployedservicedescriptor)] | A mapping from chainlet names to service descriptors. This is used to create RPC sessions to dependency chainlets. It contains only the chainlet services that are dependencies of the current chainlet. |
secrets | Mapping[str,str] | A mapping from secret names to secret values. It contains only the secrets that are listed in remote_config.assets.secret_keys of the current chainlet. |
environment | Environment|None | The environment that the chainlet is deployed in. None if the chainlet is not associated with an environment. |
get_baseten_api_key()
- Return type: str
get_service_descriptor(chainlet_name)
Parameters:
Name | Type | Description |
---|---|---|
chainlet_name | str | The name of the chainlet. |
- Return type: DeployedServiceDescriptor
class truss_chains.definitions.Environment
Bases: pydantic.BaseModel
The environment the chainlet is deployed in.
- Parameters: name (str) – The name of the environment.
class truss_chains.ChainletOptions
Bases: pydantic.BaseModel
Parameters:
Name | Type | Description |
---|---|---|
enable_b10_tracing | bool | enables baseten-internal trace data collection. This helps baseten engineers better analyze chain performance in case of issues. It is independent of a potentially user-configured tracing instrumentation. Turning this on, could add performance overhead. |
env_variables | Mapping[str,str] | static environment variables available to the deployed chainlet. |
class truss_chains.RPCOptions
Bases: pydantic.BaseModel
Options to customize RPCs to dependency chainlets.
Parameters:
Name | Type | Description |
---|---|---|
timeout_sec | int | |
retries | int | |
use_binary | bool | Whether to send data in binary format. This can give a parsing speedup and message size reduction (~25%) for numpy arrays. Use NumpyArrayField as a field type on pydantic models for integration and set this option to True . For simple text data, there is no significant benefit. |
truss_chains.mark_entrypoint
Decorator to mark a chainlet as the entrypoint of a chain.
This decorator can be applied to one chainlet in a source file and then the CLI push command simplifies: only the file, not the class within, must be specified.
Optionally a display name for the Chain (not the Chainlet) can be set (effectively
giving a custom default value for the name
arg of the CLI push command).
Example usage:
- Return type: Type[ChainletBase]
Remote Configuration
These data structures specify for each chainlet how it gets deployed remotely, e.g. dependencies and compute resources.
class truss_chains.RemoteConfig
Bases: pydantic.BaseModel
Bundles config values needed to deploy a chainlet remotely.
This is specified as a class variable for each chainlet class, e.g.:
Parameters:
Name | Type | Description |
---|---|---|
docker_image | DockerImage | |
compute | Compute | |
assets | Assets | |
name | str|None | |
options | ChainletOptions |
class truss_chains.DockerImage
Bases: pydantic.BaseModel
Configures the docker image in which a remoted chainlet is deployed.
Any paths are relative to the source file where DockerImage
is
defined and must be created with the helper function make_abs_path_here
.
This allows you for example organize chainlets in different (potentially nested)
modules and keep their requirement files right next their python source files.
Parameters:
Name | Type | Description |
---|---|---|
base_image | BasetenImage|CustomImage | The base image used by the chainlet. Other dependencies and assets are included as additional layers on top of that image. You can choose a baseten default image for a supported python version (e.g. BasetenImage.PY311 ), this will also include GPU drivers if needed, or provide a custom image (e.g. CustomImage(image="python:3.11-slim") ). |
pip_requirements_file | AbsPath|None | Path to a file containing pip requirements. The file content is naively concatenated with pip_requirements . |
pip_requirements | list[str] | A list of pip requirements to install. The items are naively concatenated with the content of the pip_requirements_file . |
apt_requirements | list[str] | A list of apt requirements to install. |
data_dir | AbsPath|None | Data from this directory is copied into the docker image and accessible to the remote chainlet at runtime. |
external_package_dirs | list[AbsPath]|None | A list of directories containing additional python packages outside the chain’s workspace dir, e.g. a shared library. This code is copied into the docker image and importable at runtime. |
class truss_chains.BasetenImage
Bases: Enum
Default images, curated by baseten, for different python versions. If a Chainlet uses GPUs, drivers will be included in the image.
Enum Member | Value |
---|---|
PY310 | py310 |
PY311 | py311 |
PY39 | py39 |
class truss_chains.CustomImage
Bases: pydantic.BaseModel
Configures the usage of a custom image hosted on dockerhub.
Parameters:
Name | Type | Description |
---|---|---|
image | str | Reference to image on dockerhub. |
python_executable_path | str|None | Absolute path to python executable (if default python is ambiguous). |
docker_auth | DockerAuthSettings|None | See corresponding truss config. |
class truss_chains.Compute
Specifies which compute resources a chainlet has in the remote deployment.
Not all combinations can be exactly satisfied by available hardware, in some cases more powerful machine types are chosen to make sure requirements are met or over-provisioned. Refer to the baseten instance reference.
Parameters:
Name | Type | Description |
---|---|---|
cpu_count | int | Minimum number of CPUs to allocate. |
memory | str | Minimum memory to allocate, e.g. “2Gi” (2 gibibytes). |
gpu | str|Accelerator|None | GPU accelerator type, e.g. “A10G”, “A100”, refer to the truss config for more choices. |
gpu_count | int | Number of GPUs to allocate. |
predict_concurrency | int|Literal[‘cpu_count’] | Number of concurrent requests a single replica of a deployed chainlet handles. |
Concurrency concepts are explained
in this guide.
It is important to understand the difference between predict_concurrency and
the concurrency target (used for autoscaling, i.e. adding or removing replicas).
Furthermore, the predict_concurrency
of a single instance is implemented in
two ways:
- Via python’s
asyncio
, ifrun_remote
is an async def. This requires thatrun_remote
yields to the event loop. - With a threadpool if it’s a synchronous function. This requires that the threads don’t have significant CPU load (due to the GIL).
class truss_chains.Assets
Specifies which assets a chainlet can access in the remote deployment.
For example, model weight caching can be used like this:
See truss caching guide for more details on caching.
Parameters:
Name | Type | Description |
---|---|---|
cached | Iterable[ModelRepo] | One or more truss_config.ModelRepo objects. |
secret_keys | Iterable[str] | Names of secrets stored on baseten, that the chainlet should have access to. You can manage secrets on baseten here. |
external_data | Iterable[ExternalDataItem] | Data to be downloaded from public URLs and made available in the deployment (via context.data_dir ). See here for more details. |
Core
General framework and helper functions.
truss_chains.push
Deploys a chain remotely (with all dependent chainlets).
Parameters:
Name | Type | Description |
---|---|---|
entrypoint | Type[ChainletBase] | The chainlet class that serves as the entrypoint to the chain. |
chain_name | str | The name of the chain. |
publish | bool | Whether to publish the chain as a published deployment (it is a draft deployment otherwise) |
promote | bool | Whether to promote the chain to be the production deployment (this implies publishing as well). |
only_generate_trusses | bool | Used for debugging purposes. If set to True, only the the underlying truss models for the chainlets are generated in /tmp/.chains_generated . |
remote | str|None | name of a remote config in .trussrc. If not provided, it will be inquired. |
environment | str|None | The name of an environment to promote deployment into. |
progress_bar | Type[progress.Progress]|None | Optional rich.progress.Progress if output is desired. |
- Returns: A chain service handle to the deployed chain.
- Return type: ChainService
class truss_chains.deployment.deployment_client.ChainService
Handle for a deployed chain.
A ChainService
is created and returned when using push
. It
bundles the individual services for each chainlet in the chain, and provides
utilities to query their status, invoke the entrypoint etc.
get_info()
Queries the statuses of all chainlets in the chain.
- Returns:
List of
DeployedChainlet
,(name, is_entrypoint, status, logs_url)
for each chainlet. - Return type: list[DeployedChainlet]
property name : str
run_remote(json)
Invokes the entrypoint with JSON data.
Parameters:
Name | Type | Description |
---|---|---|
json | JSON dict | Input data to the entrypoint |
- Returns: The JSON response.
- Return type: Any
property run_remote_url : str
URL to invoke the entrypoint.
property status_page_url : str
Link to status page on Baseten.
truss_chains.make_abs_path_here
Helper to specify file paths relative to the immediately calling module.
E.g. in you have a project structure like this:
You can now in root/sub_package/chainlet.py
point to the requirements
file like this:
This helper uses the directory of the immediately calling module as an
absolute reference point for resolving the file location. Therefore,
you MUST NOT wrap the instantiation of make_abs_path_here
into a
function (e.g. applying decorators) or use dynamic code execution.
Ok:
Not Ok:
Parameters:
Name | Type | Description |
---|---|---|
file_path | str | Absolute or relative path. |
- Return type: AbsPath
truss_chains.run_local
Context manager local debug execution of a chain.
The arguments only need to be provided if the chainlets explicitly access any the
corresponding fields of
DeploymentContext
.
Parameters:
Name | Type | Description |
---|---|---|
secrets | Mapping[str,str]|None | A dict of secrets keys and values to provide to the chainlets. |
data_dir | Path|str|None | Path to a directory with data files. |
chainlet_to_service | Mapping[str,DeployedServiceDescriptor] | A dict of chainlet names to service descriptors. |
- Return type: ContextManager[None]
Example usage (as trailing main section in a chain file):
Refer to the local debugging guide for more details.
class truss_chains.DeployedServiceDescriptor
Bases: ServiceDescriptor
Bundles values to establish an RPC session to a dependency chainlet,
specifically with StubBase
.
Parameters:
Name | Type | Description |
---|---|---|
name | str | |
predict_url | str | |
options | RPCOptions | |
predict_url | str |
class truss_chains.StubBase
Bases: BasetenSession
, ABC
Base class for stubs that invoke remote chainlets.
Extends BasetenSession
with methods for data serialization, de-serialization
and invoking other endpoints.
It is used internally for RPCs to dependency chainlets, but it can also be used in user-code for wrapping a deployed truss model into the Chains framework. It flexibly supports JSON and pydantic inputs and output. Example usage:
Parameters:
Name | Type | Description |
---|---|---|
service_descriptor | [DeployedServiceDescriptor(#class-truss-chains-deployedservicedescriptor)] | Contains the URL and other configuration. |
api_key | str | A baseten API key to authorize requests. |
classmethod from_url(predict_url, context, options=None)
Factory method, convenient to be used in chainlet’s __init__
-method.
Parameters:
Name | Type | Description |
---|---|---|
predict_url | str | URL to predict endpoint of another chain / truss model. |
context | DeploymentContext | Deployment context object, obtained in the chainlet’s __init__ . |
options | RPCOptions | RPC options, e.g. retries. |
Invocation Methods
async predict_async(inputs: PydanticModel, output_model: Type[PydanticModel]) → PydanticModel
async predict_async(inputs: JSON, output_model: Type[PydanticModel]) → PydanticModel
async predict_async(inputs: JSON) → JSON
async predict_async_stream(inputs: PydanticModel | JSON) -> AsyncIterator[bytes]
Deprecated synchronous methods:
predict_sync(inputs: PydanticModel, output_model: Type[PydanticModel]) → PydanticModel
predict_sync(inputs: JSON, output_model: Type[PydanticModel]) → PydanticModel
predict_sync(inputs: JSON) → JSON
class truss_chains.RemoteErrorDetail
Bases: pydantic.BaseModel
When a remote chainlet raises an exception, this pydantic model contains information about the error and stack trace and is included in JSON form in the error response.
Parameters:
Name | Type | Description |
---|---|---|
exception_cls_name | str | |
exception_module_name | str|None | |
exception_message | str | |
user_stack_trace | list[StackFrame] |
format()
Format the error for printing, similar to how Python formats exceptions with stack traces.
- Return type: str