truss 0.9.45__tar.gz → 0.9.45.dev101__tar.gz

{truss-0.9.45 → truss-0.9.45.dev101}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truss
-Version: 0.9.45
+Version: 0.9.45.dev101
 Summary: A seamless bridge from model development to model delivery
 Home-page: https://github.com/basetenlabs/truss
 License: MIT
@@ -17,7 +17,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: Jinja2 (>=3.1.2,<4.0.0)
 Requires-Dist: PyYAML (>=6.0)
-Requires-Dist: aiofiles (>=24.1.0,<25.0.0)
+Requires-Dist: aiohttp (>=3.10.10,<4.0.0)
 Requires-Dist: blake3 (>=0.3.3,<0.4.0)
 Requires-Dist: boto3 (>=1.34.85,<2.0.0)
 Requires-Dist: fastapi (>=0.109.1)

{truss-0.9.45 → truss-0.9.45.dev101}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.45"
+version = "0.9.45dev101"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -27,7 +27,7 @@ packages = [
 "Baseten" = "https://baseten.co"
 
 [tool.poetry.dependencies]
-aiofiles = "^24.1.0"
+aiohttp = "^3.10.10"
 blake3 = "^0.3.3"
 boto3 = "^1.34.85"
 fastapi = ">=0.109.1"
@@ -97,7 +97,6 @@ pytest = "7.2.0"
 pytest-cov = "^3.0.0"
 types-PyYAML = "^6.0.12.12"
 types-setuptools = "^69.0.0.0"
-types-aiofiles = "^24.1.0.20240626"
 
 [tool.poetry.scripts]
 truss = 'truss.cli:truss_cli'

truss-0.9.45.dev101/truss/.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

truss-0.9.45.dev101/truss/.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

truss-0.9.45.dev101/truss/.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

truss-0.9.45.dev101/truss/.pytest_cache/v/cache/lastfailed

@@ -0,0 +1 @@
+{}

truss-0.9.45.dev101/truss/.pytest_cache/v/cache/nodeids

@@ -0,0 +1,3 @@
+[
+  "tests/test_control_truss_patching.py::test_patch_added_model_dir"
+]

truss-0.9.45.dev101/truss/.pytest_cache/v/cache/stepwise

@@ -0,0 +1 @@
+[]

{truss-0.9.45 → truss-0.9.45.dev101}/truss/config/trt_llm.py

@@ -23,7 +23,6 @@ class TrussTRTLLMModel(str, Enum):
     MISTRAL = "mistral"
     DEEPSEEK = "deepseek"
     WHISPER = "whisper"
-    QWEN = "qwen"
 
 
 class TrussTRTLLMQuantizationType(str, Enum):
@@ -31,7 +30,7 @@ class TrussTRTLLMQuantizationType(str, Enum):
     WEIGHTS_ONLY_INT8 = "weights_int8"
     WEIGHTS_KV_INT8 = "weights_kv_int8"
     WEIGHTS_ONLY_INT4 = "weights_int4"
-    WEIGHTS_INT4_KV_INT8 = "weights_int4_kv_int8"
+    WEIGHTS_KV_INT4 = "weights_kv_int4"
     SMOOTH_QUANT = "smooth_quant"
     FP8 = "fp8"
     FP8_KV = "fp8_kv"
@@ -59,9 +58,10 @@ class CheckpointRepository(BaseModel):
 
 class TrussTRTLLMBuildConfiguration(BaseModel):
     base_model: TrussTRTLLMModel
-    max_seq_len: int
-    max_batch_size: Optional[int] = 256
-    max_num_tokens: Optional[int] = 8192
+    max_input_len: int
+    max_output_len: int
+    max_batch_size: int
+    max_num_tokens: Optional[int] = None
     max_beam_width: int = 1
     max_prompt_embedding_table_size: int = 0
     checkpoint_repository: CheckpointRepository
@@ -75,6 +75,7 @@ class TrussTRTLLMBuildConfiguration(BaseModel):
     plugin_configuration: TrussTRTLLMPluginConfiguration = (
         TrussTRTLLMPluginConfiguration()
     )
+    use_fused_mlp: bool = False
     kv_cache_free_gpu_mem_fraction: float = 0.9
     num_builder_gpus: Optional[int] = None
     enable_chunked_context: bool = False

{truss-0.9.45 → truss-0.9.45.dev101}/truss/constants.py

@@ -106,9 +106,17 @@ HTTP_PUBLIC_BLOB_BACKEND = "http_public"
 
 REGISTRY_BUILD_SECRET_PREFIX = "DOCKER_REGISTRY_"
 
-TRTLLM_BASE_IMAGE = "baseten/briton-server:v0.13.0_v0.0.14"
+TRTLLM_BASE_IMAGE = "baseten/briton-server:5fa9436e_v0.0.11"
 TRTLLM_PYTHON_EXECUTABLE = "/usr/bin/python3"
-BASE_TRTLLM_REQUIREMENTS = ["briton==0.3.2"]
+BASE_TRTLLM_REQUIREMENTS = [
+    "grpcio==1.62.3",
+    "grpcio-tools==1.62.3",
+    "transformers==4.44.2",
+    "truss==0.9.42rc010",
+    "outlines==0.0.46",
+    "torch==2.4.0",
+    "sentencepiece==0.2.0",
+]
 AUDIO_MODEL_TRTLLM_REQUIREMENTS = [
     "--extra-index-url https://pypi.nvidia.com",
     "tensorrt_cu12_bindings==10.2.0.post1",

{truss-0.9.45 → truss-0.9.45.dev101}/truss/contexts/image_builder/serving_image_builder.py

@@ -412,6 +412,17 @@ class ServingImageBuilder(ImageBuilder):
                     DEFAULT_BUNDLED_PACKAGES_DIR,
                 )
 
+            tensor_parallel_count = (
+                config.trt_llm.build.tensor_parallel_count  # type: ignore[union-attr]
+                if config.trt_llm.build is not None
+                else config.trt_llm.serve.tensor_parallel_count  # type: ignore[union-attr]
+            )
+
+            if tensor_parallel_count != config.resources.accelerator.count:
+                raise ValueError(
+                    "Tensor parallelism and GPU count must be the same for TRT-LLM"
+                )
+
             config.runtime.predict_concurrency = TRTLLM_PREDICT_CONCURRENCY
 
             if not is_audio_model:

{truss-0.9.45 → truss-0.9.45.dev101}/truss/local/local_config_handler.py

@@ -75,18 +75,6 @@ class LocalConfigHandler:
         bptr_data_dir.mkdir(exist_ok=True, parents=True)
         return bptr_data_dir
 
-    @staticmethod
-    def dynamic_config_path():
-        dynamic_config_dir = LocalConfigHandler.TRUSS_CONFIG_DIR / "b10_dynamic_config"
-        dynamic_config_dir.mkdir(exist_ok=True, parents=True)
-        return dynamic_config_dir
-
-    @staticmethod
-    def set_dynamic_config(key: str, value: str):
-        key_path = LocalConfigHandler.dynamic_config_path() / key
-        with key_path.open("w") as key_file:
-            key_file.write(value)
-
     @staticmethod
     def _signatures_dir_path():
         return LocalConfigHandler.TRUSS_CONFIG_DIR / "signatures"

{truss-0.9.45 → truss-0.9.45.dev101}/truss/templates/server/common/errors.py

@@ -1,7 +1,6 @@
 import contextlib
 import logging
 import sys
-import textwrap
 from http import HTTPStatus
 from types import TracebackType
 from typing import (
@@ -14,7 +13,6 @@ from typing import (
 )
 
 import fastapi
-import pydantic
 import starlette.responses
 from fastapi import HTTPException
 from fastapi.responses import JSONResponse
@@ -173,44 +171,3 @@ def intercept_exceptions(
             "Internal Server Error", exc_info=filter_traceback(model_file_name)
         )
         raise UserCodeError(str(exc))
-
-
-def _loc_to_dot_sep(loc: Tuple[Union[str, int], ...]) -> str:
-    # From https://docs.pydantic.dev/latest/errors/errors/#customize-error-messages.
-    # Chained field access is stylized with `.`-notation (corresponding to str parts)
-    # and array indexing is stylized with `[i]`-notation (corresponding to int parts).
-    # E.g. ('items', 1, 'value') -> items[1].value.
-    parts = []
-    for i, x in enumerate(loc):
-        if isinstance(x, str):
-            if i > 0:
-                parts.append(".")
-            parts.append(x)
-        elif isinstance(x, int):
-            parts.append(f"[{x}]")
-        else:
-            raise TypeError(f"Unexpected type: {x}.")
-    return "".join(parts)
-
-
-def format_pydantic_validation_error(exc: pydantic.ValidationError) -> str:
-    if exc.error_count() == 1:
-        parts = ["Input Parsing Error:"]
-    else:
-        parts = ["Input Parsing Errors:"]
-
-    try:
-        for error in exc.errors(include_url=False, include_input=False):
-            loc = _loc_to_dot_sep(error["loc"])
-            if error["msg"]:
-                msg = error["msg"]
-            else:
-                error_no_loc = dict(error)
-                error_no_loc.pop("loc")
-                msg = str(error_no_loc)
-            parts.append(textwrap.indent(f"`{loc}`: {msg}.", "  "))
-    except:  # noqa: E722
-        # Fallback in case any of the fields cannot be processed as expected.
-        return f"Input Parsing Error, {str(exc)}"
-
-    return "\n".join(parts)

{truss-0.9.45 → truss-0.9.45.dev101}/truss/templates/server/model_wrapper.py

@@ -4,7 +4,6 @@ import enum
 import importlib
 import importlib.util
 import inspect
-import json
 import logging
 import os
 import pathlib
@@ -36,7 +35,7 @@ from common.retry import retry
 from common.schema import TrussSchema
 from opentelemetry import trace
 from pydantic import BaseModel
-from shared import dynamic_config_resolver, serialization
+from shared import serialization
 from shared.lazy_data_resolver import LazyDataResolver
 from shared.secrets_resolver import SecretsResolver
 
@@ -54,7 +53,6 @@ EXTENSIONS_DIR_NAME = "extensions"
 EXTENSION_CLASS_NAME = "Extension"
 EXTENSION_FILE_NAME = "extension"
 TRT_LLM_EXTENSION_NAME = "trt_llm"
-POLL_FOR_ENVIRONMENT_UPDATES_TIMEOUT_SECS = 30
 
 
 @asynccontextmanager
@@ -193,7 +191,6 @@ class ModelDescriptor:
     predict: MethodDescriptor
     postprocess: Optional[MethodDescriptor]
     truss_schema: Optional[TrussSchema]
-    setup_environment: Optional[MethodDescriptor]
 
     @cached_property
     def skip_input_parsing(self) -> bool:
@@ -246,19 +243,11 @@ class ModelDescriptor:
         else:
             return_annotation = inspect.signature(model.predict).return_annotation
 
-        if hasattr(model, "setup_environment"):
-            setup_environment = MethodDescriptor.from_method(
-                model.setup_environment, "setup_environment"
-            )
-        else:
-            setup_environment = None
-
         return cls(
             preprocess=preprocess,
             predict=predict,
             postprocess=postprocess,
             truss_schema=TrussSchema.from_signature(parameters, return_annotation),
-            setup_environment=setup_environment,
         )
 
 
@@ -270,8 +259,6 @@ class ModelWrapper:
     _logger: logging.Logger
     _status: "ModelWrapper.Status"
     _predict_semaphore: Semaphore
-    _poll_for_environment_updates_task: Optional[asyncio.Task]
-    _environment: Optional[dict]
 
     class Status(Enum):
         NOT_READY = 0
@@ -293,8 +280,6 @@ class ModelWrapper:
                 "predict_concurrency", DEFAULT_PREDICT_CONCURRENCY
             )
         )
-        self._poll_for_environment_updates_task = None
-        self._environment = None
 
     @property
     def _model(self) -> Any:
@@ -434,9 +419,6 @@ class ModelWrapper:
 
         self._maybe_model_descriptor = ModelDescriptor.from_model(self._model)
 
-        if self._maybe_model_descriptor.setup_environment:
-            self._initialize_environment_before_load()
-
         if hasattr(self._model, "load"):
             retry(
                 self._model.load,
@@ -446,76 +428,6 @@ class ModelWrapper:
                 gap_seconds=1.0,
             )
 
-    def setup_polling_for_environment_updates(self):
-        self._poll_for_environment_updates_task = asyncio.create_task(
-            self.poll_for_environment_updates()
-        )
-
-    def _initialize_environment_before_load(self):
-        environment_str = dynamic_config_resolver.get_dynamic_config_value_sync(
-            dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY
-        )
-        if environment_str:
-            environment_json = json.loads(environment_str)
-            self._logger.info(
-                f"Executing model.setup_environment with environment: {environment_json}"
-            )
-            # TODO: Support calling an async setup_environment() here once we support async load()
-            self._model.setup_environment(environment_json)
-            self._environment = environment_json
-
-    async def setup_environment(self, environment: Optional[dict]):
-        descriptor = self.model_descriptor.setup_environment
-        if not descriptor:
-            return
-        self._logger.info(
-            f"Executing model.setup_environment with environment: {environment}"
-        )
-        if descriptor.is_async:
-            return await self._model.setup_environment(environment)
-        else:
-            return await to_thread.run_sync(self._model.setup_environment, environment)
-
-    async def poll_for_environment_updates(self) -> None:
-        last_modified_time = None
-        environment_config_filename = (
-            dynamic_config_resolver.get_dynamic_config_file_path(
-                dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY
-            )
-        )
-
-        while True:
-            # Give control back to the event loop while waiting for environment updates
-            await asyncio.sleep(POLL_FOR_ENVIRONMENT_UPDATES_TIMEOUT_SECS)
-
-            # Wait for load to finish before checking for environment updates
-            if not self.ready:
-                continue
-
-            # Skip polling if no setup_environment implementation provided
-            if not self.model_descriptor.setup_environment:
-                break
-
-            if environment_config_filename.exists():
-                try:
-                    current_mtime = os.path.getmtime(environment_config_filename)
-                    if not last_modified_time or last_modified_time != current_mtime:
-                        environment_str = await dynamic_config_resolver.get_dynamic_config_value_async(
-                            dynamic_config_resolver.ENVIRONMENT_DYNAMIC_CONFIG_KEY
-                        )
-                        if environment_str:
-                            last_modified_time = current_mtime
-                            environment_json = json.loads(environment_str)
-                            # Avoid rerunning `setup_environment` with the same environment
-                            if self._environment != environment_json:
-                                await self.setup_environment(environment_json)
-                                self._environment = environment_json
-                except Exception as e:
-                    self._logger.exception(
-                        "Exception while setting up environment: " + str(e),
-                        exc_info=errors.filter_traceback(self._model_file_name),
-                    )
-
     async def preprocess(
         self,
         inputs: serialization.InputType,

{truss-0.9.45 → truss-0.9.45.dev101}/truss/templates/server/requirements.txt

@@ -17,4 +17,3 @@ pyyaml==6.0.0
 requests==2.31.0
 uvicorn==0.24.0
 uvloop==0.19.0
-aiofiles==24.1.0

{truss-0.9.45 → truss-0.9.45.dev101}/truss/templates/server/truss_server.py

@@ -141,7 +141,7 @@ class BasetenEndpoints:
                         inputs = truss_schema.input_type.parse_obj(inputs)
                 except pydantic.ValidationError as e:
                     raise errors.InputParsingError(
-                        errors.format_pydantic_validation_error(e)
+                        f"Request Validation Error, {str(e)}"
                     ) from e
         else:
             if truss_schema:
@@ -151,7 +151,7 @@ class BasetenEndpoints:
                             inputs = truss_schema.input_type.parse_raw(body_raw)
                     except pydantic.ValidationError as e:
                         raise errors.InputParsingError(
-                            errors.format_pydantic_validation_error(e)
+                            f"Request Validation Error, {str(e)}"
                         ) from e
             else:
                 try:
@@ -281,7 +281,6 @@ class TrussServer:
         if self._setup_json_logger:
             setup_logging()
         self._model.start_load_thread()
-        self._model.setup_polling_for_environment_updates()
 
     def create_application(self):
         app = FastAPI(

truss-0.9.45.dev101/truss/templates/shared/dynamic_config_resolver.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+from typing import Optional
+
+DYNAMIC_CONFIG_MOUNT_DIR = "/etc/b10_dynamic_config"
+
+
+def get_dynamic_config_value(key: str) -> Optional[str]:
+    dynamic_config_path = Path(DYNAMIC_CONFIG_MOUNT_DIR) / key
+    if dynamic_config_path.exists() and dynamic_config_path.is_file():
+        with dynamic_config_path.open() as dynamic_config_file:
+            dynamic_config_value = dynamic_config_file.read()
+            return dynamic_config_value
+    return None