llama-stack-mirror/scripts/distro_codegen.py

#!/usr/bin/env python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

import concurrent.futures
import importlib
import subprocess
import sys
from collections.abc import Iterable
from functools import partial
from pathlib import Path

import tomlkit
from rich.progress import Progress, SpinnerColumn, TextColumn

from llama_stack.distribution.build import (
    SERVER_DEPENDENCIES,
    get_provider_dependencies,
)

REPO_ROOT = Path(__file__).parent.parent


class ChangedPathTracker:
    """Track a list of paths we may have changed."""

    def __init__(self):
        self._changed_paths = []

    def add_paths(self, *paths):
        for path in paths:
            path = str(path)
            if path not in self._changed_paths:
                self._changed_paths.append(path)

    def changed_paths(self):
        return self._changed_paths


def find_template_dirs(templates_dir: Path) -> Iterable[Path]:
    """Find immediate subdirectories in the templates folder."""
    if not templates_dir.exists():
        raise FileNotFoundError(f"Templates directory not found: {templates_dir}")

    return sorted(d for d in templates_dir.iterdir() if d.is_dir() and d.name != "__pycache__")


def process_template(template_dir: Path, progress, change_tracker: ChangedPathTracker) -> None:
    """Process a single template directory."""
    progress.print(f"Processing {template_dir.name}")

    try:
        # Import the module directly
        module_name = f"llama_stack.templates.{template_dir.name}"
        module = importlib.import_module(module_name)

        # Get and save the distribution template
        if template_func := getattr(module, "get_distribution_template", None):
            template = template_func()

            yaml_output_dir = REPO_ROOT / "llama_stack" / "templates" / template.name
            doc_output_dir = REPO_ROOT / "docs/source/distributions" / f"{template.distro_type}_distro"
            change_tracker.add_paths(yaml_output_dir, doc_output_dir)
            template.save_distribution(
                yaml_output_dir=yaml_output_dir,
                doc_output_dir=doc_output_dir,
            )
        else:
            progress.print(f"[yellow]Warning: {template_dir.name} has no get_distribution_template function")

    except Exception as e:
        progress.print(f"[red]Error processing {template_dir.name}: {str(e)}")
        raise e


def check_for_changes(change_tracker: ChangedPathTracker) -> bool:
    """Check if there are any uncommitted changes."""
    has_changes = False
    for path in change_tracker.changed_paths():
        result = subprocess.run(
            ["git", "diff", "--exit-code", path],
            cwd=REPO_ROOT,
            capture_output=True,
        )
        if result.returncode != 0:
            print(f"Change detected in '{path}'.", file=sys.stderr)
            has_changes = True
    return has_changes


def collect_template_dependencies(template_dir: Path) -> tuple[str | None, 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)
            # Combine all dependencies in order: normal deps, special deps, server deps
            all_deps = sorted(set(normal_deps + SERVER_DEPENDENCIES)) + sorted(set(special_deps))

            return template.name, all_deps
    except Exception as e:
        print("Error collecting template dependencies for", template_dir, e)
        return None, []
    return None, []


def pre_import_templates(template_dirs: list[Path]) -> None:
    # Pre-import all template modules to avoid deadlocks.
    for template_dir in template_dirs:
        module_name = f"llama_stack.templates.{template_dir.name}"
        importlib.import_module(module_name)


def generate_dependencies_files(change_tracker: ChangedPathTracker):
    templates_dir = REPO_ROOT / "llama_stack" / "templates"
    distribution_deps = {}

    for template_dir in find_template_dirs(templates_dir):
        print("template_dir", template_dir)
        name, deps = collect_template_dependencies(template_dir)
        if name:
            distribution_deps[name] = deps
        else:
            print("No template function found for", template_dir)

    # First, remove any distributions that are no longer present
    pyproject_file = REPO_ROOT / "pyproject.toml"
    change_tracker.add_paths(pyproject_file)

    # Read and parse the current pyproject.toml content
    with open(pyproject_file) as fp:
        pyproject = tomlkit.load(fp)

    # Get current optional dependencies
    current_deps = pyproject["project"]["optional-dependencies"]

    # Store ui dependencies if they exist
    ui_deps = current_deps.get("ui")

    # Remove distributions that are no longer present
    for name in list(current_deps.keys()):
        if name not in distribution_deps.keys() and name != "ui":
            del current_deps[name]

    # Now add/update the remaining distributions
    for name, deps in distribution_deps.items():
        deps_array = tomlkit.array()
        for dep in sorted(deps):
            deps_array.append(dep)
        current_deps[name] = deps_array.multiline(True)

    # Restore ui dependencies if they existed
    if ui_deps is not None:
        current_deps["ui"] = ui_deps

    # Write back to pyproject.toml
    with open(pyproject_file, "w") as fp:
        tomlkit.dump(pyproject, fp)


def main():
    templates_dir = REPO_ROOT / "llama_stack" / "templates"
    change_tracker = ChangedPathTracker()

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
    ) as progress:
        template_dirs = list(find_template_dirs(templates_dir))
        task = progress.add_task("Processing distribution templates...", total=len(template_dirs))

        pre_import_templates(template_dirs)

        # Create a partial function with the progress bar
        process_func = partial(process_template, progress=progress, change_tracker=change_tracker)

        # Process templates in parallel
        with concurrent.futures.ThreadPoolExecutor() as executor:
            # Submit all tasks and wait for completion
            list(executor.map(process_func, template_dirs))
            progress.update(task, advance=len(template_dirs))

    # TODO: generate a Containerfile for each distribution as well?
    generate_dependencies_files(change_tracker)

    if check_for_changes(change_tracker):
        print(
            "Distribution template changes detected. Please commit the changes.",
            file=sys.stderr,
        )
        sys.exit(1)

    sys.exit(0)


if __name__ == "__main__":
    main()