truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss_chains/public_api.py

@@ -0,0 +1,231 @@
+import pathlib
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    ContextManager,
+    Mapping,
+    Optional,
+    Type,
+    Union,
+    overload,
+)
+
+from truss_chains import definitions, framework
+from truss_chains.deployment import deployment_client
+
+if TYPE_CHECKING:
+    from rich import progress
+
+
+def depends_context() -> definitions.DeploymentContext:
+    """Sets a "symbolic marker" for injecting a context object at runtime.
+
+    Refer to `the docs <https://docs.baseten.co/chains/getting-started>`_ and this
+    `example chainlet <https://github.com/basetenlabs/truss/blob/main/truss-chains/truss_chains/example_chainlet.py>`_
+    for more guidance on the ``__init__``-signature of chainlets.
+
+    Warning:
+        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.
+
+
+    """
+    # The type error is silenced to because chains framework will at runtime inject
+    # a corresponding instance. Nonetheless, we want to use a type annotation here,
+    # to facilitate type inference, code-completion and type checking within the code
+    # of chainlets that depend on the other chainlet.
+    return framework.ContextDependencyMarker()  # type: ignore
+
+
+def depends(
+    chainlet_cls: Type[framework.ChainletT],
+    retries: int = 1,
+    timeout_sec: float = definitions.DEFAULT_TIMEOUT_SEC,
+    use_binary: bool = False,
+) -> framework.ChainletT:
+    """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 <https://docs.baseten.co/chains/getting-started>`_ and this
+    `example chainlet <https://github.com/basetenlabs/truss/blob/main/truss-chains/truss_chains/example_chainlet.py>`_
+    for more guidance on how make one chainlet depend on another chainlet.
+
+    Warning:
+        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.
+
+
+    Args:
+        chainlet_cls: The chainlet class of the dependency.
+        retries: The number of times to retry the remote chainlet in case of failures
+          (e.g. due to transient network issues). For streaming, retries are only made
+          if the request fails before streaming any results back. Failures mid-stream
+          not retried.
+        timeout_sec: Timeout for the HTTP request to this chainlet.
+        use_binary: 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.
+    """
+    options = definitions.RPCOptions(
+        retries=retries, timeout_sec=timeout_sec, use_binary=use_binary
+    )
+    # The type error is silenced to because chains framework will at runtime inject
+    # a corresponding instance. Nonetheless, we want to use a type annotation here,
+    # to facilitate type inference, code-completion and type checking within the code
+    # of chainlets that depend on the other chainlet.
+    return framework.ChainletDependencyMarker(chainlet_cls, options)  # type: ignore
+
+
+@overload
+def mark_entrypoint(
+    cls_or_chain_name: Type[framework.ChainletT],
+) -> Type[framework.ChainletT]: ...
+
+
+@overload
+def mark_entrypoint(
+    cls_or_chain_name: str,
+) -> Callable[[Type[framework.ChainletT]], Type[framework.ChainletT]]: ...
+
+
+def mark_entrypoint(
+    cls_or_chain_name: Union[str, Type[framework.ChainletT]],
+) -> Union[
+    Type[framework.ChainletT],
+    Callable[[Type[framework.ChainletT]], Type[framework.ChainletT]],
+]:
+    """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::
+
+        import truss_chains as chains
+
+        @chains.mark_entrypoint
+        class MyChainlet(ChainletBase):
+            ...
+
+        # OR with custom Chain name.
+        @chains.mark_entrypoint("My Chain Name")
+        class MyChainlet(ChainletBase):
+            ...
+    """
+    return framework.entrypoint(cls_or_chain_name)
+
+
+def push(
+    entrypoint: Type[definitions.ABCChainlet],
+    chain_name: str,
+    publish: bool = True,
+    promote: bool = True,
+    only_generate_trusses: bool = False,
+    remote: str = "baseten",
+    environment: Optional[str] = None,
+    progress_bar: Optional[Type["progress.Progress"]] = None,
+) -> deployment_client.BasetenChainService:
+    """
+    Deploys a chain remotely (with all dependent chainlets).
+
+    Args:
+        entrypoint: The chainlet class that serves as the entrypoint to the chain.
+        chain_name: The name of the chain.
+        publish: Whether to publish the chain as a published deployment (it is a
+          draft deployment otherwise)
+        promote: Whether to promote the chain to be the production deployment (this
+          implies publishing as well).
+        only_generate_trusses: Used for debugging purposes. If set to True, only the
+          the underlying truss models for the chainlets are generated in
+          ``/tmp/.chains_generated``.
+        remote: name of a remote config in `.trussrc`. If not provided, it will be
+          inquired.
+        environment: The name of an environment to promote deployment into.
+        progress_bar: Optional `rich.progress.Progress` if output is desired.
+
+    Returns:
+        A chain service handle to the deployed chain.
+
+    """
+    options = definitions.PushOptionsBaseten.create(
+        chain_name=chain_name,
+        publish=publish,
+        promote=promote,
+        only_generate_trusses=only_generate_trusses,
+        remote=remote,
+        environment=environment,
+    )
+    service = deployment_client.push(entrypoint, options, progress_bar=progress_bar)
+    assert isinstance(
+        service, deployment_client.BasetenChainService
+    )  # Per options above.
+    return service
+
+
+def run_local(
+    secrets: Optional[Mapping[str, str]] = None,
+    data_dir: Optional[Union[pathlib.Path, str]] = None,
+    chainlet_to_service: Optional[
+        Mapping[str, definitions.DeployedServiceDescriptor]
+    ] = None,
+) -> ContextManager[None]:
+    """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``.
+
+    Args:
+        secrets: A dict of secrets keys and values to provide to the chainlets.
+        data_dir: Path to a directory with data files.
+        chainlet_to_service: A dict of chainlet names to service descriptors.
+
+    Example usage (as trailing main section in a chain file)::
+
+        import os
+        import truss_chains as chains
+
+
+        class HelloWorld(chains.ChainletBase):
+            ...
+
+
+        if __name__ == "__main__":
+            with chains.run_local(
+                secrets={"some_token": os.environ["SOME_TOKEN"]},
+                chainlet_to_service={
+                    "SomeChainlet": chains.DeployedServiceDescriptor(
+                        name="SomeChainlet",
+                        display_name="SomeChainlet",
+                        predict_url="https://...",
+                        options=chains.RPCOptions(),
+                    )
+                },
+            ):
+                hello_world_chain = HelloWorld()
+                result = hello_world_chain.run_remote(max_value=5)
+
+            print(result)
+
+
+    Refer to the `local debugging guide <https://docs.baseten.co/chains/guide#test-a-chain-locally>`_
+    for more details.
+    """
+    data_dir = pathlib.Path(data_dir) if data_dir else None
+    return framework.run_local(secrets or {}, data_dir, chainlet_to_service or {})

truss_chains/pydantic_numpy.py

@@ -0,0 +1,131 @@
+import base64
+from typing import TYPE_CHECKING, Any, ClassVar
+
+import pydantic
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+class NumpyArrayField:
+    """Wrapper class to support numpy arrays as fields on pydantic models and provide
+    JSON or binary serialization implementations.
+
+    The JSON serialization exposes (data, shape, dtype), and the data is base-64
+    encoded which leads to ~33% overhead. A more compact serialization can be achieved
+    using ``msgpack_numpy`` (integrated in chains, if RPC-option ``use_binary`` is
+    enabled).
+
+    Usage example:
+
+    ```
+    import numpy as np
+
+    class MyModel(pydantic.BaseModel):
+        my_array: NumpyArrayField
+
+    m = MyModel(my_array=np.arange(4).reshape((2, 2)))
+    m.my_array.array += 10  # Work with the numpy array.
+    print(m)
+    # my_array=NumpyArrayField(
+    #   shape=(2, 2),
+    #   dtype=int64,
+    #   data=[[10 11] [12 13]])
+    m_json = m.model_dump_json()  # Serialize.
+    print(m_json)
+    # {"my_array":{"data_b64":"CgAAAAAAAAALAAAAAAAAAAwAAAAAAAAADQAAAAAAAAA=","shape":[2,2],"dtype":"int64"}}
+    m2 = MyModel.model_validate_json(m_json)  # De-serialize.
+    ```
+    """
+
+    data_key: ClassVar[str] = "data_b64"
+    shape_key: ClassVar[str] = "shape"
+    dtype_key: ClassVar[str] = "dtype"
+    array: "NDArray"
+
+    def __init__(self, array: "NDArray"):
+        self.array = array
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(shape={self.array.shape}, "
+            f"dtype={self.array.dtype}, data={self.array})"
+        )
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.no_info_after_validator_function(
+            cls.validate_numpy_array,
+            core_schema.any_schema(),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                cls.serialize_numpy_array, info_arg=True
+            ),
+        )
+
+    @classmethod
+    def validate_numpy_array(cls, value: Any):
+        import numpy as np
+
+        keys = {cls.data_key, cls.shape_key, cls.dtype_key}
+        if isinstance(value, dict) and keys.issubset(value):
+            try:
+                data = base64.b64decode(value[cls.data_key])
+                array = np.frombuffer(data, dtype=value[cls.dtype_key]).reshape(
+                    value[cls.shape_key]
+                )
+                return cls(array)
+            except (ValueError, TypeError) as e:
+                raise TypeError(
+                    "numpy_array_validation"
+                    f"Invalid data, shape, or dtype for NumPy array: {str(e)}"
+                )
+        if isinstance(value, np.ndarray):
+            return cls(value)
+        if isinstance(value, cls):
+            return value
+
+        raise TypeError(
+            "numpy_array_validation\n"
+            f"Expected a NumPy array or a dictionary with keys {keys}.\n"
+            f"Got:\n{value}"
+        )
+
+    @classmethod
+    def serialize_numpy_array(
+        cls, obj: "NumpyArrayField", info: core_schema.SerializationInfo
+    ):
+        if info.mode == "json":
+            return {
+                cls.data_key: base64.b64encode(obj.array.tobytes()).decode("utf-8"),
+                cls.shape_key: obj.array.shape,
+                cls.dtype_key: str(obj.array.dtype),
+            }
+        return obj.array
+
+    @classmethod
+    def __get_pydantic_json_schema__(
+        cls,
+        _core_schema: core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> JsonSchemaValue:
+        json_schema = handler(_core_schema)
+        json_schema.update(
+            {
+                "type": "object",
+                "properties": {
+                    "data": {"type": "string", "format": "byte"},
+                    "shape": {
+                        "type": "array",
+                        "items": {"type": "integer"},
+                        "minItems": 1,
+                    },
+                    "dtype": {"type": "string"},
+                },
+                "required": ["data", "shape", "dtype"],
+            }
+        )
+        return json_schema

truss_chains/reference_code/reference_chainlet.py

@@ -0,0 +1,34 @@
+import random
+
+# For more on chains, check out https://docs.baseten.co/chains/overview.
+import truss_chains as chains
+
+
+# By inhereting chains.ChainletBase, the chains framework will know to create a chainlet that hosts the RandInt class.
+class RandInt(chains.ChainletBase):
+    # run_remote must be implemented by all chainlets. This is the code that will be executed at inference time.
+    def run_remote(self, max_value: int) -> int:
+        return random.randint(1, max_value)
+
+
+# The @chains.mark_entrypoint decorator indicates that this Chainlet is the entrypoint.
+# Each chain must have exactly one entrypoint.
+@chains.mark_entrypoint
+class HelloWorld(chains.ChainletBase):
+    # chains.depends indicates that the HelloWorld chainlet depends on the RandInt Chainlet
+    # this enables the HelloWorld chainlet to call the RandInt chainlet
+    def __init__(self, rand_int=chains.depends(RandInt, retries=3)) -> None:
+        self._rand_int = rand_int
+
+    def run_remote(self, max_value: int) -> str:
+        num_repetitions = self._rand_int.run_remote(max_value)
+        return "Hello World! " * num_repetitions
+
+
+if __name__ == "__main__":
+    with chains.run_local():
+        hello_world_chain = HelloWorld()
+        result = hello_world_chain.run_remote(max_value=5)
+
+    print(result)
+    # Hello World! Hello World! Hello World!

truss_chains/reference_code/reference_model.py

@@ -0,0 +1,10 @@
+import truss_chains as chains
+
+
+class HelloWorld(chains.ModelBase):
+    def __init__(self, context: chains.DeploymentContext = chains.depends_context()):
+        self._call_count = 0
+
+    def predict(self, call_count_increment: int) -> int:
+        self._call_count += call_count_increment
+        return self._call_count

truss_chains/remote_chainlet/model_skeleton.py

@@ -0,0 +1,60 @@
+import pathlib
+from typing import Optional
+
+from truss.templates.shared import secrets_resolver
+
+from truss_chains import definitions
+from truss_chains.remote_chainlet import utils
+
+
+class TrussChainletModel:
+    _context: definitions.DeploymentContext
+    _chainlet: definitions.ABCChainlet
+
+    def __init__(
+        self,
+        config: dict,
+        data_dir: pathlib.Path,
+        secrets: secrets_resolver.Secrets,
+        # TODO: Remove the default value once all truss versions are synced up.
+        environment: Optional[dict] = None,
+    ) -> None:
+        truss_metadata: definitions.TrussMetadata = (
+            definitions.TrussMetadata.model_validate(
+                config["model_metadata"][definitions.TRUSS_CONFIG_CHAINS_KEY]
+            )
+        )
+        deployment_environment: Optional[definitions.Environment] = (
+            definitions.Environment.model_validate(environment) if environment else None
+        )
+        chainlet_to_deployed_service = utils.populate_chainlet_service_predict_urls(
+            truss_metadata.chainlet_to_service
+        )
+
+        self._context = definitions.DeploymentContext(
+            chainlet_to_service=chainlet_to_deployed_service,
+            secrets=secrets,
+            data_dir=data_dir,
+            environment=deployment_environment,
+        )
+
+    # Below illustrated code will be added by code generation.
+
+    # def load(self) -> None:
+    #     logging.info(f"Loading Chainlet `TextToNum`.")
+    #     self._chainlet = itest_chain.TextToNum(
+    #         replicator=stub.factory(TextReplicator, self._context),
+    #         side_effect=stub.factory(SideEffectOnlySubclass, self._context),
+    #     )
+    #
+    # If chainlet implements is_healthy:
+    # def is_healthy(self) -> Optional[bool]:
+    #     if hasattr(self, "_chainlet"):
+    #         return self._chainlet.is_healthy()
+    #
+    # def predict(
+    #     self, inputs: TextToNumInput, request: starlette.requests.Request
+    # ) -> TextToNumOutput:
+    #     with utils.predict_context(request):
+    #         result = self._chainlet.run_remote(**utils.pydantic_set_field_dict(inputs))
+    #     return TextToNumOutput(result)