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

truss/tests/contexts/image_builder/test_serving_image_builder.py

@@ -1,35 +1,45 @@
+import filecmp
+import os
 import time
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest.mock import patch
 
 import pytest
+from truss.base.constants import (
+    BASE_TRTLLM_REQUIREMENTS,
+    TRTLLM_BASE_IMAGE,
+    TRTLLM_PREDICT_CONCURRENCY,
+    TRTLLM_PYTHON_EXECUTABLE,
+    TRTLLM_TRUSS_DIR,
+)
+from truss.base.truss_config import ModelCache, ModelRepo, TrussConfig
 from truss.contexts.image_builder.serving_image_builder import (
     HF_ACCESS_TOKEN_FILE_NAME,
     ServingImageBuilderContext,
     get_files_to_cache,
 )
 from truss.tests.test_testing_utilities_for_other_tests import ensure_kill_all
-from truss.truss_config import ModelCache, ModelRepo, TrussConfig
-from truss.truss_handle import TrussHandle
+from truss.truss_handle.truss_handle import TrussHandle
 
 BASE_DIR = Path(__file__).parent
 
 
-def test_serving_image_dockerfile_from_user_base_image(custom_model_truss_dir):
+def test_serving_image_dockerfile_from_user_base_image(
+    test_data_path, custom_model_truss_dir
+):
     th = TrussHandle(custom_model_truss_dir)
+    # The test fixture python varies with host version, need to pin here.
+    th.update_python_version("py39")
     th.set_base_image("baseten/truss-server-base:3.9-v0.4.3", "/usr/local/bin/python3")
-
-    image_builder = ServingImageBuilderContext.run(th.spec.truss_dir)
+    builder_context = ServingImageBuilderContext
+    image_builder = builder_context.run(th.spec.truss_dir)
     with TemporaryDirectory() as tmp_dir:
         tmp_path = Path(tmp_dir)
         image_builder.prepare_image_build_dir(tmp_path)
         with open(tmp_path / "Dockerfile", "r") as f:
             gen_docker_lines = f.readlines()
-        with open(
-            f"{BASE_DIR}/../../../test_data/server.Dockerfile",
-            "r",
-        ) as f:
+        with open(test_data_path / "server.Dockerfile", "r") as f:
             server_docker_lines = f.readlines()
 
         def filter_empty_lines(lines):
@@ -43,7 +53,8 @@ def test_serving_image_dockerfile_from_user_base_image(custom_model_truss_dir):
 def test_requirements_setup_in_build_dir(custom_model_truss_dir):
     th = TrussHandle(custom_model_truss_dir)
     th.add_python_requirement("numpy")
-    image_builder = ServingImageBuilderContext.run(th.spec.truss_dir)
+    builder_context = ServingImageBuilderContext
+    image_builder = builder_context.run(th.spec.truss_dir)
 
     with TemporaryDirectory() as tmp_dir:
         tmp_path = Path(tmp_dir)
@@ -51,8 +62,10 @@ def test_requirements_setup_in_build_dir(custom_model_truss_dir):
         with open(tmp_path / "requirements.txt", "r") as f:
             requirements_content = f.read()
 
-        # We are no longer adding the base requirements because the server is installed separately.
-        assert requirements_content == "numpy\n"
+        with open(f"{BASE_DIR}/../../../templates/server/requirements.txt", "r") as f:
+            base_requirements_content = f.read()
+
+        assert requirements_content == base_requirements_content + "numpy\n"
 
 
 def flatten_cached_files(local_cache_files):
@@ -62,8 +75,7 @@ def flatten_cached_files(local_cache_files):
 def test_correct_hf_files_accessed_for_caching():
     model = "openai/whisper-small"
     config = TrussConfig(
-        python_version="py39",
-        model_cache=ModelCache(models=[ModelRepo(repo_id=model)]),
+        python_version="py39", model_cache=ModelCache(models=[ModelRepo(repo_id=model)])
     )
 
     with TemporaryDirectory() as tmp_dir:
@@ -98,8 +110,7 @@ def test_correct_gcs_files_accessed_for_caching(mock_list_bucket_files):
     model = "gs://crazy-good-new-model-7b"
 
     config = TrussConfig(
-        python_version="py39",
-        model_cache=ModelCache(models=[ModelRepo(repo_id=model)]),
+        python_version="py39", model_cache=ModelCache(models=[ModelRepo(repo_id=model)])
     )
 
     with TemporaryDirectory() as tmp_dir:
@@ -132,8 +143,7 @@ def test_correct_s3_files_accessed_for_caching(mock_list_bucket_files):
     model = "s3://crazy-good-new-model-7b"
 
     config = TrussConfig(
-        python_version="py39",
-        model_cache=ModelCache(models=[ModelRepo(repo_id=model)]),
+        python_version="py39", model_cache=ModelCache(models=[ModelRepo(repo_id=model)])
     )
 
     with TemporaryDirectory() as tmp_dir:
@@ -166,8 +176,7 @@ def test_correct_nested_gcs_files_accessed_for_caching(mock_list_bucket_files):
     model = "gs://crazy-good-new-model-7b/folder_a/folder_b"
 
     config = TrussConfig(
-        python_version="py39",
-        model_cache=ModelCache(models=[ModelRepo(repo_id=model)]),
+        python_version="py39", model_cache=ModelCache(models=[ModelRepo(repo_id=model)])
     )
 
     with TemporaryDirectory() as tmp_dir:
@@ -204,8 +213,7 @@ def test_correct_nested_s3_files_accessed_for_caching(mock_list_bucket_files):
     model = "s3://crazy-good-new-model-7b/folder_a/folder_b"
 
     config = TrussConfig(
-        python_version="py39",
-        model_cache=ModelCache(models=[ModelRepo(repo_id=model)]),
+        python_version="py39", model_cache=ModelCache(models=[ModelRepo(repo_id=model)])
     )
 
     with TemporaryDirectory() as tmp_dir:
@@ -234,24 +242,20 @@ def test_correct_nested_s3_files_accessed_for_caching(mock_list_bucket_files):
 
 
 @pytest.mark.integration
-def test_truss_server_caching_truss():
+def test_truss_server_caching_truss(test_data_path):
     with ensure_kill_all():
-        truss_root = (
-            Path(__file__).parent.parent.parent.parent.parent.resolve() / "truss"
-        )
-        truss_dir = truss_root / "test_data" / "test_truss_server_caching_truss"
+        truss_dir = test_data_path / "test_truss_server_caching_truss"
         tr = TrussHandle(truss_dir)
 
         container = tr.docker_run(
-            local_port=8090, detach=True, wait_for_server_ready=True
+            local_port=8090, detach=True, wait_for_server_ready=True, network="host"
         )
         time.sleep(15)
         assert "Downloading model.safetensors:" not in container.logs()
 
 
-def test_model_cache_dockerfile():
-    truss_root = Path(__file__).parent.parent.parent.parent.parent.resolve() / "truss"
-    truss_dir = truss_root / "test_data" / "test_truss_server_caching_truss"
+def test_model_cache_dockerfile(test_data_path):
+    truss_dir = test_data_path / "test_truss_server_caching_truss"
     tr = TrussHandle(truss_dir)
 
     builder_context = ServingImageBuilderContext
@@ -285,3 +289,42 @@ def test_ignore_files_during_build_setup(custom_model_truss_dir_with_truss_ignor
 
         assert not (build_path / ignore_folder).exists()
         assert (build_path / do_not_ignore_folder).exists()
+
+
+def test_trt_llm_build_dir(custom_model_trt_llm):
+    th = TrussHandle(custom_model_trt_llm)
+    builder_context = ServingImageBuilderContext
+    image_builder = builder_context.run(th.spec.truss_dir)
+    with TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+        image_builder.prepare_image_build_dir(tmp_path)
+        build_th = TrussHandle(tmp_path)
+
+        # Check that all files were copied
+        _assert_copied(
+            TRTLLM_TRUSS_DIR / "src", tmp_path / "server" / "extensions" / "trt_llm"
+        )
+        _assert_copied(TRTLLM_TRUSS_DIR / "packages", tmp_path / "packages")
+
+        assert (
+            build_th.spec.config.runtime.predict_concurrency
+            == TRTLLM_PREDICT_CONCURRENCY
+        )
+        assert build_th.spec.config.base_image.image == TRTLLM_BASE_IMAGE
+        assert (
+            build_th.spec.config.base_image.python_executable_path
+            == TRTLLM_PYTHON_EXECUTABLE
+        )
+        assert BASE_TRTLLM_REQUIREMENTS == build_th.spec.config.requirements
+
+
+def _assert_copied(src_path: str, dest_path: str):
+    for dirpath, dirnames, filenames in os.walk(src_path):
+        rel_path = os.path.relpath(dirpath, src_path)
+        for filename in filenames:
+            src_file = os.path.join(dirpath, filename)
+            dest_file = os.path.join(dest_path, rel_path, filename)
+            assert os.path.exists(dest_file), f"{dest_file} was not copied"
+            assert filecmp.cmp(src_file, dest_file, shallow=False), (
+                f"{src_file} and {dest_file} are not the same"
+            )

truss/tests/contexts/local_loader/test_load_local.py

@@ -1,7 +1,7 @@
+from truss.base.truss_spec import TrussSpec
 from truss.contexts.local_loader.utils import prepare_secrets
 from truss.local.local_config_handler import LocalConfigHandler
-from truss.truss_handle import TrussHandle
-from truss.truss_spec import TrussSpec
+from truss.truss_handle.truss_handle import TrussHandle
 
 
 def test_prepare_secrets(custom_model_truss_dir, tmp_path):

truss/tests/contexts/local_loader/test_truss_module_finder.py

@@ -1,8 +1,8 @@
 import tempfile
 from pathlib import Path
 
+from truss.base.truss_config import DEFAULT_BUNDLED_PACKAGES_DIR
 from truss.contexts.local_loader.truss_module_loader import truss_module_loaded
-from truss.truss_config import DEFAULT_BUNDLED_PACKAGES_DIR
 
 ORIG_MODEL_CLASS_CONTENT = """
 class Model: