truss 0.11.16__tar.gz → 0.11.17__tar.gz

{truss-0.11.16 → truss-0.11.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.16
+Version: 0.11.17
 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.16 → truss-0.11.17}/pyproject.toml

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

{truss-0.11.16 → truss-0.11.17}/truss/cli/chains_commands.py

@@ -211,6 +211,14 @@ def _create_chains_table(service) -> Tuple[rich.table.Table, List[str]]:
     default=False,
     help=common.INCLUDE_GIT_INFO_DOC,
 )
+@click.option(
+    "--disable-chain-download",
+    "disable_chain_download",
+    is_flag=True,
+    required=False,
+    default=False,
+    help="Disable downloading of pushed chain source code from the UI.",
+)
 @click.pass_context
 @common.common_options()
 def push_chain(
@@ -227,6 +235,7 @@ def push_chain(
     environment: Optional[str],
     experimental_watch_chainlet_names: Optional[str],
     include_git_info: bool = False,
+    disable_chain_download: bool = False,
 ) -> None:
     """
     Deploys a chain remotely.
@@ -284,6 +293,7 @@ def push_chain(
             environment=environment,
             include_git_info=include_git_info,
             working_dir=source.parent if source.is_file() else source.resolve(),
+            disable_chain_download=disable_chain_download,
         )
         service = deployment_client.push(
             entrypoint_cls, options, progress_bar=progress.Progress

{truss-0.11.16 → truss-0.11.17}/truss/remote/baseten/api.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Mapping, Optional
 import requests
 from pydantic import BaseModel, Field
 
+from truss.base.custom_types import SafeModel
 from truss.remote.baseten import custom_types as b10_types
 from truss.remote.baseten.auth import ApiKey, AuthService
 from truss.remote.baseten.custom_types import APIKeyCategory
@@ -13,6 +14,29 @@ from truss.remote.baseten.rest_client import RestAPIClient
 from truss.remote.baseten.utils.transfer import base64_encoded_json_str
 
 logger = logging.getLogger(__name__)
+PARAMS_INDENT = "\n                    "
+
+
+class ChainAWSCredential(SafeModel):
+    aws_access_key_id: str
+    aws_secret_access_key: str
+    aws_session_token: str
+
+
+class ChainUploadCredentials(SafeModel):
+    s3_bucket: str
+    s3_key: str
+    aws_access_key_id: str
+    aws_secret_access_key: str
+    aws_session_token: str
+
+    @property
+    def aws_credentials(self) -> ChainAWSCredential:
+        return ChainAWSCredential(
+            aws_access_key_id=self.aws_access_key_id,
+            aws_secret_access_key=self.aws_secret_access_key,
+            aws_session_token=self.aws_session_token,
+        )
 
 
 class InstanceTypeV1(BaseModel):
@@ -299,7 +323,11 @@ class BasetenApi:
         chain_name: Optional[str] = None,
         environment: Optional[str] = None,
         is_draft: bool = False,
+        original_source_artifact_s3_key: Optional[str] = None,
+        allow_truss_download: Optional[bool] = True,
     ):
+        if allow_truss_download is None:
+            allow_truss_download = True
         entrypoint_str = _chainlet_data_atomic_to_graphql_mutation(entrypoint)
 
         dependencies_str = ", ".join(
@@ -309,13 +337,28 @@ class BasetenApi:
             ]
         )
 
+        params = []
+        if chain_id:
+            params.append(f'chain_id: "{chain_id}"')
+        if chain_name:
+            params.append(f'chain_name: "{chain_name}"')
+        if environment:
+            params.append(f'environment: "{environment}"')
+        if original_source_artifact_s3_key:
+            params.append(
+                f'original_source_artifact_s3_key: "{original_source_artifact_s3_key}"'
+            )
+
+        params.append(f"is_draft: {str(is_draft).lower()}")
+        if allow_truss_download is False:
+            params.append("allow_truss_download: false")
+
+        params_str = PARAMS_INDENT.join(params)
+
         query_string = f"""
             mutation ($trussUserEnv: String) {{
                 deploy_chain_atomic(
-                    {f'chain_id: "{chain_id}"' if chain_id else ""}
-                    {f'chain_name: "{chain_name}"' if chain_name else ""}
-                    {f'environment: "{environment}"' if environment else ""}
-                    is_draft: {str(is_draft).lower()}
+                    {params_str}
                     entrypoint: {entrypoint_str}
                     dependencies: [{dependencies_str}]
                     truss_user_env: $trussUserEnv
@@ -657,8 +700,29 @@ class BasetenApi:
         return resp_json["training_projects"]
 
     def get_blob_credentials(self, blob_type: b10_types.BlobType):
+        if blob_type == b10_types.BlobType.CHAIN:
+            return self.get_chain_s3_upload_credentials()
         return self._rest_api_client.get(f"v1/blobs/credentials/{blob_type.value}")
 
+    def get_chain_s3_upload_credentials(self) -> ChainUploadCredentials:
+        """Get chain artifact credentials using GraphQL query."""
+        query = """
+        query {
+            chain_s3_upload_credentials {
+                s3_bucket
+                s3_key
+                aws_access_key_id
+                aws_secret_access_key
+                aws_session_token
+            }
+        }
+        """
+        response = self._post_graphql_query(query)
+
+        return ChainUploadCredentials.model_validate(
+            response["data"]["chain_s3_upload_credentials"]
+        )
+
     def get_training_job_metrics(
         self,
         project_id: str,

{truss-0.11.16 → truss-0.11.17}/truss/remote/baseten/core.py

@@ -8,6 +8,7 @@ from typing import IO, TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tup
 import requests
 
 from truss.base.errors import ValidationError
+from truss.util.error_utils import handle_client_error
 
 if TYPE_CHECKING:
     from rich import progress
@@ -80,6 +81,8 @@ class ChainDeploymentHandleAtomic(NamedTuple):
     chain_id: str
     chain_deployment_id: str
     is_draft: bool
+    original_source_artifact_s3_key: Optional[str] = None
+    allow_truss_download: Optional[bool] = True
 
 
 class ModelVersionHandle(NamedTuple):
@@ -127,6 +130,8 @@ def create_chain_atomic(
     is_draft: bool,
     truss_user_env: b10_types.TrussUserEnv,
     environment: Optional[str],
+    original_source_artifact_s3_key: Optional[str] = None,
+    allow_truss_download: bool = True,
 ) -> ChainDeploymentHandleAtomic:
     if environment and is_draft:
         logging.info(
@@ -149,6 +154,8 @@ def create_chain_atomic(
             chain_name=chain_name,
             is_draft=True,
             truss_user_env=truss_user_env,
+            original_source_artifact_s3_key=original_source_artifact_s3_key,
+            allow_truss_download=allow_truss_download,
         )
     elif chain_id:
         # This is the only case where promote has relevance, since
@@ -162,6 +169,8 @@ def create_chain_atomic(
                 chain_id=chain_id,
                 environment=environment,
                 truss_user_env=truss_user_env,
+                original_source_artifact_s3_key=original_source_artifact_s3_key,
+                allow_truss_download=allow_truss_download,
             )
         except ApiError as e:
             if (
@@ -182,6 +191,8 @@ def create_chain_atomic(
             dependencies=dependencies,
             chain_name=chain_name,
             truss_user_env=truss_user_env,
+            original_source_artifact_s3_key=original_source_artifact_s3_key,
+            allow_truss_download=allow_truss_download,
         )
 
     return ChainDeploymentHandleAtomic(
@@ -189,6 +200,8 @@ def create_chain_atomic(
         chain_id=res["chain_deployment"]["chain"]["id"],
         hostname=res["chain_deployment"]["chain"]["hostname"],
         is_draft=is_draft,
+        original_source_artifact_s3_key=original_source_artifact_s3_key,
+        allow_truss_download=allow_truss_download,
     )
 
 
@@ -342,6 +355,33 @@ def upload_truss(
     return s3_key
 
 
+def upload_chain_artifact(
+    api: BasetenApi,
+    serialize_file: IO,
+    progress_bar: Optional[Type["progress.Progress"]],
+) -> str:
+    """
+    Upload a chain artifact to the Baseten remote.
+
+    Args:
+        api: BasetenApi instance
+        serialize_file: File-like object containing the serialized chain artifact
+
+    Returns:
+        The S3 key of the uploaded file
+    """
+    credentials = api.get_chain_s3_upload_credentials()
+    with handle_client_error("Uploading chain source"):
+        multipart_upload_boto3(
+            serialize_file.name,
+            credentials.s3_bucket,
+            credentials.s3_key,
+            credentials.aws_credentials.model_dump(),
+            progress_bar,
+        )
+    return credentials.s3_key
+
+
 def create_truss_service(
     api: BasetenApi,
     model_name: str,

{truss-0.11.16 → truss-0.11.17}/truss/remote/baseten/custom_types.py

@@ -120,6 +120,7 @@ class TrussUserEnv(pydantic.BaseModel):
 class BlobType(Enum):
     MODEL = "model"
     TRAIN = "train"
+    CHAIN = "chain"
 
 
 class FileSummary(pydantic.BaseModel):

{truss-0.11.16 → truss-0.11.17}/truss/remote/baseten/remote.py

@@ -31,6 +31,7 @@ from truss.remote.baseten.core import (
     get_model_and_versions,
     get_prod_version_from_versions,
     get_truss_watch_state,
+    upload_chain_artifact,
     upload_truss,
     validate_truss_config_against_backend,
 )
@@ -263,9 +264,11 @@ class BasetenRemote(TrussRemote):
         entrypoint_artifact: custom_types.ChainletArtifact,
         dependency_artifacts: List[custom_types.ChainletArtifact],
         truss_user_env: b10_types.TrussUserEnv,
+        chain_root: Optional[Path] = None,
         publish: bool = False,
         environment: Optional[str] = None,
         progress_bar: Optional[Type["progress.Progress"]] = None,
+        disable_chain_download: bool = False,
     ) -> ChainDeploymentHandleAtomic:
         # If we are promoting a model to an environment after deploy, it must be published.
         # Draft models cannot be promoted.
@@ -285,6 +288,7 @@ class BasetenRemote(TrussRemote):
                 publish=publish,
                 origin=custom_types.ModelOrigin.CHAINS,
                 progress_bar=progress_bar,
+                disable_truss_download=disable_chain_download,
             )
             oracle_data = custom_types.OracleData(
                 model_name=push_data.model_name,
@@ -300,6 +304,18 @@ class BasetenRemote(TrussRemote):
                 )
             )
 
+        # Upload raw chain artifact if chain_root is provided
+        raw_chain_s3_key = None
+        if chain_root is not None:
+            logging.info("Uploading source artifact")
+            # Create a tar file from the chain root directory
+            original_source_tar = archive_dir(dir=chain_root, progress_bar=progress_bar)
+            # Upload the chain artifact to S3
+            raw_chain_s3_key = upload_chain_artifact(
+                api=self._api,
+                serialize_file=original_source_tar,
+                progress_bar=progress_bar,
+            )
         chain_deployment_handle = create_chain_atomic(
             api=self._api,
             chain_name=chain_name,
@@ -308,6 +324,8 @@ class BasetenRemote(TrussRemote):
             is_draft=not publish,
             truss_user_env=truss_user_env,
             environment=environment,
+            original_source_artifact_s3_key=raw_chain_s3_key,
+            allow_truss_download=not disable_chain_download,
         )
         logging.info("Successfully pushed to baseten. Chain is building and deploying.")
         return chain_deployment_handle

truss-0.11.17/truss/tests/cli/test_chains_cli.py

@@ -0,0 +1,100 @@
+"""Tests for truss chains CLI commands."""
+
+from unittest.mock import Mock, patch
+
+from click.testing import CliRunner
+
+from truss.cli.cli import truss_cli
+
+
+def test_chains_push_with_disable_chain_download_flag():
+    """Test that --disable-chain-download flag is properly parsed and passed through."""
+    runner = CliRunner()
+
+    mock_entrypoint_cls = Mock()
+    mock_entrypoint_cls.meta_data.chain_name = "test_chain"
+    mock_entrypoint_cls.display_name = "TestChain"
+
+    mock_service = Mock()
+    mock_service.run_remote_url = "http://test.com/run_remote"
+    mock_service.is_websocket = False
+
+    with patch(
+        "truss_chains.framework.ChainletImporter.import_target"
+    ) as mock_importer:
+        with patch("truss_chains.deployment.deployment_client.push") as mock_push:
+            mock_importer.return_value.__enter__.return_value = mock_entrypoint_cls
+            mock_push.return_value = mock_service
+
+            result = runner.invoke(
+                truss_cli,
+                [
+                    "chains",
+                    "push",
+                    "test_chain.py",
+                    "--disable-chain-download",
+                    "--remote",
+                    "test_remote",
+                    "--dryrun",
+                ],
+            )
+
+    assert result.exit_code == 0
+
+    mock_push.assert_called_once()
+    call_args = mock_push.call_args
+    options = call_args[0][1]
+
+    assert hasattr(options, "disable_chain_download")
+    assert options.disable_chain_download is True
+
+
+def test_chains_push_without_disable_chain_download_flag():
+    """Test that disable_chain_download defaults to False when flag is not provided."""
+    runner = CliRunner()
+
+    mock_entrypoint_cls = Mock()
+    mock_entrypoint_cls.meta_data.chain_name = "test_chain"
+    mock_entrypoint_cls.display_name = "TestChain"
+
+    mock_service = Mock()
+    mock_service.run_remote_url = "http://test.com/run_remote"
+    mock_service.is_websocket = False
+
+    with patch(
+        "truss_chains.framework.ChainletImporter.import_target"
+    ) as mock_importer:
+        with patch("truss_chains.deployment.deployment_client.push") as mock_push:
+            mock_importer.return_value.__enter__.return_value = mock_entrypoint_cls
+            mock_push.return_value = mock_service
+
+            result = runner.invoke(
+                truss_cli,
+                [
+                    "chains",
+                    "push",
+                    "test_chain.py",
+                    "--remote",
+                    "test_remote",
+                    "--dryrun",
+                ],
+            )
+
+    assert result.exit_code == 0
+
+    mock_push.assert_called_once()
+    call_args = mock_push.call_args
+    options = call_args[0][1]
+
+    assert hasattr(options, "disable_chain_download")
+    assert options.disable_chain_download is False
+
+
+def test_chains_push_help_includes_disable_chain_download():
+    """Test that --disable-chain-download appears in the help output."""
+    runner = CliRunner()
+
+    result = runner.invoke(truss_cli, ["chains", "push", "--help"])
+
+    assert result.exit_code == 0
+    assert "--disable-chain-download" in result.output