From 1619d37cc653cb1d9cbddcbc5627cd818b11b3e6 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Tue, 19 Nov 2024 09:54:30 -0800 Subject: [PATCH] codegen per-distro dependencies; not hooked into setup.py yet --- MANIFEST.in | 1 + distributions/dependencies.json | 177 ++++++++++++++++++++++++++ llama_stack/scripts/distro_codegen.py | 38 ++++++ 3 files changed, 216 insertions(+) create mode 100644 distributions/dependencies.json diff --git a/MANIFEST.in b/MANIFEST.in index 27cb775f7..4d1843051 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,5 @@ include requirements.txt +include distributions/dependencies.json include llama_stack/distribution/*.sh include llama_stack/cli/scripts/*.sh include llama_stack/templates/*/*.yaml diff --git a/distributions/dependencies.json b/distributions/dependencies.json new file mode 100644 index 000000000..6827af1f1 --- /dev/null +++ b/distributions/dependencies.json @@ -0,0 +1,177 @@ +{ + "together": [ + "scipy", + "scikit-learn", + "nltk", + "chardet", + "chromadb-client", + "psycopg2-binary", + "sentencepiece", + "faiss-cpu", + "blobfile", + "pandas", + "pillow", + "together", + "pypdf", + "matplotlib", + "aiosqlite", + "redis", + "transformers", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ], + "remote-vllm": [ + "scipy", + "scikit-learn", + "nltk", + "chardet", + "chromadb-client", + "psycopg2-binary", + "sentencepiece", + "faiss-cpu", + "blobfile", + "pandas", + "pillow", + "pypdf", + "matplotlib", + "openai", + "aiosqlite", + "redis", + "transformers", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ], + "fireworks": [ + "scipy", + "scikit-learn", + "nltk", + "chardet", + "chromadb-client", + "psycopg2-binary", + "sentencepiece", + "faiss-cpu", + "blobfile", + "pandas", + "pillow", + "pypdf", + "matplotlib", + "aiosqlite", + "redis", + "transformers", + "fireworks-ai", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ], + "tgi": [ + "scipy", + "scikit-learn", + "nltk", + "aiohttp", + "chardet", + "chromadb-client", + "psycopg2-binary", + "huggingface_hub", + "sentencepiece", + "faiss-cpu", + "blobfile", + "pandas", + "pillow", + "pypdf", + "matplotlib", + "aiosqlite", + "transformers", + "redis", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ], + "meta-reference-gpu": [ + "lm-format-enforcer", + "scipy", + "scikit-learn", + "nltk", + "accelerate", + "chardet", + "chromadb-client", + "psycopg2-binary", + "sentencepiece", + "zmq", + "faiss-cpu", + "torchvision", + "blobfile", + "fairscale", + "pandas", + "pillow", + "pypdf", + "matplotlib", + "transformers", + "torch", + "aiosqlite", + "redis", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ], + "ollama": [ + "scipy", + "scikit-learn", + "nltk", + "aiohttp", + "ollama", + "chardet", + "chromadb-client", + "psycopg2-binary", + "sentencepiece", + "faiss-cpu", + "blobfile", + "pandas", + "pillow", + "pypdf", + "matplotlib", + "aiosqlite", + "transformers", + "redis", + "numpy", + "tqdm", + "sentence-transformers --no-deps", + "torch --index-url https://download.pytorch.org/whl/cpu", + "aiosqlite", + "fastapi", + "fire", + "httpx", + "uvicorn" + ] +} diff --git a/llama_stack/scripts/distro_codegen.py b/llama_stack/scripts/distro_codegen.py index f0d3bb4b9..8bcf97374 100644 --- a/llama_stack/scripts/distro_codegen.py +++ b/llama_stack/scripts/distro_codegen.py @@ -6,6 +6,7 @@ import concurrent.futures import importlib +import json import subprocess import sys from functools import partial @@ -14,6 +15,11 @@ from typing import Iterator from rich.progress import Progress, SpinnerColumn, TextColumn +from llama_stack.distribution.build import ( + get_provider_dependencies, + SERVER_DEPENDENCIES, +) + REPO_ROOT = Path(__file__).parent.parent.parent @@ -67,6 +73,36 @@ def check_for_changes() -> bool: return result.returncode != 0 +def collect_template_dependencies(template_dir: Path) -> tuple[str, list[str]]: + try: + module_name = f"llama_stack.templates.{template_dir.name}" + module = importlib.import_module(module_name) + + if template_func := getattr(module, "get_distribution_template", None): + template = template_func() + normal_deps, special_deps = get_provider_dependencies(template.providers) + # Combine all dependencies in order: normal deps, special deps, server deps + all_deps = normal_deps + special_deps + SERVER_DEPENDENCIES + return template.name, all_deps + except Exception: + return None, [] + return None, [] + + +def generate_dependencies_file(): + templates_dir = REPO_ROOT / "llama_stack" / "templates" + distribution_deps = {} + + for template_dir in find_template_dirs(templates_dir): + name, deps = collect_template_dependencies(template_dir) + if name: + distribution_deps[name] = deps + + deps_file = REPO_ROOT / "distributions" / "dependencies.json" + with open(deps_file, "w") as f: + json.dump(distribution_deps, f, indent=2) + + def main(): templates_dir = REPO_ROOT / "llama_stack" / "templates" @@ -88,6 +124,8 @@ def main(): list(executor.map(process_func, template_dirs)) progress.update(task, advance=len(template_dirs)) + generate_dependencies_file() + if check_for_changes(): print( "Distribution template changes detected. Please commit the changes.",