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

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

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

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