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

truss/build/build.py

@@ -1,143 +0,0 @@
-import logging
-import os
-import sys
-from pathlib import Path
-from typing import List, Optional
-
-import yaml
-from truss.constants import CONFIG_FILE, TEMPLATES_DIR, TRUSS
-from truss.docker import kill_containers
-from truss.model_inference import infer_python_version, map_to_supported_python_version
-from truss.notebook import is_notebook_or_ipython
-from truss.truss_config import Build, TrussConfig
-from truss.truss_handle import TrussHandle
-from truss.util.path import build_truss_target_directory, copy_tree_path
-
-logger: logging.Logger = logging.getLogger(__name__)
-
-if is_notebook_or_ipython():
-    logger.setLevel(logging.INFO)
-    logger.addHandler(logging.StreamHandler(sys.stdout))
-
-
-def populate_target_directory(
-    config: TrussConfig,
-    target_directory_path: Optional[str] = None,
-    template: str = "custom",
-    populate_dirs: bool = True,
-) -> Path:
-    target_directory_path_typed = None
-    if target_directory_path is None:
-        target_directory_path_typed = build_truss_target_directory(template)
-    else:
-        target_directory_path_typed = Path(target_directory_path)
-        target_directory_path_typed.mkdir(parents=True, exist_ok=True)
-
-    if populate_dirs:
-        # Create data dir
-        (target_directory_path_typed / config.data_dir).mkdir()
-
-        # Create bundled packages dir
-        # TODO: Drop by default
-        (target_directory_path_typed / config.bundled_packages_dir).mkdir()
-
-        # Create model module dir
-        model_dir = target_directory_path_typed / config.model_module_dir
-        template_path = TEMPLATES_DIR / template
-        copy_tree_path(template_path / "model", model_dir)
-
-    # Write config
-    with (target_directory_path_typed / CONFIG_FILE).open("w") as config_file:
-        yaml.dump(config.to_dict(verbose=False), config_file)
-
-    return target_directory_path_typed
-
-
-def init(
-    target_directory: str,
-    data_files: Optional[List[str]] = None,
-    requirements_file: Optional[str] = None,
-    bundled_packages: Optional[List[str]] = None,
-    build_config: Optional[Build] = None,
-    model_name: Optional[str] = None,
-) -> TrussHandle:
-    """
-    Initialize an empty placeholder Truss. A Truss is a build context designed
-    to be built as a container locally or uploaded into a baseten serving
-    environment. This placeholder structure can be filled to represent ML
-    models.
-
-    Args:
-        target_directory: Absolute or relative path of the directory to create
-                          Truss in. The directory is created if it doesn't exist.
-    """
-    config = TrussConfig(
-        model_name=model_name,
-        python_version=map_to_supported_python_version(infer_python_version()),
-    )
-
-    if build_config:
-        config.build = build_config
-
-    target_directory_path = populate_target_directory(
-        config=config,
-        target_directory_path=target_directory,
-        populate_dirs=True,
-    )
-
-    scaf = TrussHandle(target_directory_path)
-    _update_truss_props(scaf, data_files, requirements_file, bundled_packages)
-    return scaf
-
-
-def load(truss_directory: str) -> TrussHandle:
-    """Get a handle to a Truss. A Truss is a build context designed to be built
-    as a container locally or uploaded into a model serving environment.
-
-    Args:
-        truss_directory (str): The local directory of an existing Truss
-    Returns:
-        TrussHandle
-    """
-    return TrussHandle(Path(truss_directory))
-
-
-def from_directory(*args, **kwargs):
-    logger.warn(
-        "DeprecationWarning: from_directory() is deprecated. Use load() instead."
-    )
-    return load(*args, **kwargs)
-
-
-def cleanup() -> None:
-    """
-    Cleans up .truss directory.
-    """
-    build_folder_path = Path(Path.home(), ".truss")
-    if build_folder_path.exists():
-        for obj in build_folder_path.glob("**/*"):
-            if (not obj.name == "config.yaml") and (obj.is_file()):
-                os.remove(obj)
-    return
-
-
-def _update_truss_props(
-    scaf: TrussHandle,
-    data_files: Optional[List[str]] = None,
-    requirements_file: Optional[str] = None,
-    bundled_packages: Optional[List[str]] = None,
-) -> None:
-    if data_files is not None:
-        for data_file in data_files:
-            scaf.add_data(data_file)
-
-    if bundled_packages is not None:
-        for package in bundled_packages:
-            scaf.add_bundled_package(package)
-
-    if requirements_file is not None:
-        scaf.update_requirements_from_file(requirements_file)
-
-
-def kill_all() -> None:
-    kill_containers({TRUSS: True})

truss/build/configure.py

@@ -1,63 +0,0 @@
-import json
-import logging
-import os
-import sys
-from pathlib import Path
-from typing import Optional
-
-import typer
-from truss import load
-from truss.patch.hash import directory_content_hash
-from truss.patch.signature import calc_truss_signature
-from truss.patch.truss_dir_patch_applier import TrussDirPatchApplier
-from truss.server.control.patch.types import Patch
-
-logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
-
-
-app = typer.Typer()
-
-
-@app.command()
-def configure_truss_for_build(
-    truss_dir: str,
-    build_context_path: str,
-    use_control_server: bool = False,
-    patches_path: Optional[str] = None,
-    hash_file_path: Optional[str] = None,
-    signature_file_path: Optional[str] = None,
-):
-    tr = load(truss_dir)
-
-    if patches_path is not None:
-        logging.info("Applying patches")
-        logger = logging.getLogger("patch_applier")
-        patch_applier = TrussDirPatchApplier(Path(truss_dir), logger)
-        patches = json.loads(Path(patches_path).read_text())
-        patch_applier([Patch.from_dict(patch) for patch in patches])
-
-    # Important to do this before making changes to truss, we want
-    # to capture hash of original truss.
-    if hash_file_path is not None:
-        logging.info("Recording truss hash")
-        Path(hash_file_path).write_text(directory_content_hash(Path(truss_dir)))
-
-    if signature_file_path is not None:
-        logging.info("Recording truss signature")
-        signature_str = json.dumps(calc_truss_signature(Path(truss_dir)).to_dict())
-        Path(signature_file_path).write_text(signature_str)
-
-    tr.live_reload(enable=use_control_server)
-
-    logging.debug("Setting up docker build context for truss")
-
-    # check if we have a hf_secret
-    tr.docker_build_setup(
-        Path(build_context_path), use_hf_secret="HUGGING_FACE_HUB_TOKEN" in os.environ
-    )
-    logging.info("docker build context is set up for the truss")
-
-
-if __name__ == "__main__":
-    # parse the things
-    app()

truss/cli/__init__.py

@@ -1,2 +0,0 @@
-# flake8: noqa
-from truss.cli.cli import truss_cli

truss/cli/console.py

@@ -1,5 +0,0 @@
-import rich
-
-console = rich.console.Console()
-
-error_console = rich.console.Console(stderr=True, style="bold red")

truss/cli/create.py

@@ -1,5 +0,0 @@
-from InquirerPy import inquirer
-
-
-def ask_name() -> str:
-    return inquirer.text(message="What's the name of your model?").execute()

truss/config/trt_llm.py

@@ -1,81 +0,0 @@
-import logging
-from enum import Enum
-from typing import Optional
-
-from pydantic import BaseModel, model_validator
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class TRTLLMModelArchitecture(Enum):
-    LLAMA: str = "llama"
-    MISTRAL: str = "mistral"
-
-
-class TRTLLMQuantizationType(Enum):
-    NO_QUANT: str = "no_quant"
-    WEIGHTS_ONLY_INT8: str = "weights_int8"
-    WEIGHTS_KV_INT8: str = "weights_kv_int8"
-    WEIGHTS_ONLY_INT4: str = "weights_int4"
-    WEIGHTS_KV_INT4: str = "weights_kv_int4"
-    SMOOTH_QUANT: str = "smooth_quant"
-    FP8: str = "fp8"
-    FP8_KV: str = "fp8_kv"
-
-
-class TrussTRTLLMPluginConfiguration(BaseModel):
-    multi_block_mode: bool = False
-    paged_kv_cache: bool = False
-    use_fused_mlp: bool = False
-
-
-class TrussTRTLLMBuildConfiguration(BaseModel):
-    huggingface_ckpt_repository: str
-    base_model_architecture: TRTLLMModelArchitecture
-    max_input_len: int
-    max_output_len: int
-    max_batch_size: int
-    max_beam_width: int = 1
-    max_prompt_embedding_table_size: int = 0
-    gather_all_token_logits: bool = False
-    strongly_typed: bool = False
-    quantization_type: TRTLLMQuantizationType = TRTLLMQuantizationType.NO_QUANT
-    tensor_parallel_count: int = 1
-    pipeline_parallel_count: int = 1
-    plugin_configuration: TrussTRTLLMPluginConfiguration = (
-        TrussTRTLLMPluginConfiguration()
-    )
-
-
-class TrussTRTLLMServingConfiguration(BaseModel):
-    engine_repository: str
-    tokenizer_repository: str
-    tensor_parallel_count: int = 1
-    pipeline_parallel_count: int = 1
-
-
-class TRTLLMConfiguration(BaseModel):
-    serve: Optional[TrussTRTLLMServingConfiguration] = None
-    build: Optional[TrussTRTLLMBuildConfiguration] = None
-
-    @model_validator(mode="after")
-    def check_minimum_required_configuration(self):
-        if not self.serve and not self.build:
-            raise ValueError("Either serve or build configurations must be provided")
-        if self.serve and self.build:
-            raise ValueError("Both serve and build configurations cannot be provided")
-        if self.serve is not None:
-            if (self.serve.engine_repository is None) ^ (
-                self.serve.tokenizer_repository is None
-            ):
-                raise ValueError(
-                    "Both engine_repository and tokenizer_repository must be provided"
-                )
-        return self
-
-    @property
-    def requires_build(self):
-        if self.build is not None:
-            return True
-        return False

truss/constants.py

@@ -1,61 +0,0 @@
-import os
-import pathlib
-
-TRUSS_PACKAGE_DIR = pathlib.Path(__file__).resolve().parent
-
-SKLEARN = "sklearn"
-TENSORFLOW = "tensorflow"
-KERAS = "keras"
-XGBOOST = "xgboost"
-PYTORCH = "pytorch"
-CUSTOM = "custom"
-HUGGINGFACE_TRANSFORMER = "huggingface_transformer"
-LIGHTGBM = "lightgbm"
-
-BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-CODE_DIR = pathlib.Path(BASE_DIR, "truss")
-
-TEMPLATES_DIR = pathlib.Path(CODE_DIR, "templates")
-TRTLLM_TRUSS_DIR: pathlib.Path = TEMPLATES_DIR / "trtllm"
-
-SUPPORTED_PYTHON_VERSIONS = {"3.8", "3.9", "3.10", "3.11"}
-
-
-REQUIREMENTS_TXT_FILENAME = "requirements.txt"
-USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME = "user_requirements.txt"
-SYSTEM_PACKAGES_TXT_FILENAME = "system_packages.txt"
-
-FILENAME_CONSTANTS_MAP = {
-    "config_requirements_filename": REQUIREMENTS_TXT_FILENAME,
-    "user_supplied_requirements_filename": USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
-    "system_packages_filename": SYSTEM_PACKAGES_TXT_FILENAME,
-}
-
-SERVER_DOCKERFILE_TEMPLATE_NAME = "server.Dockerfile.jinja"
-MODEL_DOCKERFILE_NAME = "Dockerfile"
-
-README_TEMPLATE_NAME = "README.md.jinja"
-MODEL_README_NAME = "README.md"
-
-CONFIG_FILE = "config.yaml"
-DOCKERFILE = "Dockerfile"
-# Used to indicate whether to associate a container with Truss
-TRUSS = "truss"
-# Used to create unique identifier based on last time truss was updated
-TRUSS_MODIFIED_TIME = "truss_modified_time"
-# Path of the Truss used to identify which Truss is being referred
-TRUSS_DIR = "truss_dir"
-TRUSS_HASH = "truss_hash"
-
-INFERENCE_SERVER_PORT = 8080
-
-HTTP_PUBLIC_BLOB_BACKEND = "http_public"
-
-REGISTRY_BUILD_SECRET_PREFIX = "DOCKER_REGISTRY_"
-
-TRTLLM_BASE_IMAGE = "baseten/trtllm-build-server:r23.12_baseten_v0.7.1_20240111"
-BASE_TRTLLM_REQUIREMENTS = [
-    "tritonclient[all]==2.42.0",
-    "transformers==4.33.1",
-    "jinja2==3.1.3",
-]

truss/model_inference.py

@@ -1,123 +0,0 @@
-import inspect
-import logging
-import sys
-from ast import ClassDef, FunctionDef
-from dataclasses import dataclass
-from typing import Any, Dict, List, Tuple
-
-logger: logging.Logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ModelBuildStageOne:
-    # the Python Class of the model
-    model_type: str
-    # the framework that the model is built in
-    model_framework: str
-
-
-def _model_class(model: Any):
-    return model.__class__
-
-
-def infer_python_version() -> str:
-    return f"py{sys.version_info.major}{sys.version_info.minor}"
-
-
-def map_to_supported_python_version(python_version: str) -> str:
-    """Map python version to truss supported python version.
-
-    Currently, it maps any versions greater than 3.9 to 3.9.
-
-    Args:
-        python_version: in the form py[major_version][minor_version] e.g. py39,
-        py310
-    """
-    python_major_version = int(python_version[2:3])
-    python_minor_version = int(python_version[3:])
-
-    if python_major_version > 3:
-        raise NotImplementedError("Only python version 3 is supported")
-
-    if python_minor_version > 11:
-        logger.info(
-            f"Mapping python version {python_major_version}.{python_minor_version}"
-            " to 3.11, the highest version that Truss currently supports."
-        )
-        return "py311"
-
-    if python_minor_version < 8:
-        logger.info(
-            f"Mapping python version {python_major_version}.{python_minor_version}"
-            " to 3.8, the lowest version that Truss currently supports."
-        )
-        return "py38"
-
-    return python_version
-
-
-def _infer_model_init_parameters(model_class: Any) -> Tuple[List, List]:
-    full_arg_spec = inspect.getfullargspec(model_class.__init__)
-    named_args = full_arg_spec.args[1:]
-    number_of_kwargs = full_arg_spec.defaults and len(full_arg_spec.defaults) or 0
-    required_args = full_arg_spec.args[1:-number_of_kwargs]
-    return named_args, required_args
-
-
-def _infer_model_init_parameters_ast(model_class_def: ClassDef) -> Tuple[List, List]:
-    named_args: List[str] = []
-    required_args: List[str] = []
-    init_model_functions = [
-        node
-        for node in model_class_def.body
-        if isinstance(node, FunctionDef) and node.name == "__init__"
-    ]
-
-    if not init_model_functions:
-        return named_args, required_args
-
-    assert (
-        len(init_model_functions) == 1
-    ), "There should only be one __init__ function in the model class"
-    init_model_function = init_model_functions[0]
-    named_args = [arg.arg for arg in init_model_function.args.args][1:]
-    number_of_defaults = len(init_model_function.args.defaults)
-    required_args = named_args[:-number_of_defaults]
-    return named_args, required_args
-
-
-def validate_provided_parameters_with_model(
-    model_class: Any, provided_parameters: Dict[str, Any]
-) -> None:
-    """
-    Validates that all provided parameters match the signature of the model.
-
-    Args:
-        model_class: The model class to validate against
-        provided_parameters: The parameters to validate
-    """
-    if type(model_class) == ClassDef:
-        named_args, required_args = _infer_model_init_parameters_ast(model_class)
-    else:
-        named_args, required_args = _infer_model_init_parameters(model_class)
-
-    # Check that there are no extra parameters
-    if not named_args:
-        return
-
-    if provided_parameters and type(provided_parameters) is not dict:
-        raise TypeError(
-            f"Provided parameters must be a dict, not {type(provided_parameters)}"
-        )
-
-    for arg in provided_parameters:
-        if arg not in named_args:
-            raise ValueError(
-                f"Provided parameter {arg} is not a valid init parameter for the model."
-            )
-
-    for arg in required_args:
-        if arg not in provided_parameters:
-            raise ValueError(
-                f"Required init parameter {arg} was not provided for this model."
-            )

truss/patch/types.py

@@ -1,30 +0,0 @@
-from dataclasses import dataclass
-from typing import Dict
-
-
-@dataclass
-class TrussSignature:
-    """Truss signature stores information for calculating patches for future
-    changes to Truss.
-
-    Currently, it stores hashes of all of the paths in the truss directory excluding the data dir,
-    and the truss config contents. Path hashes allow calculating added/updated/removes
-    paths in future trusses compared to this. Config contents allow calculating
-    config changes, such as add/update/remove of python requirements etc.
-    """
-
-    content_hashes_by_path: Dict[str, str]
-    config: str
-
-    def to_dict(self) -> dict:
-        return {
-            "content_hashes_by_path": self.content_hashes_by_path,
-            "config": self.config,
-        }
-
-    @staticmethod
-    def from_dict(d) -> "TrussSignature":
-        return TrussSignature(
-            content_hashes_by_path=d["content_hashes_by_path"],
-            config=d["config"],
-        )

truss/pytest.ini

@@ -1,7 +0,0 @@
-[pytest]
-addopts = --capture=no
-filterwarnings =
-    ignore::DeprecationWarning
-    ignore::PendingDeprecationWarning
-markers =
-    integration: mark tests for use in local integration testing only (deselect with '-m "not integration"')

truss/server/common/errors.py

@@ -1,100 +0,0 @@
-from http import HTTPStatus
-from typing import Optional
-
-from fastapi.responses import JSONResponse
-
-
-class ModelMissingError(Exception):
-    def __init__(self, path):
-        self.path = path
-
-    def __str__(self):
-        return self.path
-
-
-class InferenceError(RuntimeError):
-    def __init__(self, reason):
-        self.reason = reason
-
-    def __str__(self):
-        return self.reason
-
-
-class InvalidInput(ValueError):
-    """
-    Exception class indicating invalid input arguments.
-    HTTP Servers should return HTTP_400 (Bad Request).
-    """
-
-    def __init__(self, reason):
-        self.reason = reason
-
-    def __str__(self):
-        return self.reason
-
-
-class ModelNotFound(Exception):
-    """
-    Exception class indicating requested model does not exist.
-    HTTP Servers should return HTTP_404 (Not Found).
-    """
-
-    def __init__(self, model_name=None):
-        self.reason = f"Model with name {model_name} does not exist."
-
-    def __str__(self):
-        return self.reason
-
-
-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
-
-
-async def exception_handler(_, exc):
-    return JSONResponse(
-        status_code=HTTPStatus.INTERNAL_SERVER_ERROR, content={"error": str(exc)}
-    )
-
-
-async def invalid_input_handler(_, exc):
-    return JSONResponse(status_code=HTTPStatus.BAD_REQUEST, content={"error": str(exc)})
-
-
-async def inference_error_handler(_, exc):
-    return JSONResponse(
-        status_code=HTTPStatus.INTERNAL_SERVER_ERROR, content={"error": str(exc)}
-    )
-
-
-async def generic_exception_handler(_, exc):
-    return JSONResponse(
-        status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
-        content={"error": f"{type(exc).__name__} : {str(exc)}"},
-    )
-
-
-async def model_not_found_handler(_, exc):
-    return JSONResponse(status_code=HTTPStatus.NOT_FOUND, content={"error": str(exc)})
-
-
-async def model_not_ready_handler(_, exc):
-    return JSONResponse(
-        status_code=HTTPStatus.SERVICE_UNAVAILABLE, content={"error": str(exc)}
-    )
-
-
-async def not_implemented_error_handler(_, exc):
-    return JSONResponse(
-        status_code=HTTPStatus.NOT_IMPLEMENTED, content={"error": str(exc)}
-    )
-
-
-async def http_exception_handler(_, exc):
-    return JSONResponse(status_code=exc.status_code, content={"error": exc.detail})

truss/server/common/termination_handler_middleware.py

@@ -1,64 +0,0 @@
-import asyncio
-import signal
-from typing import Callable
-
-from fastapi import Request
-
-# This is to allow the last request's response to finish handling. There may be more
-# middlewares that the response goes through, and then there's the time for the bytes
-# to be pushed to the caller.
-DEFAULT_TERM_DELAY_SECS = 5.0
-
-
-class TerminationHandlerMiddleware:
-    """
-    This middleware allows for swiftly and safely terminating the server. It
-    listens to a set of termination signals. On receiving such a signal, it
-    first informs on the on_stop callback, then waits for currently executing
-    requests to finish, before informing on the on_term callback.
-
-    Stop means that the process to stop the server has started. As soon as
-    outstading requests go to zero after this, on_term will be called.
-
-    Term means that this is the right time to terminate the server process, no
-    outstanding requests at this point.
-
-    The caller would typically handle on_stop by stop sending more requests to
-    the FastApi server. And on_term by exiting the server process.
-    """
-
-    def __init__(
-        self,
-        on_stop: Callable[[], None],
-        on_term: Callable[[], None],
-        termination_delay_secs: float = DEFAULT_TERM_DELAY_SECS,
-    ):
-        self._outstanding_request_count = 0
-        self._on_stop = on_stop
-        self._on_term = on_term
-        self._termination_delay_secs = termination_delay_secs
-        self._stopped = False
-        for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT]:
-            signal.signal(sig, self._stop)
-
-    async def __call__(self, request: Request, call_next):
-        self._outstanding_request_count += 1
-        try:
-            response = await call_next(request)
-        finally:
-            self._outstanding_request_count -= 1
-            if self._outstanding_request_count == 0 and self._stopped:
-                # There's a delay in term to allow some time for current
-                # response flow to finish.
-                asyncio.create_task(self._term())
-        return response
-
-    def _stop(self, sig, frame):
-        self._on_stop()
-        self._stopped = True
-        if self._outstanding_request_count == 0:
-            self._on_term()
-
-    async def _term(self):
-        await asyncio.sleep(self._termination_delay_secs)
-        self._on_term()