Revert "Bundling Private Weights from GCP (#552)"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.6.2"
+version = "0.6.1"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"
@@ -41,7 +41,6 @@ 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"
@@ -60,7 +59,6 @@ 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"

truss/contexts/image_builder/cache_warmer.py

@@ -1,41 +1,25 @@
 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, 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)."
-            )
+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)."
+        )
 
 
 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/serving_image_builder.py

@@ -2,7 +2,6 @@
 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 (
@@ -28,13 +27,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, 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
@@ -82,12 +75,7 @@ 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, 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)
-
+def create_vllm_build_dir(config: TrussConfig, build_dir: Path):
     server_endpoint_config = {
         "Completions": "/v1/completions",
         "ChatCompletions": "/v1/chat/completions",
@@ -97,58 +85,13 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
 
     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"
-    )
 
-    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_content = dockerfile_template.render(hf_access_token=hf_access_token)
     dockerfile_filepath = build_dir / "Dockerfile"
     dockerfile_filepath.write_text(dockerfile_content)
 
@@ -167,47 +110,6 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
     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):
@@ -241,7 +143,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, truss_dir)
+            create_vllm_build_dir(config, build_dir)
             return
 
         data_dir = build_dir / config.data_dir  # type: ignore[operator]
@@ -273,7 +175,7 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
 
                 filtered_repo_files = list(
                     filter_repo_objects(
-                        items=list_files(repo_id, data_dir, revision=revision),
+                        items=list_repo_files(repo_id, revision=revision),
                         allow_patterns=allow_patterns,
                         ignore_patterns=ignore_patterns,
                     )

truss/templates/base.Dockerfile.jinja

@@ -44,6 +44,8 @@ RUN pip install -r requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
 {% endblock %}
 
 
+{% block cache_weights %}
+{% endblock %}
 
 
 ENV APP_HOME /app
@@ -53,10 +55,6 @@ 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/vllm/vllm.Dockerfile.jinja

@@ -2,33 +2,19 @@ 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"]