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

truss/tests/conftest.py

@@ -1,4 +1,5 @@
 import contextlib
+import copy
 import importlib
 import os
 import shutil
@@ -6,15 +7,21 @@ import subprocess
 import sys
 import time
 from pathlib import Path
-from typing import Callable, Generator
+from typing import Any, Dict
 
 import pytest
 import requests
 import yaml
-from truss.build import init
-from truss.truss_config import DEFAULT_BUNDLED_PACKAGES_DIR
-from truss.types import Example
-from truss.util.path import copy_tree_path
+
+from truss.base.custom_types import Example
+from truss.base.trt_llm_config import TrussSpecDecMode, TrussTRTLLMBatchSchedulerPolicy
+from truss.base.truss_config import DEFAULT_BUNDLED_PACKAGES_DIR, Accelerator
+from truss.contexts.image_builder.serving_image_builder import (
+    ServingImageBuilderContext,
+)
+from truss.contexts.local_loader.docker_build_emulator import DockerBuildEmulator
+from truss.truss_handle.build import init_directory
+from truss.truss_handle.truss_handle import TrussHandle
 
 CUSTOM_MODEL_CODE = """
 class Model:
@@ -203,10 +210,22 @@ class Model:
             return file.read()
 """
 
+CUSTOM_MODEL_TRT_LLM_CODE = """
+class Model:
+    def __init__(trt_llm, *args, **kwargs):
+        pass
+
+    def load(self):
+        pass
+
+    def predict(self, model_input):
+        return [1 for i in model_input]
+"""
+
 
 @pytest.fixture
-def temp_path(tmpdir):
-    yield Path(tmpdir)
+def test_data_path() -> Path:
+    return Path(__file__).parent.resolve() / "test_data"
 
 
 @pytest.fixture
@@ -215,30 +234,20 @@ def pytorch_model_init_args():
 
 
 @pytest.fixture
-def custom_model_truss_dir(tmp_path) -> Generator[Path, None, None]:
-    yield _custom_model_from_code(
-        tmp_path,
-        "custom_truss",
-        CUSTOM_MODEL_CODE,
-    )
+def custom_model_truss_dir(tmp_path) -> Path:
+    yield _custom_model_from_code(tmp_path, "custom_truss", CUSTOM_MODEL_CODE)
 
 
 @pytest.fixture
 def no_preprocess_custom_model(tmp_path):
     yield _custom_model_from_code(
-        tmp_path,
-        "my_no_preprocess_model",
-        NO_PREPROCESS_CUSTOM_MODEL_CODE,
+        tmp_path, "my_no_preprocess_model", NO_PREPROCESS_CUSTOM_MODEL_CODE
     )
 
 
 @pytest.fixture
 def long_load_model(tmp_path):
-    yield _custom_model_from_code(
-        tmp_path,
-        "long_load_model",
-        LONG_LOAD_MODEL_CODE,
-    )
+    yield _custom_model_from_code(tmp_path, "long_load_model", LONG_LOAD_MODEL_CODE)
 
 
 @pytest.fixture
@@ -258,8 +267,7 @@ def custom_model_external_data_access_tuple_fixture(tmp_path: Path):
     (tmp_path / filename).write_text(content)
     port = 9089
     proc = subprocess.Popen(
-        ["python", "-m", "http.server", str(port), "--bind", "*"],
-        cwd=tmp_path,
+        ["python", "-m", "http.server", str(port), "--bind", "*"], cwd=tmp_path
     )
     try:
         url = f"http://localhost:{port}/{filename}"
@@ -289,8 +297,7 @@ def custom_model_external_data_access_tuple_fixture_gpu(tmp_path: Path):
     (tmp_path / filename).write_text(content)
     port = 9089
     proc = subprocess.Popen(
-        ["python", "-m", "http.server", str(port), "--bind", "*"],
-        cwd=tmp_path,
+        ["python", "-m", "http.server", str(port), "--bind", "*"], cwd=tmp_path
     )
     try:
         url = f"http://localhost:{port}/{filename}"
@@ -346,27 +353,53 @@ def custom_model_with_external_package(tmp_path: Path):
 @pytest.fixture
 def no_postprocess_custom_model(tmp_path):
     yield _custom_model_from_code(
-        tmp_path,
-        "my_no_postprocess_model",
-        NO_POSTPROCESS_CUSTOM_MODEL_CODE,
+        tmp_path, "my_no_postprocess_model", NO_POSTPROCESS_CUSTOM_MODEL_CODE
     )
 
 
 @pytest.fixture
 def no_load_custom_model(tmp_path):
     yield _custom_model_from_code(
-        tmp_path,
-        "my_no_load_model",
-        NO_LOAD_CUSTOM_MODEL_CODE,
+        tmp_path, "my_no_load_model", NO_LOAD_CUSTOM_MODEL_CODE
     )
 
 
 @pytest.fixture
 def no_params_init_custom_model(tmp_path):
+    yield _custom_model_from_code(
+        tmp_path, "my_no_params_init_load_model", NO_PARAMS_INIT_CUSTOM_MODEL_CODE
+    )
+
+
+@pytest.fixture
+def custom_model_trt_llm(tmp_path):
+    def modify_handle(h: TrussHandle):
+        with _modify_yaml(h.spec.config_path) as content:
+            h.enable_gpu()
+            content["trt_llm"] = {
+                "build": {
+                    "base_model": "llama",
+                    "max_seq_len": 2048,
+                    "max_batch_size": 512,
+                    "checkpoint_repository": {
+                        "source": "HF",
+                        "repo": "meta/llama4-500B",
+                    },
+                },
+                "runtime": {
+                    "kv_cache_free_gpu_mem_fraction": 0.9,
+                    "kv_cache_host_memory_bytes": 1000,
+                    "enabled_chunked_context": True,
+                    "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.GUARANTEED_NO_EVICT.value,
+                },
+            }
+            content["resources"]["accelerator"] = "H100:1"
+
     yield _custom_model_from_code(
         tmp_path,
-        "my_no_params_init_load_model",
-        NO_PARAMS_INIT_CUSTOM_MODEL_CODE,
+        "my_trt_llm_model",
+        CUSTOM_MODEL_TRT_LLM_CODE,
+        handle_ops=modify_handle,
     )
 
 
@@ -390,23 +423,30 @@ def temp_dir(directory):
         os.chdir(current_dir)
 
 
+@pytest.fixture
+def dynamic_config_mount_dir(tmp_path, monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "truss.templates.shared.dynamic_config_resolver.DYNAMIC_CONFIG_MOUNT_DIR",
+        str(tmp_path),
+    )
+    yield
+
+
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_no_example(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post_no_example"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_with_hidden_files(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_hidden_files"
-    _ = init(str(truss_dir_path))
+    init_directory(truss_dir_path)
     (truss_dir_path / "__pycache__").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / ".git").mkdir(parents=True, exist_ok=True)
-    (truss_dir_path / "__pycache__" / "test.cpython-38.pyc").touch()
+    (truss_dir_path / "__pycache__" / "test.cpython-311.pyc").touch()
     (truss_dir_path / ".DS_Store").touch()
     (truss_dir_path / ".git" / ".test_file").touch()
     (truss_dir_path / "data" / "test_file").write_text("123456789")
@@ -416,7 +456,7 @@ def custom_model_truss_dir_with_hidden_files(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_truss_ignore(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_truss_ignore"
-    _ = init(str(truss_dir_path))
+    init_directory(truss_dir_path)
     (truss_dir_path / "random_folder_1").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / "random_folder_2").mkdir(parents=True, exist_ok=True)
     (truss_dir_path / "random_file_1.txt").touch()
@@ -436,19 +476,17 @@ def custom_model_truss_dir_with_truss_ignore(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_examples([Example("example1", {"inputs": [[0]]})])
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_examples([Example("example1", {"inputs": [[0]]})])
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_with_bundled_packages(tmp_path):
     truss_dir_path: Path = tmp_path / "custom_model_truss_dir_with_bundled_packages"
-    handle = init(str(truss_dir_path))
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_USING_BUNDLED_PACKAGE)
+    th = TrussHandle(init_directory(truss_dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_USING_BUNDLED_PACKAGE)
     packages_path = truss_dir_path / DEFAULT_BUNDLED_PACKAGES_DIR / "test_package"
     packages_path.mkdir(parents=True)
     with (packages_path / "test.py").open("w") as file:
@@ -459,11 +497,9 @@ def custom_model_truss_dir_with_bundled_packages(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_str_example(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post_str_example"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_examples(
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_examples(
         [
             Example(
                 "example1",
@@ -471,8 +507,8 @@ def custom_model_truss_dir_with_pre_and_post_str_example(tmp_path):
                     "inputs": [
                         {
                             "image_url": "https://github.com/pytorch/hub/raw/master/images/dog.jpg"
-                        },
-                    ],
+                        }
+                    ]
                 },
             )
         ]
@@ -483,57 +519,57 @@ def custom_model_truss_dir_with_pre_and_post_str_example(tmp_path):
 @pytest.fixture
 def custom_model_truss_dir_with_pre_and_post_description(tmp_path):
     dir_path = tmp_path / "custom_truss_with_pre_post"
-    handle = init(str(dir_path))
-    handle.spec.model_class_filepath.write_text(
-        CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS
-    )
-    handle.update_description("This model adds 3 to all inputs")
+    th = TrussHandle(init_directory(dir_path))
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_WITH_PRE_AND_POST_PROCESS)
+    th.update_description("This model adds 3 to all inputs")
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_for_gpu(tmp_path):
     dir_path = tmp_path / "custom_truss"
-    handle = init(str(dir_path))
-    handle.enable_gpu()
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_GPU_TESTING)
+    th = TrussHandle(init_directory(dir_path))
+    th.enable_gpu()
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_GPU_TESTING)
     yield dir_path
 
 
 @pytest.fixture
 def custom_model_truss_dir_for_secrets(tmp_path):
     dir_path = tmp_path / "custom_truss"
-    handle = init(str(dir_path))
-    handle.add_secret("secret_name", "default_secret_value")
-    handle.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_SECRETS_TESTING)
+    th = TrussHandle(init_directory(dir_path))
+    th.add_secret("secret_name", "default_secret_value")
+    th.spec.model_class_filepath.write_text(CUSTOM_MODEL_CODE_FOR_SECRETS_TESTING)
     yield dir_path
 
 
 @pytest.fixture
-def tmp_truss_dir(tmp_path, monkeypatch):
-    ROOT = Path(__file__).parent.parent.parent.resolve()
-    tmp_dir = _copy_truss_dir_to_tmp(
-        ROOT / "truss" / "test_data" / "test_truss", tmp_path
-    )
-    monkeypatch.setenv("APP_HOME", str(tmp_dir))
-    return tmp_dir
+def truss_container_fs(tmp_path, test_data_path):
+    return _build_truss_fs(test_data_path / "test_truss", tmp_path)
+
+
+@pytest.fixture
+def trt_llm_truss_container_fs(tmp_path, test_data_path):
+    return _build_truss_fs(test_data_path / "test_trt_llm_truss", tmp_path)
+
+
+@pytest.fixture
+def open_ai_container_fs(tmp_path, test_data_path):
+    return _build_truss_fs(test_data_path / "test_openai", tmp_path)
 
 
 @pytest.fixture
-def tmp_truss_control_dir(tmp_path, monkeypatch):
-    ROOT = Path(__file__).parent.parent.parent.resolve()
-    test_truss_dir = ROOT / "truss" / "test_data" / "test_truss"
+def truss_control_container_fs(tmp_path, test_data_path):
+    test_truss_dir = test_data_path / "test_truss"
     control_truss_dir = tmp_path / "control_truss"
     shutil.copytree(str(test_truss_dir), str(control_truss_dir))
     with _modify_yaml(control_truss_dir / "config.yaml") as content:
         content["live_reload"] = True
-    tmp_dir = _copy_truss_dir_to_tmp(control_truss_dir, tmp_path)
-    monkeypatch.setenv("APP_HOME", str(tmp_dir))
-    return tmp_dir
+    return _build_truss_fs(control_truss_dir, tmp_path)
 
 
 @pytest.fixture
-def patch_ping_test_server():
+def patch_ping_test_server(test_data_path):
     port = "5001"
     proc = subprocess.Popen(
         [
@@ -548,7 +584,7 @@ def patch_ping_test_server():
             "--host",
             "0.0.0.0",
         ],
-        cwd=str(Path(__file__).parent.parent / "test_data" / "patch_ping_test_server"),
+        cwd=str(test_data_path / "patch_ping_test_server"),
     )
     base_url = f"http://127.0.0.1:{port}"
     retry_secs = 10
@@ -582,16 +618,13 @@ def _pytorch_model_from_content(
 
 
 def _custom_model_from_code(
-    where_dir: Path,
-    truss_name: str,
-    model_code: str,
-    handle_ops: Callable = None,
+    where_dir: Path, truss_name: str, model_code: str, handle_ops: callable = None
 ) -> Path:
     dir_path = where_dir / truss_name
-    handle = init(str(dir_path))
+    th = TrussHandle(init_directory(dir_path))
     if handle_ops is not None:
-        handle_ops(handle)
-    handle.spec.model_class_filepath.write_text(model_code)
+        handle_ops(th)
+    th.spec.model_class_filepath.write_text(model_code)
     return dir_path
 
 
@@ -604,10 +637,7 @@ def custom_model_data_dir(tmp_path: Path):
         handle.add_data(str(data_file.resolve()))
 
     yield _custom_model_from_code(
-        tmp_path,
-        "data_dir_truss",
-        CUSTOM_MODEL_CODE,
-        handle_ops=add_data,
+        tmp_path, "data_dir_truss", CUSTOM_MODEL_CODE, handle_ops=add_data
     )
 
 
@@ -631,6 +661,18 @@ class Helpers:
         finally:
             sys.path.pop()
 
+    @staticmethod
+    @contextlib.contextmanager
+    def sys_paths(*paths: Path):
+        num_paths = len(paths)
+        try:
+            for path in paths:
+                sys.path.append(str(path))
+            yield
+        finally:
+            for _ in range(num_paths):
+                sys.path.pop()
+
     @staticmethod
     @contextlib.contextmanager
     def env_var(var: str, value: str):
@@ -648,9 +690,18 @@ def helpers():
     return Helpers()
 
 
-def _copy_truss_dir_to_tmp(truss_dir: Path, tmp_path: Path) -> Path:
-    copy_tree_path(truss_dir, tmp_path)
-    return tmp_path
+def _build_truss_fs(truss_dir: Path, tmp_path: Path) -> Path:
+    truss_fs = tmp_path / "truss_fs"
+    truss_fs.mkdir()
+    truss_build_dir = tmp_path / "truss_fs_build"
+    truss_build_dir.mkdir()
+    image_builder = ServingImageBuilderContext.run(truss_dir)
+    image_builder.prepare_image_build_dir(truss_build_dir)
+    dockerfile_path = truss_build_dir / "Dockerfile"
+
+    docker_build_emulator = DockerBuildEmulator(dockerfile_path, truss_build_dir)
+    docker_build_emulator.run(truss_fs)
+    return truss_fs
 
 
 @contextlib.contextmanager
@@ -660,3 +711,160 @@ def _modify_yaml(yaml_path: Path):
     yield content
     with yaml_path.open("w") as yaml_file:
         yaml.dump(content, yaml_file)
+
+
+@pytest.fixture
+def default_config() -> Dict[str, Any]:
+    return {
+        "build_commands": [],
+        "environment_variables": {},
+        "external_package_dirs": [],
+        "model_metadata": {},
+        "model_name": None,
+        "python_version": "py39",
+        "requirements": [],
+        "resources": {
+            "accelerator": None,
+            "cpu": "1",
+            "memory": "2Gi",
+            "use_gpu": False,
+        },
+        "secrets": {},
+        "system_packages": [],
+    }
+
+
+@pytest.fixture
+def trtllm_config(default_config) -> Dict[str, Any]:
+    trtllm_config = default_config
+    trtllm_config["resources"] = {
+        "accelerator": Accelerator.L4.value,
+        "cpu": "1",
+        "memory": "24Gi",
+        "use_gpu": True,
+        "node_count": 1,
+    }
+    trtllm_config["trt_llm"] = {
+        "build": {
+            "base_model": "llama",
+            "max_seq_len": 2048,
+            "max_batch_size": 512,
+            "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            "gather_all_token_logits": False,
+        },
+        "runtime": {},
+    }
+    return trtllm_config
+
+
+@pytest.fixture
+def deprecated_trtllm_config(default_config) -> Dict[str, Any]:
+    trtllm_config = default_config
+    trtllm_config["resources"] = {
+        "accelerator": Accelerator.L4.value,
+        "cpu": "1",
+        "memory": "24Gi",
+        "use_gpu": True,
+    }
+    trtllm_config["trt_llm"] = {
+        "build": {
+            "base_model": "llama",
+            "max_seq_len": 2048,
+            "max_batch_size": 512,
+            # start deprecated fields
+            "kv_cache_free_gpu_mem_fraction": 0.1,
+            "enable_chunked_context": True,
+            "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.MAX_UTILIZATION.value,
+            "request_default_max_tokens": 10,
+            "total_token_limit": 50,
+            # end deprecated fields
+            "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            "gather_all_token_logits": False,
+        }
+    }
+    return trtllm_config
+
+
+@pytest.fixture
+def deprecated_trtllm_config_with_runtime_existing(default_config) -> Dict[str, Any]:
+    trtllm_config = default_config
+    trtllm_config["resources"] = {
+        "accelerator": Accelerator.L4.value,
+        "cpu": "1",
+        "memory": "24Gi",
+        "use_gpu": True,
+    }
+    trtllm_config["trt_llm"] = {
+        "build": {
+            "base_model": "llama",
+            "max_seq_len": 2048,
+            "max_batch_size": 512,
+            # start deprecated fields
+            "kv_cache_free_gpu_mem_fraction": 0.1,
+            "enable_chunked_context": True,
+            "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.MAX_UTILIZATION.value,
+            "request_default_max_tokens": 10,
+            "total_token_limit": 50,
+            # end deprecated fields
+            "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            "gather_all_token_logits": False,
+        },
+        "runtime": {"total_token_limit": 100},
+    }
+    return trtllm_config
+
+
+@pytest.fixture
+def trtllm_spec_dec_config_full(trtllm_config) -> Dict[str, Any]:
+    spec_dec_config = copy.deepcopy(trtllm_config)
+    spec_dec_config["trt_llm"] = {
+        "build": {
+            "base_model": "llama",
+            "max_seq_len": 2048,
+            "max_batch_size": 512,
+            "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            "plugin_configuration": {
+                "paged_kv_cache": True,
+                "gemm_plugin": "auto",
+                "use_paged_context_fmha": True,
+            },
+            "speculator": {
+                "speculative_decoding_mode": TrussSpecDecMode.DRAFT_EXTERNAL.value,
+                "num_draft_tokens": 4,
+                "build": {
+                    "base_model": "llama",
+                    "max_seq_len": 2048,
+                    "max_batch_size": 512,
+                    "checkpoint_repository": {
+                        "source": "HF",
+                        "repo": "meta/llama4-500B",
+                    },
+                },
+            },
+        }
+    }
+    return spec_dec_config
+
+
+@pytest.fixture
+def trtllm_spec_dec_config(trtllm_config) -> Dict[str, Any]:
+    spec_dec_config = copy.deepcopy(trtllm_config)
+    spec_dec_config["trt_llm"] = {
+        "build": {
+            "base_model": "llama",
+            "max_seq_len": 2048,
+            "max_batch_size": 512,
+            "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            "plugin_configuration": {
+                "paged_kv_cache": True,
+                "gemm_plugin": "auto",
+                "use_paged_context_fmha": True,
+            },
+            "speculator": {
+                "speculative_decoding_mode": TrussSpecDecMode.DRAFT_EXTERNAL.value,
+                "num_draft_tokens": 4,
+                "checkpoint_repository": {"source": "HF", "repo": "meta/llama4-500B"},
+            },
+        }
+    }
+    return spec_dec_config