Bundling Private Weights from GCP (#552)

* working gcp bundling

* update save dir to cache_dir / {bucket}

* change toml

* update poetry.lock

* update pyproject

* bump

* vllm working

* cleaned up example

* cleanup

* bump

* model can now point to a gs bucket

* alias model_name for gcs

* bump

* fix cache_requirements

* bump

* bump

* bump

* remove cache_reqs

* remove cache_reqs

* bump

* add cache_reqs

* bump

* update serving builder

* remove spec

* revert toml

* bump

* add gc to toml

* update lock

Author: Varun Shenoy

examples/vllm-gcs/config.yaml

@@ -0,0 +1,19 @@
+build:
+  arguments:
+    endpoint: Completions
+    model: gs://llama-2-7b
+    tokenizer: hf-internal-testing/llama-tokenizer
+  model_server: VLLM
+environment_variables: {}
+external_package_dirs: []
+model_metadata: {}
+model_name: vllm llama gcs
+python_version: py39
+requirements: []
+resources:
+  accelerator: A10G
+  cpu: 500m
+  memory: 30Gi
+  use_gpu: true
+secrets: {}
+system_packages: []

examples/vllm-gcs/data/service_account.json

@@ -0,0 +1 @@
+YOUR SERVICE ACCOUNT KEY

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.6.1"
+version = "0.6.2"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -41,6 +41,7 @@ watchfiles = "^0.19.0"
 huggingface_hub = "^0.16.4"
 rich-click = "^1.6.1"
 inquirerpy = "^0.3.4"
+google-cloud-storage = "2.10.0"
 
 [tool.poetry.group.builder.dependencies]
 python = ">=3.8,<3.12"
@@ -59,6 +60,7 @@ uvicorn = "^0.21.1"
 httpx = "^0.24.1"
 psutil = "^5.9.4"
 huggingface_hub = "^0.16.4"
+google-cloud-storage = "2.10.0"
 
 [tool.poetry.dev-dependencies]
 torch = "^1.9.0"

pyproject.toml

@@ -1,25 +1,41 @@
 import sys
 from pathlib import Path
 
+from google.cloud import storage
 from huggingface_hub import hf_hub_download
 
 
-def download_file(repo_name, file_name, revision_name=None):
-    secret = None
-    secret_path = Path("/etc/secrets/hf_access_token")
-
-    if secret_path.exists():
-        secret = secret_path.read_text().strip()
-    try:
-        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)."
-        )
+def download_file(
+    repo_name, file_name, revision_name=None, key_file="/app/data/service_account.json"
+):
+    # Check if repo_name starts with "gs://"
+    if "gs://" in repo_name:
+        # Create directory if not exist
+        repo_name = repo_name.replace("gs://", "")
+        cache_dir = Path(f"/app/hf_cache/{repo_name}")
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        # Connect to GCS storage
+        try:
+            storage_client = storage.Client.from_service_account_json(key_file)
+            bucket = storage_client.bucket(repo_name)
+            blob = bucket.blob(file_name)
+            # Download the blob to a file
+            blob.download_to_filename(f"{cache_dir}/{file_name}")
+        except Exception as e:
+            raise RuntimeError(f"Failure downloading file from GCS: {e}")
+    else:
+        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)
+        except FileNotFoundError:
+            raise RuntimeError(
+                "Hugging Face repository not found (and no valid secret found for possibly private repository)."
+            )
 
 
 if __name__ == "__main__":
-    # TODO(varun): might make sense to move this check + write to a separate `prepare_cache.py` script
     file_path = Path.home() / ".cache/huggingface/hub/version.txt"
     file_path.parent.mkdir(parents=True, exist_ok=True)
 

truss/contexts/image_builder/cache_warmer.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, Optional
 
 import yaml
+from google.cloud import storage
 from huggingface_hub import list_repo_files
 from huggingface_hub.utils import filter_repo_objects
 from truss.constants import (
@@ -27,7 +28,13 @@
 )
 from truss.contexts.truss_context import TrussContext
 from truss.patch.hash import directory_content_hash
-from truss.truss_config import Build, ModelServer, TrussConfig
+from truss.truss_config import (
+    Build,
+    HuggingFaceCache,
+    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
@@ -75,7 +82,12 @@ def create_tgi_build_dir(config: TrussConfig, build_dir: Path):
     supervisord_filepath.write_text(supervisord_contents)
 
 
-def create_vllm_build_dir(config: TrussConfig, build_dir: Path):
+def create_vllm_build_dir(config: TrussConfig, build_dir: Path, truss_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]
+
+    copy_tree_path(truss_dir, build_dir)
+
     server_endpoint_config = {
         "Completions": "/v1/completions",
         "ChatCompletions": "/v1/chat/completions",
@@ -85,13 +97,58 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path):
 
     build_config: Build = config.build
     server_endpoint = server_endpoint_config[build_config.arguments.pop("endpoint")]
+
+    model_name = build_config.arguments.pop("model")
+    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])
+        build_config.arguments[
+            "model"
+        ] = f"/app/hf_cache/{model_name.replace('gs://', '')}"
+
     hf_access_token = config.secrets.get(HF_ACCESS_TOKEN_SECRET_NAME)
     dockerfile_template = read_template_from_fs(
         TEMPLATES_DIR, "vllm/vllm.Dockerfile.jinja"
     )
     nginx_template = read_template_from_fs(TEMPLATES_DIR, "vllm/proxy.conf.jinja")
+    copy_into_build_dir(
+        TEMPLATES_DIR / "cache_requirements.txt", "cache_requirements.txt"
+    )
 
-    dockerfile_content = dockerfile_template.render(hf_access_token=hf_access_token)
+    model_files = {}
+    if config.hf_cache:
+        curr_dir = Path(__file__).parent.resolve()
+        copy_into_build_dir(curr_dir / "cache_warmer.py", "cache_warmer.py")
+        for model in config.hf_cache.models:
+            repo_id = model.repo_id
+            revision = model.revision
+
+            allow_patterns = model.allow_patterns
+            ignore_patterns = model.ignore_patterns
+
+            filtered_repo_files = list(
+                filter_repo_objects(
+                    items=list_files(
+                        repo_id, truss_dir / config.data_dir, revision=revision
+                    ),
+                    allow_patterns=allow_patterns,
+                    ignore_patterns=ignore_patterns,
+                )
+            )
+            model_files[repo_id] = {
+                "files": filtered_repo_files,
+                "revision": revision,
+            }
+
+    dockerfile_content = dockerfile_template.render(
+        hf_access_token=hf_access_token,
+        models=model_files,
+        should_install_server_requirements=True,
+    )
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
 
@@ -110,6 +167,47 @@ def create_vllm_build_dir(config: TrussConfig, build_dir: Path):
     supervisord_filepath.write_text(supervisord_contents)
 
 
+def split_gs_path(gs_path):
+    # Remove the 'gs://' prefix
+    path = gs_path.replace("gs://", "")
+
+    # Split on the first slash
+    parts = path.split("/", 1)
+
+    bucket_name = parts[0]
+    prefix = parts[1] if len(parts) > 1 else ""
+
+    return bucket_name, prefix
+
+
+def list_bucket_files(bucket_name, data_dir, is_trusted=False):
+    # TODO(varun): provide support for aws s3
+
+    if is_trusted:
+        storage_client = storage.Client.from_service_account_json(
+            data_dir / "service_account.json"
+        )
+    else:
+        storage_client = storage.Client()
+    print(bucket_name.replace("gs://", ""))
+    bucket_name, prefix = split_gs_path(bucket_name)
+    blobs = storage_client.list_blobs(bucket_name, prefix=prefix)
+
+    all_objects = []
+    for blob in blobs:
+        all_objects.append(Path(blob.name).name)
+        print(Path(blob.name).name)
+    return all_objects
+
+
+def list_files(repo_id, data_dir, revision=None):
+    if repo_id.startswith(("s3://", "gs://")):
+        return list_bucket_files(repo_id, data_dir, is_trusted=True)
+    else:
+        # we assume it's a HF bucket
+        list_repo_files(repo_id, revision=revision)
+
+
 class ServingImageBuilderContext(TrussContext):
     @staticmethod
     def run(truss_dir: Path):
@@ -143,7 +241,7 @@ def prepare_image_build_dir(
             create_tgi_build_dir(config, build_dir)
             return
         elif config.build.model_server is ModelServer.VLLM:
-            create_vllm_build_dir(config, build_dir)
+            create_vllm_build_dir(config, build_dir, truss_dir)
             return
 
         data_dir = build_dir / config.data_dir  # type: ignore[operator]
@@ -175,7 +273,7 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
 
                 filtered_repo_files = list(
                     filter_repo_objects(
-                        items=list_repo_files(repo_id, revision=revision),
+                        items=list_files(repo_id, data_dir, revision=revision),
                         allow_patterns=allow_patterns,
                         ignore_patterns=ignore_patterns,
                     )

truss/contexts/image_builder/serving_image_builder.py

@@ -44,8 +44,6 @@ RUN pip install -r requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
 {% endblock %}
 
 
-{% block cache_weights %}
-{% endblock %}
 
 
 ENV APP_HOME /app
@@ -55,6 +53,10 @@ WORKDIR $APP_HOME
 {% block app_copy %}
 {% endblock %}
 
+
+{% block cache_weights %}
+{% endblock %}
+
 {% block bundled_packages_copy %}
     {%- if bundled_packages_dir_exists %}
 COPY ./{{config.bundled_packages_dir}} /packages

truss/templates/base.Dockerfile.jinja

@@ -0,0 +1,2 @@
+huggingface-hub==0.16.4
+google-cloud-storage==2.10.0

truss/templates/cache_requirements.txt

@@ -2,19 +2,33 @@ FROM baseten/vllm:latest
 
 EXPOSE 8080-9000
 
+{% if hf_access_token %}
+ENV HUGGING_FACE_HUB_TOKEN {{hf_access_token}}
+{% endif %}
+
+COPY ./data /app/data
+
+{%- if hf_cache != None %}
+COPY ./cache_warmer.py /cache_warmer.py
+
+COPY ./cache_requirements.txt /app/cache_requirements.txt
+
+RUN pip install -r /app/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
+        {% 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 %}
+{%- endif %}
+
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
         curl nginx supervisor && \
         rm -rf /var/lib/apt/lists/*
 
-
 COPY ./proxy.conf /etc/nginx/conf.d/proxy.conf
 
 RUN mkdir -p /var/log/supervisor
 COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf
 
-{% if hf_access_token %}
-ENV HUGGING_FACE_HUB_TOKEN {{hf_access_token}}
-{% endif %}
-
 ENV SERVER_START_CMD /usr/bin/supervisord
 ENTRYPOINT ["/usr/bin/supervisord"]