truss 0.11.15rc12__tar.gz → 0.11.16__tar.gz

{truss-0.11.15rc12 → truss-0.11.16}/PKG-INFO

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

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

{truss-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/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-0.11.15rc12 → truss-0.11.16}/truss/templates/docker_server/supervisord.conf.jinja

@@ -7,7 +7,7 @@ logfile_maxbytes=0           ; No size limit on logfile (since logging is disabl
 [program:model-server]
 command={{start_command}}    ; Command to start the model server (provided by Jinja variable)
 startsecs=30                 ; Wait 30 seconds before assuming the server is running
-startretries=1               ; Do not retry if server fails to start
+startretries=0               ; Do not retry if server fails to start
 autostart=true               ; Automatically start the program when supervisord starts
 autorestart=false             ; Don't restart the program
 stdout_logfile=/dev/fd/1     ; Send stdout to the first file descriptor (stdout)

truss-0.11.16/truss/tests/util/test_env_vars.py

@@ -0,0 +1,19 @@
+import os
+
+from truss.util.env_vars import modify_env_vars
+
+
+def test_modify_env_vars():
+    os.environ["API_KEY"] = "original_key"
+    os.environ["AWS_CONFIG_FILE"] = "original_config_file"
+
+    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-0.11.15rc12 → truss-0.11.16}/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.15rc12 → truss-0.11.16}/truss-chains/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-0.11.15rc12 → truss-0.11.16}/truss-chains/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-0.11.15rc12 → truss-0.11.16}/truss-chains/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-0.11.15rc12/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