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

truss/templates/server/requirements.txt

@@ -0,0 +1,21 @@
+-i https://pypi.org/simple
+
+aiocontextvars==0.2.2
+argparse==1.4.0
+cython==3.0.5
+fastapi==0.114.1
+joblib==1.2.0
+loguru==0.7.2
+msgpack-numpy==0.4.8
+msgpack==1.1.0  # Numpy/msgpack versions are finniky (1.0.2 breaks), double check when changing.
+numpy>=1.23.5
+opentelemetry-api>=1.25.0
+opentelemetry-sdk>=1.25.0
+opentelemetry-exporter-otlp>=1.25.0
+psutil==5.9.4
+python-json-logger==2.0.2
+pyyaml==6.0.0
+requests==2.31.0
+uvicorn==0.24.0
+uvloop==0.19.0
+aiofiles==24.1.0

truss/templates/server/truss_server.py

@@ -0,0 +1,447 @@
+import asyncio
+import json
+import logging
+import logging.config
+import os
+import signal
+import sys
+from http import HTTPStatus
+from pathlib import Path
+from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Optional, Union
+
+import pydantic
+import uvicorn
+import yaml
+from common import errors, tracing
+from common.schema import TrussSchema
+from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi.responses import ORJSONResponse, StreamingResponse
+from fastapi.routing import APIRoute as FastAPIRoute
+from model_wrapper import MODEL_BASENAME, MethodName, ModelWrapper
+from opentelemetry import propagate as otel_propagate
+from opentelemetry import trace
+from opentelemetry.sdk import trace as sdk_trace
+from pydantic import BaseModel
+from shared import log_config, serialization
+from shared.secrets_resolver import SecretsResolver
+from starlette.requests import ClientDisconnect
+from starlette.responses import Response
+
+if sys.version_info >= (3, 9):
+    from typing import AsyncGenerator, Generator
+else:
+    from typing_extensions import AsyncGenerator, Generator
+
+PYDANTIC_MAJOR_VERSION = int(pydantic.VERSION.split(".")[0])
+
+# [IMPORTANT] A lot of things depend on this currently, change with extreme care.
+TIMEOUT_GRACEFUL_SHUTDOWN = 120
+INFERENCE_SERVER_FAILED_FILE = Path("~/inference_server_crashed.txt").expanduser()
+
+if TYPE_CHECKING:
+    from model_wrapper import InputType, MethodDescriptor, OutputType
+
+
+async def parse_body(request: Request) -> bytes:
+    """
+    Used by FastAPI to read body in an asynchronous manner
+    """
+    try:
+        return await request.body()
+    except ClientDisconnect as exc:
+        error_message = "Client disconnected"
+        logging.error(error_message)
+        raise HTTPException(status_code=499, detail=error_message) from exc
+
+
+class BasetenEndpoints:
+    """The implementation of the model server endpoints.
+
+    Historically, we relied on the kserve server interface, which assumes that
+    multiple models are running behind a registry. As a result, some arguments to
+    to functions will rename unused except for backwards compatibility checks.
+    """
+
+    def __init__(self, model: ModelWrapper, tracer: sdk_trace.Tracer) -> None:
+        self._model = model
+        self._tracer = tracer
+
+    def _safe_lookup_model(self, model_name: str = MODEL_BASENAME) -> ModelWrapper:
+        if model_name != self._model.name:
+            raise errors.ModelMissingError(model_name)
+        return self._model
+
+    @staticmethod
+    def check_healthy(model: ModelWrapper):
+        if model.load_failed:
+            INFERENCE_SERVER_FAILED_FILE.touch()
+            os.kill(os.getpid(), signal.SIGKILL)
+
+        if not model.ready:
+            raise errors.ModelNotReady(model.name)
+
+    async def model_ready(self, model_name: str) -> dict:
+        model: ModelWrapper = self._safe_lookup_model(model_name)
+        is_healthy = await model.is_healthy()
+        if is_healthy is None:
+            self.check_healthy(model)
+        elif not is_healthy:
+            raise errors.ModelNotReady(model.name)
+
+        return {}
+
+    async def model_loaded(self, model_name: str) -> dict:
+        self.check_healthy(self._safe_lookup_model(model_name))
+
+        return {}
+
+    async def invocations_ready(self) -> Dict[str, Union[str, bool]]:
+        """
+        This method provides compatibility with Sagemaker hosting for the 'ping' endpoint.
+        """
+        if self._model is None:
+            raise errors.ModelMissingError("model")
+        self.check_healthy(self._model)
+
+        return {}
+
+    async def invocations(
+        self, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        """
+        This method provides compatibility with Sagemaker hosting for the 'invocations' endpoint.
+        """
+        return await self.predict(self._model.name, request, body_raw)
+
+    async def _parse_body(
+        self,
+        request: Request,
+        body_raw: bytes,
+        truss_schema: Optional[TrussSchema],
+        span: trace.Span,
+    ) -> "InputType":
+        if self.is_binary(request):
+            with tracing.section_as_event(span, "binary-deserialize"):
+                inputs = serialization.truss_msgpack_deserialize(body_raw)
+            if truss_schema:
+                try:
+                    with tracing.section_as_event(span, "parse-pydantic"):
+                        inputs = truss_schema.input_type.parse_obj(inputs)
+                except pydantic.ValidationError as e:
+                    raise errors.InputParsingError(
+                        errors.format_pydantic_validation_error(e)
+                    ) from e
+        else:
+            if truss_schema:
+                try:
+                    with tracing.section_as_event(span, "parse-pydantic"):
+                        inputs = truss_schema.input_type.parse_raw(body_raw)
+                except pydantic.ValidationError as e:
+                    raise errors.InputParsingError(
+                        errors.format_pydantic_validation_error(e)
+                    ) from e
+            else:
+                try:
+                    with tracing.section_as_event(span, "json-deserialize"):
+                        inputs = json.loads(body_raw)
+                except json.JSONDecodeError as e:
+                    raise errors.InputParsingError(
+                        f"Invalid JSON payload: {str(e)}"
+                    ) from e
+
+        return inputs
+
+    async def _execute_request(
+        self,
+        model: ModelWrapper,
+        method: Callable[["InputType", Request], Awaitable["OutputType"]],
+        method_name: MethodName,
+        request: Request,
+        body_raw: bytes,
+    ) -> Response:
+        """
+        Executes a predictive endpoint
+        """
+        if await request.is_disconnected():
+            msg = f"Client disconnected. Skipping `{method_name}`."
+            logging.info(msg)
+            raise ClientDisconnect(msg)
+
+        self.check_healthy(model)
+        trace_ctx = otel_propagate.extract(request.headers) or None
+        # This is the top-level span in the truss-server, so we set the context here.
+        # Nested spans "inherit" context automatically.
+        with self._tracer.start_as_current_span(
+            f"{method_name}-endpoint", context=trace_ctx
+        ) as span:
+            inputs: Optional["InputType"]
+            if model.model_descriptor.skip_input_parsing:
+                inputs = None
+            else:
+                inputs = await self._parse_body(
+                    request, body_raw, model.model_descriptor.truss_schema, span
+                )
+            with tracing.section_as_event(span, "model-call"):
+                result: "OutputType" = await method(inputs, request)
+
+            # In the case that the model returns a Generator object, return a
+            # StreamingResponse instead.
+            if isinstance(result, (AsyncGenerator, Generator)):
+                # media_type in StreamingResponse sets the Content-Type header
+                return StreamingResponse(result, media_type="application/octet-stream")
+            elif isinstance(result, Response):
+                if result.status_code >= HTTPStatus.MULTIPLE_CHOICES.value:
+                    errors.add_error_headers_to_user_response(result)
+                return result
+            return self._serialize_result(result, self.is_binary(request), span)
+
+    async def chat_completions(
+        self, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model()
+        self._raise_if_not_supported(
+            MethodName.CHAT_COMPLETIONS, model.model_descriptor.chat_completions
+        )
+
+        return await self._execute_request(
+            model=model,
+            method=model.chat_completions,
+            method_name=MethodName.CHAT_COMPLETIONS,
+            request=request,
+            body_raw=body_raw,
+        )
+
+    def _raise_if_not_supported(
+        self, method_name: MethodName, descriptor: Optional["MethodDescriptor"]
+    ):
+        if not descriptor:
+            raise HTTPException(status_code=404, detail=f"{method_name} not supported.")
+
+    async def completions(
+        self, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model()
+        self._raise_if_not_supported(
+            MethodName.COMPLETIONS, model.model_descriptor.completions
+        )
+
+        return await self._execute_request(
+            model=model,
+            method=model.completions,
+            method_name=MethodName.COMPLETIONS,
+            request=request,
+            body_raw=body_raw,
+        )
+
+    async def predict(
+        self, model_name: str, request: Request, body_raw: bytes = Depends(parse_body)
+    ) -> Response:
+        model = self._safe_lookup_model(model_name)
+
+        return await self._execute_request(
+            model=model,
+            method=model,  # We overwrote __call__ on ModelWrapper
+            method_name=MethodName.PREDICT,
+            request=request,
+            body_raw=body_raw,
+        )
+
+    def _serialize_result(
+        self, result: "OutputType", is_binary: bool, span: trace.Span
+    ) -> Response:
+        response_headers = {}
+        if is_binary:
+            if isinstance(result, BaseModel):
+                with tracing.section_as_event(span, "binary-dump"):
+                    if PYDANTIC_MAJOR_VERSION > 1:
+                        result = result.model_dump(mode="python")
+                    else:
+                        result = result.dict()
+            # If the result is not already serialize and not a pydantic model, it must
+            # be something that can be serialized with `truss_msgpack_serialize` (some
+            # dict / nested structure).
+            if not isinstance(result, bytes):
+                with tracing.section_as_event(span, "binary-serialize"):
+                    result = serialization.truss_msgpack_serialize(result)
+
+            response_headers["Content-Type"] = "application/octet-stream"
+            return Response(content=result, headers=response_headers)
+        else:
+            with tracing.section_as_event(span, "json-serialize"):
+                if isinstance(result, BaseModel):
+                    # Note: chains has a pydantic integration for numpy arrays
+                    # `NumpyArrayField`. `result.dict()`, passes through the array
+                    # object which cannot be JSON serialized.
+                    # In pydantic v2 `result.model_dump(mode="json")` could be used.
+                    # For backwards compatibility we dump directly the JSON string.
+                    content = result.json()
+                else:
+                    content = json.dumps(result, cls=serialization.DeepNumpyEncoder)
+
+                response_headers["Content-Type"] = "application/json"
+                return Response(content=content, headers=response_headers)
+
+    async def schema(self, model_name: str) -> Dict:
+        model: ModelWrapper = self._safe_lookup_model(model_name)
+        if model.model_descriptor.truss_schema is None:
+            # If there is not a TrussSchema, we return a 404.
+            if model.ready:
+                raise HTTPException(status_code=404, detail="No schema found")
+            else:
+                raise HTTPException(
+                    status_code=503,
+                    detail="Schema not available, please try again later.",
+                )
+        else:
+            return model.model_descriptor.truss_schema.serialize()
+
+    @staticmethod
+    def is_binary(request: Request):
+        return (
+            "Content-Type" in request.headers
+            and request.headers["Content-Type"] == "application/octet-stream"
+        )
+
+
+class TrussServer:
+    """This wrapper class manages creation and cleanup of uvicorn server processes
+    running the FastAPI inference server app.
+
+    TrussServer runs as a main process managing UvicornCustomServer subprocesses that
+    in turn may manage their own worker processes. Notably, this main process is kept
+    alive when running `servers_task()` because of the child uvicorn server processes'
+    main loop.
+    """
+
+    _server: Optional[uvicorn.Server]
+
+    def __init__(self, http_port: int, config_or_path: Union[str, Path, Dict]):
+        # This is run before uvicorn is up. Need explicit logging config here.
+        logging.config.dictConfig(log_config.make_log_config("INFO"))
+
+        if isinstance(config_or_path, (str, Path)):
+            with open(config_or_path, encoding="utf-8") as config_file:
+                config = yaml.safe_load(config_file)
+        else:
+            config = config_or_path
+
+        secrets = SecretsResolver.get_secrets(config)
+        tracer = tracing.get_truss_tracer(secrets, config)
+        self._http_port = http_port
+        self._config = config
+        self._model = ModelWrapper(self._config, tracer)
+        self._endpoints = BasetenEndpoints(self._model, tracer)
+        self._server = None
+
+    def cleanup(self):
+        if INFERENCE_SERVER_FAILED_FILE.exists():
+            INFERENCE_SERVER_FAILED_FILE.unlink()
+
+    def on_startup(self):
+        """
+        This method will be started inside the main process, so here is where
+        we want to setup our logging and model.
+        """
+        self.cleanup()
+        self._model.start_load_thread()
+        asyncio.create_task(self._shutdown_if_load_fails())
+        self._model.setup_polling_for_environment_updates()
+
+    async def _shutdown_if_load_fails(self):
+        while not self._model.ready:
+            await asyncio.sleep(0.5)
+            if self._model.load_failed:
+                assert self._server is not None
+                logging.info("Trying shut down after failed model load.")
+                self._server.should_exit = True
+                return
+
+    def create_application(self):
+        app = FastAPI(
+            title="Baseten Inference Server",
+            docs_url=None,
+            redoc_url=None,
+            default_response_class=ORJSONResponse,
+            on_startup=[self.on_startup],
+            routes=[
+                # liveness endpoint
+                FastAPIRoute(r"/", lambda: True),
+                # readiness endpoint
+                FastAPIRoute(
+                    r"/v1/models/{model_name}", self._endpoints.model_ready, tags=["V1"]
+                ),
+                # loaded endpoint
+                FastAPIRoute(
+                    r"/v1/models/{model_name}/loaded",
+                    self._endpoints.model_loaded,
+                    tags=["V1"],
+                ),
+                FastAPIRoute(
+                    r"/v1/models/{model_name}/schema",
+                    self._endpoints.schema,
+                    methods=["GET"],
+                    tags=["V1"],
+                ),
+                FastAPIRoute(
+                    r"/v1/models/{model_name}:predict",
+                    self._endpoints.predict,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
+                FastAPIRoute(
+                    r"/v1/models/{model_name}:predict_binary",
+                    self._endpoints.predict,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
+                # OpenAI Spec
+                FastAPIRoute(
+                    r"/v1/chat/completions",
+                    self._endpoints.chat_completions,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
+                FastAPIRoute(
+                    r"/v1/completions",
+                    self._endpoints.completions,
+                    methods=["POST"],
+                    tags=["V1"],
+                ),
+                # Endpoint aliases for Sagemaker hosting
+                FastAPIRoute(r"/ping", self._endpoints.invocations_ready),
+                FastAPIRoute(
+                    r"/invocations", self._endpoints.invocations, methods=["POST"]
+                ),
+            ],
+            exception_handlers={
+                exc: errors.exception_handler for exc in errors.HANDLED_EXCEPTIONS
+            },
+        )
+        # Above `exception_handlers` only triggers on exact exception classes.
+        # This here is a fallback to add our custom headers in all other cases.
+        app.add_exception_handler(Exception, errors.exception_handler)
+
+        return app
+
+    def start(self):
+        log_level = (
+            "DEBUG"
+            if self._config["runtime"].get("enable_debug_logs", False)
+            else "INFO"
+        )
+        cfg = uvicorn.Config(
+            self.create_application(),
+            # We hard-code the http parser as h11 (the default) in case the user has
+            # httptools installed, which does not work with our requests & version
+            # of uvicorn.
+            http="h11",
+            host="0.0.0.0",
+            port=self._http_port,
+            workers=1,
+            timeout_graceful_shutdown=TIMEOUT_GRACEFUL_SHUTDOWN,
+            log_config=log_config.make_log_config(log_level),
+        )
+        cfg.setup_event_loop()  # Call this so uvloop gets used
+        server = uvicorn.Server(config=cfg)
+        self._server = server
+        asyncio.run(server.serve())

truss/templates/server.Dockerfile.jinja

@@ -7,8 +7,9 @@
 {% block base_image_patch %}
 # If user base image is supplied in config, apply build commands from truss base image
 {% if config.base_image %}
-ENV PYTHONUNBUFFERED True
-ENV DEBIAN_FRONTEND=noninteractive
+    {%- if not config.docker_server %}
+ENV PYTHONUNBUFFERED="True"
+ENV DEBIAN_FRONTEND="noninteractive"
 
 RUN apt update && \
     apt install -y bash \
@@ -21,7 +22,11 @@ RUN apt update && \
     && apt-get clean -y \
     && rm -rf /var/lib/apt/lists/*
 
-    {%- if config.live_reload %}
+COPY ./{{base_server_requirements_filename}} {{base_server_requirements_filename}}
+RUN pip install -r {{base_server_requirements_filename}} --no-cache-dir && rm -rf /root/.cache/pip
+    {%- endif %}
+
+    {%- if config.live_reload  and not config.docker_server%}
 RUN $PYTHON_EXECUTABLE -m venv -h >/dev/null \
     || { pythonVersion=$(echo $($PYTHON_EXECUTABLE --version) | cut -d" " -f2 | cut -d"." -f1,2) \
     && add-apt-repository -y ppa:deadsnakes/ppa \
@@ -39,6 +44,10 @@ RUN ln -sf {{config.base_image.python_executable_path}} /usr/local/bin/python
 {% endblock %}
 
 {% block install_requirements %}
+    {%- if should_install_server_requirements %}
+COPY ./{{server_requirements_filename}} {{server_requirements_filename}}
+RUN pip install -r {{server_requirements_filename}} --no-cache-dir && rm -rf /root/.cache/pip
+    {%- endif %}
 {{ super() }}
 {% endblock %}
 
@@ -55,26 +64,65 @@ RUN mkdir -p {{ dst.parent }}; curl -L "{{ url }}" -o {{ dst }}
 {% endfor %}
 {%- endif  %}
 
+
+{%- if build_commands %}
+{% for command in build_commands %}
+RUN {% for secret,path in config.build.secret_to_path_mapping.items() %} --mount=type=secret,id={{secret}},target={{path}}{% endfor %} {{ command }}
+{% endfor %}
+{%- endif  %}
+
 # Copy data before code for better caching
 {%- if data_dir_exists %}
 COPY ./{{config.data_dir}} /app/data
 {%- endif %}
 
-COPY ./{{ config.model_module_dir }} /app/model
+{%- if not config.docker_server %}
+COPY ./server /app
+{%- endif %}
+
+{%- if use_local_chains_src %}
+{# This path takes precedence over site-packages. #}
+COPY ./truss_chains /app/truss_chains
+{%- endif %}
+
 COPY ./config.yaml /app/config.yaml
+    {%- if config.live_reload and not config.docker_server%}
+COPY ./control /control
+RUN python3 -m venv /control/.env \
+    && /control/.env/bin/pip3 install -r /control/requirements.txt
+    {%- endif %}
+{%- if model_dir_exists %}
+COPY ./{{ config.model_module_dir }} /app/model
+{%- endif %}
 {% endblock %}
 
+
 {% block run %}
-    {%- if config.live_reload %}
-ENV HASH_TRUSS {{truss_hash}}
-ENV CONTROL_SERVER_PORT 8080
-ENV INFERENCE_SERVER_PORT 8090
-ENV PYTHON_EXECUTABLE {{config.base_image.python_executable_path or "python3"}}
-ENV SERVER_START_CMD="{{config.base_image.python_executable_path or "python3"}} -m truss.server.control.server"
-ENTRYPOINT ["{{config.base_image.python_executable_path or "python3"}}", "-m", "truss.server.control.server"]
+    {%- if config.docker_server %}
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        curl nginx python3-pip && \
+        rm -rf /var/lib/apt/lists/*
+COPY ./docker_server_requirements.txt /app/docker_server_requirements.txt
+RUN pip install -r /app/docker_server_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
+{% set proxy_config_path = "/etc/nginx/conf.d/proxy.conf" %}
+{% set supervisor_config_path = "/etc/supervisor/supervisord.conf" %}
+{% set supervisor_log_dir = "/var/log/supervisor" %}
+{% set supervisor_server_url = "http://localhost:8080" %}
+COPY ./proxy.conf {{ proxy_config_path }}
+RUN mkdir -p {{ supervisor_log_dir }}
+COPY supervisord.conf {{ supervisor_config_path }}
+ENV SUPERVISOR_SERVER_URL="{{ supervisor_server_url }}"
+ENV SERVER_START_CMD="supervisord -c {{ supervisor_config_path }}"
+ENTRYPOINT ["supervisord", "-c", "{{ supervisor_config_path }}"]
+    {%- elif config.live_reload %}
+ENV HASH_TRUSS="{{truss_hash}}"
+ENV CONTROL_SERVER_PORT="8080"
+ENV INFERENCE_SERVER_PORT="8090"
+ENV SERVER_START_CMD="/control/.env/bin/python3 /control/control/server.py"
+ENTRYPOINT ["/control/.env/bin/python3", "/control/control/server.py"]
     {%- else %}
-ENV INFERENCE_SERVER_PORT 8080
-ENV SERVER_START_CMD="{{(config.base_image.python_executable_path or "python3") ~ " -m truss.server.inference_server"}}"
-ENTRYPOINT ["{{config.base_image.python_executable_path or "python3"}}", "-m", "truss.server.inference_server"]
+ENV INFERENCE_SERVER_PORT="8080"
+ENV SERVER_START_CMD="{{(config.base_image.python_executable_path or "python3") ~ " /app/main.py"}}"
+ENTRYPOINT ["{{config.base_image.python_executable_path or "python3"}}", "/app/main.py"]
     {%- endif %}
 {% endblock %}

truss/templates/shared/dynamic_config_resolver.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+from typing import Optional
+
+import aiofiles
+
+DYNAMIC_CONFIG_MOUNT_DIR = "/etc/b10_dynamic_config"
+ENVIRONMENT_DYNAMIC_CONFIG_KEY = "environment"
+
+
+def get_dynamic_config_value_sync(key: str) -> Optional[str]:
+    dynamic_config_path = Path(DYNAMIC_CONFIG_MOUNT_DIR) / key
+    if dynamic_config_path.exists():
+        with dynamic_config_path.open() as dynamic_config_file:
+            return dynamic_config_file.read()
+    return None
+
+
+def get_dynamic_config_file_path(key: str):
+    dynamic_config_path = Path(DYNAMIC_CONFIG_MOUNT_DIR) / key
+    return dynamic_config_path
+
+
+async def get_dynamic_config_value_async(key: str) -> Optional[str]:
+    dynamic_config_path = get_dynamic_config_file_path(key)
+    if dynamic_config_path.exists():
+        async with aiofiles.open(dynamic_config_path, "r") as dynamic_config_file:
+            return await dynamic_config_file.read()
+    return None