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

truss/truss_handle/patch/truss_dir_patch_applier.py

@@ -0,0 +1,87 @@
+import logging
+from pathlib import Path
+from typing import List
+
+from truss.base.truss_config import TrussConfig
+from truss.templates.control.control.helpers.custom_types import (
+    Action,
+    ConfigPatch,
+    EnvVarPatch,
+    ExternalDataPatch,
+    ModelCodePatch,
+    PackagePatch,
+    Patch,
+    PythonRequirementPatch,
+    SystemPackagePatch,
+)
+from truss.templates.control.control.helpers.errors import UnsupportedPatch
+from truss.templates.control.control.helpers.truss_patch.model_code_patch_applier import (
+    apply_code_patch,
+)
+from truss.templates.control.control.helpers.truss_patch.requirement_name_identifier import (
+    identify_requirement_name,
+    reqs_by_name,
+)
+from truss.templates.control.control.helpers.truss_patch.system_packages import (
+    system_packages_set,
+)
+
+
+class TrussDirPatchApplier:
+    """Applies patches to a truss directory.
+    This should be compatible with ModelContainerPatchApplier.
+
+    Note: This class imported via old_build_setup.sh.jinja in the baseten
+    repository
+    """
+
+    def __init__(self, truss_dir: Path, logger: logging.Logger) -> None:
+        self._truss_dir = truss_dir
+        self._truss_config_path = self._truss_dir / "config.yaml"
+        self._truss_config = TrussConfig.from_yaml(self._truss_config_path)
+        self._logger = logger
+
+    def __call__(self, patches: List[Patch]):
+        # Apply model code patches immediately
+        # Aggregate config patches and apply at end
+        reqs = reqs_by_name(self._truss_config.requirements)
+        pkgs = system_packages_set(self._truss_config.system_packages)
+        new_config = self._truss_config
+        for patch in patches:
+            self._logger.debug(f"Applying patch {patch.to_dict()}")
+            action = patch.body.action
+            if isinstance(patch.body, ModelCodePatch):
+                model_code_patch: ModelCodePatch = patch.body
+                model_module_dir = self._truss_dir / self._truss_config.model_module_dir
+                apply_code_patch(model_module_dir, model_code_patch, self._logger)
+            elif isinstance(patch.body, PythonRequirementPatch):
+                py_req_patch: PythonRequirementPatch = patch.body
+                req = py_req_patch.requirement
+                req_name = identify_requirement_name(req)
+                if action == Action.REMOVE:
+                    del reqs[req_name]
+                elif action == Action.ADD or Action.UPDATE:
+                    reqs[req_name] = req
+            elif isinstance(patch.body, SystemPackagePatch):
+                sys_pkg_patch: SystemPackagePatch = patch.body
+                pkg = sys_pkg_patch.package
+                if action == Action.REMOVE:
+                    pkgs.remove(pkg)
+                elif action == Action.ADD or Action.UPDATE:
+                    pkgs.add(pkg)
+            elif isinstance(patch.body, ConfigPatch):
+                new_config = TrussConfig.from_dict(patch.body.config)
+            # Each of EnvVarPatch and ExternalDataPatch can be expressed through an
+            # overwrite of the config, handled below.
+            elif isinstance(patch.body, (EnvVarPatch, ExternalDataPatch)):
+                pass
+            elif isinstance(patch.body, PackagePatch):
+                package_patch: PackagePatch = patch.body
+                package_module_dir = (
+                    self._truss_dir / self._truss_config.bundled_packages_dir
+                )
+                apply_code_patch(package_module_dir, package_patch, self._logger)
+            else:
+                raise UnsupportedPatch(f"Unknown patch type {patch.type}")
+
+        new_config.write_to_yaml_file(self._truss_config_path)

truss/{readme_generator.py → truss_handle/readme_generator.py}

@@ -1,6 +1,7 @@
 from jinja2 import Template
-from truss.constants import README_TEMPLATE_NAME, TEMPLATES_DIR
-from truss.truss_spec import TrussSpec
+
+from truss.base.constants import README_TEMPLATE_NAME, TEMPLATES_DIR
+from truss.base.truss_spec import TrussSpec
 
 
 def generate_readme(_spec: TrussSpec) -> str:

truss/{truss_gatherer.py → truss_handle/truss_gatherer.py}

@@ -1,9 +1,10 @@
 from pathlib import Path
 
 import yaml
+
 from truss.local.local_config_handler import LocalConfigHandler
-from truss.patch.hash import str_hash_str
-from truss.truss_handle import TrussHandle
+from truss.truss_handle.patch.hash import str_hash_str
+from truss.truss_handle.truss_handle import TrussHandle
 from truss.util.path import copy_file_path, copy_tree_path, remove_tree_path
 
 

truss/{truss_handle.py → truss_handle/truss_handle.py}

@@ -6,7 +6,7 @@ import sys
 import uuid
 from dataclasses import replace
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 from urllib.error import HTTPError
 
 import requests
@@ -24,19 +24,46 @@ from tenacity import (
     stop_after_delay,
     wait_fixed,
 )
-from truss.constants import (
+
+from truss.base.constants import (
     INFERENCE_SERVER_PORT,
     TRUSS,
     TRUSS_DIR,
     TRUSS_HASH,
     TRUSS_MODIFIED_TIME,
 )
+from truss.base.custom_types import Example
+from truss.base.errors import ContainerIsDownError, ContainerNotFoundError
+from truss.base.truss_config import (
+    BaseImage,
+    ExternalData,
+    ExternalDataItem,
+    TrussConfig,
+)
+from truss.base.truss_spec import TrussSpec
+from truss.base.validation import validate_secret_name
 from truss.contexts.image_builder.serving_image_builder import (
     ServingImageBuilderContext,
 )
 from truss.contexts.local_loader.load_model_local import LoadModelLocal
-from truss.decorators import proxy_to_shadow_if_scattered
-from truss.docker import (
+from truss.contexts.truss_context import TrussContext
+from truss.local.local_config_handler import LocalConfigHandler
+from truss.templates.shared.serialization import (
+    truss_msgpack_deserialize,
+    truss_msgpack_serialize,
+)
+from truss.trt_llm.validation import validate
+from truss.truss_handle.decorators import proxy_to_shadow_if_scattered
+from truss.truss_handle.patch.calc_patch import calc_truss_patch
+from truss.truss_handle.patch.custom_types import (
+    PatchDetails,
+    PatchRequest,
+    TrussSignature,
+)
+from truss.truss_handle.patch.hash import directory_content_hash
+from truss.truss_handle.patch.signature import calc_truss_signature
+from truss.truss_handle.readme_generator import generate_readme
+from truss.util.docker import (
     Docker,
     DockerStates,
     get_container_logs,
@@ -46,28 +73,13 @@ from truss.docker import (
     get_urls_from_container,
     kill_containers,
 )
-from truss.errors import ContainerIsDownError, ContainerNotFoundError
-from truss.local.local_config_handler import LocalConfigHandler
-from truss.notebook import is_notebook_or_ipython
-from truss.patch.calc_patch import calc_truss_patch
-from truss.patch.hash import directory_content_hash
-from truss.patch.signature import calc_truss_signature
-from truss.patch.types import TrussSignature
-from truss.readme_generator import generate_readme
-from truss.server.shared.serialization import (
-    truss_msgpack_deserialize,
-    truss_msgpack_serialize,
-)
-from truss.truss_config import BaseImage, ExternalData, ExternalDataItem, TrussConfig
-from truss.truss_spec import TrussSpec
-from truss.types import Example, PatchDetails, PatchRequest
+from truss.util.notebook import is_notebook_or_ipython
 from truss.util.path import (
     copy_file_path,
     copy_tree_path,
     get_max_modified_time_of_dir,
     load_trussignore_patterns,
 )
-from truss.validation import validate_secret_name
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -76,16 +88,32 @@ if is_notebook_or_ipython():
     logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
+class RunningContainer:
+    def __init__(self, container):
+        self.container = container
+
+    def logs(self):
+        from python_on_whales import docker
+
+        return docker.logs(self.container, follow=True, stream=True)
+
+    def wait(self):
+        from python_on_whales import docker
+
+        return docker.wait(self.container)
+
+
 class TrussHandle:
     def __init__(self, truss_dir: Path, validate: bool = True) -> None:
         self._truss_dir = truss_dir
-        self._spec = TrussSpec(truss_dir)
+        self._spec = TrussSpec(self._truss_dir)
         self._hash_for_mod_time: Optional[Tuple[float, str]] = None
         if validate:
             self.validate()
 
     def validate(self):
         self._validate_external_packages()
+        self._validate_extensions()
 
     @property
     def spec(self) -> TrussSpec:
@@ -149,6 +177,57 @@ class TrussHandle:
         """[Deprecated] Do not use."""
         return _docker_image_from_labels(labels)
 
+    @proxy_to_shadow_if_scattered
+    def run_python_script(self, script_path: Path, build_dir: Optional[Path] = None):
+        from python_on_whales.exceptions import DockerException
+
+        image = self.build_serving_docker_image(build_dir=build_dir)
+        secrets_mount_dir_path = _prepare_secrets_mount_dir()
+
+        envs: Dict[str, str] = {}
+        # Add bundled packages to the PYTHONPATH. Note
+        # that this is necessary to achieve the same environment as Truss Server
+        # -- this is setup that is done by Truss Server, that won't be available
+        # to the standalone script.
+        bundled_packages_path = Path("/packages")
+        envs["PYTHONPATH"] = bundled_packages_path.as_posix()
+
+        # Note that the entrypoint command should match
+        # what we use when executing Truss Server.
+        entrypoint_command = self.spec.python_executable_path or "python3"
+
+        def _docker_run(gpus: Optional[str] = None):
+            container = Docker.client().run(
+                image.id,
+                entrypoint=entrypoint_command,
+                command=["/app/script.py"],
+                detach=True,
+                mounts=[
+                    [
+                        "type=bind",
+                        f"src={str(secrets_mount_dir_path)}",
+                        "target=/secrets",
+                    ],
+                    [
+                        "type=bind",
+                        f"src={str(script_path.absolute())}",
+                        "target=/app/script.py",
+                    ],
+                ],
+                gpus=gpus,
+                envs=envs,
+                add_hosts=[("host.docker.internal", "host-gateway")],
+            )
+
+            return RunningContainer(container)
+
+        try:
+            return _docker_run("all" if self._spec.config.resources.use_gpu else None)
+        except DockerException:
+            # The reason we'd wind up here is if the Truss needs
+            # a GPU, but the host does not have one that can attach.
+            return _docker_run(None)
+
     @proxy_to_shadow_if_scattered
     def docker_run(
         self,
@@ -159,7 +238,9 @@ class TrussHandle:
         patch_ping_url: Optional[str] = None,
         wait_for_server_ready: bool = True,
         network: Optional[str] = None,
-        cache: bool = True,
+        container_name_prefix: Optional[str] = None,
+        model_server_stop_retry_override=None,
+        disable_json_logging: bool = False,
     ):
         """
         Builds a docker image and runs it as a container. For control trusses,
@@ -173,7 +254,10 @@ class TrussHandle:
             patch_ping_url:  Mostly for testing, if supplied then a live
                              reload capable truss queries for truss changes
                              by hitting this url.
-            wait_for_server_ready: If true, wait for server to pass readiness probe before returning.
+            wait_for_server_ready: If true, wait for server to pass readiness
+              probe before returning.
+            network: docker network name.
+            container_name_prefix: optional docker container name prefix.
 
         Returns:
             Container, which can be used to get information about the running,
@@ -186,7 +270,7 @@ class TrussHandle:
             container = container_if_patched
         else:
             image = self.build_serving_docker_image(
-                build_dir=build_dir, tag=tag, network=network, cache=cache
+                build_dir=build_dir, tag=tag, network=network
             )
             secrets_mount_dir_path = _prepare_secrets_mount_dir()
             publish_ports = [[local_port, INFERENCE_SERVER_PORT]]
@@ -198,6 +282,14 @@ class TrussHandle:
             envs = {}
             if patch_ping_url is not None:
                 envs["PATCH_PING_URL_TRUSS"] = patch_ping_url
+            if disable_json_logging:
+                envs["DISABLE_JSON_LOGGING"] = "true"
+
+            if container_name_prefix:
+                suffix = str(uuid.uuid4()).split("-")[0]
+                name = f"{container_name_prefix}-{suffix}"
+            else:
+                name = None
 
             def _run_docker(gpus: Optional[str] = None):
                 return Docker.client().run(
@@ -210,11 +302,23 @@ class TrussHandle:
                             "type=bind",
                             f"src={str(secrets_mount_dir_path)}",
                             "target=/secrets",
-                        ]
+                        ],
+                        [
+                            "type=bind",
+                            f"src={str(LocalConfigHandler.bptr_data_resolution_dir_path())}",
+                            "target=/bptr",
+                        ],
+                        [
+                            "type=bind",
+                            f"src={str(LocalConfigHandler.dynamic_config_path())}",
+                            "target=/etc/b10_dynamic_config",
+                            "readonly=false",
+                        ],
                     ],
                     gpus=gpus,
                     envs=envs,
                     add_hosts=[("host.docker.internal", "host-gateway")],
+                    name=name,
                 )
 
             try:
@@ -232,7 +336,12 @@ class TrussHandle:
             )
         model_base_url = f"http://localhost:{local_port}/v1/models/model"
         try:
-            wait_for_truss(model_base_url, container, wait_for_server_ready)
+            wait_for_truss(
+                model_base_url,
+                container,
+                wait_for_server_ready,
+                model_server_stop_retry_override,
+            )
         except ContainerNotFoundError as err:
             raise err
         except (ContainerIsDownError, HTTPError, ConnectionError) as err:
@@ -263,6 +372,7 @@ class TrussHandle:
         else:
             return self.server_predict(request)
 
+    # TODO(marius): can we kill this?
     def server_predict(self, request: Dict):
         """Run the prediction flow locally."""
         model = LoadModelLocal.run(self._truss_dir)
@@ -384,11 +494,7 @@ class TrussHandle:
         validate_secret_name(secret_name)
         self._update_config(
             lambda conf: replace(
-                conf,
-                secrets={
-                    **conf.secrets,
-                    secret_name: default_secret_value,
-                },
+                conf, secrets={**conf.secrets, secret_name: default_secret_value}
             )
         )
 
@@ -410,23 +516,12 @@ class TrussHandle:
             self._spec.config.external_data or ExternalData([])
         )
         new_external_data = replace(
-            current_external_data,
-            items=current_external_data.items + [item],
-        )
-        self._update_config(
-            lambda conf: replace(
-                conf,
-                external_data=new_external_data,
-            )
+            current_external_data, items=current_external_data.items + [item]
         )
+        self._update_config(lambda conf: replace(conf, external_data=new_external_data))
 
     def remove_all_external_data(self):
-        self._update_config(
-            lambda conf: replace(
-                conf,
-                external_data=None,
-            )
-        )
+        self._update_config(lambda conf: replace(conf, external_data=None))
 
     def update_requirements(self, requirements: List[str]):
         """Update requirements in truss model's config.
@@ -554,9 +649,7 @@ class TrussHandle:
 
     @proxy_to_shadow_if_scattered
     def get_serving_docker_containers_from_labels(
-        self,
-        all: bool = False,
-        labels: Optional[dict] = None,
+        self, all: bool = False, labels: Optional[dict] = None
     ) -> list:
         """Get serving docker containers, with given labels.
 
@@ -568,10 +661,7 @@ class TrussHandle:
             labels = self._get_serving_lookup_labels()
         else:
             # Make sure we're looking for serving container for this truss.
-            labels = {
-                TRUSS: True,
-                **labels,
-            }
+            labels = {TRUSS: True, **labels}
 
         return sorted(get_containers(labels, all=all), key=lambda c: c.created)
 
@@ -630,10 +720,7 @@ class TrussHandle:
                     python_executable_path=python_executable_path,
                 )
                 return replace(conf, base_image=new_base_image)
-            new_base_image = BaseImage(
-                image,
-                python_executable_path,
-            )
+            new_base_image = BaseImage(image, python_executable_path)
             return replace(conf, base_image=new_base_image)
 
         self._update_config(define_base_image_fn)
@@ -691,10 +778,7 @@ class TrussHandle:
             inferred_python_version = f"py{version_parts[0]}{version_parts[1]}"
 
         self._update_config(
-            lambda conf: replace(
-                conf,
-                python_version=inferred_python_version,
-            )
+            lambda conf: replace(conf, python_version=inferred_python_version)
         )
 
     def _control_serving_container_has_partially_applied_patch(self) -> Optional[bool]:
@@ -758,7 +842,6 @@ class TrussHandle:
         if prev_sign_str is None:
             logger.info(f"Signature not found for truss for hash {prev_truss_hash}")
             return None
-
         prev_sign = TrussSignature.from_dict(json.loads(prev_sign_str))
         ignore_patterns = truss_ignore_patterns + self._spec.hash_ignore_patterns
         patch_ops = calc_truss_patch(self._truss_dir, prev_sign, ignore_patterns)
@@ -783,7 +866,7 @@ class TrussHandle:
         gatherer and a handle to that truss is returned. These gathered trusses
         are caches and resused.
         """
-        from truss.truss_gatherer import gather
+        from truss.truss_handle.truss_gatherer import gather
 
         if not self.is_scattered():
             return self._truss_dir
@@ -847,7 +930,7 @@ class TrussHandle:
 
     def _build_image(
         self,
-        builder_context,
+        builder_context: Type[TrussContext],
         labels: Dict[str, str],
         build_dir: Optional[Path] = None,
         tag: Optional[str] = None,
@@ -855,7 +938,7 @@ class TrussHandle:
         network: Optional[str] = None,
     ):
         image = _docker_image_from_labels(labels=labels)
-        if cache and image is not None:
+        if image is not None:
             return image
 
         build_dir_path = Path(build_dir) if build_dir is not None else None
@@ -917,7 +1000,7 @@ class TrussHandle:
         )
         resp = self.patch_container(patch_request)
         if "error" in resp:
-            raise RuntimeError(f'Failed to patch control truss {resp["error"]}')
+            raise RuntimeError(f"Failed to patch control truss {resp['error']}")
         self._store_signature()
         return container
 
@@ -951,6 +1034,11 @@ class TrussHandle:
                         f"{path.resolve()} but that path does not exist."
                     )
 
+    def _validate_extensions(self):
+        # Only one extenstion right now.
+        if self._spec.config.trt_llm is not None:
+            validate(self._spec)
+
 
 def _prediction_flow(model, request: Dict):
     """This flow attempts to mimic the request life-cycle of a server"""
@@ -974,31 +1062,39 @@ def _wait_for_docker_build(container) -> None:
                 raise ContainerIsDownError(f"Container stuck in state: {state.value}.")
 
 
-@retry(
-    stop=stop_after_delay(120),
-    wait=wait_fixed(2),
-    retry=(
-        retry_if_result(lambda response: response.status_code == 503)
-        | retry_if_exception_type(exceptions.ConnectionError)
-    ),
-)
-def _wait_for_model_server(url: str) -> Response:
-    return requests.get(url)
+def _wait_for_model_server(url: str, stop=stop_after_delay(120)) -> Response:  # type: ignore[return]
+    for attempt in Retrying(
+        stop=stop,
+        wait=wait_fixed(2),
+        retry=(
+            retry_if_result(lambda response: response.status_code in [502, 503])
+            | retry_if_exception_type(exceptions.ConnectionError)
+        ),
+    ):
+        with attempt:
+            response = requests.get(url)
+            return response
 
 
 def wait_for_truss(
-    url: str, container: str, wait_for_server_ready: bool = True
+    url: str,
+    container: str,
+    wait_for_server_ready: bool = True,
+    model_server_stop_retry_override=None,
 ) -> None:
     from python_on_whales.exceptions import NoSuchContainer
 
     try:
         _wait_for_docker_build(container)
+        if wait_for_server_ready:
+            if model_server_stop_retry_override is not None:
+                _wait_for_model_server(url, stop=model_server_stop_retry_override)
+            else:
+                _wait_for_model_server(url)
     except NoSuchContainer:
         raise ContainerNotFoundError(message=f"Container {container} was not found")
     except RetryError as retry_err:
         retry_err.reraise()
-    if wait_for_server_ready:
-        _wait_for_model_server(url)
 
 
 def _prepare_secrets_mount_dir() -> Path:

truss/util/.truss_ignore

@@ -141,6 +141,9 @@ venv.bak/
 .dmypy.json
 dmypy.json
 
+# ruff
+.ruff_cache/
+
 # Pyre type checker
 .pyre/
 

truss/{docker.py → util/docker.py}

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List
 if TYPE_CHECKING:
     from python_on_whales.components.container.cli_wrapper import Container
 
-from truss.constants import TRUSS_DIR
+from truss.base.constants import TRUSS, TRUSS_DIR
 from truss.local.local_config_handler import LocalConfigHandler
 
 
@@ -57,7 +57,7 @@ def get_urls_from_container(container_details) -> Dict[int, List[str]]:
         )
 
     def urls_from_port_protocol_values(
-        port_protocol_values: List[Dict[str, str]]
+        port_protocol_values: List[Dict[str, str]],
     ) -> List[str]:
         return [url_from_port_protocol_value(v) for v in port_protocol_values]
 
@@ -115,3 +115,7 @@ def _create_label_filters(labels: Dict) -> Dict[str, Any]:
     return {
         f"label={label_key}": label_value for label_key, label_value in labels.items()
     }
+
+
+def kill_all() -> None:
+    kill_containers({TRUSS: True})

truss/util/download.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import Optional
 
 import requests
-from truss.truss_config import ExternalData
+from truss.base.truss_config import ExternalData
 
 B10CP_EXECUTABLE_NAME = "b10cp"
 BLOB_DOWNLOAD_TIMEOUT_SECS = 600  # 10 minutes
@@ -41,9 +41,7 @@ def _b10cp_path() -> Optional[str]:
 
 
 def _download_external_data_using_b10cp(
-    b10cp_path: str,
-    data_dir: Path,
-    external_data: ExternalData,
+    b10cp_path: str, data_dir: Path, external_data: ExternalData
 ):
     procs = []
     # TODO(pankaj) Limit concurrency here
@@ -56,11 +54,7 @@ def _download_external_data_using_b10cp(
         proc.wait()
 
 
-def _download_from_url_using_b10cp(
-    b10cp_path: str,
-    url: str,
-    download_to: Path,
-):
+def _download_from_url_using_b10cp(b10cp_path: str, url: str, download_to: Path):
     return subprocess.Popen(
         [
             b10cp_path,
@@ -74,18 +68,15 @@ def _download_from_url_using_b10cp(
 
 def _download_external_data_using_requests(data_dir: Path, external_data: ExternalData):
     for item in external_data.items:
-        _download_from_url_using_requests(
+        download_from_url_using_requests(
             item.url, (data_dir / item.local_data_path).resolve()
         )
 
 
-def _download_from_url_using_requests(URL: str, download_to: Path):
+def download_from_url_using_requests(URL: str, download_to: Path):
     # Streaming download to keep memory usage low
     resp = requests.get(
-        URL,
-        allow_redirects=True,
-        stream=True,
-        timeout=BLOB_DOWNLOAD_TIMEOUT_SECS,
+        URL, allow_redirects=True, stream=True, timeout=BLOB_DOWNLOAD_TIMEOUT_SECS
     )
     resp.raise_for_status()
     with download_to.open("wb") as file: