truss 0.11.14__py3-none-any.whl → 0.11.15rc10__py3-none-any.whl

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/contexts/image_builder/cache_warmer.py

@@ -11,10 +11,11 @@ from typing import Optional, Type
 import boto3
 from botocore import UNSIGNED
 from botocore.client import Config
-from botocore.exceptions import ClientError, NoCredentialsError
 from google.cloud import storage
 from huggingface_hub import hf_hub_download
 
+from truss.util.error_utils import handle_client_error
+
 B10CP_PATH_TRUSS_ENV_VAR_NAME = "B10CP_PATH_TRUSS"
 
 GCS_CREDENTIALS = "/app/data/service_account.json"
@@ -188,24 +189,14 @@ class S3File(RepositoryFile):
         if not dst_file.parent.exists():
             dst_file.parent.mkdir(parents=True)
 
-        try:
+        with handle_client_error(
+            f"accessing S3 bucket {bucket_name} for file {file_name}"
+        ):
             url = client.generate_presigned_url(
                 "get_object",
                 Params={"Bucket": bucket_name, "Key": file_name},
                 ExpiresIn=3600,
             )
-        except NoCredentialsError as nce:
-            raise RuntimeError(
-                f"No AWS credentials found\nOriginal exception: {str(nce)}"
-            )
-        except ClientError as ce:
-            raise RuntimeError(
-                f"Client error when accessing the S3 bucket (check your credentials): {str(ce)}"
-            )
-        except Exception as exc:
-            raise RuntimeError(
-                f"File not found on S3 bucket: {file_name}\nOriginal exception: {str(exc)}"
-            )
 
         download_file_using_b10cp(url, dst_file, self.file_name)
 

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/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/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/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/templates/docker_server/supervisord.conf.jinja

@@ -8,7 +8,7 @@ logfile_maxbytes=0           ; No size limit on logfile (since logging is disabl
 command={{start_command}}    ; Command to start the model server (provided by Jinja variable)
 startsecs=30                 ; Wait 30 seconds before assuming the server is running
 autostart=true               ; Automatically start the program when supervisord starts
-autorestart=true             ; Always restart the program if it exits, no matter what the exit code
+autorestart=unexpected             ; Don't restart the program if it exits with clear exit code
 stdout_logfile=/dev/fd/1     ; Send stdout to the first file descriptor (stdout)
 stdout_logfile_maxbytes=0    ; No size limit on stdout log
 redirect_stderr=true         ; Redirect stderr to stdout

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

truss/tests/remote/baseten/test_chain_upload.py

@@ -0,0 +1,285 @@
+import pathlib
+import tempfile
+from unittest.mock import Mock, patch
+
+import pytest
+
+from truss.remote.baseten import custom_types as b10_types
+from truss.remote.baseten.api import BasetenApi
+from truss.remote.baseten.core import upload_chain_artifact
+from truss.remote.baseten.remote import BasetenRemote
+
+
+@pytest.fixture
+def mock_push_data():
+    """Fixture providing mock push data for tests."""
+    mock_push_data = Mock()
+    mock_push_data.model_name = "test-model"
+    mock_push_data.s3_key = "models/test-key"
+    mock_push_data.encoded_config_str = "encoded_config"
+    mock_push_data.is_draft = False
+    mock_push_data.model_id = "model-id"
+    mock_push_data.version_name = None
+    return mock_push_data
+
+
+@pytest.fixture
+def mock_remote_context():
+    """Fixture providing mock remote and context managers for tests."""
+    api = Mock(spec=BasetenApi)
+
+    remote = BasetenRemote("https://test.baseten.co", "test-api-key")
+    remote._api = api
+
+    chain_name = "test-chain"
+    entrypoint_artifact = Mock()
+    entrypoint_artifact.truss_dir = "/path/to/truss"
+    entrypoint_artifact.display_name = "entrypoint"
+
+    dependency_artifacts = []
+    truss_user_env = Mock()
+    chain_root = pathlib.Path("/path/to/chain")
+
+    with patch.object(remote, "_prepare_push") as mock_prepare_push:
+        with patch("truss.remote.baseten.remote.truss_build.load") as mock_load:
+            mock_truss_handle = Mock()
+            mock_truss_handle.spec.config.model_name = "test-model"
+            mock_load.return_value = mock_truss_handle
+
+            yield {
+                "remote": remote,
+                "api": api,
+                "chain_name": chain_name,
+                "entrypoint_artifact": entrypoint_artifact,
+                "dependency_artifacts": dependency_artifacts,
+                "truss_user_env": truss_user_env,
+                "chain_root": chain_root,
+                "mock_prepare_push": mock_prepare_push,
+                "mock_load": mock_load,
+                "mock_truss_handle": mock_truss_handle,
+            }
+
+
+def test_get_blob_credentials_for_chain():
+    """Test that get_blob_credentials works correctly for chain blob type using GraphQL."""
+    mock_graphql_response = {
+        "data": {
+            "chain_s3_upload_credentials": {
+                "s3_bucket": "test-chain-bucket",
+                "s3_key": "chains/test-uuid/chain.tgz",
+                "aws_access_key_id": "test_access_key",
+                "aws_secret_access_key": "test_secret_key",
+                "aws_session_token": "test_session_token",
+            }
+        }
+    }
+
+    # Create a real API instance and mock the GraphQL call
+    mock_auth_service = Mock()
+    mock_auth_service.authenticate.return_value = Mock(value="test-token")
+    api = BasetenApi("https://test.baseten.co", mock_auth_service)
+    with patch.object(api, "_post_graphql_query") as mock_graphql:
+        mock_graphql.return_value = mock_graphql_response
+
+        result = api.get_chain_s3_upload_credentials()
+
+        assert result.s3_bucket == "test-chain-bucket"
+        assert result.s3_key == "chains/test-uuid/chain.tgz"
+        assert result.aws_access_key_id == "test_access_key"
+        assert result.aws_secret_access_key == "test_secret_key"
+        assert result.aws_session_token == "test_session_token"
+
+        mock_graphql.assert_called_once()
+        call_args = mock_graphql.call_args
+        assert "chain_s3_upload_credentials" in call_args[0][0]
+
+
+def test_get_blob_credentials_for_other_types_uses_rest():
+    """Test that get_blob_credentials uses REST API for non-chain blob types."""
+    mock_response = {
+        "s3_bucket": "test-bucket",
+        "s3_key": "test-key",
+        "creds": {
+            "aws_access_key_id": "test_access_key",
+            "aws_secret_access_key": "test_secret_key",
+            "aws_session_token": "test_session_token",
+        },
+    }
+
+    mock_auth_service = Mock()
+    mock_auth_service.authenticate.return_value = Mock(value="test-token")
+    api = BasetenApi("https://test.baseten.co", mock_auth_service)
+    with patch.object(api, "_rest_api_client") as mock_client, patch.object(
+        api, "_post_graphql_query"
+    ) as mock_graphql:
+        mock_client.get.return_value = mock_response
+
+        result = api.get_blob_credentials(b10_types.BlobType.MODEL)
+
+        assert result["s3_bucket"] == "test-bucket"
+        assert result["s3_key"] == "test-key"
+
+        mock_client.get.assert_called_once_with("v1/blobs/credentials/model")
+        mock_graphql.assert_not_called()
+
+
+@patch("truss.remote.baseten.core.multipart_upload_boto3")
+def test_upload_chain_artifact_function(mock_multipart_upload):
+    """Test the upload_chain_artifact function."""
+    # Mock ChainUploadCredentials object
+    mock_credentials = Mock()
+    mock_credentials.s3_bucket = "test-chain-bucket"
+    mock_credentials.s3_key = "chains/test-uuid/chain.tgz"
+    mock_credentials.aws_credentials = Mock()
+    mock_credentials.aws_credentials.model_dump.return_value = {
+        "aws_access_key_id": "test_access_key",
+        "aws_secret_access_key": "test_secret_key",
+        "aws_session_token": "test_session_token",
+    }
+
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.return_value = mock_credentials
+
+    with tempfile.NamedTemporaryFile(suffix=".tgz", delete=False) as temp_file:
+        temp_file.write(b"test chain content")
+        temp_file.flush()
+
+        result = upload_chain_artifact(api, temp_file, None)
+
+        assert result == "chains/test-uuid/chain.tgz"
+
+        api.get_chain_s3_upload_credentials.assert_called_once_with()
+
+        mock_multipart_upload.assert_called_once()
+        call_args = mock_multipart_upload.call_args
+        assert call_args[0][0] == temp_file.name  # file path
+        assert call_args[0][1] == "test-chain-bucket"  # bucket
+        assert call_args[0][2] == "chains/test-uuid/chain.tgz"  # key
+        assert call_args[0][3] == {  # credentials
+            "aws_access_key_id": "test_access_key",
+            "aws_secret_access_key": "test_secret_key",
+            "aws_session_token": "test_session_token",
+        }
+
+
+@patch("truss.remote.baseten.remote.upload_chain_artifact")
+@patch("truss.remote.baseten.remote.archive_dir")
+@patch("truss.remote.baseten.remote.create_chain_atomic")
+def test_push_chain_atomic_with_chain_upload(
+    mock_create_chain_atomic,
+    mock_archive_dir,
+    mock_upload_chain_artifact,
+    mock_push_data,
+    mock_remote_context,
+):
+    """Test that push_chain_atomic uploads raw chain artifact when chain_root is provided."""
+    mock_create_chain_atomic.return_value = Mock()
+    mock_archive_dir.return_value = Mock()
+    mock_upload_chain_artifact.return_value = "chains/test-uuid/chain.tgz"
+
+    context = mock_remote_context
+    remote = context["remote"]
+    chain_name = context["chain_name"]
+    entrypoint_artifact = context["entrypoint_artifact"]
+    dependency_artifacts = context["dependency_artifacts"]
+    truss_user_env = context["truss_user_env"]
+    chain_root = context["chain_root"]
+
+    context["mock_prepare_push"].return_value = mock_push_data
+
+    result = remote.push_chain_atomic(
+        chain_name=chain_name,
+        entrypoint_artifact=entrypoint_artifact,
+        dependency_artifacts=dependency_artifacts,
+        truss_user_env=truss_user_env,
+        chain_root=chain_root,
+        publish=True,
+    )
+    assert result == mock_create_chain_atomic.return_value
+
+    mock_archive_dir.assert_called_once_with(dir=chain_root, progress_bar=None)
+    mock_upload_chain_artifact.assert_called_once()
+
+    mock_create_chain_atomic.assert_called_once()
+
+
+@patch("truss.remote.baseten.remote.create_chain_atomic")
+def test_push_chain_atomic_without_chain_upload(
+    mock_create_chain_atomic, mock_push_data, mock_remote_context
+):
+    """Test that push_chain_atomic skips chain upload when chain_root is None."""
+    mock_create_chain_atomic.return_value = Mock()
+
+    context = mock_remote_context
+    remote = context["remote"]
+    chain_name = context["chain_name"]
+    entrypoint_artifact = context["entrypoint_artifact"]
+    dependency_artifacts = context["dependency_artifacts"]
+    truss_user_env = context["truss_user_env"]
+
+    context["mock_prepare_push"].return_value = mock_push_data
+
+    with patch("truss.remote.baseten.remote.upload_chain_artifact") as mock_upload:
+        with patch(
+            "truss.remote.baseten.core.create_tar_with_progress_bar"
+        ) as mock_tar:
+            # Call push_chain_atomic without chain_root
+            result = remote.push_chain_atomic(
+                chain_name=chain_name,
+                entrypoint_artifact=entrypoint_artifact,
+                dependency_artifacts=dependency_artifacts,
+                truss_user_env=truss_user_env,
+                chain_root=None,  # No chain root
+                publish=True,
+            )
+
+            assert result
+            # Verify chain artifact upload was NOT called
+            mock_tar.assert_not_called()
+            mock_upload.assert_not_called()
+
+            mock_create_chain_atomic.assert_called_once()
+
+
+@patch("truss.remote.baseten.core.multipart_upload_boto3")
+def test_upload_chain_artifact_error_handling(mock_multipart_upload):
+    """Test error handling in upload_chain_artifact."""
+    # Mock API to raise an exception
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.side_effect = Exception("API Error")
+
+    with tempfile.NamedTemporaryFile(suffix=".tgz") as temp_file:
+        # Should raise the exception
+        with pytest.raises(Exception, match="API Error"):
+            upload_chain_artifact(api, temp_file, None)
+
+
+def test_upload_chain_artifact_credentials_extraction():
+    """Test that credentials are properly extracted from API response."""
+    # Mock ChainUploadCredentials object
+    mock_credentials = Mock()
+    mock_credentials.s3_bucket = "test-bucket"
+    mock_credentials.s3_key = "chains/test-uuid/chain.tgz"
+    mock_credentials.aws_credentials = Mock()
+    mock_credentials.aws_credentials.model_dump.return_value = {
+        "aws_access_key_id": "access_key",
+        "aws_secret_access_key": "secret_key",
+        "aws_session_token": "session_token",
+    }
+
+    api = Mock(spec=BasetenApi)
+    api.get_chain_s3_upload_credentials.return_value = mock_credentials
+
+    with patch("truss.remote.baseten.core.multipart_upload_boto3") as mock_upload:
+        with tempfile.NamedTemporaryFile(suffix=".tgz") as temp_file:
+            upload_chain_artifact(api, temp_file, None)
+
+            call_args = mock_upload.call_args
+            credentials = call_args[0][3]
+
+            assert credentials == {
+                "aws_access_key_id": "access_key",
+                "aws_secret_access_key": "secret_key",
+                "aws_session_token": "session_token",
+            }
+            assert "extra_field" not in credentials

truss/util/error_utils.py

@@ -0,0 +1,34 @@
+from contextlib import contextmanager
+from typing import Generator
+
+from botocore.exceptions import ClientError, NoCredentialsError
+
+
+@contextmanager
+def handle_client_error(
+    operation_description: str = "AWS operation",
+) -> Generator[None, None, None]:
+    """
+    Context manager to handle common boto3 errors and convert them to RuntimeError.
+
+    Args:
+        operation_description: Description of the operation being performed for error messages
+
+    Raises:
+        RuntimeError: For NoCredentialsError, ClientError, and other exceptions
+    """
+    try:
+        yield
+    except NoCredentialsError as nce:
+        raise RuntimeError(
+            f"No AWS credentials found for {operation_description}\nOriginal exception: {str(nce)}"
+        )
+    except ClientError as ce:
+        raise RuntimeError(
+            f"AWS client error when {operation_description} (check your credentials): {str(ce)}"
+        )
+
+    except Exception as exc:
+        raise RuntimeError(
+            f"Unexpected error `{exc}` during `{operation_description}`\nOriginal exception: {str(exc)}"
+        )

{truss-0.11.14.dist-info → truss-0.11.15rc10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.14
+Version: 0.11.15rc10
 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.14.dist-info → truss-0.11.15rc10.dist-info}/RECORD

@@ -8,7 +8,7 @@ truss/base/errors.py,sha256=zDVLEvseTChdPP0oNhBBQCtQUtZJUaof5zeWMIjqz6o,691
 truss/base/trt_llm_config.py,sha256=rEtBVFg2QnNMxnaz11s5Z69dJB1w7Bpt48Wf6jSsVZI,33087
 truss/base/truss_config.py,sha256=s39Xc1e20s8IV07YLl_aVnp-uRS18ZQ2TV-3FILx4nY,28416
 truss/base/truss_spec.py,sha256=jFVF79CXoEEspl2kXBAPyi-rwISReIGTdobGpaIhwJw,5979
-truss/cli/chains_commands.py,sha256=Kpa5mCg6URAJQE2ZmZfVQFhjBHEitKT28tKiW0H6XAI,17406
+truss/cli/chains_commands.py,sha256=QijtACpuAt2O1RV_qhTNPw0jcFg-u0dX9PP-ct0t-rs,17716
 truss/cli/cli.py,sha256=PaMkuwXZflkU7sa1tEoT_Zmy-iBkEZs1m4IVqcieaeo,30367
 truss/cli/remote_cli.py,sha256=G_xCKRXzgkCmkiZJhUFfsv5YSVgde1jLA5LPQitpZgI,1905
 truss/cli/train_commands.py,sha256=CrVqWsdkmSxgi3i2sSEyiE4QdfD0Z96F2Ib-PMZJjm8,20444
@@ -34,7 +34,7 @@ truss/cli/utils/output.py,sha256=GNjU85ZAMp5BI6Yij5wYXcaAvpm_kmHV0nHNmdkMxb0,646
 truss/cli/utils/self_upgrade.py,sha256=eTJZA4Wc8uUp4Qh6viRQp6bZm--wnQp7KWe5KRRpPtg,5427
 truss/contexts/docker_build_setup.py,sha256=cF4ExZgtYvrWxvyCAaUZUvV_DB_7__MqVomUDpalvKo,3925
 truss/contexts/truss_context.py,sha256=uS6L-ACHxNk0BsJwESOHh1lA0OGGw0pb33aFKGsASj4,436
-truss/contexts/image_builder/cache_warmer.py,sha256=TGMV1Mh87n2e_dSowH0sf0rZhZraDOR-LVapZL3a5r8,7377
+truss/contexts/image_builder/cache_warmer.py,sha256=EETFAgZk7C6rQezzFxz4XqjS5LIyF7uM1VVscQt_cBA,6959
 truss/contexts/image_builder/image_builder.py,sha256=IuRgDeeoHVLzIkJvKtX3807eeqEyaroCs_KWDcIHZUg,1461
 truss/contexts/image_builder/serving_image_builder.py,sha256=1PfHtkTEdNPhSQAX8Ajk_0LN3KR2EfLKwOJsnECtKXQ,33958
 truss/contexts/image_builder/util.py,sha256=y2-CjUKv0XV-0w2sr1fUCflysDJLsoU4oPp6tvvoFnk,1203
@@ -52,12 +52,12 @@ truss/patch/truss_dir_patch_applier.py,sha256=ALnaVnu96g0kF2UmGuBFTua3lrXpwAy4sG
 truss/remote/remote_factory.py,sha256=-0gLh_yIyNDgD48Q6sR8Yo5dOMQg84lrHRvn_XR0n4s,3585
 truss/remote/truss_remote.py,sha256=TEe6h6by5-JLy7PMFsDN2QxIY5FmdIYN3bKvHHl02xM,8440
 truss/remote/baseten/__init__.py,sha256=XNqJW1zyp143XQc6-7XVwsUA_Q_ZJv_ausn1_Ohtw9Y,176
-truss/remote/baseten/api.py,sha256=5B5IXNy0v8hRHNH2ar3rldDa47kwt5s1PtKZQ9_pfmE,28263
+truss/remote/baseten/api.py,sha256=2Es2afWKnz7OlQJHIbvYAKoSrb1dn9SnsAY--uHXbTs,30210
 truss/remote/baseten/auth.py,sha256=tI7s6cI2EZgzpMIzrdbILHyGwiHDnmoKf_JBhJXT55E,776
-truss/remote/baseten/core.py,sha256=uxtmBI9RAVHu1glIEJb5Q4ccJYLeZM1Cp5Svb9W68Yw,21965
-truss/remote/baseten/custom_types.py,sha256=bYrfTzGgYr6FDoya0omyadCLSTcTc-83U2scQORyUj0,4715
+truss/remote/baseten/core.py,sha256=69utHGGFRw1ZQUobj80TSmaBgU3plnsfHZfiR15dPrY,23502
+truss/remote/baseten/custom_types.py,sha256=g7yWkE8p6uIAG5JqgfELFGHzjFLvO7vLPzbe-yl1nYs,4735
 truss/remote/baseten/error.py,sha256=3TNTwwPqZnr4NRd9Sl6SfLUQR2fz9l6akDPpOntTpzA,578
-truss/remote/baseten/remote.py,sha256=Se8AES5mk8jxa8S9fN2DSG7wnsaV7ftRjJ4Uwc_w_S0,22544
+truss/remote/baseten/remote.py,sha256=aKG1BODtrnmuRV-M8T3F3pw8oHawGwI09caKANJ19BM,23420
 truss/remote/baseten/rest_client.py,sha256=_t3CWsWARt2u0C0fDsF4rtvkkHe-lH7KXoPxWXAkKd4,1185
 truss/remote/baseten/service.py,sha256=HMaKiYbr2Mzv4BfXF9QkJ8H3Wwrq3LOMpFt9js4t0rs,5834
 truss/remote/baseten/utils/status.py,sha256=jputc9N9AHXxUuW4KOk6mcZKzQ_gOBOe5BSx9K0DxPY,1266
@@ -93,7 +93,7 @@ truss/templates/custom/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NM
 truss/templates/custom/model/model.py,sha256=J04rLxK09Pwt2F4GoKOLKL-H-CqZUdYIM-PL2CE9PoE,1079
 truss/templates/custom_python_dx/my_model.py,sha256=NG75mQ6wxzB1BYUemDFZvRLBET-UrzuUK4FuHjqI29U,910
 truss/templates/docker_server/proxy.conf.jinja,sha256=Axily8EtznvrF7mUCgy2VFY99BYRt4BycZ0p9uWfd0s,2025
-truss/templates/docker_server/supervisord.conf.jinja,sha256=dd37fwZE--cutrvOUCqEyJQQQhlp61H2IUs2huKWsSk,1808
+truss/templates/docker_server/supervisord.conf.jinja,sha256=N95c7nQvPZ5i-Oypy_7nYaV7JgcYQ25M5D4F2exS2HI,1804
 truss/templates/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/templates/server/main.py,sha256=kWXrdD8z8IpamyWxc8qcvd5ck9gM1Kz2QH5qHJCnmOQ,222
 truss/templates/server/model_wrapper.py,sha256=k75VVISwwlsx5EGb82UZsu8kCM_i6Yi3-Hd0-Kpm1yo,42055
@@ -141,6 +141,7 @@ truss/tests/test_testing_utilities_for_other_tests.py,sha256=YqIKflnd_BUMYaDBSkX
 truss/tests/test_truss_gatherer.py,sha256=bn288OEkC49YY0mhly4cAl410ktZPfElNdWwZy82WfA,1261
 truss/tests/test_truss_handle.py,sha256=-xz9VXkecXDTslmQZ-dmUmQLnvD0uumRqHS2uvGlMBA,30750
 truss/tests/test_util.py,sha256=hs1bNMkXKEdoPRx4Nw-NAEdoibR92OubZuADGmbiYsQ,1344
+truss/tests/cli/test_chains_cli.py,sha256=l9GTQrhRm9SRZn43WkMY4tdRslLmdsVyiydRPa1_Ja4,3162
 truss/tests/cli/test_cli.py,sha256=yfbVS5u1hnAmmA8mJ539vj3lhH-JVGUvC4Q_Mbort44,787
 truss/tests/cli/train/test_cache_view.py,sha256=aVRCh3atRpFbJqyYgq7N-vAW0DiKMftQ7ajUqO2ClOg,22606
 truss/tests/cli/train/test_deploy_checkpoints.py,sha256=Ndkd9YxEgDLf3zLAZYH0myFK_wkKTz0oGZ57yWQt_l8,10100
@@ -162,6 +163,7 @@ truss/tests/remote/test_truss_remote.py,sha256=Rguyrnbx5RlbPJHFfCtsRtX1czAJ9Fo0a
 truss/tests/remote/baseten/conftest.py,sha256=vNk0nfDB7XdmqatOMhjdANCWFGYM4VwSHVKlaBO2PPk,442
 truss/tests/remote/baseten/test_api.py,sha256=AKJeNsrUtTNa0QPClfEvXlBOSJ214PKp23ULehMRJOQ,15885
 truss/tests/remote/baseten/test_auth.py,sha256=ttu4bDnmwGfo3oiNut4HVGnh-QnjAefwZJctiibQJKY,669
+truss/tests/remote/baseten/test_chain_upload.py,sha256=XaaF1ocovkBYsLMJ8EpXB9FUGfQZAwu4iyOWqoVn7tc,10886
 truss/tests/remote/baseten/test_core.py,sha256=6NzJTDmoSUv6Muy1LFEYIUg10-cqw-hbLyeTSWcdNjY,26117
 truss/tests/remote/baseten/test_remote.py,sha256=y1qSPL1t7dBeYI3xMFn436fttG7wkYdAoENTz7qKObg,23634
 truss/tests/remote/baseten/test_service.py,sha256=ehbGkzzSPdLN7JHxc0O9YDPfzzKqU8OBzJGjRdw08zE,3786
@@ -340,6 +342,7 @@ truss/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/util/docker.py,sha256=6PD7kMBBrOjsdvgkuSv7JMgZbe3NoJIeGasljMm2SwA,3934
 truss/util/download.py,sha256=1lfBwzyaNLEp7SAVrBd9BX5inZpkCVp8sBnS9RNoiJA,2521
 truss/util/env_vars.py,sha256=7Bv686eER71Barrs6fNamk_TrTJGmu9yV2TxaVmupn0,1232
+truss/util/error_utils.py,sha256=aO76Vf8LMlvhM28jRJ1qzNl4E5ZyvKK4TFQq_UhbQrk,1095
 truss/util/gpu.py,sha256=YiEF_JZyzur0MDMJOebMuJBQxrHD9ApGI0aPpWdb5BU,553
 truss/util/jinja.py,sha256=7KbuYNq55I3DGtImAiCvBwR0K9-z1Jo6gMhmsy4lNZE,333
 truss/util/log_utils.py,sha256=LwSgRh2K7KFjKKqBxr-IirFxGIzHi1mUM7YEvujvHsE,1985
@@ -349,8 +352,8 @@ truss/util/requirements.py,sha256=6T4nVV_NbSl3mAEo-CAk3JFmyJ_RJD768QaR55RdUJQ,69
 truss/util/user_config.py,sha256=CvBf5oouNyfdcFXOg3HFhELVW-THiuwyOYdW3aTxdHw,9130
 truss_chains/__init__.py,sha256=QDw1YwdqMaQpz5Oltu2Eq2vzEX9fDrMoqnhtbeh60i4,1278
 truss_chains/framework.py,sha256=CS7tSegPe2Q8UUT6CDkrtSrB3utr_1QN1jTEPjrj5Ug,67519
-truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
-truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
+truss_chains/private_types.py,sha256=vdcl8FuVsL9JGIu_9K7fd2EW9Ytzoq8nfEx5pmuMKTA,9063
+truss_chains/public_api.py,sha256=civY8juJU92jSGBI7zM1qMnA7hlUdCq7L8o4IOo5meA,9722
 truss_chains/public_types.py,sha256=RPr8jgKO_F_26F7H3CpwbidL-6euoKPdFHVpEIpYqrQ,29415
 truss_chains/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/pydantic_numpy.py,sha256=MG8Ji_Inwo_JSfM2n7TPj8B-nbrBlDYsY3SOeBwD8fE,4289
@@ -358,7 +361,7 @@ truss_chains/streaming.py,sha256=DGl2LEAN67YwP7Nn9MK488KmYc4KopWmcHuE6WjyO1Q,125
 truss_chains/utils.py,sha256=LvpCG2lnN6dqPqyX3PwLH9tyjUzqQN3N4WeEFROMHak,6291
 truss_chains/deployment/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/deployment/code_gen.py,sha256=397FiSNZuW59J3Ma7N9GKGfvG_87BNFAXCIV8BW41t0,32669
-truss_chains/deployment/deployment_client.py,sha256=OoqkO3daktYzR2YsIcDvsuGfjR05X2K7QlA7wvFduzc,34208
+truss_chains/deployment/deployment_client.py,sha256=4cHuvaynVCclJ6M9pw8ukhO1E2NRKohIRxftvOfNvOE,34499
 truss_chains/reference_code/reference_chainlet.py,sha256=5feSeqGtrHDbldkfZCfX2R5YbbW0Uhc35mhaP2pXrHw,1340
 truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1hz3g3lQ00tiLo55cSM,322
 truss_chains/remote_chainlet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -371,8 +374,8 @@ truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,307
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
 truss_train/restore_from_checkpoint.py,sha256=8hdPm-WSgkt74HDPjvCjZMBpvA9MwtoYsxVjOoa7BaM,1176
-truss-0.11.14.dist-info/METADATA,sha256=3F8SLMvpqzqYhOb4WB0Ngy_eFRBS_AJqhHtfy4j73DQ,6678
-truss-0.11.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.14.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.14.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.14.dist-info/RECORD,,
+truss-0.11.15rc10.dist-info/METADATA,sha256=eFZQXzoDyrBqITGt5jfqUKBdFWLaL92ZxYIPjAd1b58,6682
+truss-0.11.15rc10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.15rc10.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.15rc10.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.15rc10.dist-info/RECORD,,

truss_chains/deployment/deployment_client.py

@@ -516,14 +516,21 @@ def _create_baseten_chain(
 
     _create_chains_secret_if_missing(remote_provider)
 
+    # Get chain root for raw chain artifact upload
+    chain_root = None
+    if entrypoint_descriptor is not None:
+        chain_root = _get_chain_root(entrypoint_descriptor.chainlet_cls)
+
     chain_deployment_handle = remote_provider.push_chain_atomic(
         baseten_options.chain_name,
         entrypoint_artifact,
         dependency_artifacts,
         truss_user_env,
+        chain_root=chain_root,
         publish=baseten_options.publish,
         environment=baseten_options.environment,
         progress_bar=progress_bar,
+        disable_chain_download=baseten_options.disable_chain_download,
     )
     return BasetenChainService(
         baseten_options.chain_name,

truss_chains/private_types.py

@@ -265,6 +265,7 @@ class PushOptionsBaseten(PushOptions):
     environment: Optional[str]
     include_git_info: bool
     working_dir: pathlib.Path
+    disable_chain_download: bool = False
 
     @classmethod
     def create(
@@ -277,6 +278,7 @@ class PushOptionsBaseten(PushOptions):
         include_git_info: bool,
         working_dir: pathlib.Path,
         environment: Optional[str] = None,
+        disable_chain_download: bool = False,
     ) -> "PushOptionsBaseten":
         if promote and not environment:
             environment = PRODUCTION_ENVIRONMENT_NAME
@@ -290,6 +292,7 @@ class PushOptionsBaseten(PushOptions):
             environment=environment,
             include_git_info=include_git_info,
             working_dir=working_dir,
+            disable_chain_download=disable_chain_download,
         )
 
 

truss_chains/public_api.py

@@ -151,6 +151,7 @@ def push(
     environment: Optional[str] = None,
     progress_bar: Optional[Type["progress.Progress"]] = None,
     include_git_info: bool = False,
+    disable_chain_download: bool = False,
 ) -> deployment_client.BasetenChainService:
     """
     Deploys a chain remotely (with all dependent chainlets).
@@ -172,6 +173,7 @@ def push(
         include_git_info: Whether to attach git versioning info (sha, branch, tag) to
           deployments made from within a git repo. If set to True in `.trussrc`, it
           will always be attached.
+        disable_chain_download: Disable downloading of pushed chain source code from the UI.
 
     Returns:
         A chain service handle to the deployed chain.
@@ -186,6 +188,7 @@ def push(
         environment=environment,
         include_git_info=include_git_info,
         working_dir=pathlib.Path(inspect.getfile(entrypoint)).parent,
+        disable_chain_download=disable_chain_download,
     )
     service = deployment_client.push(entrypoint, options, progress_bar=progress_bar)
     assert isinstance(