truss 0.11.15rc13__py3-none-any.whl → 0.11.16__py3-none-any.whl

truss/cli/chains_commands.py

@@ -211,14 +211,6 @@ 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(
@@ -235,7 +227,6 @@ 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.
@@ -293,7 +284,6 @@ 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,11 +11,10 @@ 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"
@@ -189,14 +188,24 @@ class S3File(RepositoryFile):
         if not dst_file.parent.exists():
             dst_file.parent.mkdir(parents=True)
 
-        with handle_client_error(
-            f"accessing S3 bucket {bucket_name} for file {file_name}"
-        ):
+        try:
             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,7 +5,6 @@ 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
@@ -14,29 +13,6 @@ 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):
@@ -323,11 +299,7 @@ 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(
@@ -337,28 +309,13 @@ 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(
-                    {params_str}
+                    {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()}
                     entrypoint: {entrypoint_str}
                     dependencies: [{dependencies_str}]
                     truss_user_env: $trussUserEnv
@@ -700,29 +657,8 @@ 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,7 +8,6 @@ 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
@@ -81,8 +80,6 @@ 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):
@@ -130,8 +127,6 @@ 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(
@@ -154,8 +149,6 @@ 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
@@ -169,8 +162,6 @@ 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 (
@@ -191,8 +182,6 @@ 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(
@@ -200,8 +189,6 @@ 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,
     )
 
 
@@ -355,33 +342,6 @@ 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,7 +120,6 @@ class TrussUserEnv(pydantic.BaseModel):
 class BlobType(Enum):
     MODEL = "model"
     TRAIN = "train"
-    CHAIN = "chain"
 
 
 class FileSummary(pydantic.BaseModel):

truss/remote/baseten/remote.py

@@ -31,7 +31,6 @@ 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,
 )
@@ -264,11 +263,9 @@ 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.
@@ -288,7 +285,6 @@ 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,
@@ -304,18 +300,6 @@ 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,
@@ -324,8 +308,6 @@ 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/remote/baseten/utils/transfer.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Optional, Type
 import boto3
 from boto3.s3.transfer import TransferConfig
 
-from truss.util.env_vars import override_env_vars
+from truss.util.env_vars import modify_env_vars
 
 if TYPE_CHECKING:
     from rich import progress
@@ -26,7 +26,10 @@ def multipart_upload_boto3(
 ) -> None:
     # In the CLI flow, ignore any local ~/.aws/config files,
     # which can interfere with uploading the Truss to S3.
-    with override_env_vars({"AWS_CONFIG_FILE": ""}):
+    aws_env_vars = set(
+        env_var for env_var in os.environ.keys() if env_var.startswith("AWS_")
+    )
+    with modify_env_vars(deletions=aws_env_vars):
         s3_resource = boto3.resource("s3", **credentials)
         filesize = os.stat(file_path).st_size
 

truss/tests/util/test_env_vars.py

@@ -1,14 +1,19 @@
 import os
 
-from truss.util.env_vars import override_env_vars
+from truss.util.env_vars import modify_env_vars
 
 
-def test_override_env_vars():
+def test_modify_env_vars():
     os.environ["API_KEY"] = "original_key"
+    os.environ["AWS_CONFIG_FILE"] = "original_config_file"
 
-    with override_env_vars({"API_KEY": "new_key", "DEBUG": "true"}):
+    with modify_env_vars(
+        overrides={"API_KEY": "new_key", "DEBUG": "true"}, deletions={"AWS_CONFIG_FILE"}
+    ):
         assert os.environ["API_KEY"] == "new_key"
         assert os.environ["DEBUG"] == "true"
+        assert "AWS_CONFIG_FILE" not in os.environ
 
     assert os.environ["API_KEY"] == "original_key"
     assert "DEBUG" not in os.environ
+    assert os.environ["AWS_CONFIG_FILE"] == "original_config_file"

truss/util/env_vars.py

@@ -1,32 +1,43 @@
 import os
-from typing import Dict, Optional
+from typing import Dict, Optional, Set
 
 
-class override_env_vars:
+class modify_env_vars:
     """A context manager for temporarily overwriting environment variables.
 
     Usage:
-        with override_env_vars({'API_KEY': 'test_key', 'DEBUG': 'true'}):
+        with modify_env_vars(overrides={'API_KEY': 'test_key', 'DEBUG': 'true'}, deletions={'AWS_CONFIG_FILE'}):
             # Environment variables are modified here
             ...
         # Original environment is restored here
     """
 
-    def __init__(self, env_vars: Dict[str, str]):
+    def __init__(
+        self,
+        overrides: Optional[Dict[str, str]] = None,
+        deletions: Optional[Set[str]] = None,
+    ):
         """
         Args:
-            env_vars: Dictionary of environment variables to set
+            overrides: Dictionary of environment variables to set
+            deletions: Set of environment variables to delete
         """
-        self.env_vars = env_vars
+        self.overrides: Dict[str, str] = overrides or dict()
+        self.deletions: Set[str] = deletions or set()
         self.original_vars: Dict[str, Optional[str]] = {}
 
     def __enter__(self):
-        for key in self.env_vars:
+        all_keys = set(self.overrides.keys()) | self.deletions
+        for key in all_keys:
             self.original_vars[key] = os.environ.get(key)
 
-        for key, value in self.env_vars.items():
+        for key, value in self.overrides.items():
             os.environ[key] = value
 
+        for key in self.deletions:
+            if key in os.environ:
+                del os.environ[key]
+
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):

{truss-0.11.15rc13.dist-info → truss-0.11.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.15rc13
+Version: 0.11.16
 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.15rc13.dist-info → truss-0.11.16.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=QijtACpuAt2O1RV_qhTNPw0jcFg-u0dX9PP-ct0t-rs,17716
+truss/cli/chains_commands.py,sha256=Kpa5mCg6URAJQE2ZmZfVQFhjBHEitKT28tKiW0H6XAI,17406
 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=EETFAgZk7C6rQezzFxz4XqjS5LIyF7uM1VVscQt_cBA,6959
+truss/contexts/image_builder/cache_warmer.py,sha256=TGMV1Mh87n2e_dSowH0sf0rZhZraDOR-LVapZL3a5r8,7377
 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,18 +52,18 @@ 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=2Es2afWKnz7OlQJHIbvYAKoSrb1dn9SnsAY--uHXbTs,30210
+truss/remote/baseten/api.py,sha256=5B5IXNy0v8hRHNH2ar3rldDa47kwt5s1PtKZQ9_pfmE,28263
 truss/remote/baseten/auth.py,sha256=tI7s6cI2EZgzpMIzrdbILHyGwiHDnmoKf_JBhJXT55E,776
-truss/remote/baseten/core.py,sha256=69utHGGFRw1ZQUobj80TSmaBgU3plnsfHZfiR15dPrY,23502
-truss/remote/baseten/custom_types.py,sha256=g7yWkE8p6uIAG5JqgfELFGHzjFLvO7vLPzbe-yl1nYs,4735
+truss/remote/baseten/core.py,sha256=uxtmBI9RAVHu1glIEJb5Q4ccJYLeZM1Cp5Svb9W68Yw,21965
+truss/remote/baseten/custom_types.py,sha256=bYrfTzGgYr6FDoya0omyadCLSTcTc-83U2scQORyUj0,4715
 truss/remote/baseten/error.py,sha256=3TNTwwPqZnr4NRd9Sl6SfLUQR2fz9l6akDPpOntTpzA,578
-truss/remote/baseten/remote.py,sha256=aKG1BODtrnmuRV-M8T3F3pw8oHawGwI09caKANJ19BM,23420
+truss/remote/baseten/remote.py,sha256=Se8AES5mk8jxa8S9fN2DSG7wnsaV7ftRjJ4Uwc_w_S0,22544
 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
 truss/remote/baseten/utils/tar.py,sha256=pMUv--YkwXDngUx1WUOK-KmAIKMcOg2E-CD5x4heh3s,2514
 truss/remote/baseten/utils/time.py,sha256=Ry9GMjYnbIGYVIGwtmv4V8ljWjvdcaCf5NOQzlNeGxI,397
-truss/remote/baseten/utils/transfer.py,sha256=d3VptuQb6M1nyS6kz0BAfeOYDLkMKUjatJXpY-mp-As,1548
+truss/remote/baseten/utils/transfer.py,sha256=vMI-Dcd_HaRkqVjWI02y3eK9DjiqgG0ULayS0KxIIvA,1652
 truss/templates/README.md.jinja,sha256=N7CJdyldZuJamj5jLh47le0hFBdu9irVsTBqoxhPNPQ,2476
 truss/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/templates/base.Dockerfile.jinja,sha256=tdMmK5TeiQuYbz4gqbACM3R-l-mazqL9tAZtJ4sxC4g,5331
@@ -141,7 +141,6 @@ 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
@@ -163,7 +162,6 @@ 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
@@ -318,7 +316,7 @@ truss/tests/test_data/test_truss_with_error/packages/helpers_2.py,sha256=q_UpVfX
 truss/tests/trt_llm/test_trt_llm_config.py,sha256=lNQ4EEkOsiT17KvnvW1snCeEBd7K_cl9_Y0dko3qpn8,8505
 truss/tests/trt_llm/test_validation.py,sha256=dmax2EHxRfqxJvWzV8uubkTef50833KBBHw-WkHufL8,2120
 truss/tests/util/test_config_checks.py,sha256=aoZF_Q-eRd3qz5wjUqa8Cr_7qF2SxodXbBIY_DBuFWg,522
-truss/tests/util/test_env_vars.py,sha256=hthgB1mU0bJb1H4Jugc-0khArlLZ3x6tLE82cDaa-J0,390
+truss/tests/util/test_env_vars.py,sha256=kz5FlynWvXhNR9bhf-xC-0cXuyYyAOvtKLLxKLDUGf4,616
 truss/tests/util/test_path.py,sha256=YfW3-IM_7iRsdR1Cb26KB1BkDsG_53_BUGBzoxY2Nog,7408
 truss/trt_llm/config_checks.py,sha256=Efxb5l7vRNveDglse78untq9V-IgtLxejk3_8JKxN5I,4671
 truss/trt_llm/validation.py,sha256=cse-EnmuHmRpwBuSc3IvmSnl-StSQCIFM1nssgnaRUQ,1848
@@ -341,8 +339,7 @@ truss/util/.truss_ignore,sha256=jpQA9ou-r_JEIcEHsUqGLHhir_m3d4IPGNyzKXtS-2g,3131
 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/env_vars.py,sha256=PmKsXdN-PX2-_xk9XcdHTFuRRWiFaMw2iNUKxE8B1Ro,1671
 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
@@ -352,8 +349,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=vdcl8FuVsL9JGIu_9K7fd2EW9Ytzoq8nfEx5pmuMKTA,9063
-truss_chains/public_api.py,sha256=civY8juJU92jSGBI7zM1qMnA7hlUdCq7L8o4IOo5meA,9722
+truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
+truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
 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
@@ -361,7 +358,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=4cHuvaynVCclJ6M9pw8ukhO1E2NRKohIRxftvOfNvOE,34499
+truss_chains/deployment/deployment_client.py,sha256=OoqkO3daktYzR2YsIcDvsuGfjR05X2K7QlA7wvFduzc,34208
 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
@@ -374,8 +371,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.15rc13.dist-info/METADATA,sha256=wxGVIems9WcjdUZMHCH-ZEnmWT9RG0Cb4NYQAd9Xv84,6682
-truss-0.11.15rc13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.15rc13.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.15rc13.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.15rc13.dist-info/RECORD,,
+truss-0.11.16.dist-info/METADATA,sha256=ykIbKzLWO3FqVipKHGs7Ad8yNkqGXoHm74AchWFVgHA,6678
+truss-0.11.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.16.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.16.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.16.dist-info/RECORD,,

truss_chains/deployment/deployment_client.py

@@ -516,21 +516,14 @@ 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,7 +265,6 @@ class PushOptionsBaseten(PushOptions):
     environment: Optional[str]
     include_git_info: bool
     working_dir: pathlib.Path
-    disable_chain_download: bool = False
 
     @classmethod
     def create(
@@ -278,7 +277,6 @@ 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
@@ -292,7 +290,6 @@ 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,7 +151,6 @@ 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).
@@ -173,7 +172,6 @@ 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.
@@ -188,7 +186,6 @@ 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(

truss/tests/cli/test_chains_cli.py

@@ -1,100 +0,0 @@
-"""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

@@ -1,285 +0,0 @@
-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

@@ -1,34 +0,0 @@
-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)}"
-        )