diff --git a/bptr_generator/Dockerfile.bptr b/bptr_generator/Dockerfile.bptr new file mode 100644 index 000000000..28eeb138e --- /dev/null +++ b/bptr_generator/Dockerfile.bptr @@ -0,0 +1,20 @@ +FROM python:3.11-slim + +WORKDIR /app + +# Install curl and other dependencies +RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/* + +COPY bptr_requirements.txt . +RUN pip install --no-cache-dir -r bptr_requirements.txt + +# Install truss-transfer-cli +RUN curl -sSL --fail --retry 5 --retry-delay 2 -o /usr/local/bin/truss-transfer-cli https://github.com/basetenlabs/truss/releases/download/v0.11.13rc3/truss-transfer-cli-v0.11.13rc3-linux-x86_64-unknown-linux-musl && \ + chmod +x /usr/local/bin/truss-transfer-cli + +COPY bptr_manifest_generator.py /usr/local/bin/bptr_manifest_generator.py +COPY bptr_init_script.sh /usr/local/bin/bptr_init_script.sh + +RUN chmod +x /usr/local/bin/bptr_init_script.sh + +ENTRYPOINT ["/usr/local/bin/bptr_init_script.sh"] diff --git a/bptr_generator/bptr_init_script.sh b/bptr_generator/bptr_init_script.sh new file mode 100755 index 000000000..74dc31c44 --- /dev/null +++ b/bptr_generator/bptr_init_script.sh @@ -0,0 +1,30 @@ +#!/bin/bash +set -e + +BPTR_DIR="${BPTR_DIR:-/cache/model/bptr}" +MODEL_CACHE_DIR="${MODEL_CACHE_DIR:-/cache/model/model_cache}" +APP_MODEL_CACHE="${APP_MODEL_CACHE:-/app/model_cache}" +MODEL_CACHE_CONFIG="${MODEL_CACHE_CONFIG:-/app/model_cache_config.json}" + +mkdir -p "$BPTR_DIR" "$MODEL_CACHE_DIR" +ln -sf "$BPTR_DIR" /static-bptr +ln -sf "$BPTR_DIR" /bptr + +python /usr/local/bin/bptr_manifest_generator.py \ + --model-cache "$MODEL_CACHE_CONFIG" \ + --output "$BPTR_DIR/static-bptr-manifest.json" \ + --model-path "$MODEL_CACHE_DIR" + +ln -sf "$BPTR_DIR/static-bptr-manifest.json" "$BPTR_DIR/bptr-manifest.json" +ln -sf "$BPTR_DIR/static-bptr-manifest.json" "$BPTR_DIR/bptr-manifest" + +if [ ! -f "$BPTR_DIR/static-bptr-manifest.json" ]; then + echo "ERROR: Manifest not found" >&2 + ls -la "$BPTR_DIR" || true + exit 1 +fi + +truss-transfer-cli + +mkdir -p "$APP_MODEL_CACHE" +ln -sf "$MODEL_CACHE_DIR"/* "$APP_MODEL_CACHE"/ 2>/dev/null || true diff --git a/bptr_generator/bptr_manifest_generator.py b/bptr_generator/bptr_manifest_generator.py new file mode 100644 index 000000000..2992fcb9d --- /dev/null +++ b/bptr_generator/bptr_manifest_generator.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +import argparse +import json +import logging +import sys +from pathlib import Path + +import yaml +from truss_transfer import PyModelRepo, create_basetenpointer_from_models + +logging.basicConfig(level=logging.INFO, format="%(message)s") + + +def load_truss_config(config_path: Path) -> dict: + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found: {config_path}") + with open(config_path, "r") as f: + config = yaml.safe_load(f) + if "model_cache" not in config: + raise ValueError("No 'model_cache' found in config.yaml") + return config + + +def load_model_cache_config(config_path: Path) -> list[dict]: + """Load model_cache config directly from JSON file""" + if not config_path.exists(): + raise FileNotFoundError(f"Model cache config file not found: {config_path}") + with open(config_path, "r") as f: + model_cache = json.load(f) + if not isinstance(model_cache, list): + raise ValueError("Model cache config must be a list of model configurations") + return model_cache + + +def generate_bptr_manifest_from_model_cache( + model_cache: list[dict], + output_path: Path, + model_path: str = "/cache/model/model_cache", +): + """Generate bptr-manifest directly from model_cache configuration""" + if not any(model.get("use_volume", False) for model in model_cache): + raise ValueError("No v2 models found (use_volume=True)") + if not all( + model.get("volume_folder") + for model in model_cache + if model.get("use_volume", False) + ): + raise ValueError("All v2 models must have volume_folder") + + py_models = [ + PyModelRepo( + repo_id=model["repo_id"], + revision=model.get("revision"), + runtime_secret_name=model.get("runtime_secret_name", "hf_access_token"), + allow_patterns=model.get("allow_patterns"), + ignore_patterns=model.get("ignore_patterns"), + volume_folder=model.get("volume_folder"), + kind=model.get("kind", "hf"), + ) + for model in model_cache + if model.get("use_volume", False) + ] + + basetenpointer_json = create_basetenpointer_from_models( + models=py_models, model_path=model_path + ) + bptr_py = json.loads(basetenpointer_json)["pointers"] + logging.info(f"created ({len(bptr_py)}) Basetenpointer") + logging.info(f"pointers json: {basetenpointer_json}") + + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + f.write(basetenpointer_json) + + +def generate_bptr_manifest( + config: dict, output_path: Path, model_path: str = "/cache/model/model_cache" +): + """Legacy function for backward compatibility""" + model_cache = config.get("model_cache", []) + generate_bptr_manifest_from_model_cache(model_cache, output_path, model_path) + + +def main(): + parser = argparse.ArgumentParser( + description="Generate bptr-manifest from truss config.yaml or model_cache config" + ) + parser.add_argument( + "--config", "-c", type=Path, help="Path to truss config.yaml file" + ) + parser.add_argument( + "--model-cache", "-m", type=Path, help="Path to model_cache JSON config file" + ) + parser.add_argument( + "--output", + "-o", + type=Path, + required=True, + help="Output path for bptr-manifest file", + ) + parser.add_argument( + "--model-path", + type=str, + default="/cache/model/model_cache", + help="Base path for model files in manifest", + ) + args = parser.parse_args() + + if not args.config and not args.model_cache: + parser.error("Either --config or --model-cache must be specified") + if args.config and args.model_cache: + parser.error("Cannot specify both --config and --model-cache") + + try: + if args.model_cache: + model_cache = load_model_cache_config(args.model_cache) + generate_bptr_manifest_from_model_cache( + model_cache, args.output, args.model_path + ) + else: + config = load_truss_config(args.config) + generate_bptr_manifest_from_model_cache( + config.get("model_cache", []), args.output, args.model_path + ) + except Exception as e: + logging.error(f"Error generating bptr-manifest: {e}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/bptr_generator/bptr_requirements.txt b/bptr_generator/bptr_requirements.txt new file mode 100644 index 000000000..762c34d48 --- /dev/null +++ b/bptr_generator/bptr_requirements.txt @@ -0,0 +1,3 @@ +# Requirements for bptr_manifest_generator standalone app +truss-transfer>=0.0.37,<0.0.40 +PyYAML>=6.0