Merge pull request #652 from basetenlabs/bump-version-0.7.3

Release 0.7.3

Author: squidarth

docs/guides/secrets.mdx

@@ -0,0 +1,55 @@
+---
+title: How to use secrets
+description: "A guide to using secrets securely in your ML models"
+---
+
+Your model server may need to use access tokens, API keys, passwords, or other secret values. Truss gives you everything you need to use secrets securely.
+
+## Setting secrets in `config.yaml`
+
+If your model needs a secret, first add its name in `config.yaml` with a placeholder value:
+
+```yaml config.yaml
+secrets:
+  hf_access_token: null
+```
+
+<Warning>
+Never set the actual value of a secret in the `config.yaml` file. Only put secret values in secure places, like the Baseten workspace secret manager.
+</Warning>
+
+## Using secrets in `model.py`
+
+Secrets are passed to your `Model` class as a keyword argument in `init`. They can be accessed with:
+
+```py model/model.py
+def __init__(self, **kwargs):
+    self._secrets = kwargs["secrets"]
+```
+
+You can then use the `self._secrets` dictionary in the `load` and `predict` functions:
+
+```py model/model.py
+def load(self):
+    self._model = pipeline(
+        "fill-mask",
+        model="baseten/docs-example-gated-model",
+        use_auth_token=self._secrets["hf_access_token"]
+    )
+```
+
+## Storing secrets on your remote
+
+On your remote host, such as your Baseten account, store both the secret name and value before deploying your model. On Baseten, you can add secrets to your workspace on the [secrets workspace settings page](https://app.baseten.co/settings/secrets).
+
+Make sure to use the same name (case sensitive) as used in the Truss on the remote.
+
+## Deploying with secrets
+
+For additional security, models don't have access to secrets by default. To deploy a model and give it access to secrets, pass the `--trusted` flag during `truss push` as follows:
+
+```sh
+truss push --trusted
+```
+
+Your model will be deployed with access to secrets stored on your remote.

docs/mint.json

@@ -64,6 +64,12 @@
         "examples/performance/vllm-server"
       ]
     },
+    {
+      "group": "Guides",
+      "pages": [
+        "guides/secrets"
+      ]
+    },
     {
       "group": "Remotes",
       "pages": [

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.7.2"
+version = "0.7.3"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/cli/cli.py

@@ -433,7 +433,7 @@ def push(
 
         click.echo(draft_model_text)
 
-    logs_url = remote_provider.get_remote_logs_url(model_name)  # type: ignore[attr-defined]
+    logs_url = remote_provider.get_remote_logs_url(model_name, publish)  # type: ignore[attr-defined]
     rich.print(f"🪵  View logs for your deployment at {logs_url}")
     should_open_logs = inquirer.confirm(
         message="🗂  Open logs in a new tab?", default=True

truss/contexts/image_builder/cache_warmer.py

@@ -63,7 +63,12 @@ def download_file(
         secret_path = Path("/etc/secrets/hf-access-token")
         secret = secret_path.read_text().strip() if secret_path.exists() else None
         try:
-            hf_hub_download(repo_name, file_name, revision=revision_name, token=secret)
+            hf_hub_download(
+                repo_name,
+                file_name,
+                revision=revision_name,
+                token=secret,
+            )
         except FileNotFoundError:
             raise RuntimeError(
                 "Hugging Face repository not found (and no valid secret found for possibly private repository)."

truss/contexts/image_builder/serving_image_builder.py

@@ -1,9 +1,9 @@
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import yaml
 from google.cloud import storage
-from huggingface_hub import list_repo_files
+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,
@@ -29,13 +29,7 @@
 )
 from truss.contexts.truss_context import TrussContext
 from truss.patch.hash import directory_content_hash
-from truss.truss_config import (
-    Build,
-    HuggingFaceCache,
-    HuggingFaceModel,
-    ModelServer,
-    TrussConfig,
-)
+from truss.truss_config import Build, HuggingFaceModel, ModelServer, TrussConfig
 from truss.truss_spec import TrussSpec
 from truss.util.download import download_external_data
 from truss.util.jinja import read_template_from_fs
@@ -133,33 +127,43 @@ def list_files(repo_id, data_dir, revision=None):
         return list_repo_files(repo_id, revision=revision)
 
 
-def update_config_and_gather_files(
-    config: TrussConfig, truss_dir: Path, build_dir: Path, server_name: str
-):
+def update_model_key(config: TrussConfig) -> str:
+    server_name = config.build.model_server
+
+    if server_name == ModelServer.TGI:
+        return "model_id"
+    elif server_name == ModelServer.VLLM:
+        return "model"
+
+    raise ValueError(
+        f"Invalid server name (must be `TGI` or `VLLM`, not `{server_name}`)."
+    )
+
+
+def update_model_name(config: TrussConfig, model_key: str) -> str:
+    if model_key not in config.build.arguments:
+        # We should definitely just use the same key across both vLLM and TGI
+        raise KeyError(
+            "Key for model missing in config or incorrect key used. Use `model` for VLLM and `model_id` for TGI."
+        )
+    model_name = config.build.arguments[model_key]
+    if "gs://" in model_name:
+        # if we are pulling from a gs bucket, we want to alias it as a part of the cache
+        model_to_cache = HuggingFaceModel(model_name)
+        config.hf_cache.models.append(model_to_cache)
+
+        config.build.arguments[
+            model_key
+        ] = f"/app/hf_cache/{model_name.replace('gs://', '')}"
+    return model_name
+
+
+def get_files_to_cache(config: TrussConfig, truss_dir: Path, build_dir: Path):
     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]
 
-    if server_name == "TGI":
-        model_key = "model_id"
-    elif server_name == "vLLM":
-        model_key = "model"
-
-    if server_name != "TrussServer":
-        model_name = config.build.arguments[model_key]
-        if "gs://" in model_name:
-            # if we are pulling from a gs bucket, we want to alias it as a part of the cache
-            model_to_cache = {"repo_id": model_name}
-            if config.hf_cache:
-                config.hf_cache.models.append(
-                    HuggingFaceModel.from_dict(model_to_cache)
-                )
-            else:
-                config.hf_cache = HuggingFaceCache.from_list([model_to_cache])
-            config.build.arguments[
-                model_key
-            ] = f"/app/hf_cache/{model_name.replace('gs://', '')}"
-
     model_files = {}
+    cached_files: List[str] = []
     if config.hf_cache:
         curr_dir = Path(__file__).parent.resolve()
         copy_into_build_dir(curr_dir / "cache_warmer.py", "cache_warmer.py")
@@ -180,18 +184,52 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
                 )
             )
 
-            if "gs://" in repo_id:
-                repo_id, _ = split_gs_path(repo_id)
-                repo_id = f"gs://{repo_id}"
+            cached_files = fetch_files_to_cache(
+                cached_files, repo_id, filtered_repo_files
+            )
+
+            model_files[repo_id] = {"files": filtered_repo_files, "revision": revision}
 
-            model_files[repo_id] = {
-                "files": filtered_repo_files,
-                "revision": revision,
-            }
     copy_into_build_dir(
         TEMPLATES_DIR / "cache_requirements.txt", "cache_requirements.txt"
     )
-    return model_files
+    return model_files, cached_files
+
+
+def fetch_files_to_cache(cached_files: list, repo_id: str, filtered_repo_files: list):
+    if "gs://" in repo_id:
+        bucket_name, _ = split_gs_path(repo_id)
+        repo_id = f"gs://{bucket_name}"
+
+        for filename in filtered_repo_files:
+            cached_files.append(f"/app/hf_cache/{bucket_name}/{filename}")
+    else:
+        repo_folder_name = f"models--{repo_id.replace('/', '--')}"
+        for filename in filtered_repo_files:
+            hf_url = hf_hub_url(repo_id, filename)
+            hf_file_metadata = get_hf_file_metadata(hf_url)
+
+            cached_files.append(f"{repo_folder_name}/blobs/{hf_file_metadata.etag}")
+
+        # snapshots is just a set of folders with symlinks -- we can copy the entire thing separately
+        cached_files.append(f"{repo_folder_name}/snapshots/")
+
+        # refs just has files with revision commit hashes
+        cached_files.append(f"{repo_folder_name}/refs/")
+
+        cached_files.append("version.txt")
+
+    return cached_files
+
+
+def update_config_and_gather_files(
+    config: TrussConfig, truss_dir: Path, build_dir: Path
+):
+    if config.build.model_server != ModelServer.TrussServer:
+        model_key = update_model_key(config)
+        update_model_name(config, model_key)
+
+    return get_files_to_cache(config, truss_dir, build_dir)
 
 
 def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
@@ -200,19 +238,24 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path):
     if not build_dir.exists():
         build_dir.mkdir(parents=True)
 
-    model_files = update_config_and_gather_files(config, truss_dir, build_dir, "TGI")
+    model_files, cached_file_paths = update_config_and_gather_files(
+        config, truss_dir, build_dir
+    )
 
     hf_access_token = config.secrets.get(HF_ACCESS_TOKEN_SECRET_NAME)
     dockerfile_template = read_template_from_fs(
         TEMPLATES_DIR, "tgi/tgi.Dockerfile.jinja"
     )
 
     data_dir = build_dir / "data"
+    credentials_file = data_dir / "service_account.json"
     dockerfile_content = dockerfile_template.render(
         hf_access_token=hf_access_token,
         models=model_files,
         hf_cache=config.hf_cache,
-        data_dir_exists=Path(data_dir).exists(),
+        data_dir_exists=data_dir.exists(),
+        credentials_exists=credentials_file.exists(),
+        cached_files=cached_file_paths,
     )
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
@@ -247,7 +290,9 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path)
     build_config: Build = config.build
     server_endpoint = server_endpoint_config[build_config.arguments.pop("endpoint")]
 
-    model_files = update_config_and_gather_files(config, truss_dir, build_dir, "vLLM")
+    model_files, cached_file_paths = update_config_and_gather_files(
+        config, truss_dir, build_dir
+    )
 
     hf_access_token = config.secrets.get(HF_ACCESS_TOKEN_SECRET_NAME)
     dockerfile_template = read_template_from_fs(
@@ -256,12 +301,15 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_dir: Path)
     nginx_template = read_template_from_fs(TEMPLATES_DIR, "vllm/proxy.conf.jinja")
 
     data_dir = build_dir / "data"
+    credentials_file = data_dir / "service_account.json"
     dockerfile_content = dockerfile_template.render(
         hf_access_token=hf_access_token,
         models=model_files,
         should_install_server_requirements=True,
         hf_cache=config.hf_cache,
         data_dir_exists=data_dir.exists(),
+        credentials_exists=credentials_file.exists(),
+        cached_files=cached_file_paths,
     )
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
@@ -336,8 +384,8 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
         download_external_data(self._spec.external_data, data_dir)
 
         # Download from HuggingFace
-        model_files = update_config_and_gather_files(
-            config, truss_dir, build_dir, server_name="TrussServer"
+        model_files, cached_files = update_config_and_gather_files(
+            config, truss_dir, build_dir
         )
 
         # Copy inference server code
@@ -391,7 +439,11 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
         (build_dir / SYSTEM_PACKAGES_TXT_FILENAME).write_text(spec.system_packages_txt)
 
         self._render_dockerfile(
-            build_dir, should_install_server_requirements, model_files, use_hf_secret
+            build_dir,
+            should_install_server_requirements,
+            model_files,
+            use_hf_secret,
+            cached_files,
         )
 
     def _render_dockerfile(
@@ -400,10 +452,12 @@ def _render_dockerfile(
         should_install_server_requirements: bool,
         model_files: Dict[str, Any],
         use_hf_secret: bool,
+        cached_files: List[str],
     ):
         config = self._spec.config
         data_dir = build_dir / config.data_dir
         bundled_packages_dir = build_dir / config.bundled_packages_dir
+        credentials_file = data_dir / "service_account.json"
         dockerfile_template = read_template_from_fs(
             TEMPLATES_DIR, SERVER_DOCKERFILE_TEMPLATE_NAME
         )
@@ -437,6 +491,9 @@ def _render_dockerfile(
             truss_hash=directory_content_hash(self._truss_dir),
             models=model_files,
             use_hf_secret=use_hf_secret,
+            cached_files=cached_files,
+            credentials_exists=credentials_file.exists(),
+            hf_cache=len(config.hf_cache.models) > 0,
         )
         docker_file_path = build_dir / MODEL_DOCKERFILE_NAME
         docker_file_path.write_text(dockerfile_contents)

truss/remote/baseten/remote.py

@@ -183,6 +183,13 @@ def patch(
         dev_version = get_dev_version_info(self._api, model_name)  # type: ignore
         truss_hash = dev_version.get("truss_hash", None)
         truss_signature = dev_version.get("truss_signature", None)
+        if not (truss_hash and truss_signature):
+            logger.error(
+                """Failed to inspect a running remote deployment to watch for changes.
+Ensure that there exists a running remote deployment before attempting to watch for changes
+            """
+            )
+            return
         LocalConfigHandler.add_signature(truss_hash, truss_signature)
         try:
             patch_request = truss_handle.calc_patch(truss_hash)

truss/templates/base.Dockerfile.jinja

@@ -1,5 +1,5 @@
 ARG PYVERSION={{config.python_version}}
-FROM {{base_image_name_and_tag}}
+FROM {{base_image_name_and_tag}} as truss_server
 
 ENV PYTHON_EXECUTABLE {{ config.base_image.python_executable_path or 'python3' }}
 

truss/templates/cache.Dockerfile.jinja

@@ -0,0 +1,22 @@
+FROM python:3.11-slim as cache_warmer
+
+RUN mkdir -p /app/hf_cache
+WORKDIR /app
+
+{% if hf_access_token %}
+ENV HUGGING_FACE_HUB_TOKEN {{hf_access_token}}
+{% endif %}
+{%- if credentials_exists %}
+COPY ./data/service_account.json /app/data/service_account.json
+{%- endif %}
+
+RUN apt-get -y update; apt-get -y install curl; curl -s https://baseten-public.s3.us-west-2.amazonaws.com/bin/b10cp-5fe8dc7da-linux-amd64 -o /app/b10cp; chmod +x /app/b10cp
+ENV B10CP_PATH_TRUSS /app/b10cp
+COPY ./cache_requirements.txt /app/cache_requirements.txt
+RUN pip install -r /app/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
+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 %}
+        {% endfor %}
+{% endfor %}