truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/tests/templates/core/server/test_lazy_data_resolver_v2.py

@@ -0,0 +1,79 @@
+import json
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+from truss.templates.shared.lazy_data_resolver import (
+    TRUSS_TRANSFER_AVAILABLE,
+    LazyDataResolverV2,
+)
+
+LAZY_DATA_RESOLVER_PATH = Path("/bptr/bptr-manifest")
+TARGET_FILE = Path("nested/config.json")
+
+
+def write_bptr_manifest_to_file(expiration_timestamp: int = 2683764059):
+    bptr_manifest = {
+        "pointers": [
+            {
+                "resolution": {
+                    "url": "https://raw.githubusercontent.com/basetenlabs/truss/00e01b679afbe353b0b2fe4de6b138d912bb7167/.circleci/config.yml",
+                    "expiration_timestamp": expiration_timestamp,
+                },
+                "uid": "8c6b2f215f0333437cdc3fe7c79be0c802847d2f2a0ccdc0bb251814e63cf375",
+                "file_name": TARGET_FILE.as_posix(),
+                "hashtype": "blake3",
+                "hash": "8c6b2f215f0333437cdc3fe7c79be0c802847d2f2a0ccdc0bb251814e63cf375",
+                "size": 1482,
+            }
+        ]
+    }
+    # write to LAZY_DATA_RESOLVER_PATH
+    with open(LAZY_DATA_RESOLVER_PATH, "w") as f:
+        json.dump(bptr_manifest, f)
+
+
+@pytest.mark.skipif(not TRUSS_TRANSFER_AVAILABLE, reason="Truss Transfer not available")
+def test_lazy_data_resolver_v2():
+    # truss_transfer reads from LAZY_DATA_RESOLVER_PATH
+    if LAZY_DATA_RESOLVER_PATH.exists():
+        LAZY_DATA_RESOLVER_PATH.unlink()
+    with pytest.raises(Exception):
+        # LAZY_DATA_RESOLVER_PATH does not exist
+        # should raise an exception
+        LazyDataResolverV2(Path("/tmp")).fetch()
+
+    try:
+        LAZY_DATA_RESOLVER_PATH.mkdir(parents=True, exist_ok=True)
+    except Exception as e:
+        pytest.skip(
+            f"Unable to create {LAZY_DATA_RESOLVER_PATH} due to missing os permissions: {e}"
+        )
+
+    # without LAZY_DATA_RESOLVER_PATH -> does not create folder / file
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        resolver = LazyDataResolverV2(data_dir).fetch()
+        assert not (data_dir / TARGET_FILE).exists()
+
+    # with LAZY_DATA_RESOLVER_PATH -> fetches data
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        write_bptr_manifest_to_file()
+        resolver = LazyDataResolverV2(data_dir).fetch()
+        resolver.fetch()
+        assert (data_dir / TARGET_FILE).exists()
+        assert (data_dir / TARGET_FILE).stat().st_size == 1482
+
+    # with expired LAZY_DATA_RESOLVER_PATH -> raises exception
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_dir = Path(tempdir)
+        write_bptr_manifest_to_file(expiration_timestamp=int(time.time()) - 1)
+        resolver = LazyDataResolverV2(data_dir).fetch()
+        with pytest.raises(Exception):
+            resolver.fetch()
+
+
+if __name__ == "__main__":
+    test_lazy_data_resolver_v2()

truss/tests/{server → templates}/core/server/test_secrets_resolver.py

@@ -2,7 +2,7 @@ import os
 from contextlib import contextmanager
 from pathlib import Path
 
-from truss.server.shared.secrets_resolver import SecretsResolver
+from truss.templates.shared.secrets_resolver import SecretsResolver
 
 CONFIG = {"secrets": {"secret_key": "default_secret_value"}}
 

truss/tests/{server → templates/server}/common/test_retry.py

@@ -1,8 +1,8 @@
-from typing import Any, Callable
+from typing import Any
 from unittest.mock import Mock
 
 import pytest
-from truss.server.common.retry import retry
+from truss.templates.server.common.retry import retry
 
 
 class FailForCallCount:
@@ -20,7 +20,7 @@ class FailForCallCount:
         return self._call_count
 
 
-def fail_for_call_count(count: int) -> Callable:
+def fail_for_call_count(count: int) -> callable:
     call_count = 0
 
     def inner():

truss/tests/templates/server/test_model_wrapper.py

@@ -0,0 +1,248 @@
+import importlib
+import os
+import sys
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+from unittest.mock import MagicMock, Mock, patch
+
+import opentelemetry.sdk.trace as sdk_trace
+import pytest
+import yaml
+from starlette.requests import Request
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"
+
+
+@pytest.fixture
+def app_path(truss_container_fs: Path, helpers: Any):
+    truss_container_app_path = truss_container_fs / "app"
+    model_file_content = """
+class Model:
+    def __init__(self):
+        self.load_count = 0
+
+    def load(self):
+        self.load_count += 1
+        if self.load_count <= 2:
+            raise RuntimeError('Simulated error')
+
+    def predict(self, request):
+        return request
+    """
+    with helpers.file_content(
+        truss_container_app_path / "model" / "model.py", model_file_content
+    ), helpers.sys_path(truss_container_app_path):
+        yield truss_container_app_path
+
+
+@pytest.mark.anyio
+async def test_model_wrapper_load_error_once(app_path):
+    if "model_wrapper" in sys.modules:
+        model_wrapper_module = sys.modules["model_wrapper"]
+        importlib.reload(model_wrapper_module)
+    else:
+        model_wrapper_module = importlib.import_module("model_wrapper")
+    model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+    config = yaml.safe_load((app_path / "config.yaml").read_text())
+    os.chdir(app_path)
+    model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+    model_wrapper.load()
+    # Allow load thread to execute
+    time.sleep(1)
+    output = await model_wrapper.predict({}, MagicMock(spec=Request))
+    assert output == {}
+    assert model_wrapper._model.load_count == 2
+
+
+def test_model_wrapper_load_error_more_than_allowed(app_path, helpers):
+    with helpers.env_var("NUM_LOAD_RETRIES_TRUSS", "0"):
+        if "model_wrapper" in sys.modules:
+            model_wrapper_module = sys.modules["model_wrapper"]
+            importlib.reload(model_wrapper_module)
+        else:
+            model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+        os.chdir(app_path)
+        model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+        model_wrapper.load()
+        # Allow load thread to execute
+        time.sleep(1)
+        assert model_wrapper.load_failed
+
+
+@pytest.mark.anyio
+@pytest.mark.integration
+async def test_model_wrapper_streaming_timeout(app_path):
+    if "model_wrapper" in sys.modules:
+        model_wrapper_module = sys.modules["model_wrapper"]
+        importlib.reload(model_wrapper_module)
+    else:
+        model_wrapper_module = importlib.import_module("model_wrapper")
+    model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+
+    # Create an instance of ModelWrapper with streaming_read_timeout set to 5 seconds
+    config = yaml.safe_load((app_path / "config.yaml").read_text())
+    config["runtime"]["streaming_read_timeout"] = 5
+    model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+    model_wrapper.load()
+    assert model_wrapper._config.get("runtime").get("streaming_read_timeout") == 5
+
+
+@pytest.mark.anyio
+async def test_trt_llm_truss_init_extension(trt_llm_truss_container_fs, helpers):
+    app_path = trt_llm_truss_container_fs / "app"
+    packages_path = trt_llm_truss_container_fs / "packages"
+    with _clear_model_load_modules(), helpers.sys_paths(app_path, packages_path):
+        model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+        mock_extension = Mock()
+        mock_extension.load = Mock()
+        with patch.object(
+            model_wrapper_module, "_init_extension", return_value=mock_extension
+        ) as mock_init_extension:
+            model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+            model_wrapper.load()
+            called_with_specific_extension = any(
+                call_args[0][0] == "trt_llm"
+                for call_args in mock_init_extension.call_args_list
+            )
+            assert called_with_specific_extension, (
+                "Expected extension_name was not called"
+            )
+
+
+@pytest.mark.anyio
+async def test_trt_llm_truss_predict(trt_llm_truss_container_fs, helpers):
+    app_path = trt_llm_truss_container_fs / "app"
+    packages_path = trt_llm_truss_container_fs / "packages"
+    with _clear_model_load_modules(), helpers.sys_paths(
+        app_path, packages_path
+    ), _change_directory(app_path):
+        model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+
+        expected_predict_response = "test"
+        mock_predict_called = False
+
+        async def mock_predict(return_value, request):
+            nonlocal mock_predict_called
+            mock_predict_called = True
+            return expected_predict_response
+
+        mock_engine = Mock(predict=mock_predict)
+        mock_extension = Mock()
+        mock_extension.load = Mock()
+        mock_extension.model_args = Mock(return_value={"engine": mock_engine})
+        with patch.object(
+            model_wrapper_module, "_init_extension", return_value=mock_extension
+        ):
+            model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+            model_wrapper.load()
+            resp = await model_wrapper.predict({}, MagicMock(spec=Request))
+            mock_extension.load.assert_called()
+            mock_extension.model_args.assert_called()
+            assert mock_predict_called
+            assert resp == expected_predict_response
+
+
+@pytest.mark.anyio
+async def test_trt_llm_truss_missing_model_py(trt_llm_truss_container_fs, helpers):
+    app_path = trt_llm_truss_container_fs / "app"
+    (app_path / "model" / "model.py").unlink()
+
+    packages_path = trt_llm_truss_container_fs / "packages"
+    with _clear_model_load_modules(), helpers.sys_paths(
+        app_path, packages_path
+    ), _change_directory(app_path):
+        model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+
+        expected_predict_response = "test"
+        mock_predict_called = False
+
+        async def mock_predict(return_value, request: Request):
+            nonlocal mock_predict_called
+            mock_predict_called = True
+            return expected_predict_response
+
+        mock_engine = Mock(predict=mock_predict, spec=["predict"])
+        mock_extension = Mock()
+        mock_extension.load = Mock()
+        mock_extension.model_override = Mock(return_value=mock_engine)
+        with patch.object(
+            model_wrapper_module, "_init_extension", return_value=mock_extension
+        ):
+            model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+            model_wrapper.load()
+            resp = await model_wrapper.predict({}, MagicMock(spec=Request))
+            mock_extension.load.assert_called()
+            mock_extension.model_override.assert_called()
+            assert mock_predict_called
+            assert resp == expected_predict_response
+
+
+@pytest.mark.anyio
+async def test_open_ai_completion_endpoints(open_ai_container_fs, helpers):
+    app_path = open_ai_container_fs / "app"
+    with _clear_model_load_modules(), helpers.sys_paths(app_path), _change_directory(
+        app_path
+    ):
+        model_wrapper_module = importlib.import_module("model_wrapper")
+        model_wrapper_class = getattr(model_wrapper_module, "ModelWrapper")
+        config = yaml.safe_load((app_path / "config.yaml").read_text())
+
+        model_wrapper = model_wrapper_class(config, sdk_trace.NoOpTracer())
+        model_wrapper.load()
+
+        mock_req = MagicMock(spec=Request)
+        predict_resp = await model_wrapper.predict({}, mock_req)
+        assert predict_resp == "predict"
+
+        completions_resp = await model_wrapper.completions({}, mock_req)
+        assert completions_resp == "completions"
+
+        chat_completions_resp = await model_wrapper.chat_completions({}, mock_req)
+        assert chat_completions_resp == "chat_completions"
+
+
+@contextmanager
+def _change_directory(new_directory: Path):
+    original_directory = os.getcwd()
+    os.chdir(str(new_directory))
+    try:
+        yield
+    finally:
+        os.chdir(original_directory)
+
+
+@contextmanager
+def _clear_model_load_modules():
+    """Clear dangling references to model and model_wrapper modules
+
+    We do this before to clear any debris from before, and after to clean up
+    after self. This is meant for cases where we simulate running a truss model
+    in process, where these modules are loaded dyamically.
+    """
+    # These are left over by TrussModuleLoader used by local prediction tests.
+    # TODO(pankaj) Find a way for TrussModuleLoader to clean up after itself.
+    _remove_model_load_sys_modules()
+    yield
+    _remove_model_load_sys_modules()
+
+
+def _remove_model_load_sys_modules():
+    if "model" in sys.modules:
+        del sys.modules["model"]
+    if "model.model" in sys.modules:
+        del sys.modules["model.model"]
+    if "model_wrapper" in sys.modules:
+        del sys.modules["model_wrapper"]

truss/tests/{server → templates/server}/test_schema.py

@@ -2,7 +2,7 @@ import inspect
 from typing import AsyncGenerator, Awaitable, Generator, Union
 
 from pydantic import BaseModel
-from truss.server.common.schema import TrussSchema
+from truss.templates.server.common.schema import TrussSchema
 
 
 class ModelInput(BaseModel):
@@ -181,8 +181,7 @@ def test_truss_schema_union_sync():
 def test_truss_schema_union_async():
     class Model:
         async def predict(
-            self,
-            request: ModelInput,
+            self, request: ModelInput
         ) -> Union[Awaitable[ModelOutput], AsyncGenerator[str, None]]:
             if request.stream:
 
@@ -208,8 +207,7 @@ def test_truss_schema_union_async():
 def test_truss_schema_union_async_non_pydantic():
     class Model:
         async def predict(
-            self,
-            request: ModelInput,
+            self, request: ModelInput
         ) -> Union[Awaitable[str], AsyncGenerator[str, None]]:
             return "hello"
 

truss/tests/{server/core/server/common → templates/server}/test_truss_server.py

@@ -8,18 +8,21 @@ from multiprocessing import Process
 from pathlib import Path
 
 import pytest
-import yaml
-from truss.server.common.truss_server import TrussServer
 
 
 @pytest.mark.integration
-def test_truss_server_termination(tmp_truss_dir):
+def test_truss_server_termination(truss_container_fs):
     port = 10123
 
     def start_truss_server(stdout_capture_file_path):
         sys.stdout = open(stdout_capture_file_path, "w")
-        config = yaml.safe_load((tmp_truss_dir / "config.yaml").read_text())
-        server = TrussServer(http_port=port, config=config)
+        app_path = truss_container_fs / "app"
+        sys.path.append(str(app_path))
+        os.chdir(app_path)
+
+        from truss_server import TrussServer
+
+        server = TrussServer(http_port=port, config_or_path=app_path / "config.yaml")
         server.start()
 
     stdout_capture_file = tempfile.NamedTemporaryFile()

truss/tests/test_build.py

@@ -1,63 +1,20 @@
-from pathlib import Path
-
-from truss.build import init
-from truss.truss_spec import TrussSpec
+from truss.base.truss_spec import TrussSpec
+from truss.truss_handle.build import init_directory, load
+from truss_chains.deployment import code_gen
 
 
 def test_truss_init(tmp_path):
-    dir_name = str(tmp_path)
-    init(dir_name)
-    spec = TrussSpec(Path(dir_name))
+    spec = TrussSpec(init_directory(tmp_path))
     assert spec.model_module_dir.exists()
     assert spec.data_dir.exists()
     assert spec.truss_dir == tmp_path
     assert spec.config_path.exists()
 
 
-def test_truss_init_with_data_file_and_requirements_file_and_bundled_packages(
-    tmp_path,
-):
-    dir_path = tmp_path / "truss"
-    dir_name = str(dir_path)
-
-    # Init data files
-    data_path = tmp_path / "data.txt"
-    with data_path.open("w") as data_file:
-        data_file.write("test")
-
-    # Init requirements file
-    req_file_path = tmp_path / "requirements.txt"
-    requirements = [
-        "tensorflow==2.3.1",
-        "uvicorn==0.12.2",
-    ]
-    with req_file_path.open("w") as req_file:
-        for req in requirements:
-            req_file.write(f"{req}\n")
+def test_truss_init_with_python_dx(tmp_path):
+    init_directory(tmp_path, model_name="Test Model Name", python_config=True)
 
-    # init bundled packages
-    packages_path = tmp_path / "dep_pkg"
-    packages_path.mkdir()
-    packages_path_file_py = packages_path / "file.py"
-    packages_path_init_py = packages_path / "__init__.py"
-    pkg_files = [packages_path_init_py, packages_path_file_py]
-    for pkg_file in pkg_files:
-        with pkg_file.open("w") as fh:
-            fh.write("test")
+    generated_truss_dir = code_gen.gen_truss_model_from_source(tmp_path / "my_model.py")
+    truss_handle = load(generated_truss_dir)
 
-    init(
-        dir_name,
-        data_files=[str(data_path)],
-        requirements_file=str(req_file_path),
-        bundled_packages=[str(packages_path)],
-    )
-    spec = TrussSpec(Path(dir_name))
-    assert spec.model_module_dir.exists()
-    assert spec.truss_dir == dir_path
-    assert spec.config_path.exists()
-    assert spec.data_dir.exists()
-    assert spec.bundled_packages_dir.exists()
-    assert (spec.data_dir / "data.txt").exists()
-    assert spec.requirements == requirements
-    assert (spec.bundled_packages_dir / "dep_pkg" / "__init__.py").exists()
-    assert (spec.bundled_packages_dir / "dep_pkg" / "file.py").exists()
+    assert truss_handle.spec.config.model_name == "Test Model Name"