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

truss/tests/util/test_config_checks.py

@@ -0,0 +1,40 @@
+from unittest.mock import patch
+
+import pytest
+from truss.base.constants import TRTLLM_MIN_MEMORY_REQUEST_GI
+from truss.trt_llm.config_checks import (
+    is_missing_secrets_for_trt_llm_builder,
+    memory_updated_for_trt_llm_builder,
+)
+from truss.truss_handle.truss_handle import TrussHandle
+
+
+@patch("truss.trt_llm.config_checks._is_model_public")
+@pytest.mark.parametrize(
+    "has_secret, is_model_public, expected_result",
+    [
+        (False, False, True),
+        (False, True, False),
+        (True, False, False),
+        (True, True, False),
+    ],
+)
+def test_is_missing_secrets_for_trt_llm_builder(
+    _is_model_public_mock,
+    has_secret,
+    is_model_public,
+    expected_result,
+    custom_model_trt_llm,
+):
+    _is_model_public_mock.return_value = is_model_public
+    handle = TrussHandle(custom_model_trt_llm)
+    if has_secret:
+        handle.add_secret("hf_access_token")
+    assert is_missing_secrets_for_trt_llm_builder(handle) == expected_result
+
+
+def test_check_and_update_memory_for_trt_llm_builder(custom_model_trt_llm):
+    handle = TrussHandle(custom_model_trt_llm)
+    assert memory_updated_for_trt_llm_builder(handle)
+    assert handle.spec.memory == f"{TRTLLM_MIN_MEMORY_REQUEST_GI}Gi"
+    assert handle.spec.memory_in_bytes == TRTLLM_MIN_MEMORY_REQUEST_GI * 1024**3

truss/tests/util/test_env_vars.py

@@ -0,0 +1,14 @@
+import os
+
+from truss.util.env_vars import override_env_vars
+
+
+def test_override_env_vars():
+    os.environ["API_KEY"] = "original_key"
+
+    with override_env_vars({"API_KEY": "new_key", "DEBUG": "true"}):
+        assert os.environ["API_KEY"] == "new_key"
+        assert os.environ["DEBUG"] == "true"
+
+    assert os.environ["API_KEY"] == "original_key"
+    assert "DEBUG" not in os.environ

truss/tests/util/test_path.py

@@ -3,10 +3,10 @@ import tempfile
 import time
 from pathlib import Path
 
-from truss import load
 from truss.contexts.image_builder.serving_image_builder import (
     ServingImageBuilderContext,
 )
+from truss.truss_handle.build import load
 from truss.util import path
 
 
@@ -43,7 +43,7 @@ def test_copy_tree_path_with_hidden_files(custom_model_truss_dir_with_hidden_fil
     with path.given_or_temporary_dir() as dir:
         path.copy_tree_path(custom_model_truss_dir_with_hidden_files, dir)
 
-        assert not (dir / "__pycache__" / "test.cpython-38.pyc").exists()
+        assert not (dir / "__pycache__" / "test.cpython-311.pyc").exists()
         assert not (dir / ".DS_Store").exists()
         assert not (dir / ".git").exists()
         assert (dir / "model").exists()
@@ -56,7 +56,7 @@ def test_is_ignored(custom_model_truss_dir_with_hidden_files):
     assert path.is_ignored(
         custom_model_truss_dir_with_hidden_files
         / "__pycache__"
-        / "test.cpython-38.pyc",
+        / "test.cpython-311.pyc",
         patterns,
     )
 
@@ -116,9 +116,7 @@ def test_is_ignored_with_base_dir(custom_model_truss_dir_with_hidden_files):
     )
 
 
-def test_ignored_files_in_docker_context(
-    custom_model_truss_dir_with_hidden_files,
-):
+def test_ignored_files_in_docker_context(custom_model_truss_dir_with_hidden_files):
     tr = load(custom_model_truss_dir_with_hidden_files)
 
     with path.given_or_temporary_dir() as dir:
@@ -127,7 +125,7 @@ def test_ignored_files_in_docker_context(
 
         assert dir.exists()
 
-        assert not (dir / "__pycache__" / "test.cpython-38.pyc").exists()
+        assert not (dir / "__pycache__" / "test.cpython-311.pyc").exists()
         assert not (dir / ".DS_Store").exists()
         assert not (dir / ".git").exists()
         assert (dir / "model").exists()
@@ -135,7 +133,7 @@ def test_ignored_files_in_docker_context(
         assert (
             custom_model_truss_dir_with_hidden_files
             / "__pycache__"
-            / "test.cpython-38.pyc"
+            / "test.cpython-311.pyc"
         ).exists()
         assert (custom_model_truss_dir_with_hidden_files / ".DS_Store").exists()
         assert (custom_model_truss_dir_with_hidden_files / ".git").exists()
@@ -159,13 +157,7 @@ def test_copy_tree_path_with_truss_ignore(custom_model_truss_dir_with_truss_igno
 
 
 def test_get_ignored_relative_paths():
-    ignore_patterns = [
-        ".mypy_cache/",
-        "venv/",
-        "*.tmp",
-        ".git",
-        "data/*",
-    ]
+    ignore_patterns = [".mypy_cache/", "venv/", "*.tmp", ".git", "data/*"]
 
     root_relative_paths = {
         ".mypy_cache/should_ignore.json",
@@ -174,7 +166,7 @@ def test_get_ignored_relative_paths():
         ".git",
         ".git/HEAD",
         "data/should_ignore.txt",
-        "data.txt" "config.yaml",
+        "data.txtconfig.yaml",
         "model/model.py",
     }
 
@@ -192,12 +184,7 @@ def test_get_ignored_relative_paths():
 
 
 def test_get_ignored_relative_paths_from_root(custom_model_truss_dir_with_hidden_files):
-    ignore_patterns = [
-        "__pycache__",
-        ".DS_Store",
-        ".git",
-        "data/*",
-    ]
+    ignore_patterns = ["__pycache__", ".DS_Store", ".git", "data/*"]
 
     unignored_relative_paths = path.get_unignored_relative_paths_from_root(
         custom_model_truss_dir_with_hidden_files, ignore_patterns
@@ -220,7 +207,7 @@ def test_get_ignored_relative_paths_from_root(custom_model_truss_dir_with_hidden
     )
     ignored_relative_paths_strs = {
         "__pycache__",
-        "__pycache__/test.cpython-38.pyc",
+        "__pycache__/test.cpython-311.pyc",
         ".DS_Store",
         ".git",
         ".git/.test_file",

truss/trt_llm/config_checks.py

@@ -0,0 +1,43 @@
+import requests
+from truss.base.constants import (
+    HF_ACCESS_TOKEN_KEY,
+    HF_MODELS_API_URL,
+    TRTLLM_MIN_MEMORY_REQUEST_GI,
+)
+from truss.base.trt_llm_config import CheckpointSource
+from truss.truss_handle.truss_handle import TrussHandle
+
+
+def is_missing_secrets_for_trt_llm_builder(tr: TrussHandle) -> bool:
+    for trt_llm_build_config in tr.spec.config.parsed_trt_llm_build_configs:
+        source = trt_llm_build_config.checkpoint_repository.source
+        hf_model_id = trt_llm_build_config.checkpoint_repository.repo
+        if (
+            source == CheckpointSource.HF
+            and HF_ACCESS_TOKEN_KEY not in tr.spec.secrets
+            and not _is_model_public(hf_model_id)
+        ):
+            return True
+    return False
+
+
+def memory_updated_for_trt_llm_builder(tr: TrussHandle) -> bool:
+    if uses_trt_llm_builder(tr):
+        if tr.spec.memory_in_bytes < TRTLLM_MIN_MEMORY_REQUEST_GI * 1024**3:
+            tr.spec.config.resources.memory = f"{TRTLLM_MIN_MEMORY_REQUEST_GI}Gi"
+            tr.spec.config.write_to_yaml_file(tr.spec.config_path, verbose=False)
+            return True
+    return False
+
+
+def _is_model_public(model_id: str) -> bool:
+    """Validate that a huggingface hub model is public.
+
+    The hf hub API will return 401 when trying to access a private or gated model without auth.
+    """
+    response = requests.get(f"{HF_MODELS_API_URL}/{model_id}")
+    return response.status_code == 200
+
+
+def uses_trt_llm_builder(tr: TrussHandle) -> bool:
+    return tr.spec.config.trt_llm is not None

truss/trt_llm/validation.py

@@ -0,0 +1,42 @@
+import ast
+
+from truss.base.errors import ValidationError
+from truss.base.truss_spec import TrussSpec
+
+
+def validate(truss_spec: TrussSpec):
+    model_class_filepath = truss_spec.model_class_filepath
+    if not model_class_filepath.exists():
+        # It's ok if model class file is not provided,
+        # trt_llm will generate one.
+        return
+
+    source = model_class_filepath.read_text()
+    _verify_has_class_init_arg(source, truss_spec.model_class_name, "trt_llm")
+
+
+def _verify_has_class_init_arg(source: str, class_name: str, arg_name: str):
+    tree = ast.parse(source)
+    model_class_init_found = False
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
+            for child in node.body:
+                if child.name == "__init__":  # type: ignore[attr-defined]
+                    model_class_init_found = True
+                    arg_found = False
+                    for arg in child.args.args:  # type: ignore[attr-defined]
+                        if arg.arg == arg_name:
+                            arg_found = True
+                    if not arg_found:
+                        raise ValidationError(
+                            (
+                                "Model class `__init__` method is required to have `trt_llm` as an argument. Please add that argument.\n  "
+                                "Or if you want to use the automatically generated model class then remove the `model.py` file.\n "
+                                "Refer to https://docs.baseten.co/performance/engine-builder-customization for details on engine object usage."
+                            )
+                        )
+
+    if not model_class_init_found:
+        raise ValidationError(
+            "Model class does not have an `__init__` method; when using `trt_llm` it is required"
+        )

truss/truss_handle/build.py

@@ -0,0 +1,122 @@
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Optional, Union
+
+import yaml
+
+from truss.base.constants import (
+    CONFIG_FILE,
+    PYTHON_DX_CUSTOM_TEMPLATE_DIR,
+    TRADITIONAL_CUSTOM_TEMPLATE_DIR,
+)
+from truss.base.truss_config import (
+    Build,
+    TrussConfig,
+    map_local_to_supported_python_version,
+)
+from truss.truss_handle.truss_handle import TrussHandle
+from truss.util.notebook import is_notebook_or_ipython
+from truss.util.path import copy_tree_path
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+if is_notebook_or_ipython():
+    logger.setLevel(logging.INFO)
+    logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def _populate_traditional_target_directory(config: TrussConfig, dir_path: Path):
+    # Create data dir
+    (dir_path / config.data_dir).mkdir()
+
+    # Create bundled packages dir
+    # TODO: Drop by default
+    (dir_path / config.bundled_packages_dir).mkdir()
+
+    # Create model module dir
+    model_dir = dir_path / config.model_module_dir
+    copy_tree_path(TRADITIONAL_CUSTOM_TEMPLATE_DIR / "model", model_dir)
+
+    # Write config
+    with (dir_path / CONFIG_FILE).open("w") as config_file:
+        yaml.dump(config.to_dict(verbose=False), config_file)
+
+
+def _populate_python_dx_target_directory(config: TrussConfig, dir_path: Path):
+    copy_tree_path(PYTHON_DX_CUSTOM_TEMPLATE_DIR, dir_path)
+
+    # Hack: We want to place the user provided model name into generated code. Until
+    # this gets more complicated, we rely on a brittle string replace. Eventually, we
+    # can consider moving to jinja templates.
+    model_file_path = dir_path / "my_model.py"
+    with open(model_file_path, "r") as f:
+        content = f.read()
+
+    assert config.model_name is not None
+    content = content.replace("{{ MODEL_NAME }}", config.model_name)
+    with open(model_file_path, "w") as f:
+        f.write(content)
+
+
+def init_directory(
+    target_directory: Union[str, Path],
+    build_config: Optional[Build] = None,
+    model_name: Optional[str] = None,
+    python_config: bool = False,
+) -> Path:
+    """
+    Initialize an empty placeholder Truss. A Truss is a build context designed
+    to be built as a container locally or uploaded into a baseten serving
+    environment. This placeholder structure can be filled to represent ML
+    models.
+
+    Args:
+        target_directory: Absolute or relative path of the directory to create
+                          Truss in. The directory is created if it doesn't exist.
+    """
+    config = TrussConfig(
+        model_name=model_name, python_version=map_local_to_supported_python_version()
+    )
+
+    if build_config:
+        config.build = build_config
+
+    target_directory_path = Path(target_directory)
+    target_directory_path.mkdir(parents=True, exist_ok=True)
+
+    if not python_config:
+        _populate_traditional_target_directory(
+            config=config, dir_path=target_directory_path
+        )
+    else:
+        _populate_python_dx_target_directory(
+            config=config, dir_path=target_directory_path
+        )
+
+    return target_directory_path
+
+
+def load(truss_directory: Union[str, Path]) -> TrussHandle:
+    """Get a handle to a Truss. A Truss is a build context designed to be built
+    as a container locally or uploaded into a model serving environment.
+
+    Args:
+        truss_directory (str | Path): The local directory of an existing Truss
+    Returns:
+        TrussHandle
+    """
+    return TrussHandle(Path(truss_directory))
+
+
+def cleanup() -> None:
+    """
+    Cleans up .truss directory.
+    """
+    build_folder_path = Path(Path.home(), ".truss")
+    if build_folder_path.exists():
+        for obj in build_folder_path.glob("**/*"):
+            if (not obj.name == "config.yaml") and (obj.is_file()):
+                os.remove(obj)
+    return

truss/{decorators.py → truss_handle/decorators.py}

@@ -1,6 +1,6 @@
 def proxy_to_shadow_if_scattered(func):
     def wrapper(*args, **kwargs):
-        from truss.truss_handle import TrussHandle
+        from truss.truss_handle.truss_handle import TrussHandle
 
         truss_handle = args[0]
         if not truss_handle.is_scattered():