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

truss/{remote → cli}/remote_cli.py

@@ -23,9 +23,7 @@ def inquire_remote_config() -> RemoteConfig:
     # can do so manually in the .trussrc file.
     remote_url = "https://app.baseten.co"
     api_key = inquirer.secret(
-        message="🤫 Quietly paste your API_KEY:",
-        qmark="",
-        validate=NonEmptyValidator(),
+        message="🤫 Quietly paste your API_KEY:", qmark="", validate=NonEmptyValidator()
     ).execute()
 
     return RemoteConfig(
@@ -56,7 +54,4 @@ def inquire_remote_name(available_remotes: List[str]) -> str:
 
 
 def inquire_model_name() -> str:
-    return inquirer.text(
-        "📦 Name this model:",
-        qmark="",
-    ).execute()
+    return inquirer.text("📦 Name this model:", qmark="").execute()

truss/contexts/docker_build_setup.py

@@ -0,0 +1,67 @@
+import json
+import logging
+import os
+import pathlib
+import sys
+
+import click
+from truss.patch.hash import directory_content_hash
+from truss.patch.truss_dir_patch_applier import TrussDirPatchApplier
+from truss.templates.control.control.helpers.custom_types import Patch
+from truss.truss_handle import truss_handle
+from truss.truss_handle.patch import signature
+
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+
+working_dir = pathlib.Path("/")
+
+TRUSS_SRC_DIR = working_dir / "build/model_scaffold"
+TRUSS_HASH_FILE = working_dir / "scaffold/truss_hash"
+TRUSS_SIGNATURE_FILE = working_dir / "scaffold/truss_signature"
+TRUSS_BUILD_CONTEXT_DIR = working_dir / "build/context"
+
+
+@click.command()
+@click.option("--truss_type", required=True)
+def docker_build_setup(truss_type: str) -> None:
+    """
+    Prepares source and asset files in a build directory (build context), on which a
+    docker build command can be run.
+
+    This is to be run for remote builds on baseten.
+    Local builds use `TrussHandle.build_serving_docker_image`.
+    """
+    logging.info("Loading truss")
+    tr = truss_handle.TrussHandle(TRUSS_SRC_DIR)
+    logging.info("Truss is loaded")
+
+    if patches_dir := os.environ.get("PATCHES_DIR"):
+        logging.info("Applying patches")
+        logger = logging.getLogger("patch_applier")
+        patch_applier = TrussDirPatchApplier(TRUSS_SRC_DIR, logger)
+        patches = json.loads(pathlib.Path(patches_dir).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.
+    logging.info("Recording truss hash")
+    TRUSS_HASH_FILE.write_text(directory_content_hash(TRUSS_SRC_DIR))
+
+    logging.info("Recording truss signature.")
+    sign = signature.calc_truss_signature(TRUSS_SRC_DIR)
+    TRUSS_SIGNATURE_FILE.write_text(json.dumps(sign.to_dict()))
+
+    if truss_type == "server_control":
+        tr.live_reload(enable=True)
+    else:
+        tr.live_reload(enable=False)
+
+    # check if we have a hf_secret
+    tr.docker_build_setup(
+        TRUSS_BUILD_CONTEXT_DIR, use_hf_secret="HUGGING_FACE_HUB_TOKEN" in os.environ
+    )
+    print("Docker build context is set up for the truss.")
+
+
+if __name__ == "__main__":
+    docker_build_setup()

truss/contexts/image_builder/cache_warmer.py

@@ -25,11 +25,7 @@ def _b10cp_path() -> Optional[str]:
     return os.environ.get(B10CP_PATH_TRUSS_ENV_VAR_NAME)
 
 
-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,
@@ -157,9 +153,7 @@ class GCSFile(RepositoryFile):
 
         if is_private:
             url = blob.generate_signed_url(
-                version="v4",
-                expiration=datetime.timedelta(minutes=15),
-                method="GET",
+                version="v4", expiration=datetime.timedelta(minutes=15), method="GET"
             )
         else:
             base_url = "https://storage.googleapis.com"

truss/contexts/image_builder/image_builder.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Optional
 
-from truss.docker import Docker
+from truss.util.docker import Docker
 from truss.util.path import given_or_temporary_dir
 
 

truss/contexts/image_builder/serving_image_builder.py

@@ -6,24 +6,51 @@ from pathlib import Path
 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 (
+from truss.base import constants
+from truss.base.constants import (
+    BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
     BASE_TRTLLM_REQUIREMENTS,
+    BEI_MAX_CONCURRENCY_TARGET_REQUESTS,
+    BEI_REQUIRED_MAX_NUM_TOKENS,
+    BEI_TRTLLM_BASE_IMAGE,
+    BEI_TRTLLM_CLIENT_BATCH_SIZE,
+    BEI_TRTLLM_PYTHON_EXECUTABLE,
+    CHAINS_CODE_DIR,
+    CONTROL_SERVER_CODE_DIR,
+    DOCKER_SERVER_TEMPLATES_DIR,
     FILENAME_CONSTANTS_MAP,
+    MAX_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE,
+    MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE,
     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_PREDICT_CONCURRENCY,
+    TRTLLM_PYTHON_EXECUTABLE,
     TRTLLM_TRUSS_DIR,
-    TRUSS_PACKAGE_DIR,
+    TRUSSLESS_MAX_PAYLOAD_SIZE,
     USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
 )
+from truss.base.trt_llm_config import TRTLLMConfiguration, TrussTRTLLMModel
+from truss.base.truss_config import (
+    DEFAULT_BUNDLED_PACKAGES_DIR,
+    BaseImage,
+    DockerServer,
+    TrussConfig,
+)
+from truss.base.truss_spec import TrussSpec
 from truss.contexts.image_builder.cache_warmer import (
     AWSCredentials,
     parse_s3_credentials_file,
@@ -37,9 +64,7 @@ from truss.contexts.image_builder.util import (
     truss_base_image_tag,
 )
 from truss.contexts.truss_context import TrussContext
-from truss.patch.hash import directory_content_hash
-from truss.truss_config import BaseImage, ModelServer, TrussConfig
-from truss.truss_spec import TrussSpec
+from truss.truss_handle.patch.hash import directory_content_hash
 from truss.util.jinja import read_template_from_fs
 from truss.util.path import (
     build_truss_target_directory,
@@ -48,12 +73,16 @@ from truss.util.path import (
     load_trussignore_patterns,
 )
 
+BUILD_SERVER_DIR_NAME = "server"
+BUILD_CONTROL_SERVER_DIR_NAME = "control"
+BUILD_SERVER_EXTENSIONS_PATH = "extensions"
+BUILD_CHAINS_DIR_NAME = "truss_chains"
+
 CONFIG_FILE = "config.yaml"
 USER_TRUSS_IGNORE_FILE = ".truss_ignore"
 GCS_CREDENTIALS = "service_account.json"
 S3_CREDENTIALS = "s3_credentials.json"
 
-HF_ACCESS_TOKEN_SECRET_NAME = "hf_access_token"
 HF_ACCESS_TOKEN_FILE_NAME = "hf-access-token"
 
 CLOUD_BUCKET_CACHE = Path("/app/model_cache/")
@@ -131,7 +160,7 @@ class GCSCache(RemoteCache):
 
 
 class S3Cache(RemoteCache):
-    def list_files(self, revision=None):
+    def list_files(self, revision=None) -> List[str]:
         s3_credentials_file = self.data_dir / S3_CREDENTIALS
 
         if s3_credentials_file.exists():
@@ -279,6 +308,49 @@ def update_config_and_gather_files(
     return get_files_to_cache(config, truss_dir, build_dir)
 
 
+def generate_docker_server_nginx_config(build_dir, config):
+    nginx_template = read_template_from_fs(
+        DOCKER_SERVER_TEMPLATES_DIR, "proxy.conf.jinja"
+    )
+
+    assert config.docker_server.predict_endpoint is not None, (
+        "docker_server.predict_endpoint is required to use custom server"
+    )
+    assert config.docker_server.server_port is not None, (
+        "docker_server.server_port is required to use custom server"
+    )
+    assert config.docker_server.readiness_endpoint is not None, (
+        "docker_server.readiness_endpoint is required to use custom server"
+    )
+    assert config.docker_server.liveness_endpoint is not None, (
+        "docker_server.liveness_endpoint is required to use custom server"
+    )
+
+    nginx_content = nginx_template.render(
+        server_endpoint=config.docker_server.predict_endpoint,
+        readiness_endpoint=config.docker_server.readiness_endpoint,
+        liveness_endpoint=config.docker_server.liveness_endpoint,
+        server_port=config.docker_server.server_port,
+        client_max_body_size=TRUSSLESS_MAX_PAYLOAD_SIZE,
+    )
+    nginx_filepath = build_dir / "proxy.conf"
+    nginx_filepath.write_text(nginx_content)
+
+
+def generate_docker_server_supervisord_config(build_dir, config):
+    supervisord_template = read_template_from_fs(
+        DOCKER_SERVER_TEMPLATES_DIR, "supervisord.conf.jinja"
+    )
+    assert config.docker_server.start_command is not None, (
+        "docker_server.start_command is required to use custom server"
+    )
+    supervisord_contents = supervisord_template.render(
+        start_command=config.docker_server.start_command
+    )
+    supervisord_filepath = build_dir / "supervisord.conf"
+    supervisord_filepath.write_text(supervisord_contents)
+
+
 class ServingImageBuilderContext(TrussContext):
     @staticmethod
     def run(truss_dir: Path):
@@ -294,6 +366,99 @@ class ServingImageBuilder(ImageBuilder):
     def default_tag(self):
         return f"{self._spec.model_framework_name}-model:latest"
 
+    def _copy_into_build_dir(
+        self, from_path: Path, build_dir: Path, path_in_build_dir: str
+    ):
+        copy_tree_or_file(from_path, build_dir / path_in_build_dir)  # type: ignore[operator]
+
+    def prepare_trtllm_bei_encoder_build_dir(self, build_dir: Path):
+        """prepares the build directory for a trtllm ENCODER model to launch a Baseten Embeddings Inference (BEI) server"""
+        config = self._spec.config
+        assert (
+            config.trt_llm
+            and config.trt_llm.build
+            and config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER
+        ), (
+            "prepare_trtllm_bei_encoder_build_dir should only be called for ENCODER tensorrt-llm model"
+        )
+        # TRTLLM has performance degradation with batch size >> 32, so we limit the runtime settings
+        # runtime batch size may not be higher than what the build settings of the model allow
+        # to 32 even if the engine.rank0 allows for higher batch_size
+        runtime_max_batch_size = min(config.trt_llm.build.max_batch_size, 32)
+        # make sure the user gets good performance, enforcing max_num_tokens here and in engine-builder
+        runtime_max_batch_tokens = max(
+            config.trt_llm.build.max_num_tokens, BEI_REQUIRED_MAX_NUM_TOKENS
+        )
+        port = 7997
+        start_command = " ".join(
+            [
+                "truss-transfer-cli && text-embeddings-router",
+                f"--port {port}",
+                # assert the max_batch_size is within trt-engine limits
+                f"--max-batch-requests {runtime_max_batch_size}",
+                # assert the max_num_tokens is within trt-engine limits
+                f"--max-batch-tokens {runtime_max_batch_tokens}",
+                # how many sentences can be in a single json payload.
+                # limited default to improve request based autoscaling.
+                f"--max-client-batch-size {BEI_TRTLLM_CLIENT_BATCH_SIZE}",
+                # how many concurrent requests can be handled by the server until 429 is returned.
+                # limited by https://docs.baseten.co/performance/concurrency#concurrency-target
+                # 2048 is a safe max value for the server
+                f"--max-concurrent-requests {BEI_MAX_CONCURRENCY_TARGET_REQUESTS}",
+            ]
+        )
+        self._spec.config.docker_server = DockerServer(
+            start_command=f"/bin/sh -c '{start_command}'",
+            server_port=port,
+            predict_endpoint="/v1/embeddings",
+            readiness_endpoint="/health",
+            liveness_endpoint="/health",
+        )
+        copy_tree_path(DOCKER_SERVER_TEMPLATES_DIR, build_dir, ignore_patterns=[])
+
+        config.base_image = BaseImage(
+            image=BEI_TRTLLM_BASE_IMAGE,
+            python_executable_path=BEI_TRTLLM_PYTHON_EXECUTABLE,
+        )
+
+    def prepare_trtllm_decoder_build_dir(self, build_dir: Path):
+        """prepares the build directory for a trtllm decoder-like models to launch BRITON server"""
+        config = self._spec.config
+        assert (
+            config.trt_llm
+            and config.trt_llm.build
+            and config.trt_llm.build.base_model != TrussTRTLLMModel.ENCODER
+        ), (
+            "prepare_trtllm_decoder_build_dir should only be called for decoder tensorrt-llm model"
+        )
+
+        # trt_llm is treated as an extension at model run time.
+        self._copy_into_build_dir(
+            TRTLLM_TRUSS_DIR / "src",
+            build_dir,
+            f"{BUILD_SERVER_DIR_NAME}/{BUILD_SERVER_EXTENSIONS_PATH}/trt_llm",
+        )
+        # TODO(pankaj) Do this differently. This is not ideal, user
+        # supplied code in bundled packages can conflict with those from
+        # the trtllm extension. We don't want to put this in the build
+        # directory directly either because of chances of conflict there
+        # as well and the noise it can create there. We need to find a
+        # new place that's made available in model's pythonpath. This is
+        # a bigger lift and feels overkill right now. Worth revisiting
+        # if we come across cases of actual conflicts.
+        self._copy_into_build_dir(
+            TRTLLM_TRUSS_DIR / DEFAULT_BUNDLED_PACKAGES_DIR,
+            build_dir,
+            DEFAULT_BUNDLED_PACKAGES_DIR,
+        )
+
+        config.runtime.predict_concurrency = TRTLLM_PREDICT_CONCURRENCY
+
+        config.base_image = BaseImage(
+            image=TRTLLM_BASE_IMAGE, python_executable_path=TRTLLM_PYTHON_EXECUTABLE
+        )
+        config.requirements.extend(BASE_TRTLLM_REQUIREMENTS)
+
     def prepare_image_build_dir(
         self, build_dir: Optional[Path] = None, use_hf_secret: bool = False
     ):
@@ -310,17 +475,6 @@ class ServingImageBuilder(ImageBuilder):
 
         data_dir = build_dir / config.data_dir  # type: ignore[operator]
 
-        def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
-            # using default ignore patterns ignores the `build` dir in truss
-            copy_tree_or_file(from_path, build_dir / path_in_build_dir, ignore_files=False)  # 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(
@@ -329,36 +483,37 @@ class ServingImageBuilder(ImageBuilder):
 
         # Copy over truss
         copy_tree_path(truss_dir, build_dir, ignore_patterns=truss_ignore_patterns)
-        # Copy over template truss for TRT-LLM (we overwrite the model and packages dir)
-        if config.build.model_server is ModelServer.TRT_LLM:
-            copy_tree_path(TRTLLM_TRUSS_DIR, build_dir, ignore_patterns=[])
-
-            # Check to see if TP and GPU count are the same
-            # TODO(Abu): Consolidate these config parameters so that we don't have to
-            # keep truss + template in sync if we change th einterface
-            if "tensor_parallel_count" in config.build.arguments:
-                if (
-                    config.build.arguments["tensor_parallel_count"]
-                    != config.resources.accelerator.count
-                ):
-                    raise ValueError(
-                        "Tensor parallelism and GPU count must be the same for TRT-LLM"
-                    )
-
-            config.base_image = BaseImage(
-                image=TRTLLM_BASE_IMAGE, python_executable_path="/usr/bin/python3"
+        if (
+            isinstance(config.trt_llm, TRTLLMConfiguration)
+            and config.trt_llm.build is not None
+        ):
+            if config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER:
+                # Run the specific encoder build
+                self.prepare_trtllm_bei_encoder_build_dir(build_dir=build_dir)
+            else:
+                self.prepare_trtllm_decoder_build_dir(build_dir=build_dir)
+
+        if config.docker_server is not None:
+            self._copy_into_build_dir(
+                TEMPLATES_DIR / "docker_server_requirements.txt",
+                build_dir,
+                "docker_server_requirements.txt",
             )
-            config.requirements.extend(BASE_TRTLLM_REQUIREMENTS)
+
+            generate_docker_server_nginx_config(build_dir, config)
+
+            generate_docker_server_supervisord_config(build_dir, config)
+
+        # 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:
             for ext_file in self._spec.external_data.items:
                 external_data_files.append(
-                    (
-                        ext_file.url,
-                        (data_dir / ext_file.local_data_path).resolve(),
-                    )
+                    (ext_file.url, (data_dir / ext_file.local_data_path).resolve())
                 )
 
         # Download from HuggingFace
@@ -366,12 +521,74 @@ class ServingImageBuilder(ImageBuilder):
             config, truss_dir, build_dir
         )
 
-        user_provided_python_requirements = (
-            spec.requirements_txt if spec.requirements else ""
+        # Copy inference server code
+        self._copy_into_build_dir(SERVER_CODE_DIR, build_dir, BUILD_SERVER_DIR_NAME)
+        self._copy_into_build_dir(
+            SHARED_SERVING_AND_TRAINING_CODE_DIR,
+            build_dir,
+            BUILD_SERVER_DIR_NAME + "/" + SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME,
+        )
+
+        # Copy control server code
+        if config.live_reload:
+            self._copy_into_build_dir(
+                CONTROL_SERVER_CODE_DIR, build_dir, BUILD_CONTROL_SERVER_DIR_NAME
+            )
+            self._copy_into_build_dir(
+                SHARED_SERVING_AND_TRAINING_CODE_DIR,
+                build_dir,
+                BUILD_CONTROL_SERVER_DIR_NAME
+                + "/control/"
+                + SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME,
+            )
+
+        if config.use_local_chains_src:
+            self._copy_into_build_dir(CHAINS_CODE_DIR, build_dir, BUILD_CHAINS_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:
+            self._copy_into_build_dir(
+                base_truss_server_reqs_filepath,
+                build_dir,
+                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:
+            self._copy_into_build_dir(
+                server_reqs_filepath, build_dir, SERVER_REQUIREMENTS_TXT_FILENAME
+            )
+
+        with open(base_truss_server_reqs_filepath, "r") as f:
+            base_server_requirements = f.read()
+
+        if config.docker_server:
+            # when docker server is enabled, no need to install truss requirements
+            # only install user-provided python requirements
+            user_provided_python_requirements = spec.requirements_txt
+        else:
+            # 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.
+            # TODO (BT-10217): above reasoning leads to inconsistencies. To get consistent
+            #  images tentatively add server requirements always. This whole point needs
+            #  more thought and potentially a re-design.
+            user_provided_python_requirements = (
+                base_server_requirements + spec.requirements_txt
+                if spec.requirements
+                else base_server_requirements
+            )
         if spec.requirements_file is not None:
-            copy_into_build_dir(
+            self._copy_into_build_dir(
                 truss_dir / spec.requirements_file,
+                build_dir,
                 USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
             )
         (build_dir / REQUIREMENTS_TXT_FILENAME).write_text(
@@ -381,24 +598,28 @@ class ServingImageBuilder(ImageBuilder):
 
         self._render_dockerfile(
             build_dir,
+            should_install_server_requirements,
             model_files,
             use_hf_secret,
             cached_files,
             external_data_files,
+            self._spec.build_commands,
         )
 
     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],
         external_data_files: List[Tuple[str, str]],
+        build_commands: List[str],
     ):
         config = self._spec.config
         data_dir = build_dir / config.data_dir
+        model_dir = build_dir / config.model_module_dir
         bundled_packages_dir = build_dir / config.bundled_packages_dir
-
         dockerfile_template = read_template_from_fs(
             TEMPLATES_DIR, SERVER_DOCKERFILE_TEMPLATE_NAME
         )
@@ -423,9 +644,31 @@ class ServingImageBuilder(ImageBuilder):
             build_dir / USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME
         )
 
-        hf_access_token = config.secrets.get(HF_ACCESS_TOKEN_SECRET_NAME)
+        max_supported_python_version_in_custom_base_image = (
+            MAX_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE
+        )
+        min_supported_python_version_in_custom_base_image = (
+            MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE
+        )
+        max_supported_python_minor_version_in_custom_base_image = (
+            MAX_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE.split(".")[1]
+        )
+        min_supported_python_minor_version_in_custom_base_image = (
+            MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE.split(".")[1]
+        )
+        supported_python_major_version_in_custom_base_image = (
+            MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE.split(".")[0]
+        )
+
+        hf_access_token = config.secrets.get(constants.HF_ACCESS_TOKEN_KEY)
         dockerfile_contents = dockerfile_template.render(
+            should_install_server_requirements=should_install_server_requirements,
             base_image_name_and_tag=base_image_name_and_tag,
+            max_supported_python_version_in_custom_base_image=max_supported_python_version_in_custom_base_image,
+            min_supported_python_version_in_custom_base_image=min_supported_python_version_in_custom_base_image,
+            max_supported_python_minor_version_in_custom_base_image=max_supported_python_minor_version_in_custom_base_image,
+            min_supported_python_minor_version_in_custom_base_image=min_supported_python_minor_version_in_custom_base_image,
+            supported_python_major_version_in_custom_base_image=supported_python_major_version_in_custom_base_image,
             should_install_system_requirements=should_install_system_requirements,
             should_install_requirements=should_install_python_requirements,
             should_install_user_requirements_file=should_install_user_requirements_file,
@@ -433,6 +676,7 @@ class ServingImageBuilder(ImageBuilder):
             python_version=python_version,
             live_reload=config.live_reload,
             data_dir_exists=data_dir.exists(),
+            model_dir_exists=model_dir.exists(),
             bundled_packages_dir_exists=bundled_packages_dir.exists(),
             truss_hash=directory_content_hash(
                 self._truss_dir, self._spec.hash_ignore_patterns
@@ -445,6 +689,8 @@ class ServingImageBuilder(ImageBuilder):
             hf_access_token=hf_access_token,
             hf_access_token_file_name=HF_ACCESS_TOKEN_FILE_NAME,
             external_data_files=external_data_files,
+            build_commands=build_commands,
+            use_local_chains_src=config.use_local_chains_src,
             **FILENAME_CONSTANTS_MAP,
         )
         docker_file_path = build_dir / MODEL_DOCKERFILE_NAME

truss/contexts/image_builder/util.py

@@ -36,9 +36,7 @@ def truss_base_image_name(job_type: str) -> str:
 
 
 def truss_base_image_tag(
-    python_version: str,
-    use_gpu: bool,
-    version_tag: Optional[str] = None,
+    python_version: str, use_gpu: bool, version_tag: Optional[str] = None
 ) -> str:
     if version_tag is None:
         version_tag = f"v{__version__}"

truss/contexts/local_loader/docker_build_emulator.py

@@ -0,0 +1,58 @@
+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

@@ -1,14 +1,14 @@
 import inspect
 from pathlib import Path
 
+from truss.base.truss_spec import TrussSpec
 from truss.contexts.local_loader.truss_module_loader import truss_module_loaded
 from truss.contexts.local_loader.utils import (
     prepare_secrets,
     signature_accepts_keyword_arg,
 )
 from truss.contexts.truss_context import TrussContext
-from truss.server.common.patches import apply_patches
-from truss.truss_spec import TrussSpec
+from truss.templates.server.common.patches import apply_patches
 
 
 class LoadModelLocal(TrussContext):