Revert "Refactor truss.server out of templates to run as normal pyt…

…hon program (#829)"

This reverts commit 8b1bf5d.

pyproject.toml

@@ -1,14 +1,13 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.2rc0"
+version = "0.9.3"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
 authors = ["Pankaj Gupta <pankaj@baseten.co>", "Phil Howes <phil@baseten.co>"]
 include = ["*.txt", "*.Dockerfile", "*.md"]
 repository = "https://github.com/basetenlabs/truss"
 keywords = ["MLOps", "AI", "Model Serving", "Model Deployment", "Machine Learning"]
-packages = [{include = "truss", format = "wheel"}]
 
 [tool.poetry.urls]
 "Homepage" = "https://truss.baseten.co"
@@ -43,7 +42,6 @@ inquirerpy = "^0.3.4"
 google-cloud-storage = "2.10.0"
 loguru = ">=0.7.2"
 uvloop = "^0.19.0"
-pathspec = ">=0.9.0"
 
 
 [tool.poetry.group.builder.dependencies]
@@ -65,7 +63,6 @@ huggingface_hub = ">=0.19.4"
 google-cloud-storage = "2.10.0"
 boto3 = "^1.26.157"
 loguru = ">=0.7.2"
-pathspec = ">=0.9.0"
 
 [tool.poetry.dev-dependencies]
 ipython = "^7.16"

truss/constants.py

@@ -1,7 +1,6 @@
 import os
 import pathlib
-
-TRUSS_PACKAGE_DIR = pathlib.Path(__file__).resolve().parent
+from typing import Set
 
 SKLEARN = "sklearn"
 TENSORFLOW = "tensorflow"
@@ -16,18 +15,32 @@
 CODE_DIR = pathlib.Path(BASE_DIR, "truss")
 
 TEMPLATES_DIR = pathlib.Path(CODE_DIR, "templates")
+SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "server"
+TRITON_SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "triton"
 TRTLLM_TRUSS_DIR: pathlib.Path = TEMPLATES_DIR / "trtllm"
+SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME = "shared"
+SHARED_SERVING_AND_TRAINING_CODE_DIR: pathlib.Path = (
+    TEMPLATES_DIR / SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME
+)
+CONTROL_SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "control"
 
 SUPPORTED_PYTHON_VERSIONS = {"3.8", "3.9", "3.10", "3.11"}
 
 
+# Alias for TEMPLATES_DIR
+SERVING_DIR: pathlib.Path = TEMPLATES_DIR
+
 REQUIREMENTS_TXT_FILENAME = "requirements.txt"
 USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME = "user_requirements.txt"
+BASE_SERVER_REQUIREMENTS_TXT_FILENAME = "base_server_requirements.txt"
+SERVER_REQUIREMENTS_TXT_FILENAME = "server_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,
+    "base_server_requirements_filename": BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
+    "server_requirements_filename": SERVER_REQUIREMENTS_TXT_FILENAME,
     "system_packages_filename": SYSTEM_PACKAGES_TXT_FILENAME,
 }
 
@@ -47,6 +60,38 @@
 TRUSS_DIR = "truss_dir"
 TRUSS_HASH = "truss_hash"
 
+HUGGINGFACE_TRANSFORMER_MODULE_NAME: Set[str] = set({})
+
+# list from https://scikit-learn.org/stable/developers/advanced_installation.html
+SKLEARN_REQ_MODULE_NAMES: Set[str] = {
+    "numpy",
+    "scipy",
+    "joblib",
+    "scikit-learn",
+    "threadpoolctl",
+}
+
+XGBOOST_REQ_MODULE_NAMES: Set[str] = {"xgboost"}
+
+# list from https://www.tensorflow.org/install/pip
+# if problematic, lets look to https://www.tensorflow.org/install/source
+TENSORFLOW_REQ_MODULE_NAMES: Set[str] = {
+    "tensorflow",
+}
+
+LIGHTGBM_REQ_MODULE_NAMES: Set[str] = {
+    "lightgbm",
+}
+
+# list from https://pytorch.org/get-started/locally/
+PYTORCH_REQ_MODULE_NAMES: Set[str] = {
+    "torch",
+    "torchvision",
+    "torchaudio",
+}
+
+MLFLOW_REQ_MODULE_NAMES: Set[str] = {"mlflow"}
+
 INFERENCE_SERVER_PORT = 8080
 
 HTTP_PUBLIC_BLOB_BACKEND = "http_public"

truss/contexts/image_builder/serving_image_builder.py

@@ -6,22 +6,28 @@
 from typing import Any, Dict, List, Optional, Tuple, Type
 
 import boto3
+import yaml
 from botocore import UNSIGNED
 from botocore.client import Config
 from google.cloud import storage
 from huggingface_hub import get_hf_file_metadata, hf_hub_url, list_repo_files
 from huggingface_hub.utils import filter_repo_objects
 from truss.constants import (
+    BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
     BASE_TRTLLM_REQUIREMENTS,
+    CONTROL_SERVER_CODE_DIR,
     FILENAME_CONSTANTS_MAP,
     MODEL_DOCKERFILE_NAME,
     REQUIREMENTS_TXT_FILENAME,
+    SERVER_CODE_DIR,
     SERVER_DOCKERFILE_TEMPLATE_NAME,
+    SERVER_REQUIREMENTS_TXT_FILENAME,
+    SHARED_SERVING_AND_TRAINING_CODE_DIR,
+    SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME,
     SYSTEM_PACKAGES_TXT_FILENAME,
     TEMPLATES_DIR,
     TRTLLM_BASE_IMAGE,
     TRTLLM_TRUSS_DIR,
-    TRUSS_PACKAGE_DIR,
     USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
 )
 from truss.contexts.image_builder.cache_warmer import (
@@ -48,6 +54,9 @@
     load_trussignore_patterns,
 )
 
+BUILD_SERVER_DIR_NAME = "server"
+BUILD_CONTROL_SERVER_DIR_NAME = "control"
+
 CONFIG_FILE = "config.yaml"
 USER_TRUSS_IGNORE_FILE = ".truss_ignore"
 GCS_CREDENTIALS = "service_account.json"
@@ -313,13 +322,6 @@ def prepare_image_build_dir(
         def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
             copy_tree_or_file(from_path, build_dir / path_in_build_dir)  # type: ignore[operator]
 
-        # Copy truss package from the context builder image to build dir
-        copy_into_build_dir(TRUSS_PACKAGE_DIR, "./truss")
-        copy_into_build_dir(
-            TRUSS_PACKAGE_DIR.parent / "pyproject.toml", "./pyproject.toml"
-        )
-        copy_into_build_dir(TRUSS_PACKAGE_DIR.parent / "README.md", "./README.md")
-
         truss_ignore_patterns = []
         if (truss_dir / USER_TRUSS_IGNORE_FILE).exists():
             truss_ignore_patterns = load_trussignore_patterns(
@@ -349,6 +351,10 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
             )
             config.requirements.extend(BASE_TRTLLM_REQUIREMENTS)
 
+        # Override config.yml
+        with (build_dir / CONFIG_FILE).open("w") as config_file:
+            yaml.dump(config.to_dict(verbose=True), config_file)
+
         external_data_files: list = []
         data_dir = Path("/app/data/")
         if self._spec.external_data is not None:
@@ -365,8 +371,50 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
             config, truss_dir, build_dir
         )
 
+        # Copy inference server code
+        copy_into_build_dir(SERVER_CODE_DIR, BUILD_SERVER_DIR_NAME)
+        copy_into_build_dir(
+            SHARED_SERVING_AND_TRAINING_CODE_DIR,
+            BUILD_SERVER_DIR_NAME + "/" + SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME,
+        )
+
+        # Copy control server code
+        if config.live_reload:
+            copy_into_build_dir(CONTROL_SERVER_CODE_DIR, BUILD_CONTROL_SERVER_DIR_NAME)
+            copy_into_build_dir(
+                SHARED_SERVING_AND_TRAINING_CODE_DIR,
+                BUILD_CONTROL_SERVER_DIR_NAME
+                + "/control/"
+                + SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME,
+            )
+
+        # Copy base TrussServer requirements if supplied custom base image
+        base_truss_server_reqs_filepath = SERVER_CODE_DIR / REQUIREMENTS_TXT_FILENAME
+        if config.base_image:
+            copy_into_build_dir(
+                base_truss_server_reqs_filepath, BASE_SERVER_REQUIREMENTS_TXT_FILENAME
+            )
+
+        # Copy model framework specific requirements file
+        server_reqs_filepath = (
+            TEMPLATES_DIR / model_framework_name / REQUIREMENTS_TXT_FILENAME
+        )
+        should_install_server_requirements = file_is_not_empty(server_reqs_filepath)
+        if should_install_server_requirements:
+            copy_into_build_dir(server_reqs_filepath, SERVER_REQUIREMENTS_TXT_FILENAME)
+
+        with open(base_truss_server_reqs_filepath, "r") as f:
+            base_server_requirements = f.read()
+
+        # If the user has provided python requirements,
+        # append the truss server requirements, so that any conflicts
+        # are detected and cause a build failure. If there are no
+        # requirements provided, we just pass an empty string,
+        # as there's no need to install anything.
         user_provided_python_requirements = (
-            spec.requirements_txt if spec.requirements else ""
+            base_server_requirements + spec.requirements_txt
+            if spec.requirements
+            else ""
         )
         if spec.requirements_file is not None:
             copy_into_build_dir(
@@ -380,6 +428,7 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
 
         self._render_dockerfile(
             build_dir,
+            should_install_server_requirements,
             model_files,
             use_hf_secret,
             cached_files,
@@ -389,6 +438,7 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
     def _render_dockerfile(
         self,
         build_dir: Path,
+        should_install_server_requirements: bool,
         model_files: Dict[str, Any],
         use_hf_secret: bool,
         cached_files: List[str],
@@ -424,6 +474,7 @@ def _render_dockerfile(
 
         hf_access_token = config.secrets.get(HF_ACCESS_TOKEN_SECRET_NAME)
         dockerfile_contents = dockerfile_template.render(
+            should_install_server_requirements=should_install_server_requirements,
             base_image_name_and_tag=base_image_name_and_tag,
             should_install_system_requirements=should_install_system_requirements,
             should_install_requirements=should_install_python_requirements,

truss/contexts/local_loader/docker_build_emulator.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List
+
+from truss.util.path import copy_tree_or_file
+
+
+@dataclass
+class DockerBuildEmulatorResult:
+    workdir: Path = field(default_factory=lambda: Path("/"))
+    env: Dict = field(default_factory=dict)
+    entrypoint: List = field(default_factory=list)
+
+
+class DockerBuildEmulator:
+    """Emulates Docker Builds
+
+    As running docker builds is expensive, this class emulates the docker build
+    by parsing the docker file and applying certain commands to create an
+    appropriate enviroment in a directory to simulate the root of the file system.
+
+    Support COPY, ENV, ENTRYPOINT, WORKDIR commands. All other commands are ignored.
+    """
+
+    def __init__(
+        self,
+        dockerfile_path: Path,
+        context_dir: Path,
+    ) -> None:
+        import dockerfile
+
+        self._commands = dockerfile.parse_file(str(dockerfile_path))
+        self._context_dir = context_dir
+
+    def run(self, fs_root_dir: Path) -> DockerBuildEmulatorResult:
+        def _resolve_env(key: str) -> str:
+            if key.startswith("$"):
+                key = key.replace("$", "", 1)
+                v = result.env[key]
+                return v
+            return key
+
+        def _resolve_values(keys: List[str]) -> List[str]:
+            return list(map(_resolve_env, keys))
+
+        result = DockerBuildEmulatorResult()
+        for cmd in self._commands:
+            if cmd.cmd not in ["ENV", "ENTRYPOINT", "COPY", "WORKDIR"]:
+                continue
+            values = _resolve_values(cmd.value)
+            if cmd.cmd == "ENV":
+                result.env[values[0]] = values[1]
+            if cmd.cmd == "ENTRYPOINT":
+                result.entrypoint = list(values)
+            if cmd.cmd == "COPY":
+                src, dst = values
+                src = src.replace("./", "", 1)
+                dst = dst.replace("/", "", 1)
+                copy_tree_or_file(self._context_dir / src, fs_root_dir / dst)
+            if cmd.cmd == "WORKDIR":
+                result.workdir = result.workdir / values[0]
+        return result

truss/contexts/local_loader/load_model_local.py

@@ -7,7 +7,7 @@
     signature_accepts_keyword_arg,
 )
 from truss.contexts.truss_context import TrussContext
-from truss.server.common.patches import apply_patches
+from truss.templates.server.common.patches import apply_patches
 from truss.truss_spec import TrussSpec
 
 

truss/patch/calc_patch.py

@@ -6,9 +6,13 @@
 from truss.constants import CONFIG_FILE
 from truss.patch.hash import file_content_hash_str
 from truss.patch.types import TrussSignature
-from truss.server.control.patch.requirement_name_identifier import reqs_by_name
-from truss.server.control.patch.system_packages import system_packages_set
-from truss.server.control.patch.types import (
+from truss.templates.control.control.helpers.truss_patch.requirement_name_identifier import (
+    reqs_by_name,
+)
+from truss.templates.control.control.helpers.truss_patch.system_packages import (
+    system_packages_set,
+)
+from truss.templates.control.control.helpers.types import (
     Action,
     ConfigPatch,
     DataPatch,

truss/patch/local_truss_patch_applier.py

@@ -3,9 +3,11 @@
 from pathlib import Path
 from typing import List
 
-from truss.server.control.errors import UnsupportedPatch
-from truss.server.control.patch.model_code_patch_applier import apply_code_patch
-from truss.server.control.patch.types import (
+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.types import (
     Action,
     ModelCodePatch,
     Patch,

truss/patch/truss_dir_patch_applier.py

@@ -2,14 +2,18 @@
 from pathlib import Path
 from typing import List
 
-from truss.server.control.errors import UnsupportedPatch
-from truss.server.control.patch.model_code_patch_applier import apply_code_patch
-from truss.server.control.patch.requirement_name_identifier import (
+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.server.control.patch.system_packages import system_packages_set
-from truss.server.control.patch.types import (
+from truss.templates.control.control.helpers.truss_patch.system_packages import (
+    system_packages_set,
+)
+from truss.templates.control.control.helpers.types import (
     Action,
     ConfigPatch,
     EnvVarPatch,