codegen per-distro dependencies; not hooked into setup.py yet

This commit is contained in:
Ashwin Bharambe 2024-11-19 09:54:30 -08:00
parent 5e4ac1b7c1
commit 1619d37cc6
3 changed files with 216 additions and 0 deletions

View file

@ -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

View file

@ -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"
]
}

View file

@ -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.",