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

truss/templates/shared/lazy_data_resolver.py

@@ -0,0 +1,164 @@
+import logging
+import os
+import shutil
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import pydantic
+import requests
+import yaml
+
+try:
+    from shared.util import BLOB_DOWNLOAD_TIMEOUT_SECS
+except ModuleNotFoundError:
+    from truss.templates.shared.util import BLOB_DOWNLOAD_TIMEOUT_SECS
+
+try:
+    import truss_transfer
+
+    TRUSS_TRANSFER_AVAILABLE = True
+except ImportError:
+    TRUSS_TRANSFER_AVAILABLE = False
+
+LAZY_DATA_RESOLVER_PATH = Path("/bptr/bptr-manifest")
+NUM_WORKERS = 4
+CACHE_DIR = Path("/cache/org/artifacts")
+BASETEN_FS_ENABLED_ENV_VAR = "BASETEN_FS_ENABLED"
+
+logger = logging.getLogger(__name__)
+
+
+class Resolution(pydantic.BaseModel):
+    url: str
+    expiration_timestamp: int
+
+
+class BasetenPointer(pydantic.BaseModel):
+    """Specification for lazy data resolution for download of large files, similar to Git LFS pointers"""
+
+    resolution: Resolution
+    uid: str
+    file_name: str
+    hashtype: str
+    hash: str
+    size: int
+
+
+class BasetenPointerManifest(pydantic.BaseModel):
+    pointers: List[BasetenPointer]
+
+
+class LazyDataResolver:
+    """Deprecation warning: This class is deprecated and will be removed in a future release.
+
+    Please use LazyDataResolverV2 instead (using the `truss_transfer` package).
+    """
+
+    def __init__(self, data_dir: Path):
+        self._data_dir: Path = data_dir
+        self._bptr_resolution: Dict[str, Tuple[str, str, int]] = _read_bptr_resolution()
+        self._resolution_done = False
+        self._uses_b10_cache = (
+            os.environ.get(BASETEN_FS_ENABLED_ENV_VAR, "False") == "True"
+        )
+
+    def cached_download_from_url_using_requests(
+        self, URL: str, hash: str, file_name: str, size: int
+    ):
+        """Download object from URL, attempt to write to cache and symlink to data directory if applicable, data directory otherwise.
+        In case of failure, write to data directory
+        """
+        if self._uses_b10_cache:
+            file_path = CACHE_DIR / hash
+            if file_path.exists():
+                try:
+                    os.symlink(file_path, self._data_dir / file_name)
+                    return
+                except FileExistsError:
+                    # symlink may already exist if the inference server was restarted
+                    return
+
+        # Streaming download to keep memory usage low
+        resp = requests.get(
+            URL, allow_redirects=True, stream=True, timeout=BLOB_DOWNLOAD_TIMEOUT_SECS
+        )
+        resp.raise_for_status()
+
+        if self._uses_b10_cache:
+            try:
+                # Check whether the cache has sufficient space to store the file
+                cache_free_space = shutil.disk_usage(CACHE_DIR).free
+                if cache_free_space < size:
+                    raise OSError(
+                        f"Cache directory does not have sufficient space to save file {file_name}. Free space in cache: {cache_free_space}, file size: {size}"
+                    )
+
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                with file_path.open("wb") as file:
+                    shutil.copyfileobj(resp.raw, file)
+                # symlink to data directory
+                os.symlink(file_path, self._data_dir / file_name)
+                return
+            except FileExistsError:
+                # symlink may already exist if the inference server was restarted
+                return
+            except OSError as e:
+                logger.debug(
+                    "Failed to save artifact to cache dir, saving to data dir instead. Error: %s",
+                    e,
+                )
+                # Cache likely has no space left on device, break to download to data dir as fallback
+                pass
+
+        file_path = self._data_dir / file_name
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        with file_path.open("wb") as file:
+            shutil.copyfileobj(resp.raw, file)
+
+    def fetch(self):
+        if self._resolution_done:
+            return
+
+        with ThreadPoolExecutor(NUM_WORKERS) as executor:
+            futures = {}
+            for file_name, (resolved_url, hash, size) in self._bptr_resolution.items():
+                futures[
+                    executor.submit(
+                        self.cached_download_from_url_using_requests,
+                        resolved_url,
+                        hash,
+                        file_name,
+                        size,
+                    )
+                ] = file_name
+            for future in as_completed(futures):
+                if future.exception():
+                    file_name = futures[future]
+                    raise RuntimeError(f"Download failure for file {file_name}")
+        self._resolution_done = True
+
+
+class LazyDataResolverV2:
+    def __init__(self, data_dir: Path):
+        self._data_dir: Path = data_dir
+
+    def fetch(self):
+        truss_transfer.lazy_data_resolve(str(self._data_dir))
+
+
+def _read_bptr_resolution() -> Dict[str, Tuple[str, str, int]]:
+    if not LAZY_DATA_RESOLVER_PATH.is_file():
+        return {}
+    bptr_manifest = BasetenPointerManifest(
+        **yaml.safe_load(LAZY_DATA_RESOLVER_PATH.read_text())
+    )
+    resolution_map = {}
+    for bptr in bptr_manifest.pointers:
+        if bptr.resolution.expiration_timestamp < int(
+            datetime.now(timezone.utc).timestamp()
+        ):
+            raise RuntimeError("Baseten pointer lazy data resolution has expired")
+        resolution_map[bptr.file_name] = bptr.resolution.url, bptr.hash, bptr.size
+    return resolution_map

truss/templates/shared/log_config.py

@@ -0,0 +1,125 @@
+import logging
+import os
+import urllib.parse
+from typing import Any, Mapping
+
+from pythonjsonlogger import jsonlogger
+
+LOCAL_DATE_FORMAT = "%H:%M:%S"
+
+
+def _disable_json_logging() -> bool:
+    return bool(os.environ.get("DISABLE_JSON_LOGGING"))
+
+
+class _HealthCheckFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        excluded_paths = {
+            "GET / ",
+            "GET /v1/models/model ",
+            "GET /v1/models/model/loaded ",
+        }
+        msg = record.getMessage()
+        return not any(path in msg for path in excluded_paths)
+
+
+class _AccessJsonFormatter(jsonlogger.JsonFormatter):
+    def format(self, record: logging.LogRecord) -> str:
+        # Uvicorn sets record.msg = '%s - "%s %s HTTP/%s" %d' and
+        # record.args = (addr, method, path, version, status).
+        # Python's logging system resolves final
+        # record.message = record.msg % record.args unless we override record.msg.
+        if record.name == "uvicorn.access" and record.args and len(record.args) == 5:
+            client_addr, method, raw_path, version, status = record.args
+            path_decoded = urllib.parse.unquote(str(raw_path))
+            new_message = (
+                f"Handled request from {client_addr} - {method} "
+                f"{path_decoded} HTTP/{version} {status}"
+            )
+            record.msg = new_message
+            record.args = ()  # Ensure Python doesn't reapply the old format string
+        return super().format(record)
+
+
+class _AccessFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        if record.name == "uvicorn.access" and record.args and len(record.args) == 5:
+            client_addr, method, raw_path, version, status = record.args
+            path_decoded = urllib.parse.unquote(str(raw_path))
+            new_message = (
+                f"Handled request from  {client_addr} - {method} "
+                f"{path_decoded} HTTP/{version} {status}"
+            )
+            record.msg = new_message
+            record.args = ()
+        return super().format(record)
+
+
+def make_log_config(log_level: str) -> Mapping[str, Any]:
+    # Warning: `ModelWrapper` depends on correctly setup `uvicorn` logger,
+    # if you change/remove that logger, make sure `ModelWrapper` has a suitable
+    # alternative logger that is also correctly setup in the load thread.
+    formatters = (
+        {
+            "default_formatter": {
+                "format": "%(asctime)s.%(msecs)04d %(levelname)s %(message)s",
+                "datefmt": LOCAL_DATE_FORMAT,
+            },
+            "access_formatter": {
+                "()": _AccessFormatter,
+                "format": "%(asctime)s.%(msecs)04d %(levelname)s %(message)s",
+                "datefmt": LOCAL_DATE_FORMAT,
+            },
+        }
+        if _disable_json_logging()
+        else {
+            "default_formatter": {
+                "()": jsonlogger.JsonFormatter,
+                "format": "%(asctime)s %(levelname)s %(message)s",
+            },
+            "access_formatter": {
+                "()": _AccessJsonFormatter,
+                "format": "%(asctime)s %(levelname)s %(message)s",
+            },
+        }
+    )
+
+    log_config = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "filters": {"health_check_filter": {"()": _HealthCheckFilter}},
+        "formatters": formatters,
+        "handlers": {
+            "default_handler": {
+                "formatter": "default_formatter",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stderr",
+            },
+            "access_handler": {
+                "formatter": "access_formatter",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",
+            },
+        },
+        "loggers": {
+            "uvicorn": {
+                "handlers": ["default_handler"],
+                "level": log_level,
+                "propagate": False,
+            },
+            "uvicorn.error": {
+                "handlers": ["default_handler"],
+                "level": "INFO",
+                "propagate": False,
+            },
+            "uvicorn.access": {
+                "handlers": ["access_handler"],
+                "level": "INFO",
+                "propagate": False,
+                "filters": ["health_check_filter"],
+            },
+        },
+        # Catch-all for module loggers
+        "root": {"handlers": ["default_handler"], "level": log_level},
+    }
+    return log_config

truss/{server → templates}/shared/secrets_resolver.py

@@ -3,7 +3,7 @@ from collections.abc import Mapping
 from pathlib import Path
 from typing import Dict, Optional
 
-SECRETS_DOC_LINK = "https://truss.baseten.co/docs/using-secrets"
+SECRETS_DOC_LINK = "https://truss.baseten.co/guides/secrets"
 
 
 class SecretNotFound(Exception):
@@ -62,7 +62,6 @@ def _secret_missing_error_message(key: str) -> str:
     return f"""
 Secret '{key}' not found. Please ensure that:
   * Secret '{key}' is defined in the 'secrets' section of the Truss config file
-  * The model was pushed with the --trusted flag
   * Secret '{key}' is defined in the secret manager
  Read more about secrets here: {SECRETS_DOC_LINK}.
     """

truss/{server → templates}/shared/serialization.py

@@ -2,11 +2,33 @@ import json
 import uuid
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
-from typing import Any, Callable, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+JSONType = Union[str, int, float, bool, None, List["JSONType"], Dict[str, "JSONType"]]
+MsgPackType = Union[
+    str,
+    int,
+    float,
+    bool,
+    None,
+    date,
+    Decimal,
+    datetime,
+    time,
+    timedelta,
+    uuid.UUID,
+    "NDArray",
+    List["MsgPackType"],
+    Dict[str, "MsgPackType"],
+]
 
 
 # mostly cribbed from django.core.serializer.DjangoJSONEncoder
-def truss_msgpack_encoder(
+def _truss_msgpack_encoder(
     obj: Union[Decimal, date, time, timedelta, uuid.UUID, Dict],
     chain: Optional[Callable] = None,
 ) -> Dict:
@@ -36,7 +58,7 @@ def truss_msgpack_encoder(
         return obj if chain is None else chain(obj)
 
 
-def truss_msgpack_decoder(obj: Any, chain=None):
+def _truss_msgpack_decoder(obj: Any, chain=None):
     try:
         if b"__dt_datetime_iso__" in obj:
             return datetime.fromisoformat(obj[b"data"])
@@ -58,7 +80,7 @@ def truss_msgpack_decoder(obj: Any, chain=None):
 
 
 # this json object is JSONType + np.array + datetime
-def is_truss_serializable(obj) -> bool:
+def is_truss_serializable(obj: Any) -> bool:
     import numpy as np
 
     # basic JSON types
@@ -72,21 +94,21 @@ def is_truss_serializable(obj) -> bool:
         return False
 
 
-def truss_msgpack_serialize(obj):
+def truss_msgpack_serialize(obj: MsgPackType) -> bytes:
     import msgpack
     import msgpack_numpy as mp_np
 
     return msgpack.packb(
-        obj, default=lambda x: truss_msgpack_encoder(x, chain=mp_np.encode)
+        obj, default=lambda x: _truss_msgpack_encoder(x, chain=mp_np.encode)
     )
 
 
-def truss_msgpack_deserialize(obj):
+def truss_msgpack_deserialize(data: bytes) -> MsgPackType:
     import msgpack
     import msgpack_numpy as mp_np
 
     return msgpack.unpackb(
-        obj, object_hook=lambda x: truss_msgpack_decoder(x, chain=mp_np.decode)
+        data, object_hook=lambda x: _truss_msgpack_decoder(x, chain=mp_np.decode)
     )
 
 
@@ -101,4 +123,4 @@ class DeepNumpyEncoder(json.JSONEncoder):
         elif isinstance(obj, np.ndarray):
             return obj.tolist()
         else:
-            return super(DeepNumpyEncoder, self).default(obj)
+            return super().default(obj)

truss/{server → templates}/shared/util.py

@@ -1,10 +1,11 @@
 import multiprocessing
 import os
 import sys
-from typing import Callable, Dict, List, Mapping, TypeVar
+from typing import List
 
 import psutil
 
+BLOB_DOWNLOAD_TIMEOUT_SECS = 600  # 10 minutes
 # number of seconds to wait for truss server child processes before sending kill signal
 CHILD_PROCESS_WAIT_TIMEOUT_SECONDS = 120
 
@@ -12,9 +13,7 @@ CHILD_PROCESS_WAIT_TIMEOUT_SECONDS = 120
 def model_supports_predict_proba(model: object) -> bool:
     if not hasattr(model, "predict_proba"):
         return False
-    if hasattr(
-        model, "_check_proba"
-    ):  # noqa eg Support Vector Machines *can* predict proba if they made certain choices while training
+    if hasattr(model, "_check_proba"):  # noqa eg Support Vector Machines *can* predict proba if they made certain choices while training
         try:
             model._check_proba()
             return True
@@ -76,12 +75,3 @@ def kill_child_processes(parent_pid: int):
     )
     for process in alive:
         process.kill()
-
-
-X = TypeVar("X")
-Y = TypeVar("Y")
-Z = TypeVar("Z")
-
-
-def transform_keys(d: Mapping[X, Z], fn: Callable[[X], Y]) -> Dict[Y, Z]:
-    return {fn(key): value for key, value in d.items()}

truss/templates/trtllm-audio/model/model.py

@@ -0,0 +1,49 @@
+import base64
+
+import requests
+import sigint_patch
+
+# We patch TensorRT-LLM v0.11 to fix an issue where they set a SIGINT handler
+# which breaks because we do not use the main thread.
+sigint_patch.patch()
+from whisper_trt import WhisperModel
+from whisper_trt.custom_types import WhisperResult
+
+
+class Model:
+    def __init__(self, **kwargs):
+        self._data_dir = kwargs["data_dir"]
+        self._secrets = kwargs["secrets"]
+        self._model = None
+
+    def load(self):
+        self._model = WhisperModel(str(self._data_dir), max_queue_time=0.050)
+
+    def preprocess(self, request: dict):
+        audio_base64 = request.get("audio")
+        audio_url = request.get("url")
+
+        if audio_base64 and audio_url:
+            return {
+                "error": "Only a base64 audio file OR a URL can be passed to the API, not both of them."
+            }
+        if not audio_base64 and not audio_url:
+            return {
+                "error": "Please provide either an audio file in base64 string format or a URL to an audio file."
+            }
+
+        binary_data = None
+
+        if audio_base64:
+            binary_data = base64.b64decode(audio_base64.encode("utf-8"))
+        elif audio_url:
+            resp = requests.get(audio_url)
+            binary_data = resp.content
+        return binary_data, request
+
+    async def predict(self, preprocessed_request) -> WhisperResult:
+        binary_data, request = preprocessed_request
+        waveform = self._model.preprocess_audio(binary_data)
+        return await self._model.transcribe(
+            waveform, language="english", timestamps=True
+        )

truss/templates/trtllm-audio/packages/sigint_patch.py

@@ -0,0 +1,14 @@
+import fileinput
+import sys
+
+MODULE_FILE_PATH = "/usr/local/lib/python3.10/dist-packages/tensorrt_llm/hlapi/utils.py"
+
+
+def patch():
+    search_text = "signal.signal(signal.SIGINT, sigint_handler)"
+
+    with fileinput.FileInput(MODULE_FILE_PATH, inplace=True) as file:
+        for line in file:
+            if search_text in line:
+                line = "    # " + line.lstrip()
+            sys.stdout.write(line)