Refactor truss.server out of templates to run as normal python prog…

…ram (#829) * [RFC][POC] Run truss server as normal python process * WIP: control server * Working draft server * Fix tests * WIP: continue fixing tests; issue with patch typing * Fix some mypy issues * Clean up some more imports * Move test files * Move some things around for readabliity * Cleanup everything except fs in tests * Only 4 tests left * Revert changes to test data * Fix all unit tests * Drop unnecessary config dump * Fix some integration tests; move packages dir * Clean up logging and patches * Fix packages path for TRT * Fix issue holding control server integration tests from passing * Fix rest of integration tests * Fix unit test and drop test_durations * Fix last integration test by passing APP_HOME and adding to path
basetenlabs · Feb 26, 2024 · 8b1bf5d · 8b1bf5d
1 parent bfcd86b
commit 8b1bf5d
Show file tree

Hide file tree

Showing 69 changed files with 266 additions and 600 deletions.
diff --git a/.test_durations b/.test_durations
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,13 +1,14 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.2"
+version = "0.9.2rc0"
 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"

diff --git a/truss/constants.py b/truss/constants.py
@@ -1,6 +1,7 @@
 import os
 import pathlib
-from typing import Set
+
+TRUSS_PACKAGE_DIR = pathlib.Path(__file__).resolve().parent
 
 SKLEARN = "sklearn"
 TENSORFLOW = "tensorflow"
@@ -15,7 +16,6 @@
 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"
 TRAINING_JOB_WRAPPER_CODE_DIR_NAME = "training"
@@ -26,26 +26,18 @@
 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"
 TRAINING_REQUIREMENTS_TXT_FILENAME = "training_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,
     "training_requirements_filename": TRAINING_REQUIREMENTS_TXT_FILENAME,
     "system_packages_filename": SYSTEM_PACKAGES_TXT_FILENAME,
 }
@@ -70,38 +62,6 @@
 TRAINING_TRUSS_HASH = "training_truss_hash"
 TRAINING_LABEL = "training"
 
-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
 
 TRAINING_VARIABLES_FILENAME = "variables.yaml"

diff --git a/truss/contexts/image_builder/serving_image_builder.py b/truss/contexts/image_builder/serving_image_builder.py
@@ -13,22 +13,19 @@
 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,
     TRITON_SERVER_CODE_DIR,
     TRTLLM_BASE_IMAGE,
     TRTLLM_TRUSS_DIR,
+    TRUSS_PACKAGE_DIR,
     USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
 )
 from truss.contexts.image_builder.cache_warmer import (
@@ -55,9 +52,6 @@
     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"
@@ -509,6 +503,13 @@ 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(
@@ -538,10 +539,6 @@ 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:
@@ -558,50 +555,8 @@ 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 = (
-            base_server_requirements + spec.requirements_txt
-            if spec.requirements
-            else ""
+            spec.requirements_txt if spec.requirements else ""
         )
         if spec.requirements_file is not None:
             copy_into_build_dir(
@@ -615,7 +570,6 @@ 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,
@@ -625,7 +579,6 @@ 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],
@@ -661,7 +614,6 @@ 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,