truss 0.11.6rc102__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/tests/templates/control/control/test_server.py

@@ -1,30 +1,15 @@
 import os
 import socket
-import sys
 from contextlib import contextmanager
-from pathlib import Path
 from typing import Dict, List
 from unittest.mock import AsyncMock, patch
 
 import httpx
 import pytest
 
-from truss.truss_handle.patch.custom_types import PatchRequest
-
-# Needed to simulate the set up on the model docker container
-sys.path.append(
-    str(
-        Path(__file__).parent.parent.parent.parent.parent
-        / "templates"
-        / "control"
-        / "control"
-    )
-)
+from truss.tests.templates.control.control.conftest import setup_control_imports
 
-sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent / "templates"))
-sys.path.append(
-    str(Path(__file__).parent.parent.parent.parent.parent / "templates" / "shared")
-)
+setup_control_imports()
 
 from truss.templates.control.control.application import create_app  # noqa
 from truss.templates.control.control.helpers.custom_types import (  # noqa
@@ -34,6 +19,7 @@ from truss.templates.control.control.helpers.custom_types import (  # noqa
     PatchType,
     PythonRequirementPatch,
 )
+from truss.truss_handle.patch.custom_types import PatchRequest
 
 
 @pytest.fixture
@@ -67,7 +53,7 @@ def app(truss_container_fs, truss_original_hash, ports):
                 "control_server_port": ports["control_server_port"],
                 "inference_server_port": ports["inference_server_port"],
                 "oversee_inference_server": False,
-                "pip_path": "pip",
+                "uv_path": "uv",
             }
         )
         inference_server_controller = control_app.state.inference_server_controller
@@ -273,14 +259,12 @@ async def test_retries(client, app):
         ]
     )
 
-    with (
-        patch("endpoints.INFERENCE_SERVER_START_WAIT_SECS", new=4),
-        pytest.raises(httpx.RemoteProtocolError),
-    ):
-        await client.get("/v1/models/model")
+    with patch("endpoints.INFERENCE_SERVER_START_WAIT_SECS", new=4):
+        resp = await client.get("/v1/models/model")
 
-    # We should have made 5 attempts
+    # We should have made 5 attempts, and then surfaced a 500
     assert app.state.proxy_client.send.call_count == 5
+    assert resp.status_code == 500
 
 
 async def _verify_apply_patch_success(client, patch: Patch):

truss/tests/templates/control/control/test_server_integration.py

@@ -35,10 +35,12 @@ def _start_truss_server(
             f"http://localhost:{patch_ping_server_port}"
         )
     sys.stdout = open(stdout_capture_file_path, "w")
+
+    # NB(nikhil): Insert paths at the beginning to ensure we search there first.
     app_path = truss_control_container_fs / "app"
-    sys.path.append(str(app_path))
+    sys.path.insert(0, str(app_path))
     control_path = truss_control_container_fs / "control" / "control"
-    sys.path.append(str(control_path))
+    sys.path.insert(0, str(control_path))
 
     from server import ControlServer
 

truss/tests/test_config.py

@@ -23,6 +23,7 @@ from truss.base.truss_config import (
     HTTPOptions,
     ModelCache,
     ModelRepo,
+    ModelRepoCacheInternal,
     Resources,
     Runtime,
     TransportKind,
@@ -292,7 +293,10 @@ def test_cache_internal_with_models(default_config):
     config = TrussConfig(
         python_version="py39",
         cache_internal=CacheInternal(
-            [ModelRepo(repo_id="test/model"), ModelRepo(repo_id="test/model2")]
+            [
+                ModelRepoCacheInternal(repo_id="test/model"),
+                ModelRepoCacheInternal(repo_id="test/model2"),
+            ]
         ),
     )
     new_config = default_config
@@ -305,11 +309,12 @@ def test_cache_internal_with_models(default_config):
 
 def test_huggingface_cache_single_model_default_revision(default_config):
     config = TrussConfig(
-        python_version="py39", model_cache=ModelCache([ModelRepo(repo_id="test/model")])
+        python_version="py39",
+        model_cache=ModelCache([ModelRepo(repo_id="test/model", use_volume=False)]),
     )
 
     new_config = default_config
-    new_config["model_cache"] = [{"repo_id": "test/model"}]
+    new_config["model_cache"] = [{"repo_id": "test/model", "use_volume": False}]
 
     assert new_config == config.to_dict(verbose=False)
     assert config.to_dict(verbose=True)["model_cache"][0].get("revision") is None
@@ -319,7 +324,9 @@ def test_huggingface_cache_single_model_non_default_revision_v1():
     config = TrussConfig(
         python_version="py39",
         requirements=[],
-        model_cache=ModelCache([ModelRepo(repo_id="test/model", revision="not-main")]),
+        model_cache=ModelCache(
+            [ModelRepo(repo_id="test/model", revision="not-main", use_volume=False)]
+        ),
     )
 
     assert config.to_dict(verbose=False)["model_cache"][0].get("revision") == "not-main"
@@ -330,16 +337,16 @@ def test_huggingface_cache_multiple_models_default_revision(default_config):
         python_version="py39",
         model_cache=ModelCache(
             [
-                ModelRepo(repo_id="test/model1", revision="main"),
-                ModelRepo(repo_id="test/model2"),
+                ModelRepo(repo_id="test/model1", revision="main", use_volume=False),
+                ModelRepo(repo_id="test/model2", use_volume=False),
             ]
         ),
     )
 
     new_config = default_config
     new_config["model_cache"] = [
-        {"repo_id": "test/model1", "revision": "main"},
-        {"repo_id": "test/model2"},
+        {"repo_id": "test/model1", "revision": "main", "use_volume": False},
+        {"repo_id": "test/model2", "use_volume": False},
     ]
 
     assert new_config == config.to_dict(verbose=False)
@@ -355,16 +362,18 @@ def test_huggingface_cache_multiple_models_mixed_revision(default_config):
         python_version="py39",
         model_cache=ModelCache(
             [
-                ModelRepo(repo_id="test/model1"),
-                ModelRepo(repo_id="test/model2", revision="not-main2"),
+                ModelRepo(repo_id="test/model1", use_volume=False),
+                ModelRepo(
+                    repo_id="test/model2", revision="not-main2", use_volume=False
+                ),
             ]
         ),
     )
 
     new_config = default_config
     new_config["model_cache"] = [
-        {"repo_id": "test/model1"},
-        {"repo_id": "test/model2", "revision": "not-main2"},
+        {"repo_id": "test/model1", "use_volume": False},
+        {"repo_id": "test/model2", "revision": "not-main2", "use_volume": False},
     ]
 
     assert new_config == config.to_dict(verbose=False)

truss/tests/test_data/server.Dockerfile

@@ -17,7 +17,8 @@ RUN /usr/local/bin/python3 -c "import sys; \
     || { echo "ERROR: Supplied base image does not have 3.8 <= python <= 3.13"; exit 1; }
 RUN if ! command -v uv >/dev/null 2>&1; then \
     command -v curl >/dev/null 2>&1 || (apt update && apt install -y curl) && \
-    curl -LsSf --retry 5 --retry-delay 5 https://astral.sh/uv/0.7.19/install.sh | sh; \
+    curl -LsSf --retry 5 --retry-delay 5 https://astral.sh/uv/0.8.22/install.sh | sh && \
+    test -x ${HOME}/.local/bin/uv; \
 fi
 ENV PATH=${PATH}:${HOME}/.local/bin
 ENV PYTHONUNBUFFERED="True"
@@ -35,6 +36,7 @@ COPY --chown= ./data ${APP_HOME}/data
 COPY --chown= ./server ${APP_HOME}
 COPY --chown= ./config.yaml ${APP_HOME}/config.yaml
 COPY --chown= ./model ${APP_HOME}/model
+RUN mkdir -p /packages
 COPY --chown= ./packages /packages
 ENV INFERENCE_SERVER_PORT="8080"
 ENV SERVER_START_CMD="/usr/local/bin/python3 /app/main.py"

truss/tests/test_data/test_build_commands_truss/config.yaml

@@ -0,0 +1,14 @@
+environment_variables: {}
+external_package_dirs: []
+model_metadata: {}
+model_name: Test Build Commands
+python_version: py39
+resources:
+  accelerator: null
+  cpu: '1'
+  memory: 2Gi
+  use_gpu: false
+secrets: {}
+system_packages: []
+build_commands:
+  - sed -i 's/TEST_FIRST_VALUE/TEST_SECOND_VALUE/g' /packages/constants/constants.py

truss/tests/test_data/test_build_commands_truss/model/model.py

@@ -0,0 +1,12 @@
+from constants.constants import TEST_KEY
+
+
+class Model:
+    def __init__(self, **kwargs):
+        pass
+
+    def load(self):
+        pass
+
+    def predict(self, input) -> str:
+        return TEST_KEY

truss/tests/test_data/test_build_commands_truss/packages/constants/constants.py

@@ -0,0 +1 @@
+TEST_KEY = "TEST_FIRST_VALUE"

truss/tests/test_data/test_truss_server_model_cache_v1/config.yaml

@@ -8,6 +8,7 @@ requirements:
   - torch
 model_cache:
   - repo_id: julien-c/EsperBERTo-small
+    use_volume: false
     ignore_patterns:
       - "*.bin"
       - "*.msgpack"

truss/tests/test_model_inference.py

@@ -2099,3 +2099,16 @@ async def test_websocket_ping_timeout_behavior(caplog):
 
         # We wait 3 seconds, so there should be ~3 PING/PONGS
         assert 2 <= caplog.text.count("PING") <= 4
+
+
+@pytest.mark.integration
+def test_build_commands_on_model_files(test_data_path):
+    with ensure_kill_all():
+        truss_dir = test_data_path / "test_build_commands_truss"
+        tr = TrussHandle(truss_dir)
+        container, urls = tr.docker_run_for_test()
+        time.sleep(3)  # Sleeping to allow the load to finish
+
+        response = requests.post(urls.predict_url, json={})
+        assert response.status_code == 200
+        assert response.json() == "TEST_SECOND_VALUE"

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/util/error_utils.py

@@ -0,0 +1,37 @@
+import logging
+from contextlib import contextmanager
+from typing import Generator
+
+from botocore.exceptions import ClientError, NoCredentialsError
+
+logger = logging.getLogger(__name__)
+
+
+@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.6rc102.dist-info → truss-0.11.24rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.6rc102
+Version: 0.11.24rc2
 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
@@ -37,7 +37,7 @@ Requires-Dist: rich<14,>=13.4.2
 Requires-Dist: ruff>=0.4.8
 Requires-Dist: tenacity>=8.0.1
 Requires-Dist: tomlkit>=0.13.2
-Requires-Dist: truss-transfer<0.0.35,>=0.0.30
+Requires-Dist: truss-transfer<0.0.40,>=0.0.37
 Requires-Dist: watchfiles<0.20,>=0.19.0
 Description-Content-Type: text/markdown