Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Aug 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 0 additions & 19 deletions examples/vllm-gcs/config.yaml

This file was deleted.

1 change: 0 additions & 1 deletion examples/vllm-gcs/data/service_account.json

This file was deleted.

230 changes: 26 additions & 204 deletions poetry.lock

Large diffs are not rendered by default.

4 changes: 1 addition & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down Expand Up @@ -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"
Expand All @@ -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"
Expand Down
42 changes: 13 additions & 29 deletions truss/contexts/image_builder/cache_warmer.py
Original file line number Diff line number Diff line change
@@ -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)

Expand Down
108 changes: 5 additions & 103 deletions truss/contexts/image_builder/serving_image_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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)

Expand All @@ -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):
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
)
Expand Down
6 changes: 2 additions & 4 deletions truss/templates/base.Dockerfile.jinja
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
2 changes: 0 additions & 2 deletions truss/templates/cache_requirements.txt

This file was deleted.

24 changes: 5 additions & 19 deletions truss/templates/vllm/vllm.Dockerfile.jinja
Original file line number Diff line number Diff line change
Expand Up @@ -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"]