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

truss/templates/control/control/helpers/truss_patch/requirement_name_identifier.py

@@ -0,0 +1,66 @@
+import re
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+from urllib.parse import parse_qs, urlparse
+
+from packaging import requirements  # type: ignore
+
+URL_PATTERN = re.compile(r"^(https?|git|svn|hg|bzr)(\+|:\/\/)")
+
+
+def is_url_based_requirement(req: str) -> bool:
+    return bool(URL_PATTERN.match(req.strip()))
+
+
+def get_egg_tag(req: str) -> Optional[List[str]]:
+    if is_url_based_requirement(req):
+        parsed_url = urlparse(req)
+        fragments = parse_qs(parsed_url.fragment)
+        return fragments.get("egg", None)
+    return None
+
+
+def reqs_by_name(reqs: List[str]) -> Dict[str, str]:
+    return {identify_requirement_name(req): req for req in reqs if req.strip()}
+
+
+def identify_requirement_name(req: str) -> str:
+    try:
+        # parse as a url requirement
+        req = req.strip()
+        if is_url_based_requirement(req):
+            parsed_url = urlparse(req)
+            # Identify the package by it's unique components. We can't reliably id a package name
+            # via URL, so we fallback on this instead.
+            vcs_options = ["git", "svn", "hg", "bzr"]
+            for vcs in vcs_options:
+                if req.startswith(vcs + "+"):
+                    return f"{vcs}+{parsed_url.netloc}{parsed_url.path.split('@')[0]}"
+
+            return req
+
+        parsed_req = requirements.Requirement(req)
+        return parsed_req.name  # type: ignore
+    except (requirements.InvalidRequirement, ValueError):
+        # default to the whole line if we can't parse it.
+        return req.strip()
+
+
+@dataclass
+class RequirementMeta:
+    requirement: str
+    name: str
+    is_url_based_requirement: bool
+    egg_tag: Optional[List[str]]
+
+    def __hash__(self):
+        return hash(self.name)
+
+    @staticmethod
+    def from_req(req: str) -> "RequirementMeta":
+        return RequirementMeta(
+            req,
+            identify_requirement_name(req),
+            is_url_based_requirement(req),
+            get_egg_tag(req),
+        )

truss/{server → templates/control}/control/server.py

@@ -1,14 +1,20 @@
 import asyncio
 import os
-from pathlib import Path
 
 import uvicorn
-from truss.server.control.application import create_app
+from application import create_app
 
 CONTROL_SERVER_PORT = int(os.environ.get("CONTROL_SERVER_PORT", "8080"))
 INFERENCE_SERVER_PORT = int(os.environ.get("INFERENCE_SERVER_PORT", "8090"))
 
 
+def _identify_python_executable_path() -> str:
+    if "PYTHON_EXECUTABLE" in os.environ:
+        return os.environ["PYTHON_EXECUTABLE"]
+
+    raise RuntimeError("Unable to find python, make sure it's installed.")
+
+
 class ControlServer:
     def __init__(
         self,
@@ -29,8 +35,7 @@ class ControlServer:
                 "inference_server_home": self._inf_serv_home,
                 "inference_server_process_args": [
                     self._python_executable_path,
-                    "-m",
-                    "truss.server.inference_server",
+                    f"{self._inf_serv_home}/main.py",
                 ],
                 "control_server_host": "0.0.0.0",
                 "control_server_port": self._control_server_port,
@@ -59,8 +64,8 @@ class ControlServer:
 
 if __name__ == "__main__":
     control_server = ControlServer(
-        python_executable_path=os.environ.get("PYTHON_EXECUTABLE", default="python3"),
-        inf_serv_home=os.environ.get("APP_HOME", default=str(Path.cwd())),
+        python_executable_path=_identify_python_executable_path(),
+        inf_serv_home=os.environ["APP_HOME"],
         control_server_port=CONTROL_SERVER_PORT,
         inference_server_port=INFERENCE_SERVER_PORT,
     )

truss/templates/control/requirements.txt

@@ -0,0 +1,9 @@
+dataclasses-json==0.5.7
+truss==0.9.14  # TODO(marius/TaT): remove after TaT.
+fastapi==0.114.1
+uvicorn==0.24.0
+uvloop==0.19.0
+tenacity==8.1.0
+httpx==0.27.0
+python-json-logger==2.0.2
+loguru==0.7.2

truss/templates/custom_python_dx/my_model.py

@@ -0,0 +1,28 @@
+"""
+The `Model` class is an interface between the ML model that you're packaging and the model
+server that you're running it on.
+
+The main method to implement here is:
+* `predict`: runs every time the model server is called. Include any logic for model
+  inference and return the model output.
+
+See https://truss.baseten.co/quickstart for more.
+"""
+
+import truss_chains as baseten
+
+
+class Model(baseten.ModelBase):
+    # Configure resources for your model here.
+    remote_config: baseten.RemoteConfig = baseten.RemoteConfig(name="{{ MODEL_NAME }}")
+
+    def __init__(
+        self, context: baseten.DeploymentContext = baseten.depends_context()
+    ) -> None:
+        # Access secrets via optional `context` variable.
+        # Load model here and assign to self._model.
+        self._model = None
+
+    def predict(self, input_field: int) -> int:
+        # Run model inference here
+        return input_field

truss/templates/docker_server/proxy.conf.jinja

@@ -0,0 +1,42 @@
+server {
+    # We use the proxy_read_timeout directive here (instead of proxy_send_timeout) as it sets the timeout for reading a response from the proxied server vs. setting a timeout for sending a request to the proxied server.
+    listen 8080;
+    client_max_body_size {{client_max_body_size}};
+
+    # Liveness endpoint override
+    location = / {
+        proxy_redirect off;
+        proxy_read_timeout 300s;
+
+        rewrite ^/$ {{liveness_endpoint}} break;
+
+        proxy_pass http://127.0.0.1:{{server_port}};
+    }
+    # Readiness endpoint override
+    location ~ ^/v1/models/model$ {
+        proxy_redirect off;
+        proxy_read_timeout 300s;
+
+        rewrite ^/v1/models/model$ {{readiness_endpoint}} break;
+
+        proxy_pass http://127.0.0.1:{{server_port}};
+    }
+    # Predict
+    location ~ ^/v1/models/model:predict$ {
+        proxy_redirect off;
+        proxy_read_timeout 300s;
+
+        rewrite ^/v1/models/model:predict$ {{server_endpoint}} break;
+
+        proxy_pass http://127.0.0.1:{{server_port}};
+    }
+
+    # Forward all other paths
+    location / {
+        proxy_redirect off;
+        proxy_read_timeout 300s;
+
+        proxy_pass http://127.0.0.1:{{server_port}};
+    }
+
+}

truss/templates/docker_server/supervisord.conf.jinja

@@ -0,0 +1,27 @@
+[supervisord]
+nodaemon=true                ; Run supervisord in the foreground (useful for containers)
+logfile=/dev/null            ; Disable logging to file (send logs to /dev/null)
+logfile_maxbytes=0           ; No size limit on logfile (since logging is disabled)
+
+[program:model-server]
+command={{start_command}}    ; Command to start the model server (provided by Jinja variable)
+startsecs=30                 ; Wait 30 seconds before assuming the server is running
+autostart=true               ; Automatically start the program when supervisord starts
+autorestart=true             ; Always restart the program if it exits, no matter what the exit code
+stdout_logfile=/dev/fd/1     ; Send stdout to the first file descriptor (stdout)
+stdout_logfile_maxbytes=0    ; No size limit on stdout log
+redirect_stderr=true         ; Redirect stderr to stdout
+
+[program:nginx]
+command=nginx -g "daemon off;" ; Command to start nginx without daemonizing (keeps it in the foreground)
+startsecs=0                   ; Assume nginx starts immediately
+autostart=true                ; Automatically start nginx when supervisord starts
+autorestart=true              ; Always restart the program if it exits, no matter what the exit code
+stdout_logfile=/dev/fd/1      ; Send nginx stdout to the first file descriptor (stdout)
+stdout_logfile_maxbytes=0     ; No size limit on stdout log
+redirect_stderr=true          ; Redirect nginx stderr to stdout
+
+[eventlistener:quit_on_failure]
+events=PROCESS_STATE_FATAL    ; Listen for fatal process state events
+command=sh -c 'echo "READY"; read line; kill -15 1; echo "RESULT 2";'
+                              ; Stop supervisord if a fatal event occurs (by sending SIGTERM to PID 1)

truss/templates/docker_server_requirements.txt

@@ -0,0 +1 @@
+supervisor==4.2.5

truss/templates/server/common/errors.py

@@ -0,0 +1,231 @@
+import contextlib
+import logging
+import sys
+import textwrap
+from http import HTTPStatus
+from types import TracebackType
+from typing import Generator, Mapping, Optional, Tuple, Type, Union
+
+import fastapi
+import pydantic
+import starlette.responses
+from fastapi.responses import JSONResponse
+
+# See https://github.com/basetenlabs/baseten/blob/master/docs/Error-Propagation.md
+_TRUSS_SERVER_SERVICE_ID = 4
+_BASETEN_UNEXPECTED_ERROR = 500
+_BASETEN_DOWNSTREAM_ERROR_CODE = 600
+_BASETEN_CLIENT_ERROR_CODE = 700
+
+
+class ModelMissingError(Exception):
+    def __init__(self, path):
+        self.path = path
+
+    def __str__(self):
+        return self.path
+
+
+class ModelNotReady(RuntimeError):
+    def __init__(self, model_name: str, detail: Optional[str] = None):
+        self.model_name = model_name
+        self.error_msg = f"Model with name {self.model_name} is not ready."
+        if detail:
+            self.error_msg = self.error_msg + " " + detail
+
+    def __str__(self):
+        return self.error_msg
+
+
+class InputParsingError(ValueError):
+    pass
+
+
+class UserCodeError(Exception):
+    pass
+
+
+class ModelDefinitionError(TypeError):
+    """When the user-defined truss model does not meet the contract."""
+
+
+def _make_baseten_error_headers(error_code: int) -> Mapping[str, str]:
+    return {
+        "X-BASETEN-ERROR-SOURCE": f"{_TRUSS_SERVER_SERVICE_ID:02}",
+        "X-BASETEN-ERROR-CODE": f"{error_code:03}",
+    }
+
+
+def add_error_headers_to_user_response(response: starlette.responses.Response) -> None:
+    response.headers.update(_make_baseten_error_headers(_BASETEN_CLIENT_ERROR_CODE))
+
+
+def _make_baseten_response(
+    http_status: int, info: Union[str, Exception], baseten_error_code: int
+) -> fastapi.Response:
+    msg = str(info) if isinstance(info, Exception) else info
+    return JSONResponse(
+        status_code=http_status,
+        content={"error": msg},
+        headers=_make_baseten_error_headers(baseten_error_code),
+    )
+
+
+async def exception_handler(_: fastapi.Request, exc: Exception) -> fastapi.Response:
+    if isinstance(exc, ModelMissingError):
+        return _make_baseten_response(
+            HTTPStatus.NOT_FOUND.value, exc, _BASETEN_DOWNSTREAM_ERROR_CODE
+        )
+    if isinstance(exc, ModelNotReady):
+        return _make_baseten_response(
+            HTTPStatus.SERVICE_UNAVAILABLE.value, exc, _BASETEN_DOWNSTREAM_ERROR_CODE
+        )
+    if isinstance(exc, InputParsingError):
+        return _make_baseten_response(
+            HTTPStatus.BAD_REQUEST.value, exc, _BASETEN_CLIENT_ERROR_CODE
+        )
+    if isinstance(exc, ModelDefinitionError):
+        return _make_baseten_response(
+            HTTPStatus.PRECONDITION_FAILED.value,
+            f"{type(exc).__name__}: {str(exc)}",
+            _BASETEN_DOWNSTREAM_ERROR_CODE,
+        )
+    if isinstance(exc, UserCodeError):
+        return _make_baseten_response(
+            HTTPStatus.INTERNAL_SERVER_ERROR.value,
+            "Internal Server Error",
+            _BASETEN_DOWNSTREAM_ERROR_CODE,
+        )
+    if isinstance(exc, fastapi.HTTPException):
+        # This is a pass through, but additionally adds our custom error headers.
+        return _make_baseten_response(
+            exc.status_code, exc.detail, _BASETEN_DOWNSTREAM_ERROR_CODE
+        )
+    # Fallback case.
+    return _make_baseten_response(
+        HTTPStatus.INTERNAL_SERVER_ERROR.value,
+        f"Unhandled exception: {type(exc).__name__}: {str(exc)}",
+        _BASETEN_UNEXPECTED_ERROR,
+    )
+
+
+HANDLED_EXCEPTIONS = {
+    ModelMissingError,
+    ModelNotReady,
+    NotImplementedError,
+    InputParsingError,
+    UserCodeError,
+    ModelDefinitionError,
+    fastapi.HTTPException,
+}
+
+
+def filter_traceback(
+    model_file_name: str,
+) -> Union[
+    Tuple[Type[BaseException], BaseException, TracebackType], Tuple[None, None, None]
+]:
+    exc_type, exc_value, tb = sys.exc_info()
+    if tb is None:
+        return exc_type, exc_value, tb  # type: ignore[return-value]
+
+    # Store the last occurrence of the traceback matching the condition
+    last_matching_tb: Optional[TracebackType] = None
+    current_tb: Optional[TracebackType] = tb
+
+    while current_tb is not None:
+        filename = current_tb.tb_frame.f_code.co_filename
+        if filename.endswith(model_file_name):
+            last_matching_tb = current_tb
+        current_tb = current_tb.tb_next
+
+    # If a match was found, truncate the traceback at the last occurrence
+    if last_matching_tb is not None:
+        return exc_type, exc_value, last_matching_tb  # type: ignore[return-value]
+
+    # If `model_file_name` not found, return the original exception info
+    return exc_type, exc_value, tb  # type: ignore[return-value]
+
+
+@contextlib.contextmanager
+def intercept_exceptions(
+    logger: logging.Logger, model_file_name: str
+) -> Generator[None, None, None]:
+    try:
+        yield
+    # Note that logger.error logs the stacktrace, such that the user can
+    # debug this error from the logs.
+    except fastapi.HTTPException as e:
+        # TODO: we try to avoid any dependency of the truss server on chains, but for
+        #  the purpose of getting readable chained-stack traces in the server logs,
+        #  we have to add a special-case here.
+        if "user_stack_trace" in e.detail:
+            try:
+                from truss_chains import definitions
+
+                chains_error = definitions.RemoteErrorDetail.model_validate(e.detail)
+                # The formatted error contains a (potentially chained) stack trace
+                # with all framework code removed, see
+                # truss_chains/remote_chainlet/utils.py::response_raise_errors.
+                logger.error(f"Chainlet raised Exception:\n{chains_error.format()}")
+            except:  # If we cannot import chains or parse the error.
+                logger.error(
+                    "Model raised HTTPException",
+                    exc_info=filter_traceback(model_file_name),
+                )
+                raise
+            # If error was extracted successfully, the customized stack trace is
+            # already printed above, so we raise with a clear traceback.
+            e.__traceback__ = None
+            raise e from None
+
+        logger.error(
+            "Model raised HTTPException", exc_info=filter_traceback(model_file_name)
+        )
+        raise
+    except Exception as exc:
+        logger.error(
+            "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/{server → templates/server}/common/patches/whisper/patch.py

@@ -7,6 +7,7 @@ to fail. This patch makes the download more defensive, addressing this issue. tq
 works when the Content-Length is None.
 Open PR (https://github.com/openai/whisper/pull/1366) - once merged we can remove the patch.
 """
+
 import os
 from typing import Union
 

truss/{server/common/patches/__init__.py → templates/server/common/patches.py}

@@ -2,8 +2,6 @@ import importlib.util
 import logging
 from pathlib import Path
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
@@ -12,7 +10,7 @@ def apply_patches(enabled: bool, requirements: list):
     Apply patches to certain functions. The patches are contained in each patch module under 'patches' directory.
     If a patch cannot be applied, it logs the name of the function and the exception details.
     """
-    PATCHES_DIR = Path(__file__).parent
+    PATCHES_DIR = Path(__file__).parent / "patches"
     if not enabled:
         return
     for requirement in requirements:

truss/{server → templates/server}/common/retry.py

@@ -2,6 +2,7 @@ import time
 from typing import Callable
 
 
+# TODO: replace with tenacity.
 def retry(
     fn: Callable,
     count: int,

truss/{server → templates/server}/common/schema.py

@@ -2,12 +2,15 @@ from types import MappingProxyType
 from typing import (
     Any,
     AsyncGenerator,
+    AsyncIterator,
     Awaitable,
     Generator,
+    Iterator,
     List,
     Optional,
     Type,
     Union,
+    cast,
     get_args,
     get_origin,
 )
@@ -61,10 +64,7 @@ class TrussSchema(BaseModel):
 
 def _parse_input_type(input_parameters: MappingProxyType) -> Optional[type]:
     parameter_types = list(input_parameters.values())
-
-    if len(parameter_types) > 1:
-        return None
-
+    # In `ArgConfig.from_signature` the arguments are validated.
     input_type = parameter_types[0].annotation
 
     if _annotation_is_pydantic_model(input_type):
@@ -85,7 +85,7 @@ def _annotation_is_pydantic_model(annotation: Any) -> bool:
 
 def _parse_output_type(output_annotation: Any) -> Optional[OutputType]:
     """
-    Therea are 4 possible cases for output_annotation:
+    There are 4 possible cases for output_annotation:
     1. Data object -- represented by a Pydantic BaseModel
     2. Streaming -- represented by a Generator or AsyncGenerator
     3. Async -- represented by an Awaitable
@@ -119,7 +119,7 @@ def _parse_output_type(output_annotation: Any) -> Optional[OutputType]:
 def _is_generator_type(annotation: Any) -> bool:
     base_type = get_origin(annotation)
     return isinstance(base_type, type) and issubclass(
-        base_type, (Generator, AsyncGenerator)
+        base_type, (Generator, AsyncGenerator, Iterator, AsyncIterator)
     )
 
 
@@ -158,11 +158,13 @@ def _extract_pydantic_base_models(union_args: tuple) -> List[Type[BaseModel]]:
     2. Union[Awaitable[PydanticBaseModel], AsyncGenerator]
     So for Awaitables, we need to extract the base class from the Awaitable type
     """
+
     return [
-        retrieve_base_class_from_awaitable(arg) if _is_awaitable_type(arg) else arg
-        for arg in union_args
+        cast(Type[BaseModel], retrieve_base_class_from_awaitable(arg))
         if _is_awaitable_type(arg)
-        or (isinstance(arg, type) and issubclass(arg, BaseModel))
+        else arg
+        for arg in union_args
+        if _is_awaitable_type(arg) or _annotation_is_pydantic_model(arg)
     ]