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

truss/server/common/truss_server.py

@@ -1,389 +0,0 @@
-import asyncio
-import json
-import logging
-import multiprocessing
-import os
-import signal
-import socket
-import sys
-import time
-from collections.abc import Generator
-from pathlib import Path
-from typing import AsyncGenerator, Dict, List, Optional, Union
-
-import truss.server.common.errors as errors
-import truss.server.shared.util as utils
-import uvicorn
-from fastapi import Depends, FastAPI, HTTPException, Request
-from fastapi.responses import ORJSONResponse, StreamingResponse
-from fastapi.routing import APIRoute as FastAPIRoute
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.requests import ClientDisconnect
-from starlette.responses import Response
-from truss.server.common.termination_handler_middleware import (
-    TerminationHandlerMiddleware,
-)
-from truss.server.model_wrapper import ModelWrapper
-from truss.server.shared.logging import setup_logging
-from truss.server.shared.serialization import (
-    DeepNumpyEncoder,
-    truss_msgpack_deserialize,
-    truss_msgpack_serialize,
-)
-
-# [IMPORTANT] A lot of things depend on this currently.
-# Please consider the following when increasing this:
-# 1. Self-termination on model load fail.
-# 2. Graceful termination.
-DEFAULT_NUM_WORKERS = 1
-DEFAULT_NUM_SERVER_PROCESSES = 1
-WORKER_TERMINATION_TIMEOUT_SECS = 120.0
-WORKER_TERMINATION_CHECK_INTERVAL_SECS = 0.5
-
-
-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
-
-
-FORMAT = "%(asctime)s.%(msecs)03d %(name)s %(levelname)s [%(funcName)s():%(lineno)s] %(message)s"
-DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
-INFERENCE_SERVER_FAILED_FILE = Path("~/inference_server_crashed.txt").expanduser()
-logging.basicConfig(level=logging.INFO, format=FORMAT, datefmt=DATE_FORMAT)
-
-
-class UvicornCustomServer(multiprocessing.Process):
-    def __init__(
-        self, config: uvicorn.Config, sockets: Optional[List[socket.socket]] = None
-    ):
-        super().__init__()
-        self.sockets = sockets
-        self.config = config
-
-    def stop(self):
-        self.terminate()
-
-    def run(self):
-        server = uvicorn.Server(config=self.config)
-        asyncio.run(server.serve(sockets=self.sockets))
-
-
-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) -> None:
-        self._model = model
-
-    def _safe_lookup_model(self, model_name: str) -> 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[str, Union[str, bool]]:
-        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 predict(
-        self, model_name: str, request: Request, body_raw: bytes = Depends(parse_body)
-    ) -> Response:
-        """
-        This method calls the user-provided predict method
-        """
-        model: ModelWrapper = self._safe_lookup_model(model_name)
-
-        self.check_healthy(model)
-
-        body: Dict
-        if self.is_binary(request):
-            body = truss_msgpack_deserialize(body_raw)
-        else:
-            try:
-                body = json.loads(body_raw)
-            except json.JSONDecodeError as e:
-                error_message = f"Invalid JSON payload: {str(e)}"
-                logging.error(error_message)
-                raise HTTPException(status_code=400, detail=error_message)
-
-        # calls ModelWrapper.__call__, which runs validate, preprocess, predict, and postprocess
-        response: Union[Dict, Generator] = await model(
-            body,
-            headers=utils.transform_keys(request.headers, lambda key: key.lower()),
-        )
-
-        # In the case that the model returns a Generator object, return a
-        # StreamingResponse instead.
-        if isinstance(response, (AsyncGenerator, Generator)):
-            # media_type in StreamingResponse sets the Content-Type header
-            return StreamingResponse(response, media_type="application/octet-stream")
-
-        response_headers = {}
-        if self.is_binary(request):
-            response_headers["Content-Type"] = "application/octet-stream"
-            return Response(
-                content=truss_msgpack_serialize(response), headers=response_headers
-            )
-        else:
-            response_headers["Content-Type"] = "application/json"
-            return Response(
-                content=json.dumps(response, cls=DeepNumpyEncoder),
-                headers=response_headers,
-            )
-
-    async def schema(self, model_name: str) -> Dict:
-        model: ModelWrapper = self._safe_lookup_model(model_name)
-
-        if model.truss_schema is None:
-            print("DID NOT FIND MODEL SCHEMA")
-            # 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.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.
-
-    """
-
-    def __init__(
-        self,
-        http_port: int,
-        config: Dict,
-        setup_json_logger: bool = True,
-    ):
-        self.http_port = http_port
-        self._config = config
-        self._model = ModelWrapper(self._config)
-        self._endpoints = BasetenEndpoints(self._model)
-        self._setup_json_logger = setup_json_logger
-
-    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()
-
-        if self._setup_json_logger:
-            setup_logging()
-
-        self._model.start_load()
-
-    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"]
-                ),
-                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"],
-                ),
-                # Endpoint aliases for Sagemaker hosting
-                FastAPIRoute(r"/ping", self._endpoints.invocations_ready),
-                FastAPIRoute(
-                    r"/invocations",
-                    self._endpoints.invocations,
-                    methods=["POST"],
-                ),
-            ],
-            exception_handlers={
-                errors.InferenceError: errors.inference_error_handler,
-                errors.ModelNotFound: errors.model_not_found_handler,
-                errors.ModelNotReady: errors.model_not_ready_handler,
-                NotImplementedError: errors.not_implemented_error_handler,
-                HTTPException: errors.http_exception_handler,
-                Exception: errors.generic_exception_handler,
-            },
-        )
-
-        def exit_self():
-            # Note that this kills the current process, the worker process, not
-            # the main truss_server process.
-            utils.kill_child_processes(os.getpid())
-            sys.exit()
-
-        termination_handler_middleware = TerminationHandlerMiddleware(
-            on_stop=lambda: None,
-            on_term=exit_self,
-        )
-        app.add_middleware(BaseHTTPMiddleware, dispatch=termination_handler_middleware)
-        return app
-
-    def start(self):
-        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=DEFAULT_NUM_WORKERS,
-            log_config={
-                "version": 1,
-                "formatters": {
-                    "default": {
-                        "()": "uvicorn.logging.DefaultFormatter",
-                        "datefmt": DATE_FORMAT,
-                        "fmt": "%(asctime)s.%(msecs)03d %(name)s %(levelprefix)s %(message)s",
-                        "use_colors": None,
-                    },
-                    "access": {
-                        "()": "uvicorn.logging.AccessFormatter",
-                        "datefmt": DATE_FORMAT,
-                        "fmt": "%(asctime)s.%(msecs)03d %(name)s %(levelprefix)s %(client_addr)s %(process)s - "
-                        '"%(request_line)s" %(status_code)s',
-                        # noqa: E501
-                    },
-                },
-                "handlers": {
-                    "default": {
-                        "formatter": "default",
-                        "class": "logging.StreamHandler",
-                        "stream": "ext://sys.stderr",
-                    },
-                    "access": {
-                        "formatter": "access",
-                        "class": "logging.StreamHandler",
-                        "stream": "ext://sys.stdout",
-                    },
-                },
-                "loggers": {
-                    "uvicorn": {"handlers": ["default"], "level": "INFO"},
-                    "uvicorn.error": {"level": "INFO"},
-                    "uvicorn.access": {
-                        "handlers": ["access"],
-                        "level": "INFO",
-                        "propagate": False,
-                    },
-                },
-            },
-        )
-
-        # Call this so uvloop gets used
-        cfg.setup_event_loop()
-
-        async def serve():
-            serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            serversocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            serversocket.bind((cfg.host, cfg.port))
-            serversocket.listen(5)
-
-            num_server_procs = self._config.get("runtime", {}).get(
-                "num_workers", DEFAULT_NUM_SERVER_PROCESSES
-            )
-            logging.info(f"starting {num_server_procs} uvicorn server processes")
-            servers: List[UvicornCustomServer] = []
-            for _ in range(num_server_procs):
-                server = UvicornCustomServer(config=cfg, sockets=[serversocket])
-                server.start()
-                servers.append(server)
-
-            def stop_servers():
-                # Send stop signal, then wait for all to exit
-                for server in servers:
-                    # Sends term signal to the process, which should be handled
-                    # by the termination handler.
-                    server.stop()
-
-                termination_check_attempts = int(
-                    WORKER_TERMINATION_TIMEOUT_SECS
-                    / WORKER_TERMINATION_CHECK_INTERVAL_SECS
-                )
-                for _ in range(termination_check_attempts):
-                    time.sleep(WORKER_TERMINATION_CHECK_INTERVAL_SECS)
-                    if utils.all_processes_dead(servers):
-                        return
-
-            for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT]:
-                signal.signal(sig, lambda sig, frame: stop_servers())
-
-        async def servers_task():
-            servers = [serve()]
-            await asyncio.gather(*servers)
-
-        asyncio.run(servers_task())

truss/server/control/patch/model_code_patch_applier.py

@@ -1,46 +0,0 @@
-import logging
-import os
-from pathlib import Path
-from typing import Union
-
-from truss.server.control.patch.types import Action, ModelCodePatch, PackagePatch
-
-
-def apply_code_patch(
-    relative_dir: Path,
-    patch: Union[ModelCodePatch, PackagePatch],
-    logger: logging.Logger,
-):
-    logger.debug(f"Applying code patch {patch.to_dict()}")
-    filepath: Path = relative_dir / patch.path
-    action = patch.action
-
-    if action in [Action.ADD, Action.UPDATE]:
-        filepath.parent.mkdir(parents=True, exist_ok=True)
-        action_log = "Adding" if action == Action.ADD else "Updating"
-        logger.info(f"{action_log} file {filepath}")
-        with filepath.open("w") as file:
-            content = patch.content
-            if content is None:
-                raise ValueError(
-                    "Invalid patch: content of a file update patch should not be None."
-                )
-            file.write(content)
-
-    elif action == Action.REMOVE:
-        if not filepath.exists():
-            logger.warning(f"Could not delete file {filepath}: not found.")
-        elif filepath.is_file():
-            logger.info(f"Deleting file {filepath}")
-            filepath.unlink()
-            # attempt to recursively remove potentially empty directories, if applicable
-            # os.removedirs raises OSError with errno 39 when this process encounters a non-empty dir
-            try:
-                os.removedirs(filepath.parent)
-            except OSError as e:
-                if e.errno == 39:  # Directory not empty
-                    pass
-                else:
-                    raise
-    else:
-        raise ValueError(f"Unknown patch action {action}")

truss/server/control/patch/requirement_name_identifier.py

@@ -1,17 +0,0 @@
-from typing import Dict, List
-
-import pkg_resources
-from pkg_resources.extern.packaging import requirements  # type: ignore
-
-
-def identify_requirement_name(req: str) -> str:
-    try:
-        parsed_req = pkg_resources.Requirement.parse(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()
-
-
-def reqs_by_name(reqs: List[str]) -> Dict[str, str]:
-    return {identify_requirement_name(req): req for req in reqs}

truss/server/inference_server.py

@@ -1,29 +0,0 @@
-import os
-from pathlib import Path
-from typing import Dict
-
-from truss.server.common.truss_server import TrussServer
-from truss.server.shared.logging import setup_logging
-from truss.truss_config import TrussConfig
-
-CONFIG_FILE = "config.yaml"
-
-setup_logging()
-
-
-class ConfiguredTrussServer:
-    _config: Dict
-    _port: int
-
-    def __init__(self, config_path: str, port: int):
-        self._port = port
-        self._config = TrussConfig.from_yaml(Path(config_path)).to_dict(verbose=True)
-
-    def start(self):
-        server = TrussServer(http_port=self._port, config=self._config)
-        server.start()
-
-
-if __name__ == "__main__":
-    env_port = int(os.environ.get("INFERENCE_SERVER_PORT", "8080"))
-    ConfiguredTrussServer(CONFIG_FILE, env_port).start()