truss 0.11.1rc14__tar.gz → 0.11.2__tar.gz

{truss-0.11.1rc14 → truss-0.11.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc14
+Version: 0.11.2
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.1rc14 → truss-0.11.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.1rc14"
+version = "0.11.2"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.11.1rc14 → truss-0.11.2}/truss/base/constants.py

@@ -85,3 +85,5 @@ OPENAI_NON_COMPATIBLE_TAG = "force-legacy-api-non-openai-compatible"  # deprecat
 PRODUCTION_ENVIRONMENT_NAME = "production"
 
 TRUSS_BASE_IMAGE_NAME = "baseten/truss-server-base"
+
+DEFAULT_TRAINING_CHECKPOINT_FOLDER = "/tmp/loaded_checkpoints"

{truss-0.11.1rc14 → truss-0.11.2}/truss/cli/chains_commands.py

@@ -43,18 +43,31 @@ def _load_example_chainlet_code() -> str:
     return source
 
 
-def _make_chains_curl_snippet(run_remote_url: str, environment: Optional[str]) -> str:
+def _make_chains_curl_snippet(
+    run_remote_url: str, environment: Optional[str], is_websocket: bool = False
+) -> str:
     if environment:
         idx = run_remote_url.find("deployment")
         if idx != -1:
             run_remote_url = (
                 run_remote_url[:idx] + f"environments/{environment}/run_remote"
             )
-    return (
-        f"curl -X POST '{run_remote_url}' \\\n"
-        '    -H "Authorization: Api-Key $BASETEN_API_KEY" \\\n'
-        "    -d '<JSON_INPUT>'"
-    )
+
+    if is_websocket:
+        # Replace 'run_remote' with 'websocket' for websocket endpoints
+        websocket_url = run_remote_url.replace("run_remote", "websocket").replace(
+            "https", "wss"
+        )
+        return (
+            f'websocat -H="Authorization: Api-Key $BASETEN_API_KEY" \\\n'
+            f"    {websocket_url}"
+        )
+    else:
+        return (
+            f"curl -X POST '{run_remote_url}' \\\n"
+            '    -H "Authorization: Api-Key $BASETEN_API_KEY" \\\n'
+            "    -d '<JSON_INPUT>'"
+        )
 
 
 def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]:
@@ -281,7 +294,7 @@ def push_chain(
 
     assert isinstance(service, deployment_client.BasetenChainService)
     curl_snippet = _make_chains_curl_snippet(
-        service.run_remote_url, options.environment
+        service.run_remote_url, options.environment, service.is_websocket
     )
 
     table, statuses = _create_chains_table(service)

{truss-0.11.1rc14 → truss-0.11.2}/truss/cli/train/deploy_checkpoints/deploy_checkpoints.py

@@ -296,18 +296,6 @@ def _get_checkpoint_ids_to_deploy(
     return checkpoint_ids
 
 
-def _select_single_checkpoint(checkpoint_id_options: List[str]) -> List[str]:
-    """Select a single checkpoint using interactive prompt."""
-    checkpoint_id = inquirer.select(
-        message="Select the checkpoints to deploy:", choices=checkpoint_id_options
-    ).execute()
-
-    if not checkpoint_id:
-        raise click.UsageError("A checkpoint must be selected.")
-
-    return [checkpoint_id]
-
-
 def _select_multiple_checkpoints(checkpoint_id_options: List[str]) -> List[str]:
     """Select multiple checkpoints using interactive checkbox."""
     checkpoint_ids = inquirer.checkbox(

{truss-0.11.1rc14 → truss-0.11.2}/truss/templates/control/requirements.txt

@@ -6,7 +6,7 @@ loguru>=0.7.2
 python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
-truss==0.11.1rc14
+truss==0.11.1
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.1rc14 → truss-0.11.2}/truss-chains/truss_chains/deployment/code_gen.py

@@ -544,7 +544,11 @@ async def websocket(self, websocket: fastapi.WebSocket) -> None:
         )"""
     return _Source(
         src=src,
-        imports={"import fastapi", "from truss_chains.remote_chainlet import utils"},
+        imports={
+            "import fastapi",
+            "from starlette.websockets import WebSocketState",
+            "from truss_chains.remote_chainlet import utils",
+        },
     )
 
 

{truss-0.11.1rc14 → truss-0.11.2}/truss-chains/truss_chains/deployment/deployment_client.py

@@ -138,7 +138,12 @@ class ChainService(abc.ABC):
 
 def _generate_chainlet_artifacts(
     options: private_types.PushOptions, entrypoint: Type[private_types.ABCChainlet]
-) -> tuple[b10_types.ChainletArtifact, list[b10_types.ChainletArtifact], bool]:
+) -> tuple[
+    b10_types.ChainletArtifact,
+    list[b10_types.ChainletArtifact],
+    bool,
+    Optional[private_types.ChainletAPIDescriptor],
+]:
     chain_root = _get_chain_root(entrypoint)
     entrypoint_artifact: Optional[b10_types.ChainletArtifact] = None
     dependency_artifacts: list[b10_types.ChainletArtifact] = []
@@ -192,7 +197,19 @@ def _generate_chainlet_artifacts(
 
     assert entrypoint_artifact is not None
 
-    return entrypoint_artifact, dependency_artifacts, has_engine_builder_chainlets
+    # Find the entrypoint descriptor
+    entrypoint_descriptor = None
+    for chainlet_descriptor in _get_ordered_dependencies([entrypoint]):
+        if chainlet_descriptor.chainlet_cls == entrypoint:
+            entrypoint_descriptor = chainlet_descriptor
+            break
+
+    return (
+        entrypoint_artifact,
+        dependency_artifacts,
+        has_engine_builder_chainlets,
+        entrypoint_descriptor,
+    )
 
 
 @framework.raise_validation_errors_before
@@ -201,9 +218,12 @@ def push(
     options: private_types.PushOptions,
     progress_bar: Optional[Type["progress.Progress"]] = None,
 ) -> Optional[ChainService]:
-    entrypoint_artifact, dependency_artifacts, has_engine_builder_chainlets = (
-        _generate_chainlet_artifacts(options, entrypoint)
-    )
+    (
+        entrypoint_artifact,
+        dependency_artifacts,
+        has_engine_builder_chainlets,
+        entrypoint_descriptor,
+    ) = _generate_chainlet_artifacts(options, entrypoint)
     if options.only_generate_trusses:
         return None
     if isinstance(options, private_types.PushOptionsBaseten):
@@ -213,7 +233,11 @@ def push(
                 "not supportd, push with `--publish`."
             )
         return _create_baseten_chain(
-            options, entrypoint_artifact, dependency_artifacts, progress_bar
+            options,
+            entrypoint_artifact,
+            dependency_artifacts,
+            progress_bar,
+            entrypoint_descriptor,
         )
     elif isinstance(options, private_types.PushOptionsLocalDocker):
         if has_engine_builder_chainlets:
@@ -369,16 +393,19 @@ def _create_docker_chain(
 class BasetenChainService(ChainService):
     _chain_deployment_handle: b10_core.ChainDeploymentHandleAtomic
     _remote: b10_remote.BasetenRemote
+    _entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor]
 
     def __init__(
         self,
         name: str,
         chain_deployment_handle: b10_core.ChainDeploymentHandleAtomic,
         remote: b10_remote.BasetenRemote,
+        entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor] = None,
     ) -> None:
         super().__init__(name)
         self._chain_deployment_handle = chain_deployment_handle
         self._remote = remote
+        self._entrypoint_descriptor = entrypoint_descriptor
 
     @property
     def run_remote_url(self) -> str:
@@ -393,6 +420,13 @@ class BasetenChainService(ChainService):
             is_draft=handle.is_draft,
         )
 
+    @property
+    def is_websocket(self) -> bool:
+        """Check if the entrypoint uses websockets."""
+        if self._entrypoint_descriptor is None:
+            return False
+        return self._entrypoint_descriptor.endpoint.is_websocket
+
     def run_remote(self, json_data: Dict) -> Any:
         """Invokes the entrypoint with JSON data.
 
@@ -462,6 +496,7 @@ def _create_baseten_chain(
     entrypoint_artifact: b10_types.ChainletArtifact,
     dependency_artifacts: list[b10_types.ChainletArtifact],
     progress_bar: Optional[Type["progress.Progress"]],
+    entrypoint_descriptor: Optional[private_types.ChainletAPIDescriptor] = None,
 ):
     logging.info(
         f"Pushing Chain '{baseten_options.chain_name}' to Baseten "
@@ -491,7 +526,10 @@ def _create_baseten_chain(
         progress_bar=progress_bar,
     )
     return BasetenChainService(
-        baseten_options.chain_name, chain_deployment_handle, remote_provider
+        baseten_options.chain_name,
+        chain_deployment_handle,
+        remote_provider,
+        entrypoint_descriptor,
     )
 
 

{truss-0.11.1rc14 → truss-0.11.2}/truss-chains/truss_chains/public_types.py

@@ -482,9 +482,11 @@ class WebSocketProtocol(Protocol):
     async def send_bytes(self, data: bytes) -> None: ...
     async def send_json(self, data: Any) -> None: ...
 
-    def iter_text(self) -> AsyncIterator[str]: ...
-    def iter_bytes(self) -> AsyncIterator[bytes]: ...
-    def iter_json(self) -> AsyncIterator[Any]: ...
+    async def iter_text(self) -> AsyncIterator[str]: ...
+    async def iter_bytes(self) -> AsyncIterator[bytes]: ...
+    async def iter_json(self) -> AsyncIterator[Any]: ...
+
+    def is_connected(self) -> bool: ...
 
 
 class EngineBuilderLLMInput(pydantic.BaseModel):

{truss-0.11.1rc14 → truss-0.11.2}/truss-chains/truss_chains/remote_chainlet/utils.py

@@ -39,7 +39,7 @@ if TYPE_CHECKING:
 try:
     import prometheus_client
 except ImportError:
-    logging.warning("Optional `prometheus_client` is not installed. ")
+    logging.warning("Optional `prometheus_client` is not installed.")
 
     class _NoOpMetric:
         def labels(self, *args: object, **kwargs: object) -> "_NoOpMetric":
@@ -65,6 +65,26 @@ except ImportError:
             return _NoOpMetric()
 
 
+try:
+    from fastapi import WebSocketDisconnect
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketDisconnect, in case local environment doesn't have
+    # fastapi.
+    class WebSocketDisconnect(Exception):  # type: ignore[no-redef]
+        def __init__(self, code: int, reason: Optional[str] = None):
+            super().__init__()
+            self.code = code
+            self.reason = reason
+
+
+try:
+    from starlette.websockets import WebSocketState
+except ImportError:
+    # NB(nikhil): Stub implementation of WebSocketState, in case local environment doesn't have starlette.
+    class WebSocketState:  # type: ignore[no-redef]
+        CONNECTED = "connected"
+
+
 T = TypeVar("T")
 
 _LockT = TypeVar("_LockT", bound=Union[threading.Lock, asyncio.Lock])
@@ -588,17 +608,12 @@ class WebsocketWrapperFastAPI:
         await self._websocket.close(code=code, reason=reason)
 
     async def receive(self) -> Union[str, bytes]:
-        try:
-            import fastapi
-        except ImportError:
-            raise utils.make_optional_import_error("fastapi")
-
         message = await self._websocket.receive()
 
         if message.get("type") == "websocket.disconnect":
             # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
             # way of detecting that the client disconnected.
-            raise fastapi.WebSocketDisconnect(message["code"], message.get("reason"))
+            raise WebSocketDisconnect(message["code"], message.get("reason"))
         elif message.get("text"):
             return typing.cast(str, message["text"])
         else:
@@ -623,13 +638,19 @@ class WebsocketWrapperFastAPI:
         await self._websocket.send_json(data)
 
     async def iter_text(self) -> AsyncIterator[str]:
-        while True:
-            yield await self.receive_text()
+        return self._websocket.iter_text()
 
     async def iter_bytes(self) -> AsyncIterator[bytes]:
-        while True:
-            yield await self.receive_bytes()
+        return self._websocket.iter_bytes()
 
     async def iter_json(self) -> AsyncIterator[Any]:
-        while True:
-            yield await self.receive_json()
+        return self._websocket.iter_json()
+
+    async def is_connected(self) -> bool:
+        # NB(nikhil): This isn't a foolproof mechanism for detecting whether a websocket
+        # connection is actually alive, ping/pong messages are best suited for that. However,
+        # as a heuristic to determine if a message is safe to send, this can do a pretty good job.
+        return (
+            self._websocket.application_state == WebSocketState.CONNECTED
+            and self._websocket.client_state == WebSocketState.CONNECTED
+        )

{truss-0.11.1rc14 → truss-0.11.2}/truss-train/truss_train/__init__.py

@@ -1,5 +1,6 @@
 from truss_train.definitions import (
     AWSIAMDockerAuth,
+    BasetenCheckpoint,
     CacheConfig,
     CheckpointingConfig,
     CheckpointList,
@@ -10,6 +11,7 @@ from truss_train.definitions import (
     FullCheckpoint,
     GCPServiceAccountJSONDockerAuth,
     Image,
+    LoadCheckpointConfig,
     LoRACheckpoint,
     LoRADetails,
     ModelWeightsFormat,
@@ -36,6 +38,8 @@ __all__ = [
     "CacheConfig",
     "AWSIAMDockerAuth",
     "GCPServiceAccountJSONDockerAuth",
+    "LoadCheckpointConfig",
+    "BasetenCheckpoint",
     "DockerAuth",
     "Image",
 ]

{truss-0.11.1rc14 → truss-0.11.2}/truss-train/truss_train/definitions.py

@@ -1,11 +1,11 @@
 import enum
 from abc import ABC
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 import pydantic
 from pydantic import field_validator, model_validator
 
-from truss.base import custom_types, truss_config
+from truss.base import constants, custom_types, truss_config
 
 DEFAULT_LORA_RANK = 16
 
@@ -56,6 +56,50 @@ class Compute(custom_types.SafeModelNoExtra):
         )
 
 
+class _CheckpointBase(custom_types.SafeModelNoExtra):
+    typ: str
+
+
+class _BasetenLatestCheckpoint(_CheckpointBase):
+    job_id: Optional[str] = None
+    project_name: Optional[str] = None
+    typ: Literal["baseten_latest_checkpoint"] = "baseten_latest_checkpoint"
+
+
+class _BasetenNamedCheckpoint(_CheckpointBase):
+    checkpoint_name: str
+    job_id: Optional[str]
+    project_name: Optional[str]
+    typ: Literal["baseten_named_checkpoint"] = "baseten_named_checkpoint"
+
+
+class BasetenCheckpoint:
+    @staticmethod
+    def from_latest_checkpoint(
+        project_name: Optional[str] = None, job_id: Optional[str] = None
+    ) -> _BasetenLatestCheckpoint:
+        return _BasetenLatestCheckpoint(project_name=project_name, job_id=job_id)
+
+    @classmethod
+    def from_named_checkpoint(
+        cls,
+        checkpoint_name: str,
+        project_name: Optional[str] = None,
+        job_id: Optional[str] = None,
+    ) -> _BasetenNamedCheckpoint:
+        return _BasetenNamedCheckpoint(
+            checkpoint_name=checkpoint_name, project_name=project_name, job_id=job_id
+        )
+
+
+class LoadCheckpointConfig(custom_types.SafeModelNoExtra):
+    enabled: bool = False
+    checkpoints: List[Union[_BasetenLatestCheckpoint, _BasetenNamedCheckpoint]] = [
+        _BasetenLatestCheckpoint()
+    ]
+    download_folder: str = constants.DEFAULT_TRAINING_CHECKPOINT_FOLDER
+
+
 class CheckpointingConfig(custom_types.SafeModelNoExtra):
     enabled: bool = False
     checkpoint_path: Optional[str] = None
@@ -72,6 +116,7 @@ class Runtime(custom_types.SafeModelNoExtra):
     environment_variables: Dict[str, Union[str, SecretReference]] = {}
     enable_cache: Optional[bool] = None
     checkpointing_config: CheckpointingConfig = CheckpointingConfig()
+    load_checkpoint_config: Optional[LoadCheckpointConfig] = None
     cache_config: Optional[CacheConfig] = None
 
     @model_validator(mode="before")

truss-0.11.2/truss-train/truss_train/restore_from_checkpoint.py

@@ -0,0 +1,42 @@
+from truss_train import (
+    BasetenCheckpoint,
+    CheckpointingConfig,
+    Image,
+    LoadCheckpointConfig,
+    Runtime,
+    TrainingJob,
+    TrainingProject,
+)
+
+load_checkpoint_config = LoadCheckpointConfig(enabled=True)
+
+load_from_most_recent_checkpoint = BasetenCheckpoint.from_latest_checkpoint()
+
+load_most_recent_checkpoint = BasetenCheckpoint.from_latest_checkpoint(
+    job_id="lqz4pw4",  # Optional
+    project_name="first-project",  # Optional
+)
+
+load_from_named_checkpoint = BasetenCheckpoint.from_named_checkpoint(
+    checkpoint_name="checkpoint-24",
+    project_name="first-project",  # Optional
+    job_id="lqz4pw4",  # Optional
+)
+
+load_checkpoint_config = LoadCheckpointConfig(
+    enabled=True,
+    download_folder="/tmp/custom_location",  # default is None -> default path set by server-side
+    checkpoints=[load_from_most_recent_checkpoint, load_from_named_checkpoint],
+)
+
+checkpointing_config = CheckpointingConfig(enabled=True)
+
+job = TrainingJob(
+    image=Image(base_image="ghcr.io/baseten-ai/truss-train-base:latest"),
+    runtime=Runtime(
+        checkpointing_config=checkpointing_config,
+        load_checkpoint_config=load_checkpoint_config,
+    ),
+)
+
+project = TrainingProject(name="new-project", job=job)