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

truss/templates/trtllm/packages/triton_client.py

@@ -1,150 +0,0 @@
-import json
-import os
-import subprocess
-import time
-from pathlib import Path
-from typing import AsyncGenerator, Optional
-
-import tritonclient.grpc.aio as grpcclient
-import tritonclient.http as httpclient
-from constants import (
-    ENTRYPOINT_MODEL_NAME,
-    GRPC_SERVICE_PORT,
-    TENSORRT_LLM_MODEL_REPOSITORY_PATH,
-)
-from schema import ModelInput
-from utils import download_engine, prepare_model_repository
-
-
-class TritonServer:
-    def __init__(self, grpc_port: int = 8001, http_port: int = 8003):
-        self.grpc_port = grpc_port
-        self.http_port = http_port
-        self._server_process = None
-
-    def create_model_repository(
-        self,
-        truss_data_dir: Path,
-        engine_repository_path: Optional[str] = None,
-        huggingface_auth_token: Optional[str] = None,
-    ) -> None:
-        if engine_repository_path:
-            download_engine(
-                engine_repository=engine_repository_path,
-                fp=truss_data_dir,
-                auth_token=huggingface_auth_token,
-            )
-        prepare_model_repository(truss_data_dir)
-        return
-
-    def start(self, tensor_parallelism: int = 1, env: dict = {}) -> None:
-        def build_server_start_command() -> list:
-            """
-            Triton Inference Server has different startup commands depending on the tensor
-            parallelism (TP) configuration. This function starts the server with the appropriate
-            command.
-            """
-            base_command = [
-                "tritonserver",
-                f"--model-repository={TENSORRT_LLM_MODEL_REPOSITORY_PATH}",
-                f"--grpc-port={str(self.grpc_port)}",
-                f"--http-port={str(self.http_port)}",
-                # "--log-verbose=1",
-            ]
-
-            if tensor_parallelism == 1:
-                return base_command
-
-            mpirun_command = ["mpirun", "--allow-run-as-root"]
-            mpi_commands = []
-            for i in range(tensor_parallelism):
-                mpi_command = [
-                    "-n=1",
-                    *base_command,
-                    "--disable-auto-complete-config",
-                    f"--backend-config=python,shm-region-prefix-name=prefix{str(i)}_",
-                ]
-                mpi_commands.append(" ".join(mpi_command))
-
-            combined_mpi_commands = " : ".join(mpi_commands)
-            return mpirun_command + [combined_mpi_commands]
-
-        command = build_server_start_command()
-        self._server_process = subprocess.Popen(  # type: ignore
-            command,
-            env={**os.environ, **env},
-        )
-        while not self.is_alive and not self.is_ready:
-            time.sleep(2)
-        return
-
-    def stop(self):
-        if self._server_process:
-            if self.is_server_ready:
-                self._server_process.kill()
-            self._server_process = None
-        return
-
-    @property
-    def is_alive(self) -> bool:
-        try:
-            http_client = httpclient.InferenceServerClient(
-                url=f"localhost:{self.http_port}", verbose=False
-            )
-            return http_client.is_server_live()
-        except ConnectionRefusedError:
-            return False
-
-    @property
-    def is_ready(self) -> bool:
-        try:
-            http_client = httpclient.InferenceServerClient(
-                url=f"localhost:{self.http_port}", verbose=False
-            )
-            return http_client.is_model_ready(model_name=ENTRYPOINT_MODEL_NAME)
-        except ConnectionRefusedError:
-            return False
-
-
-class TritonClient:
-    def __init__(self, grpc_service_port: int = GRPC_SERVICE_PORT):
-        self.grpc_service_port = grpc_service_port
-        self._grpc_client = None
-
-    def start_grpc_stream(self) -> grpcclient.InferenceServerClient:
-        if self._grpc_client:
-            return self._grpc_client
-
-        self._grpc_client = grpcclient.InferenceServerClient(
-            url=f"localhost:{self.grpc_service_port}", verbose=False
-        )
-        return self._grpc_client
-
-    async def infer(
-        self, model_input: ModelInput, model_name="ensemble"
-    ) -> AsyncGenerator[str, None]:
-        grpc_client_instance = self.start_grpc_stream()
-        inputs = model_input.to_tensors()
-
-        async def input_generator():
-            yield {
-                "model_name": model_name,
-                "inputs": inputs,
-                "request_id": model_input.request_id,
-            }
-
-        response_iterator = grpc_client_instance.stream_infer(
-            inputs_iterator=input_generator(),
-        )
-
-        try:
-            async for response in response_iterator:
-                result, error = response
-                if result:
-                    result = result.as_numpy("text_output")
-                    yield result[0].decode("utf-8")
-                else:
-                    yield json.dumps({"status": "error", "message": error.message()})
-
-        except grpcclient.InferenceServerException as e:
-            print(f"InferenceServerException: {e}")

truss/templates/trtllm/packages/utils.py

@@ -1,43 +0,0 @@
-from pathlib import Path
-
-from constants import TENSORRT_LLM_MODEL_REPOSITORY_PATH
-from huggingface_hub import snapshot_download
-
-
-def move_all_files(src: Path, dest: Path) -> None:
-    """
-    Moves all files from `src` to `dest` recursively.
-    """
-    for item in src.iterdir():
-        dest_item = dest / item.name
-        if item.is_dir():
-            dest_item.mkdir(parents=True, exist_ok=True)
-            move_all_files(item, dest_item)
-        else:
-            item.rename(dest_item)
-
-
-def prepare_model_repository(data_dir: Path) -> None:
-    # Ensure the destination directory exists
-    dest_dir = TENSORRT_LLM_MODEL_REPOSITORY_PATH / "tensorrt_llm" / "1"
-    dest_dir.mkdir(parents=True, exist_ok=True)
-
-    # Ensure empty version directory for `ensemble` model exists
-    ensemble_dir = TENSORRT_LLM_MODEL_REPOSITORY_PATH / "ensemble" / "1"
-    ensemble_dir.mkdir(parents=True, exist_ok=True)
-
-    # Move all files and directories from data_dir to dest_dir
-    move_all_files(data_dir, dest_dir)
-
-
-def download_engine(engine_repository: str, fp: Path, auth_token=None):
-    """
-    Downloads the specified engine from Hugging Face Hub.
-    """
-    snapshot_download(
-        engine_repository,
-        local_dir=fp,
-        local_dir_use_symlinks=False,
-        max_workers=4,
-        **({"use_auth_token": auth_token} if auth_token is not None else {}),
-    )

truss/test_data/context_builder_image_test/test.py

@@ -1,4 +0,0 @@
-from truss import init
-
-th = init("test_truss")
-th.docker_build_setup()

truss/test_data/happy.ipynb

@@ -1,54 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "id": "7a4aad7e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "    import truss\n",
-    "    from sklearn.linear_model import LogisticRegression\n",
-    "    import numpy as np\n",
-    "    import pandas as pd\n",
-    "    import random\n",
-    "    import tempfile\n",
-    "    \n",
-    "    # 10 rows of 10 features\n",
-    "    features = random.choices(range(10), k=10)\n",
-    "    values = np.random.randint(0, 100, size=(10, 10))\n",
-    "    \n",
-    "    df = pd.DataFrame(values, columns=features)\n",
-    "\n",
-    "    lr = LogisticRegression(random_state=0, max_iter=1000)\n",
-    "    lr.fit(df.values, df.iloc[:, 0])\n",
-    "    \n",
-    "    with tempfile.TemporaryDirectory() as tmp:\n",
-    "        tr = truss.create(lr, target_directory=tmp)\n",
-    "        \n",
-    "        assert truss.load(tmp).spec.yaml_string == tr.spec.yaml_string"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.9.9"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

truss/test_data/model_load_failure_test/config.yaml

@@ -1,2 +0,0 @@
-model_name: Test Loaf Failure
-python_version: py39

truss/test_data/test_concurrency_truss/config.yaml

@@ -1,2 +0,0 @@
-runtime:
-  predict_concurrency: 2

truss/test_data/test_streaming_async_generator_truss/config.yaml

@@ -1,2 +0,0 @@
-model_name: Test Streaming Async Generator
-python_version: py39

truss/test_data/test_streaming_truss/config.yaml

@@ -1,3 +0,0 @@
-
-model_name: Test Streaming
-python_version: py39

truss/test_data/test_truss/config.yaml

@@ -1,2 +0,0 @@
-model_name: Test
-python_version: py39

truss/tests/server/common/test_termination_handler_middleware.py

@@ -1,93 +0,0 @@
-import multiprocessing
-import tempfile
-import time
-from pathlib import Path
-from typing import Awaitable, Callable, List
-
-import pytest
-from truss.server.common.termination_handler_middleware import (
-    TerminationHandlerMiddleware,
-)
-
-
-async def noop(*args, **kwargs):
-    return
-
-
-@pytest.mark.integration
-def test_termination_sequence_no_pending_requests(tmp_path):
-    # Create middleware in separate process, on sending term signal to process,
-    # it should print the right messages.
-    def main_coro_gen(middleware: TerminationHandlerMiddleware):
-        import asyncio
-
-        async def main(*args, **kwargs):
-            await middleware(1, call_next=noop)
-            await asyncio.sleep(1)
-            print("should not print due to termination")
-
-        return main()
-
-    _verify_term(main_coro_gen, ["stopped", "terminated"])
-
-
-@pytest.mark.integration
-def test_termination_sequence_with_pending_requests(tmp_path):
-    def main_coro_gen(middleware: TerminationHandlerMiddleware):
-        import asyncio
-
-        async def main(*args, **kwargs):
-            async def call_next(req):
-                await asyncio.sleep(1.0)
-                return "call_next_called"
-
-            resp = await middleware(1, call_next=call_next)
-            print(f"call_next response: {resp}")
-            await asyncio.sleep(1)
-            print("should not print due to termination")
-
-        return main()
-
-    _verify_term(
-        main_coro_gen,
-        [
-            "stopped",
-            "call_next response: call_next_called",
-            "terminated",
-        ],
-    )
-
-
-def _verify_term(
-    main_coro_gen: Callable[[TerminationHandlerMiddleware], Awaitable],
-    expected_lines: List[str],
-):
-    def run(stdout_capture_file_path):
-        import asyncio
-        import os
-        import signal
-        import sys
-
-        sys.stdout = open(stdout_capture_file_path, "w")
-
-        def term():
-            print("terminated", flush=True)
-            os.kill(os.getpid(), signal.SIGKILL)
-
-        middleware = TerminationHandlerMiddleware(
-            on_stop=lambda: print("stopped", flush=True),
-            on_term=term,
-            termination_delay_secs=0.1,
-        )
-        asyncio.run(main_coro_gen(middleware))
-
-    stdout_capture_file = tempfile.NamedTemporaryFile()
-    proc = multiprocessing.Process(target=run, args=(stdout_capture_file.name,))
-    proc.start()
-    time.sleep(1)
-    proc.terminate()
-    proc.join(timeout=6.0)
-    with Path(stdout_capture_file.name).open() as file:
-        lines = [line.strip() for line in file]
-
-    assert lines == expected_lines

truss/tests/server/control/test_model_container_patch_applier.py

@@ -1,203 +0,0 @@
-import os
-from unittest import mock
-
-import pytest
-from truss.server.control.patch.model_container_patch_applier import (
-    ModelContainerPatchApplier,
-)
-from truss.server.control.patch.types import (
-    Action,
-    ConfigPatch,
-    EnvVarPatch,
-    ExternalDataPatch,
-    ModelCodePatch,
-    PackagePatch,
-    Patch,
-    PatchType,
-)
-from truss.truss_config import TrussConfig
-
-
-@pytest.fixture
-def patch_applier(tmp_truss_dir):
-    return ModelContainerPatchApplier(tmp_truss_dir, mock.Mock())
-
-
-def test_patch_applier_model_code_patch_add(
-    patch_applier: ModelContainerPatchApplier, tmp_truss_dir
-):
-    patch = Patch(
-        type=PatchType.MODEL_CODE,
-        body=ModelCodePatch(
-            action=Action.ADD,
-            path="dummy",
-            content="",
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    assert (tmp_truss_dir / "model" / "dummy").exists()
-
-
-def test_patch_applier_model_code_patch_remove(
-    patch_applier: ModelContainerPatchApplier, tmp_truss_dir
-):
-    patch = Patch(
-        type=PatchType.MODEL_CODE,
-        body=ModelCodePatch(
-            action=Action.REMOVE,
-            path="model.py",
-        ),
-    )
-    assert (tmp_truss_dir / "model" / "model.py").exists()
-    patch_applier(patch, os.environ.copy())
-    assert not (tmp_truss_dir / "model" / "model.py").exists()
-
-
-def test_patch_applier_model_code_patch_update(
-    patch_applier: ModelContainerPatchApplier, tmp_truss_dir
-):
-    new_model_file_content = """
-    class Model:
-        pass
-    """
-    patch = Patch(
-        type=PatchType.MODEL_CODE,
-        body=ModelCodePatch(
-            action=Action.UPDATE,
-            path="model.py",
-            content=new_model_file_content,
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    assert (tmp_truss_dir / "model" / "model.py").read_text() == new_model_file_content
-
-
-def test_patch_applier_package_patch_add(
-    patch_applier: ModelContainerPatchApplier, tmp_truss_dir
-):
-    patch = Patch(
-        type=PatchType.PACKAGE,
-        body=PackagePatch(
-            action=Action.ADD,
-            path="test_package/test.py",
-            content="foobar",
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    assert (tmp_truss_dir / "packages" / "test_package" / "test.py").exists()
-
-
-def test_patch_applier_package_patch_remove(
-    patch_applier: ModelContainerPatchApplier,
-    tmp_truss_dir,
-):
-    patch = Patch(
-        type=PatchType.PACKAGE,
-        body=PackagePatch(
-            action=Action.REMOVE,
-            path="test_package/test.py",
-        ),
-    )
-    assert (tmp_truss_dir / "packages" / "test_package" / "test.py").exists()
-    patch_applier(patch, os.environ.copy())
-    assert not (tmp_truss_dir / "packages" / "test_package" / "test.py").exists()
-
-
-def test_patch_applier_package_patch_update(
-    patch_applier: ModelContainerPatchApplier,
-    tmp_truss_dir,
-):
-    new_package_content = """X = 2"""
-    patch = Patch(
-        type=PatchType.PACKAGE,
-        body=PackagePatch(
-            action=Action.UPDATE,
-            path="test_package/test.py",
-            content=new_package_content,
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    assert (
-        tmp_truss_dir / "packages" / "test_package" / "test.py"
-    ).read_text() == new_package_content
-
-
-def test_patch_applier_config_patch_update(
-    patch_applier: ModelContainerPatchApplier, tmp_truss_dir
-):
-    new_config_dict = {"model_name": "foobar"}
-    patch = Patch(
-        type=PatchType.CONFIG,
-        body=ConfigPatch(
-            action=Action.UPDATE,
-            config=new_config_dict,
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    new_config = TrussConfig.from_yaml(tmp_truss_dir / "config.yaml")
-    assert new_config.model_name == "foobar"
-
-
-def test_patch_applier_env_var_patch_update(
-    patch_applier: ModelContainerPatchApplier,
-):
-    env_var_dict = {"FOO": "BAR"}
-    patch = Patch(
-        type=PatchType.ENVIRONMENT_VARIABLE,
-        body=EnvVarPatch(
-            action=Action.UPDATE,
-            item={"FOO": "BAR-PATCHED"},
-        ),
-    )
-    patch_applier(patch, env_var_dict)
-    assert env_var_dict["FOO"] == "BAR-PATCHED"
-
-
-def test_patch_applier_env_var_patch_add(
-    patch_applier: ModelContainerPatchApplier,
-):
-    env_var_dict = {"FOO": "BAR"}
-    patch = Patch(
-        type=PatchType.ENVIRONMENT_VARIABLE,
-        body=EnvVarPatch(
-            action=Action.ADD,
-            item={"BAR": "FOO"},
-        ),
-    )
-    patch_applier(patch, env_var_dict)
-    assert env_var_dict["FOO"] == "BAR"
-    assert env_var_dict["BAR"] == "FOO"
-
-
-def test_patch_applier_env_var_patch_remove(
-    patch_applier: ModelContainerPatchApplier,
-):
-    env_var_dict = {"FOO": "BAR"}
-    patch = Patch(
-        type=PatchType.ENVIRONMENT_VARIABLE,
-        body=EnvVarPatch(
-            action=Action.REMOVE,
-            item={"FOO": "BAR"},
-        ),
-    )
-    patch_applier(patch, env_var_dict)
-    with pytest.raises(KeyError):
-        _ = env_var_dict["FOO"]
-
-
-def test_patch_applier_external_data_patch_add(
-    patch_applier: ModelContainerPatchApplier,
-    tmp_truss_dir,
-):
-    patch = Patch(
-        type=PatchType.EXTERNAL_DATA,
-        body=ExternalDataPatch(
-            action=Action.ADD,
-            item={
-                "url": "https://raw.githubusercontent.com/basetenlabs/truss/main/docs/favicon.svg",
-                "local_data_path": "truss_icon",
-            },
-        ),
-    )
-    patch_applier(patch, os.environ.copy())
-    assert (tmp_truss_dir / "data" / "truss_icon").exists()

truss/tests/server/core/server/common/test_util.py

@@ -1,19 +0,0 @@
-# This file doesn't test anything, but provides utilities for testing.
-from unittest import mock
-
-
-def model_supports_predict_proba():
-    mock_not_predict_proba = mock.Mock(name="mock_not_predict_proba")
-    mock_not_predict_proba.predict_proba.return_value = False
-
-    mock_check_proba = mock.Mock(name="mock_check_proba")
-    mock_check_proba.predict_proba.return_value = True
-    mock_check_proba._check_proba.return_value = True
-
-    mock_not_check_proba = mock.Mock(name="mock_not_check_proba")
-    mock_not_check_proba.predict_proba.return_value = True
-    mock_not_check_proba._check_proba.side_effect = AttributeError
-
-    assert not model_supports_predict_proba(mock_not_predict_proba)
-    assert model_supports_predict_proba(mock_check_proba)
-    assert not model_supports_predict_proba(mock_not_check_proba)

truss/tests/server/test_model_wrapper.py

@@ -1,87 +0,0 @@
-import importlib
-import sys
-import time
-from pathlib import Path
-from typing import Any
-
-import pytest
-import yaml
-
-
-@pytest.fixture
-def app_path(tmp_truss_dir: Path, helpers: Any):
-    truss_container_app_path = tmp_truss_dir
-    model_file_content = """
-class Model:
-    def __init__(self):
-        self.load_count = 0
-    def load(self):
-        self.load_count += 1
-        if self.load_count <= 2:
-            raise RuntimeError('Simulated error')
-    def predict(self, request):
-        return request
-    """
-    with helpers.file_content(
-        truss_container_app_path / "model" / "model.py",
-        model_file_content,
-    ), helpers.sys_path(truss_container_app_path):
-        yield truss_container_app_path
-
-
-# TODO: Make this test work
-@pytest.mark.skip(
-    reason="Succeeds when tests in this file are run alone, but fails with the whole suit"
-)
-def test_model_wrapper_load_error_once(app_path):
-    if "model_wrapper" in sys.modules:
-        model_wrapper_module = sys.modules["model_wrapper"]
-        importlib.reload(model_wrapper_module)
-    else:
-        model_wrapper_module = importlib.import_module("model_wrapper")
-    model_wraper_class = getattr(model_wrapper_module, "ModelWrapper")
-    config = yaml.safe_load((app_path / "config.yaml").read_text())
-    model_wrapper = model_wraper_class(config)
-    model_wrapper.load()
-    # Allow load thread to execute
-    time.sleep(1)
-    output = model_wrapper.predict({})
-    assert output == {}
-    assert model_wrapper._model.load_count == 3
-
-
-# TODO: Make this test work
-@pytest.mark.skip(
-    reason="Succeeds when tests in this file are run alone, but fails with the whole suit"
-)
-def test_model_wrapper_load_error_more_than_allowed(app_path, helpers):
-    with helpers.env_var("NUM_LOAD_RETRIES_TRUSS", "0"):
-        if "model_wrapper" in sys.modules:
-            model_wrapper_module = sys.modules["model_wrapper"]
-            importlib.reload(model_wrapper_module)
-        else:
-            model_wrapper_module = importlib.import_module("model_wrapper")
-        model_wraper_class = getattr(model_wrapper_module, "ModelWrapper")
-        config = yaml.safe_load((app_path / "config.yaml").read_text())
-        model_wrapper = model_wraper_class(config)
-        model_wrapper.load()
-        # Allow load thread to execute
-        time.sleep(1)
-        assert model_wrapper.load_failed()
-
-
-@pytest.mark.integration
-async def test_model_wrapper_streaming_timeout(app_path):
-    if "model_wrapper" in sys.modules:
-        model_wrapper_module = sys.modules["model_wrapper"]
-        importlib.reload(model_wrapper_module)
-    else:
-        model_wrapper_module = importlib.import_module("model_wrapper")
-    model_wraper_class = getattr(model_wrapper_module, "ModelWrapper")
-
-    # Create an instance of ModelWrapper with streaming_read_timeout set to 5 seconds
-    config = yaml.safe_load((app_path / "config.yaml").read_text())
-    config["runtime"]["streaming_read_timeout"] = 5
-    model_wrapper = model_wraper_class(config)
-    model_wrapper.load()
-    assert model_wrapper._config.get("runtime").get("streaming_read_timeout") == 5

truss/util/data_structures.py

@@ -1,16 +0,0 @@
-from typing import Callable, Dict, Optional, TypeVar
-
-X = TypeVar("X")
-Y = TypeVar("Y")
-Z = TypeVar("Z")
-
-
-def transform_optional(x: Optional[X], fn: Callable[[X], Optional[Y]]) -> Optional[Y]:
-    if x is None:
-        return None
-
-    return fn(x)
-
-
-def transform_keys(d: Dict[X, Z], fn: Callable[[X], Y]) -> Dict[Y, Z]:
-    return {fn(key): value for key, value in d.items()}