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

truss/tests/test_config.py

@@ -1,21 +1,28 @@
+import copy
 import tempfile
+from contextlib import nullcontext as does_not_raise
 from pathlib import Path
 
 import pytest
 import yaml
-from truss.truss_config import (
+
+from truss.base.custom_types import ModelFrameworkType
+from truss.base.trt_llm_config import TrussTRTLLMQuantizationType
+from truss.base.truss_config import (
     DEFAULT_CPU,
     DEFAULT_MEMORY,
     DEFAULT_USE_GPU,
     Accelerator,
     AcceleratorSpec,
     BaseImage,
+    DockerAuthSettings,
+    DockerAuthType,
     ModelCache,
     ModelRepo,
     Resources,
     TrussConfig,
 )
-from truss.types import ModelFrameworkType
+from truss.truss_handle.truss_handle import TrussHandle
 
 
 @pytest.mark.parametrize(
@@ -71,6 +78,17 @@ from truss.types import ModelFrameworkType
                 "accelerator": "A10G:4",
             },
         ),
+        (
+            {"node_count": 2},
+            Resources(node_count=2),
+            {
+                "cpu": DEFAULT_CPU,
+                "memory": DEFAULT_MEMORY,
+                "use_gpu": False,
+                "accelerator": None,
+                "node_count": 2,
+            },
+        ),
     ],
 )
 def test_parse_resources(input_dict, expect_resources, output_dict):
@@ -87,6 +105,8 @@ def test_parse_resources(input_dict, expect_resources, output_dict):
         ("A10G:4", AcceleratorSpec(Accelerator.A10G, 4)),
         ("A100:8", AcceleratorSpec(Accelerator.A100, 8)),
         ("H100", AcceleratorSpec(Accelerator.H100, 1)),
+        ("H200", AcceleratorSpec(Accelerator.H200, 1)),
+        ("H100_40GB", AcceleratorSpec(Accelerator.H100_40GB, 1)),
     ],
 )
 def test_acc_spec_from_str(input_str, expected_acc):
@@ -99,10 +119,7 @@ def test_acc_spec_from_str(input_str, expected_acc):
         (
             {},
             BaseImage(),
-            {
-                "image": "",
-                "python_executable_path": "",
-            },
+            {"image": "", "python_executable_path": "", "docker_auth": None},
         ),
         (
             {"image": "custom_base_image", "python_executable_path": "/path/python"},
@@ -110,6 +127,36 @@ def test_acc_spec_from_str(input_str, expected_acc):
             {
                 "image": "custom_base_image",
                 "python_executable_path": "/path/python",
+                "docker_auth": None,
+            },
+        ),
+        (
+            {
+                "image": "custom_base_image",
+                "python_executable_path": "/path/python",
+                "docker_auth": {
+                    "auth_method": "GCP_SERVICE_ACCOUNT_JSON",
+                    "secret_name": "some-secret-name",
+                    "registry": "some-docker-registry",
+                },
+            },
+            BaseImage(
+                image="custom_base_image",
+                python_executable_path="/path/python",
+                docker_auth=DockerAuthSettings(
+                    auth_method=DockerAuthType.GCP_SERVICE_ACCOUNT_JSON,
+                    secret_name="some-secret-name",
+                    registry="some-docker-registry",
+                ),
+            ),
+            {
+                "image": "custom_base_image",
+                "python_executable_path": "/path/python",
+                "docker_auth": {
+                    "auth_method": "GCP_SERVICE_ACCOUNT_JSON",
+                    "secret_name": "some-secret-name",
+                    "registry": "some-docker-registry",
+                },
             },
         ),
     ],
@@ -120,33 +167,11 @@ def test_parse_base_image(input_dict, expect_base_image, output_dict):
     assert parsed_result.to_dict() == output_dict
 
 
-def generate_default_config():
-    config = {
-        "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": [],
-    }
-    return config
-
-
 def test_default_config_not_crowded_end_to_end():
-    config = TrussConfig(
-        python_version="py39",
-        requirements=[],
-    )
+    config = TrussConfig(python_version="py39", requirements=[])
 
-    config_yaml = """environment_variables: {}
+    config_yaml = """build_commands: []
+environment_variables: {}
 external_package_dirs: []
 model_metadata: {}
 model_name: null
@@ -168,14 +193,12 @@ system_packages: []
     "model_framework",
     [ModelFrameworkType.CUSTOM, ModelFrameworkType.SKLEARN, ModelFrameworkType.PYTORCH],
 )
-def test_model_framework(model_framework):
+def test_model_framework(model_framework, default_config):
     config = TrussConfig(
-        python_version="py39",
-        requirements=[],
-        model_framework=model_framework,
+        python_version="py39", requirements=[], model_framework=model_framework
     )
 
-    new_config = generate_default_config()
+    new_config = default_config
     if model_framework == ModelFrameworkType.CUSTOM:
         assert new_config == config.to_dict(verbose=False)
     else:
@@ -199,19 +222,15 @@ def test_null_hf_cache_key():
     assert config.model_cache == ModelCache.from_list([])
 
 
-def test_huggingface_cache_single_model_default_revision():
+def test_huggingface_cache_single_model_default_revision(default_config):
     config = TrussConfig(
         python_version="py39",
         requirements=[],
         model_cache=ModelCache(models=[ModelRepo("test/model")]),
     )
 
-    new_config = generate_default_config()
-    new_config["model_cache"] = [
-        {
-            "repo_id": "test/model",
-        }
-    ]
+    new_config = default_config
+    new_config["model_cache"] = [{"repo_id": "test/model"}]
 
     assert new_config == config.to_dict(verbose=False)
     assert config.to_dict(verbose=True)["model_cache"][0].get("revision") is None
@@ -227,24 +246,19 @@ def test_huggingface_cache_single_model_non_default_revision():
     assert config.to_dict(verbose=False)["model_cache"][0].get("revision") == "not-main"
 
 
-def test_huggingface_cache_multiple_models_default_revision():
+def test_huggingface_cache_multiple_models_default_revision(default_config):
     config = TrussConfig(
         python_version="py39",
         requirements=[],
         model_cache=ModelCache(
-            models=[
-                ModelRepo("test/model1", "main"),
-                ModelRepo("test/model2"),
-            ]
+            models=[ModelRepo("test/model1", "main"), ModelRepo("test/model2")]
         ),
     )
 
-    new_config = generate_default_config()
+    new_config = default_config
     new_config["model_cache"] = [
         {"repo_id": "test/model1", "revision": "main"},
-        {
-            "repo_id": "test/model2",
-        },
+        {"repo_id": "test/model2"},
     ]
 
     assert new_config == config.to_dict(verbose=False)
@@ -252,23 +266,18 @@ def test_huggingface_cache_multiple_models_default_revision():
     assert config.to_dict(verbose=True)["model_cache"][1].get("revision") is None
 
 
-def test_huggingface_cache_multiple_models_mixed_revision():
+def test_huggingface_cache_multiple_models_mixed_revision(default_config):
     config = TrussConfig(
         python_version="py39",
         requirements=[],
         model_cache=ModelCache(
-            models=[
-                ModelRepo("test/model1"),
-                ModelRepo("test/model2", "not-main2"),
-            ]
+            models=[ModelRepo("test/model1"), ModelRepo("test/model2", "not-main2")]
         ),
     )
 
-    new_config = generate_default_config()
+    new_config = default_config
     new_config["model_cache"] = [
-        {
-            "repo_id": "test/model1",
-        },
+        {"repo_id": "test/model1"},
         {"repo_id": "test/model2", "revision": "not-main2"},
     ]
 
@@ -277,37 +286,287 @@ def test_huggingface_cache_multiple_models_mixed_revision():
     assert config.to_dict(verbose=True)["model_cache"][1].get("revision") == "not-main2"
 
 
-def test_empty_config():
+def test_empty_config(default_config):
     config = TrussConfig()
-    new_config = generate_default_config()
+    new_config = default_config
 
     assert new_config == config.to_dict(verbose=False)
 
 
 def test_from_yaml():
-    yaml_path = Path("test.yaml")
     data = {"description": "this is a test"}
-    with yaml_path.open("w") as yaml_file:
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
         yaml.safe_dump(data, yaml_file)
 
-    result = TrussConfig.from_yaml(yaml_path)
-
-    assert result.description == "this is a test"
+        result = TrussConfig.from_yaml(yaml_path)
 
-    yaml_path.unlink()
+        assert result.description == "this is a test"
 
 
 def test_from_yaml_empty():
-    yaml_path = Path("test.yaml")
     data = {}
-    with yaml_path.open("w") as yaml_file:
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
         yaml.safe_dump(data, yaml_file)
 
-    result = TrussConfig.from_yaml(yaml_path)
+        result = TrussConfig.from_yaml(yaml_path)
 
-    # test some attributes (should be default)
-    assert result.description is None
-    assert result.spec_version == "2.0"
-    assert result.bundled_packages_dir == "packages"
+        # test some attributes (should be default)
+        assert result.description is None
+        assert result.spec_version == "2.0"
+        assert result.bundled_packages_dir == "packages"
+
+
+def test_from_yaml_secrets_as_list():
+    data = {"description": "this is a test", "secrets": ["foo", "bar"]}
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(data, yaml_file)
 
-    yaml_path.unlink()
+        with pytest.raises(ValueError):
+            TrussConfig.from_yaml(yaml_path)
+
+
+def test_from_yaml_python_version():
+    invalid_py_version_data = {
+        "description": "this is a test",
+        "python_version": "py37",
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(invalid_py_version_data, yaml_file)
+
+        with pytest.raises(ValueError):
+            TrussConfig.from_yaml(yaml_path)
+
+    valid_py_version_data = {"description": "this is a test", "python_version": "py39"}
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(valid_py_version_data, yaml_file)
+
+        result = TrussConfig.from_yaml(yaml_path)
+        assert result.python_version == "py39"
+
+
+def test_from_yaml_environment_variables():
+    data = {
+        "description": "this is a test",
+        "environment_variables": {"foo": "bar", "bool": True, "int": 0},
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(data, yaml_file)
+
+        result = TrussConfig.from_yaml(yaml_path)
+        assert result.environment_variables == {
+            "foo": "bar",
+            "bool": "true",
+            "int": "0",
+        }
+
+
+def test_secret_to_path_mapping_correct_type(default_config):
+    data = {
+        "description": "this is a test",
+        "build": {"secret_to_path_mapping": {"foo": "/bar"}},
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(data, yaml_file)
+
+        truss_config = TrussConfig.from_yaml(yaml_path)
+        assert truss_config.build.secret_to_path_mapping == {"foo": "/bar"}
+
+
+def test_secret_to_path_mapping_invalid_secret_name(default_config):
+    data = {
+        "description": "this is a test",
+        "build": {"secret_to_path_mapping": {"foo_bar": "/bar"}},
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(data, yaml_file)
+
+        with pytest.raises(ValueError):
+            TrussConfig.from_yaml(yaml_path)
+
+
+def test_secret_to_path_mapping_incorrect_type(default_config):
+    data = {
+        "description": "this is a test",
+        "build": {"secret_to_path_mapping": ["something else"]},
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(data, yaml_file)
+
+        with pytest.raises(ValueError):
+            TrussConfig.from_yaml(yaml_path)
+
+
+def test_max_beam_width_check(trtllm_config):
+    trtllm_config["trt_llm"]["build"]["max_beam_width"] = 2
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+
+
+def test_plugin_paged_context_fmha_check(trtllm_config):
+    trtllm_config["trt_llm"]["build"]["plugin_configuration"] = {
+        "paged_kv_cache": False,
+        "use_paged_context_fmha": True,
+        "use_fp8_context_fmha": False,
+    }
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+
+
+@pytest.mark.parametrize(
+    "repo",
+    [
+        "./llama-3.1-8b",
+        "../my-model-is-in-parent-directory",
+        "~/.huggingface/my--model--cache/model",
+        "foo.git",
+        "datasets/foo/bar",
+        ".repo_idother..repo..id",
+    ],
+)
+def test_invalid_hf_repo(trtllm_config, repo):
+    trtllm_config["trt_llm"]["build"]["checkpoint_repository"]["source"] = "HF"
+    trtllm_config["trt_llm"]["build"]["checkpoint_repository"]["repo"] = repo
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+
+
+def test_plugin_paged_fp8_context_fmha_check(trtllm_config):
+    trtllm_config["trt_llm"]["build"]["plugin_configuration"] = {
+        "paged_kv_cache": False,
+        "use_paged_context_fmha": False,
+        "use_fp8_context_fmha": True,
+    }
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+    trtllm_config["trt_llm"]["build"]["plugin_configuration"] = {
+        "paged_kv_cache": True,
+        "use_paged_context_fmha": False,
+        "use_fp8_context_fmha": True,
+    }
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+
+
+def test_fp8_context_fmha_check_kv_dtype(trtllm_config):
+    trtllm_config["trt_llm"]["build"]["plugin_configuration"] = {
+        "paged_kv_cache": True,
+        "use_paged_context_fmha": True,
+        "use_fp8_context_fmha": True,
+    }
+    trtllm_config["trt_llm"]["build"]["quantization_type"] = (
+        TrussTRTLLMQuantizationType.FP8_KV.value
+    )
+    TrussConfig.from_dict(trtllm_config)
+
+    del trtllm_config["trt_llm"]["build"]["quantization_type"]
+    with pytest.raises(ValueError):
+        TrussConfig.from_dict(trtllm_config)
+
+
+@pytest.mark.parametrize("verbose, expect_equal", [(False, True), (True, False)])
+def test_to_dict_trtllm(
+    verbose,
+    expect_equal,
+    trtllm_config,
+    trtllm_spec_dec_config,
+    trtllm_spec_dec_config_full,
+):
+    assert (
+        TrussConfig.from_dict(trtllm_config).to_dict(verbose=verbose) == trtllm_config
+    ) == expect_equal
+    assert (
+        TrussConfig.from_dict(trtllm_spec_dec_config).to_dict(verbose=verbose)
+        == trtllm_spec_dec_config
+    ) == expect_equal
+    assert (
+        TrussConfig.from_dict(trtllm_spec_dec_config_full).to_dict(verbose=verbose)
+        == trtllm_spec_dec_config_full
+    ) == expect_equal
+
+
+@pytest.mark.parametrize("should_raise", [False, True])
+def test_from_dict_spec_dec_trt_llm(should_raise, trtllm_spec_dec_config):
+    test_config = copy.deepcopy(trtllm_spec_dec_config)
+    if should_raise:
+        test_config["trt_llm"]["build"]["speculator"]["speculative_decoding_mode"] = (
+            None
+        )
+        with pytest.raises(ValueError):
+            TrussConfig.from_dict(test_config)
+        test_config["trt_llm"]["build"]["speculator"]["checkpoint_repository"] = None
+        with pytest.raises(ValueError):
+            TrussConfig.from_dict(test_config)
+        test_config["trt_llm"]["build"]["speculator"]["checkpoint_repository"] = (
+            trtllm_spec_dec_config["trt_llm"]["build"]["speculator"][
+                "checkpoint_repository"
+            ]
+        )
+        test_config["trt_llm"]["build"]["plugin_configuration"][
+            "use_paged_context_fmha"
+        ] = False
+        with pytest.raises(ValueError):
+            TrussConfig.from_dict(test_config)
+        test_config["trt_llm"]["build"]["plugin_configuration"][
+            "use_paged_context_fmha"
+        ] = True
+        test_config["trt_llm"]["build"]["speculator"]["speculative_decoding_mode"] = (
+            trtllm_spec_dec_config["trt_llm"]["build"]["speculator"][
+                "speculative_decoding_mode"
+            ]
+        )
+        test_config["trt_llm"]["build"]["speculator"]["num_draft_tokens"] = None
+        with pytest.raises(ValueError):
+            TrussConfig.from_dict(test_config)
+    else:
+        TrussConfig.from_dict(trtllm_spec_dec_config)
+
+
+def test_from_yaml_invalid_requirements_configuration():
+    invalid_requirements = {
+        "requirements_file": "requirements.txt",
+        "requirements": ["requests"],
+    }
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as yaml_file:
+        yaml_path = Path(yaml_file.name)
+        yaml.safe_dump(invalid_requirements, yaml_file)
+
+        with pytest.raises(ValueError):
+            TrussConfig.from_yaml(yaml_path)
+
+
+@pytest.mark.parametrize(
+    "quant_format, accelerator, expectation",
+    [
+        (TrussTRTLLMQuantizationType.NO_QUANT, Accelerator.A100, does_not_raise()),
+        (TrussTRTLLMQuantizationType.FP8, Accelerator.H100, does_not_raise()),
+        (TrussTRTLLMQuantizationType.FP8_KV, Accelerator.H100_40GB, does_not_raise()),
+        (
+            TrussTRTLLMQuantizationType.WEIGHTS_ONLY_INT8,
+            Accelerator.A100,
+            pytest.raises(ValueError),
+        ),
+        (TrussTRTLLMQuantizationType.FP8, Accelerator.A100, pytest.raises(ValueError)),
+        (
+            TrussTRTLLMQuantizationType.FP8_KV,
+            Accelerator.A100,
+            pytest.raises(ValueError),
+        ),
+    ],
+)
+def test_validate_quant_format_and_accelerator_for_trt_llm_builder(
+    quant_format, accelerator, expectation, custom_model_trt_llm
+):
+    config = TrussHandle(custom_model_trt_llm).spec.config
+    config.trt_llm.build.quantization_type = quant_format
+    config.resources.accelerator.accelerator = accelerator
+    with expectation:
+        TrussConfig.from_dict(config.to_dict())

truss/tests/test_context_builder_image.py

@@ -5,12 +5,10 @@ import pytest
 
 
 @pytest.mark.integration
-def test_build_docker_image():
+def test_build_docker_image(test_data_path):
     root_path = Path(__file__).parent.parent.parent
     root = str(root_path)
-    context_builder_image_test_dir = str(
-        root_path / "truss" / "test_data" / "context_builder_image_test"
-    )
+    context_builder_image_test_dir = str(test_data_path / "context_builder_image_test")
 
     subprocess.run(
         [
@@ -44,11 +42,5 @@ def test_build_docker_image():
 
     # This will throw if building docker build context fails
     subprocess.run(
-        [
-            "docker",
-            "run",
-            "baseten/truss-context-builder-test",
-        ],
-        check=True,
-        cwd=root,
+        ["docker", "run", "baseten/truss-context-builder-test"], check=True, cwd=root
     )

truss/tests/test_control_truss_patching.py

@@ -1,9 +1,10 @@
 from dataclasses import replace
 from pathlib import Path
-from typing import Tuple
 
 import pytest
-from truss.constants import SUPPORTED_PYTHON_VERSIONS
+
+from truss.base.constants import SUPPORTED_PYTHON_VERSIONS
+from truss.base.truss_config import ExternalDataItem
 from truss.local.local_config_handler import LocalConfigHandler
 from truss.tests.test_testing_utilities_for_other_tests import ensure_kill_all
 from truss.tests.test_truss_handle import (
@@ -12,9 +13,8 @@ from truss.tests.test_truss_handle import (
     verify_system_package_installed_on_container,
     verify_system_requirement_not_installed_on_container,
 )
-from truss.truss_config import ExternalDataItem
-from truss.truss_gatherer import calc_shadow_truss_dirname
-from truss.truss_handle import TrussHandle
+from truss.truss_handle.truss_gatherer import calc_shadow_truss_dirname
+from truss.truss_handle.truss_handle import TrussHandle
 
 
 def current_num_docker_images(th: TrussHandle) -> int:
@@ -24,7 +24,7 @@ def current_num_docker_images(th: TrussHandle) -> int:
 @pytest.fixture
 def control_model_handle_tag_tuple(
     custom_model_control,
-) -> Tuple[Path, TrussHandle, str]:
+) -> tuple[Path, TrussHandle, str]:
     th = TrussHandle(custom_model_control)
     tag = "test-docker-custom-model-control-tag:0.0.1"
     return (custom_model_control, th, tag)
@@ -410,12 +410,7 @@ class Model:
                 )
             ],
         )
-        th._update_config(
-            lambda conf: replace(
-                conf,
-                external_data=new_external_data,
-            )
-        )
+        th._update_config(lambda conf: replace(conf, external_data=new_external_data))
         result = th.docker_predict([], tag=tag, network="host")
         assert result == content and orig_num_truss_images == current_num_docker_images(
             th

truss/tests/test_custom_server.py

@@ -0,0 +1,38 @@
+import pytest
+import requests
+from tenacity import stop_after_attempt
+
+from truss.local.local_config_handler import LocalConfigHandler
+from truss.tests.test_testing_utilities_for_other_tests import ensure_kill_all
+from truss.truss_handle.truss_handle import TrussHandle
+
+
+@pytest.mark.integration
+def test_custom_server_truss(test_data_path):
+    with ensure_kill_all():
+        print("Running test_custom_server_truss")
+        truss_dir = test_data_path / "test_custom_server_truss"
+        print(f"truss_dir: {truss_dir}")
+        tr = TrussHandle(truss_dir)
+        print("Setting secret")
+        LocalConfigHandler.set_secret("hf_access_token", "123")
+        try:
+            print("Starting container")
+            _ = tr.docker_run(
+                local_port=8090,
+                detach=True,
+                wait_for_server_ready=True,
+                model_server_stop_retry_override=stop_after_attempt(3),
+            )
+        except Exception as e:
+            raise Exception(f"Failed to start container: {e}")
+        truss_server_addr = "http://localhost:8090"
+        full_url = f"{truss_server_addr}/v1/models/model:predict"
+
+        response = requests.post(full_url, json={})
+        assert response.status_code == 200
+        assert response.json() == {
+            "message": "Hello World",
+            "is_env_var_passed": True,
+            "is_secret_mounted": True,
+        }

truss/tests/test_data/context_builder_image_test/test.py

@@ -0,0 +1,3 @@
+from truss.base import truss_config
+
+print(truss_config)