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

truss/{truss_config.py → base/truss_config.py}

@@ -1,21 +1,30 @@
 import logging
+import os
+import sys
 from dataclasses import _MISSING_TYPE, dataclass, field, fields
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, TypeVar
 
 import yaml
-from truss.config.trt_llm import TRTLLMConfiguration
-from truss.constants import HTTP_PUBLIC_BLOB_BACKEND
-from truss.errors import ValidationError
-from truss.types import ModelFrameworkType
-from truss.util.data_structures import transform_optional
-from truss.validation import (
+
+from truss.base.constants import HTTP_PUBLIC_BLOB_BACKEND
+from truss.base.custom_types import ModelFrameworkType
+from truss.base.errors import ValidationError
+from truss.base.trt_llm_config import (
+    TRTLLMConfiguration,
+    TrussTRTLLMBuildConfiguration,
+    TrussTRTLLMQuantizationType,
+)
+from truss.base.validation import (
     validate_cpu_spec,
     validate_memory_spec,
+    validate_node_count,
     validate_python_executable_path,
     validate_secret_name,
+    validate_secret_to_path_mapping,
 )
+from truss.util.requirements import parse_requirement_string
 
 DEFAULT_MODEL_FRAMEWORK_TYPE = ModelFrameworkType.CUSTOM
 DEFAULT_MODEL_TYPE = "Model"
@@ -30,17 +39,27 @@ DEFAULT_DATA_DIRECTORY = "data"
 DEFAULT_EXAMPLES_FILENAME = "examples.yaml"
 DEFAULT_SPEC_VERSION = "2.0"
 DEFAULT_PREDICT_CONCURRENCY = 1
-DEFAULT_NUM_WORKERS = 1
 DEFAULT_STREAMING_RESPONSE_READ_TIMEOUT = 60
-
+DEFAULT_ENABLE_TRACING_DATA = False  # This should be in sync with tracing.py.
 DEFAULT_CPU = "1"
 DEFAULT_MEMORY = "2Gi"
 DEFAULT_USE_GPU = False
 
 DEFAULT_BLOB_BACKEND = HTTP_PUBLIC_BLOB_BACKEND
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+VALID_PYTHON_VERSIONS = ["py38", "py39", "py310", "py311"]
+
+X = TypeVar("X")
+Y = TypeVar("Y")
+
+
+def transform_optional(x: Optional[X], fn: Callable[[X], Optional[Y]]) -> Optional[Y]:
+    if x is None:
+        return None
+
+    return fn(x)
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -51,6 +70,8 @@ class Accelerator(Enum):
     V100 = "V100"
     A100 = "A100"
     H100 = "H100"
+    H200 = "H200"
+    H100_40GB = "H100_40GB"
 
 
 @dataclass
@@ -133,31 +154,66 @@ class ModelCache:
         return [model.to_dict(verbose=verbose) for model in self.models]
 
 
+@dataclass
+class HealthChecks:
+    restart_check_delay_seconds: Optional[int] = None
+    restart_threshold_seconds: Optional[int] = None
+    stop_traffic_threshold_seconds: Optional[int] = None
+
+    @staticmethod
+    def from_dict(d):
+        return HealthChecks(
+            restart_check_delay_seconds=d.get("restart_check_delay_seconds"),
+            restart_threshold_seconds=d.get("restart_threshold_seconds"),
+            stop_traffic_threshold_seconds=d.get("stop_traffic_threshold_seconds"),
+        )
+
+    def to_dict(self):
+        return {
+            "restart_check_delay_seconds": self.restart_check_delay_seconds,
+            "restart_threshold_seconds": self.restart_threshold_seconds,
+            "stop_traffic_threshold_seconds": self.stop_traffic_threshold_seconds,
+        }
+
+
 @dataclass
 class Runtime:
     predict_concurrency: int = DEFAULT_PREDICT_CONCURRENCY
-    num_workers: int = DEFAULT_NUM_WORKERS
     streaming_read_timeout: int = DEFAULT_STREAMING_RESPONSE_READ_TIMEOUT
+    enable_tracing_data: bool = DEFAULT_ENABLE_TRACING_DATA
+    enable_debug_logs: bool = False
+    health_checks: HealthChecks = field(default_factory=HealthChecks)
 
     @staticmethod
     def from_dict(d):
         predict_concurrency = d.get("predict_concurrency", DEFAULT_PREDICT_CONCURRENCY)
-        num_workers = d.get("num_workers", DEFAULT_NUM_WORKERS)
+        num_workers = d.get("num_workers", 1)
+        if num_workers != 1:
+            raise ValueError(
+                "After truss 0.9.49 only 1 worker per server is allowed. "
+                "For concurrency utilize asyncio, autoscaling replicas "
+                "and as a last resort thread/process pools inside the "
+                "truss model."
+            )
         streaming_read_timeout = d.get(
             "streaming_read_timeout", DEFAULT_STREAMING_RESPONSE_READ_TIMEOUT
         )
+        enable_tracing_data = d.get("enable_tracing_data", DEFAULT_ENABLE_TRACING_DATA)
+        health_checks = HealthChecks.from_dict(d.get("health_checks", {}))
 
         return Runtime(
             predict_concurrency=predict_concurrency,
-            num_workers=num_workers,
             streaming_read_timeout=streaming_read_timeout,
+            enable_tracing_data=enable_tracing_data,
+            health_checks=health_checks,
         )
 
     def to_dict(self):
         return {
             "predict_concurrency": self.predict_concurrency,
-            "num_workers": self.num_workers,
             "streaming_read_timeout": self.streaming_read_timeout,
+            "enable_tracing_data": self.enable_tracing_data,
+            "health_checks": self.health_checks.to_dict(),
         }
 
 
@@ -176,15 +232,23 @@ class ModelServer(Enum):
 class Build:
     model_server: ModelServer = ModelServer.TrussServer
     arguments: Dict = field(default_factory=dict)
+    secret_to_path_mapping: Dict = field(default_factory=dict)
 
     @staticmethod
     def from_dict(d):
         model_server = ModelServer[d.get("model_server", "TrussServer")]
         arguments = d.get("arguments", {})
+        secret_to_path_mapping = d.get("secret_to_path_mapping", {})
+        validate_secret_to_path_mapping(secret_to_path_mapping)
+        if not isinstance(secret_to_path_mapping, dict):
+            raise ValueError(
+                "Please pass a valid mapping for `secret_to_path_mapping`."
+            )
 
         return Build(
             model_server=model_server,
             arguments=arguments,
+            secret_to_path_mapping=secret_to_path_mapping,
         )
 
     def to_dict(self):
@@ -197,6 +261,7 @@ class Resources:
     memory: str = DEFAULT_MEMORY
     use_gpu: bool = DEFAULT_USE_GPU
     accelerator: AcceleratorSpec = field(default_factory=AcceleratorSpec)
+    node_count: Optional[int] = None
 
     @staticmethod
     def from_dict(d):
@@ -209,20 +274,26 @@ class Resources:
         if accelerator.accelerator is not None:
             use_gpu = True
 
-        return Resources(
-            cpu=cpu,
-            memory=memory,
-            use_gpu=use_gpu,
-            accelerator=accelerator,
-        )
+        r = Resources(cpu=cpu, memory=memory, use_gpu=use_gpu, accelerator=accelerator)
+
+        # only add node_count if not None. This helps keep
+        # config generated by truss init concise.
+        node_count = d.get("node_count")
+        validate_node_count(node_count)
+        r.node_count = node_count
+
+        return r
 
     def to_dict(self):
-        return {
+        d = {
             "cpu": self.cpu,
             "memory": self.memory,
             "use_gpu": self.use_gpu,
             "accelerator": self.accelerator.to_str(),
         }
+        if self.node_count is not None:
+            d["node_count"] = self.node_count
+        return d
 
 
 @dataclass
@@ -296,25 +367,112 @@ class ExternalData:
         return [item.to_dict() for item in self.items]
 
 
+class DockerAuthType(Enum):
+    """
+    This enum will express all of the types of registry
+    authentication we support.
+    """
+
+    GCP_SERVICE_ACCOUNT_JSON = "GCP_SERVICE_ACCOUNT_JSON"
+
+
+@dataclass
+class DockerAuthSettings:
+    """
+    Provides information about how to authenticate to the docker registry containing
+    the custom base image.
+    """
+
+    auth_method: DockerAuthType
+    secret_name: str
+    registry: Optional[str] = ""
+
+    @staticmethod
+    def from_dict(d: Dict[str, str]):
+        auth_method = d.get("auth_method")
+        secret_name = d.get("secret_name")
+
+        if auth_method:
+            # Capitalize the auth method so that we support this field passed
+            # as "gcs_service_account".
+            auth_method = auth_method.upper()
+
+        if (
+            not secret_name
+            or not auth_method
+            or auth_method not in [auth_type.value for auth_type in DockerAuthType]
+        ):
+            raise ValueError("Please provide a `secret_name`, and valid `auth_method`")
+
+        return DockerAuthSettings(
+            auth_method=DockerAuthType[auth_method],
+            secret_name=secret_name,
+            registry=d.get("registry"),
+        )
+
+    def to_dict(self):
+        return {
+            "auth_method": self.auth_method.value,
+            "secret_name": self.secret_name,
+            "registry": self.registry,
+        }
+
+
 @dataclass
 class BaseImage:
     image: str = ""
     python_executable_path: str = ""
+    docker_auth: Optional[DockerAuthSettings] = None
 
     @staticmethod
     def from_dict(d):
         image = d.get("image", "")
         python_executable_path = d.get("python_executable_path", "")
+        docker_auth = d.get("docker_auth")
         validate_python_executable_path(python_executable_path)
         return BaseImage(
             image=image,
             python_executable_path=python_executable_path,
+            docker_auth=(
+                DockerAuthSettings.from_dict(docker_auth) if docker_auth else None
+            ),
         )
 
     def to_dict(self):
         return {
             "image": self.image,
             "python_executable_path": self.python_executable_path,
+            "docker_auth": transform_optional(
+                self.docker_auth, lambda docker_auth: docker_auth.to_dict()
+            ),
+        }
+
+
+@dataclass
+class DockerServer:
+    start_command: str
+    server_port: int
+    predict_endpoint: str
+    readiness_endpoint: str
+    liveness_endpoint: str
+
+    @staticmethod
+    def from_dict(d) -> "DockerServer":
+        return DockerServer(
+            start_command=d.get("start_command"),
+            server_port=d.get("server_port"),
+            predict_endpoint=d.get("predict_endpoint"),
+            readiness_endpoint=d.get("readiness_endpoint"),
+            liveness_endpoint=d.get("liveness_endpoint"),
+        )
+
+    def to_dict(self):
+        return {
+            "start_command": self.start_command,
+            "server_port": self.server_port,
+            "readiness_endpoint": self.readiness_endpoint,
+            "liveness_endpoint": self.liveness_endpoint,
+            "predict_endpoint": self.predict_endpoint,
         }
 
 
@@ -375,6 +533,7 @@ class TrussConfig:
               memory: 14Gi
               use_gpu: true
               accelerator: A10G
+              node_count: 2
             ```
         secrets (Dict[str, str]):
             <Warning>
@@ -434,16 +593,25 @@ class TrussConfig:
     # spec_version is a version string
     spec_version: str = DEFAULT_SPEC_VERSION
     base_image: Optional[BaseImage] = None
+    docker_server: Optional[DockerServer] = None
     model_cache: ModelCache = field(default_factory=ModelCache)
     trt_llm: Optional[TRTLLMConfiguration] = None
+    build_commands: List[str] = field(default_factory=list)
+    use_local_chains_src: bool = False
 
     @property
     def canonical_python_version(self) -> str:
-        return {
-            "py39": "3.9",
-            "py38": "3.8",
-            "py37": "3.7",
-        }[self.python_version]
+        return {"py311": "3.11", "py310": "3.10", "py39": "3.9", "py38": "3.8"}[
+            self.python_version
+        ]
+
+    @property
+    def parsed_trt_llm_build_configs(self) -> List[TrussTRTLLMBuildConfiguration]:
+        if self.trt_llm:
+            if self.trt_llm.build.speculator and self.trt_llm.build.speculator.build:
+                return [self.trt_llm.build, self.trt_llm.build.speculator.build]
+            return [self.trt_llm.build]
+        return []
 
     @staticmethod
     def from_dict(d):
@@ -464,7 +632,7 @@ class TrussConfig:
             requirements_file=d.get("requirements_file", None),
             requirements=d.get("requirements", []),
             system_packages=d.get("system_packages", []),
-            environment_variables=d.get("environment_variables", {}),
+            environment_variables=_handle_env_vars(d.get("environment_variables", {})),
             resources=Resources.from_dict(d.get("resources", {})),
             runtime=Runtime.from_dict(d.get("runtime", {})),
             build=Build.from_dict(d.get("build", {})),
@@ -483,19 +651,49 @@ class TrussConfig:
                 d.get("external_data"), ExternalData.from_list
             ),
             base_image=transform_optional(d.get("base_image"), BaseImage.from_dict),
+            docker_server=transform_optional(
+                d.get("docker_server"), DockerServer.from_dict
+            ),
             model_cache=transform_optional(
-                d.get("model_cache") or d.get("hf_cache") or [],
+                d.get("model_cache") or d.get("hf_cache") or [],  # type: ignore
                 ModelCache.from_list,
             ),
             trt_llm=transform_optional(
-                d.get("trt_llm"), lambda x: TRTLLMConfiguration(**x)
+                d.get("trt_llm"), lambda x: (TRTLLMConfiguration(**x))
             ),
+            build_commands=d.get("build_commands", []),
+            use_local_chains_src=d.get("use_local_chains_src", False),
         )
         config.validate()
         return config
 
+    def load_requirements_from_file(self, truss_dir: Path) -> List[str]:
+        if self.requirements_file:
+            requirements_path = truss_dir / self.requirements_file
+            try:
+                requirements = []
+                with open(requirements_path) as f:
+                    for line in f.readlines():
+                        parsed_line = parse_requirement_string(line)
+                        if parsed_line:
+                            requirements.append(parsed_line)
+                return requirements
+            except Exception as e:
+                logger.exception(
+                    f"failed to read requirements file: {self.requirements_file}"
+                )
+                raise e
+        return []
+
+    @staticmethod
+    def load_requirements_file_from_filepath(yaml_path: Path) -> List[str]:
+        config = TrussConfig.from_yaml(yaml_path)
+        return config.load_requirements_from_file(yaml_path.parent)
+
     @staticmethod
     def from_yaml(yaml_path: Path):
+        if not os.path.isfile(yaml_path):
+            raise ValueError(f"Expected a truss configuration file at {yaml_path}")
         with yaml_path.open() as yaml_file:
             raw_data = yaml.safe_load(yaml_file) or {}
             if "hf_cache" in raw_data:
@@ -516,10 +714,68 @@ class TrussConfig:
     def clone(self):
         return TrussConfig.from_dict(self.to_dict())
 
+    def _validate_trt_llm_config(self) -> None:
+        if self.trt_llm:
+            if (
+                self.trt_llm.build.quantization_type
+                is TrussTRTLLMQuantizationType.WEIGHTS_ONLY_INT8
+                and self.resources.accelerator.accelerator is Accelerator.A100
+            ):
+                raise ValueError(
+                    "Weight only int8 quantization on A100 accelerators is not currently supported"
+                )
+            elif self.trt_llm.build.quantization_type in [
+                TrussTRTLLMQuantizationType.FP8,
+                TrussTRTLLMQuantizationType.FP8_KV,
+            ] and self.resources.accelerator.accelerator not in [
+                Accelerator.H100,
+                Accelerator.H100_40GB,
+                Accelerator.L4,
+            ]:
+                raise ValueError(
+                    "FP8 quantization is only supported on L4 and H100 accelerators"
+                )
+            tensor_parallel_count = self.trt_llm.build.tensor_parallel_count
+
+            if tensor_parallel_count != self.resources.accelerator.count:
+                raise ValueError(
+                    "Tensor parallelism and GPU count must be the same for TRT-LLM"
+                )
+
     def validate(self):
+        if self.python_version not in VALID_PYTHON_VERSIONS:
+            raise ValueError(
+                f"Please ensure that `python_version` is one of {VALID_PYTHON_VERSIONS}"
+            )
+
+        if not isinstance(self.secrets, dict):
+            raise ValueError(
+                "Please ensure that `secrets` is a mapping of the form:\n"
+                "```\n"
+                "secrets:\n"
+                '  secret1: "some default value"\n'
+                '  secret2: "some other default value"\n'
+                "```"
+            )
         for secret_name in self.secrets:
             validate_secret_name(secret_name)
 
+        if self.requirements and self.requirements_file:
+            raise ValueError(
+                "Please ensure that only one of `requirements` and `requirements_file` is specified"
+            )
+        self._validate_trt_llm_config()
+
+
+def _handle_env_vars(env_vars: Dict[str, Any]) -> Dict[str, str]:
+    new_env_vars = {}
+    for k, v in env_vars.items():
+        if isinstance(v, bool):
+            new_env_vars[k] = str(v).lower()
+        else:
+            new_env_vars[k] = str(v)
+    return new_env_vars
+
 
 DATACLASS_TO_REQ_KEYS_MAP = {
     Resources: {"accelerator", "cpu", "memory", "use_gpu"},
@@ -535,6 +791,7 @@ DATACLASS_TO_REQ_KEYS_MAP = {
         "resources",
         "secrets",
         "system_packages",
+        "build_commands",
     },
     BaseImage: {"image", "python_executable_path"},
 }
@@ -581,9 +838,65 @@ def obj_to_dict(obj, verbose: bool = False):
                 )
             elif isinstance(field_curr_value, TRTLLMConfiguration):
                 d["trt_llm"] = transform_optional(
-                    field_curr_value, lambda data: data.model_dump(mode="json")
+                    field_curr_value, lambda data: data.to_json_dict(verbose=verbose)
+                )
+            elif isinstance(field_curr_value, BaseImage):
+                d["base_image"] = transform_optional(
+                    field_curr_value, lambda data: data.to_dict()
+                )
+            elif isinstance(field_curr_value, DockerServer):
+                d["docker_server"] = transform_optional(
+                    field_curr_value, lambda data: data.to_dict()
+                )
+            elif isinstance(field_curr_value, DockerAuthSettings):
+                d["docker_auth"] = transform_optional(
+                    field_curr_value, lambda data: data.to_dict()
+                )
+            elif isinstance(field_curr_value, HealthChecks):
+                d["health_checks"] = transform_optional(
+                    field_curr_value, lambda data: data.to_dict()
                 )
             else:
                 d[field_name] = field_curr_value
 
     return d
+
+
+# TODO(marius): consolidate this with config/validation:
+def _infer_python_version() -> str:
+    return f"py{sys.version_info.major}{sys.version_info.minor}"
+
+
+def map_local_to_supported_python_version() -> str:
+    return map_to_supported_python_version(_infer_python_version())
+
+
+def map_to_supported_python_version(python_version: str) -> str:
+    """Map python version to truss supported python version.
+
+    Currently, it maps any versions greater than 3.11 to 3.11.
+
+    Args:
+        python_version: in the form py[major_version][minor_version] e.g. py39,
+          py310.
+    """
+    python_major_version = int(python_version[2:3])
+    python_minor_version = int(python_version[3:])
+
+    if python_major_version != 3:
+        raise NotImplementedError("Only python version 3 is supported")
+
+    if python_minor_version > 11:
+        logger.info(
+            f"Mapping python version {python_major_version}.{python_minor_version}"
+            " to 3.11, the highest version that Truss currently supports."
+        )
+        return "py311"
+
+    if python_minor_version < 8:
+        raise ValueError(
+            f"Mapping python version {python_major_version}.{python_minor_version}"
+            " to 3.8, the lowest version that Truss currently supports."
+        )
+
+    return python_version

truss/{truss_spec.py → base/truss_spec.py}

@@ -3,10 +3,12 @@ from pathlib import Path
 from typing import Dict, List, Optional
 
 import yaml
-from truss.constants import CONFIG_FILE
-from truss.errors import ValidationError
-from truss.truss_config import ExternalData, ModelServer, TrussConfig
-from truss.types import Example, ModelFrameworkType
+
+from truss.base.constants import CONFIG_FILE
+from truss.base.custom_types import Example, ModelFrameworkType
+from truss.base.errors import ValidationError
+from truss.base.truss_config import ExternalData, ModelServer, TrussConfig
+from truss.base.validation import validate_memory_spec
 
 
 class TrussSpec:
@@ -34,6 +36,10 @@ class TrussSpec:
     def external_data(self) -> Optional[ExternalData]:
         return self._config.external_data
 
+    @property
+    def build_commands(self) -> List[str]:
+        return self._config.build_commands
+
     @property
     def model_module_dir(self) -> Path:
         return self._truss_dir / self._config.model_module_dir
@@ -90,6 +96,10 @@ class TrussSpec:
     def memory(self) -> str:
         return self._config.resources.memory
 
+    @property
+    def memory_in_bytes(self) -> int:
+        return validate_memory_spec(self.memory)
+
     @property
     def use_gpu(self) -> str:
         return self._config.resources.use_gpu
@@ -180,8 +190,12 @@ class TrussSpec:
         return self._config.base_image.image
 
     @property
-    def python_executable_path(self) -> str:
-        return self._config.base_image.python_executable_path
+    def python_executable_path(self) -> Optional[str]:
+        return (
+            self._config.base_image.python_executable_path
+            if self._config.base_image
+            else None
+        )
 
     @property
     def apply_library_patches(self) -> bool: