Release 0.7.7 (#665)

* Always upgrade when pip installing during patch (#659) * Always upgrade when pip installing. * Bump pyproject. * Private `hf_cache` fix (#663) * add fix * bump pyproj * update cache to read secret from correct path * clean up pyproject.toml * add test to check for mount * template out hf token file name * Bump version to 0.7.7 --------- Co-authored-by: Sidharth Shanker <sid.shanker@baseten.co> Co-authored-by: Varun Shenoy <vnshenoy@stanford.edu>
basetenlabs · Sep 15, 2023 · e7485f3 · e7485f3
1 parent ac16c0c
commit e7485f3
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 6 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.7.6"
+version = "0.7.7"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

diff --git a/truss/contexts/image_builder/serving_image_builder.py b/truss/contexts/image_builder/serving_image_builder.py
@@ -45,6 +45,7 @@
 CONFIG_FILE = "config.yaml"
 
 HF_ACCESS_TOKEN_SECRET_NAME = "hf_access_token"
+HF_ACCESS_TOKEN_FILE_NAME = "hf-access-token"
 
 
 def create_triton_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
@@ -234,7 +235,9 @@ def update_config_and_gather_files(
     return get_files_to_cache(config, truss_dir, build_dir)
 
 
-def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
+def create_tgi_build_dir(
+    config: TrussConfig, build_dir: Path, truss_dir: Path, use_hf_secret: bool
+):
     copy_tree_path(truss_dir, build_dir)
 
     if not build_dir.exists():
@@ -258,6 +261,8 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
         data_dir_exists=data_dir.exists(),
         credentials_exists=credentials_file.exists(),
         cached_files=cached_file_paths,
+        use_hf_secret=use_hf_secret,
+        hf_access_token_file_name=HF_ACCESS_TOKEN_FILE_NAME,
     )
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
@@ -279,7 +284,9 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
     supervisord_filepath.write_text(supervisord_contents)
 
 
-def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
+def create_vllm_build_dir(
+    config: TrussConfig, build_dir: Path, truss_dir: Path, use_hf_secret
+):
     copy_tree_path(truss_dir, build_dir)
 
     server_endpoint_config = {
@@ -312,6 +319,8 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path)
         data_dir_exists=data_dir.exists(),
         credentials_exists=credentials_file.exists(),
         cached_files=cached_file_paths,
+        use_hf_secret=use_hf_secret,
+        hf_access_token_file_name=HF_ACCESS_TOKEN_FILE_NAME,
     )
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
@@ -361,10 +370,10 @@ def prepare_image_build_dir(
             build_dir = build_truss_target_directory(model_framework_name)
 
         if config.build.model_server is ModelServer.TGI:
-            create_tgi_build_dir(config, build_dir, truss_dir)
+            create_tgi_build_dir(config, build_dir, truss_dir, use_hf_secret)
             return
         elif config.build.model_server is ModelServer.VLLM:
-            create_vllm_build_dir(config, build_dir, truss_dir)
+            create_vllm_build_dir(config, build_dir, truss_dir, use_hf_secret)
             return
         elif config.build.model_server is ModelServer.TRITON:
             create_triton_build_dir(config, build_dir, truss_dir)
@@ -481,6 +490,7 @@ def _render_dockerfile(
             build_dir / REQUIREMENTS_TXT_FILENAME
         )
 
+        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,
@@ -497,6 +507,8 @@ def _render_dockerfile(
             cached_files=cached_files,
             credentials_exists=credentials_file.exists(),
             hf_cache=len(config.hf_cache.models) > 0,
+            hf_access_token=hf_access_token,
+            hf_access_token_file_name=HF_ACCESS_TOKEN_FILE_NAME,
         )
         docker_file_path = build_dir / MODEL_DOCKERFILE_NAME
         docker_file_path.write_text(dockerfile_contents)
diff --git a/truss/templates/cache.Dockerfile.jinja b/truss/templates/cache.Dockerfile.jinja
@@ -17,6 +17,6 @@ RUN pip install -r /app/cache_requirements.txt --no-cache-dir && rm -rf /root/.c
 COPY ./cache_warmer.py /cache_warmer.py
 {% for repo, hf_dir in models.items() %}
         {% for file in hf_dir.files %}
-RUN python3 /cache_warmer.py {{file}} {{repo}} {% if hf_dir.revision != None %}{{hf_dir.revision}}{% endif %}
+{{ "RUN --mount=type=secret,id=" + hf_access_token_file_name + ",dst=/etc/secrets/" + hf_access_token_file_name if use_hf_secret else "RUN" }} python3 /cache_warmer.py {{file}} {{repo}} {% if hf_dir.revision != None %}{{hf_dir.revision}}{% endif %}
         {% endfor %}
 {% endfor %}
diff --git a/truss/templates/control/control/helpers/truss_patch/model_container_patch_applier.py b/truss/templates/control/control/helpers/truss_patch/model_container_patch_applier.py
@@ -106,6 +106,7 @@ def _apply_python_requirement_patch(
                     self._pip_path,
                     "install",
                     python_requirement_patch.requirement,
+                    "--upgrade",
                 ],
                 check=True,
             )

diff --git a/truss/tests/contexts/image_builder/test_serving_image_builder.py b/truss/tests/contexts/image_builder/test_serving_image_builder.py
@@ -5,6 +5,7 @@
 
 import pytest
 from truss.contexts.image_builder.serving_image_builder import (
+    HF_ACCESS_TOKEN_FILE_NAME,
     ServingImageBuilderContext,
     get_files_to_cache,
     update_model_key,
@@ -243,3 +244,20 @@ def test_truss_server_caching_truss():
         )
         time.sleep(15)
         assert "Downloading model.safetensors:" not in container.logs()
+
+
+def test_hf_cache_dockerfile():
+    truss_root = Path(__file__).parent.parent.parent.parent.parent.resolve() / "truss"
+    truss_dir = truss_root / "test_data" / "test_truss_server_caching_truss"
+    tr = TrussHandle(truss_dir)
+
+    builder_context = ServingImageBuilderContext
+    image_builder = builder_context.run(tr.spec.truss_dir)
+
+    secret_mount = f"RUN --mount=type=secret,id={HF_ACCESS_TOKEN_FILE_NAME}"
+    with TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+        image_builder.prepare_image_build_dir(tmp_path, use_hf_secret=True)
+        with open(tmp_path / "Dockerfile", "r") as f:
+            gen_docker_file = f.read()
+            assert secret_mount in gen_docker_file