diff --git a/pyproject.toml b/pyproject.toml index 737266603..da8a64c1e 100644 --- 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 index 3fc4aa882..98bc54c4e 100644 --- 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 index c60ab7264..edc44c61f 100644 --- 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 index 84dc2d05f..75e1015fe 100644 --- 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 index 2c3522ff4..bac5ccb20 100644 --- 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