chore: consolidate scripts under ./scripts directory (#1646)

scripts/distro_codegen.py

@@ -0,0 +1,159 @@
+#!/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 json
+import subprocess
+import sys
+from functools import partial
+from pathlib import Path
+from typing import Iterable
+
+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.providers)
+            # 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:
+        return None, []
+    return None, []
+
+
+def generate_dependencies_file(change_tracker: ChangedPathTracker):
+    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"
+    change_tracker.add_paths(deps_file)
+    with open(deps_file, "w") as f:
+        f.write(json.dumps(distribution_deps, indent=2) + "\n")
+
+
+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))
+
+        # 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))
+
+    generate_dependencies_file(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()

scripts/gen-changelog.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #

scripts/generate_prompt_format.py

@@ -0,0 +1,73 @@
+#!/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.
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# top-level folder for each specific model found within the models/ directory at
+# the top-level of this source tree.
+
+import importlib
+from pathlib import Path
+
+import fire
+
+from llama_stack.models.llama.sku_list import resolve_model
+from llama_stack.providers.inline.inference.meta_reference.config import MetaReferenceInferenceConfig
+from llama_stack.providers.inline.inference.meta_reference.llama3.generation import Llama3
+
+THIS_DIR = Path(__file__).parent.resolve()
+
+
+def run_main(
+    model_id: str,
+    checkpoint_dir: str,
+    module_name: str,
+    output_path: str,
+):
+    module = importlib.import_module(module_name)
+    assert hasattr(module, "usecases"), f"Module {module_name} missing usecases function"
+
+    config = MetaReferenceInferenceConfig(
+        model=model_id,
+        max_seq_len=512,
+        max_batch_size=1,
+        checkpoint_dir=checkpoint_dir,
+    )
+    llama_model = resolve_model(model_id)
+    if not llama_model:
+        raise ValueError(f"Model {model_id} not found")
+    generator = Llama3.build(
+        config=config,
+        model_id=model_id,
+        llama_model=llama_model,
+    )
+
+    use_cases = module.usecases()
+    text = ""
+    for u in use_cases:
+        if isinstance(u, str):
+            use_case_text = f"\n{u}\n"
+        else:
+            use_case_text = u.to_text(generator)
+
+        text += use_case_text
+        print(use_case_text)
+
+    text += "Thank You!\n"
+
+    with open(output_path, "w") as f:
+        f.write(text)
+
+
+def main():
+    fire.Fire(run_main)
+
+
+if __name__ == "__main__":
+    main()

scripts/run_client_sdk_tests.py

@@ -0,0 +1,66 @@
+#!/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 argparse
+import os
+from pathlib import Path
+
+import pytest
+
+"""
+Script for running api on AsyncLlamaStackAsLibraryClient with templates
+
+Assuming directory structure:
+- llama-stack
+    - scripts
+    - tests
+        - api
+
+Example command:
+
+cd llama-stack
+EXPORT TOGETHER_API_KEY=<..>
+EXPORT FIREWORKS_API_KEY=<..>
+./scripts/run_client_sdk_tests.py --templates together fireworks --report
+"""
+
+REPO_ROOT = Path(__file__).parent.parent
+CLIENT_SDK_TESTS_RELATIVE_PATH = "tests/api/"
+
+
+def main(parser: argparse.ArgumentParser):
+    args = parser.parse_args()
+    templates_dir = REPO_ROOT / "llama_stack" / "templates"
+    user_specified_templates = [templates_dir / t for t in args.templates] if args.templates else []
+    for d in templates_dir.iterdir():
+        if d.is_dir() and d.name != "__pycache__":
+            template_configs = list(d.rglob("run.yaml"))
+            if len(template_configs) == 0:
+                continue
+            config = template_configs[0]
+            if user_specified_templates:
+                if not any(config.parent == t for t in user_specified_templates):
+                    continue
+            os.environ["LLAMA_STACK_CONFIG"] = str(config)
+            pytest_args = "--report" if args.report else ""
+            pytest.main(
+                [
+                    pytest_args,
+                    "-s",
+                    "-v",
+                    str(REPO_ROOT / CLIENT_SDK_TESTS_RELATIVE_PATH),
+                ]
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="llama_test",
+    )
+    parser.add_argument("--templates", nargs="+")
+    parser.add_argument("--report", action="store_true")
+    main(parser)