diff --git a/src/llama_stack/cli/stack/list_stacks.py b/src/llama_stack/cli/stack/list_stacks.py index 2ea0fdeea..ae59ba911 100644 --- a/src/llama_stack/cli/stack/list_stacks.py +++ b/src/llama_stack/cli/stack/list_stacks.py @@ -9,48 +9,69 @@ from pathlib import Path from llama_stack.cli.subcommand import Subcommand from llama_stack.cli.table import print_table +from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR class StackListBuilds(Subcommand): - """List built stacks in .llama/distributions directory""" + """List available distributions (both built-in and custom)""" def __init__(self, subparsers: argparse._SubParsersAction): super().__init__() self.parser = subparsers.add_parser( "list", prog="llama stack list", - description="list the build stacks", + description="list available distributions", formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) self._add_arguments() self.parser.set_defaults(func=self._list_stack_command) - def _get_distribution_dirs(self) -> dict[str, Path]: - """Return a dictionary of distribution names and their paths""" - distributions = {} - dist_dir = Path.home() / ".llama" / "distributions" + def _get_distribution_dirs(self) -> dict[str, tuple[Path, str]]: + """Return a dictionary of distribution names and their paths with source type + + Returns: + dict mapping distro name to (path, source_type) where source_type is 'built-in' or 'custom' + """ + distributions = {} + + # Get built-in distributions from source code + distro_dir = Path(__file__).parent.parent.parent / "distributions" + if distro_dir.exists(): + for stack_dir in distro_dir.iterdir(): + if stack_dir.is_dir() and not stack_dir.name.startswith(".") and not stack_dir.name.startswith("__"): + distributions[stack_dir.name] = (stack_dir, "built-in") + + # Get custom/run distributions from ~/.llama/distributions + # These override built-in ones if they have the same name + if DISTRIBS_BASE_DIR.exists(): + for stack_dir in DISTRIBS_BASE_DIR.iterdir(): + if stack_dir.is_dir() and not stack_dir.name.startswith("."): + # Clean up the name (remove llamastack- prefix if present) + name = stack_dir.name.replace("llamastack-", "") + distributions[name] = (stack_dir, "custom") - if dist_dir.exists(): - for stack_dir in dist_dir.iterdir(): - if stack_dir.is_dir(): - distributions[stack_dir.name] = stack_dir return distributions def _list_stack_command(self, args: argparse.Namespace) -> None: distributions = self._get_distribution_dirs() if not distributions: - print("No stacks found in ~/.llama/distributions") + print("No distributions found") return - headers = ["Stack Name", "Path"] - headers.extend(["Build Config", "Run Config"]) + headers = ["Stack Name", "Source", "Path", "Build Config", "Run Config"] rows = [] - for name, path in distributions.items(): - row = [name, str(path)] + for name, (path, source_type) in sorted(distributions.items()): + row = [name, source_type, str(path)] # Check for build and run config files - build_config = "Yes" if (path / f"{name}-build.yaml").exists() else "No" - run_config = "Yes" if (path / f"{name}-run.yaml").exists() else "No" + # For built-in distributions, configs are named build.yaml and run.yaml + # For custom distributions, configs are named {name}-build.yaml and {name}-run.yaml + if source_type == "built-in": + build_config = "Yes" if (path / "build.yaml").exists() else "No" + run_config = "Yes" if (path / "run.yaml").exists() else "No" + else: + build_config = "Yes" if (path / f"{name}-build.yaml").exists() else "No" + run_config = "Yes" if (path / f"{name}-run.yaml").exists() else "No" row.extend([build_config, run_config]) rows.append(row) print_table(rows, headers, separate_rows=True) diff --git a/tests/unit/distribution/test_stack_list.py b/tests/unit/distribution/test_stack_list.py new file mode 100644 index 000000000..725ce3410 --- /dev/null +++ b/tests/unit/distribution/test_stack_list.py @@ -0,0 +1,130 @@ +# 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. + +"""Tests for the llama stack list command.""" + +import argparse +from unittest.mock import MagicMock, patch + +import pytest + +from llama_stack.cli.stack.list_stacks import StackListBuilds + + +@pytest.fixture +def list_stacks_command(): + """Create a StackListBuilds instance for testing.""" + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + return StackListBuilds(subparsers) + + +@pytest.fixture +def mock_distribs_base_dir(tmp_path): + """Create a mock DISTRIBS_BASE_DIR with some custom distributions.""" + custom_dir = tmp_path / "distributions" + custom_dir.mkdir(parents=True, exist_ok=True) + + # Create a custom distribution + starter_custom = custom_dir / "starter" + starter_custom.mkdir() + (starter_custom / "starter-build.yaml").write_text("# build config") + (starter_custom / "starter-run.yaml").write_text("# run config") + + return custom_dir + + +@pytest.fixture +def mock_distro_dir(tmp_path): + """Create a mock distributions directory with built-in distributions.""" + distro_dir = tmp_path / "src" / "llama_stack" / "distributions" + distro_dir.mkdir(parents=True, exist_ok=True) + + # Create some built-in distributions + for distro_name in ["starter", "nvidia", "dell"]: + distro_path = distro_dir / distro_name + distro_path.mkdir() + (distro_path / "build.yaml").write_text("# build config") + (distro_path / "run.yaml").write_text("# run config") + + return distro_dir + + +def create_path_mock(builtin_dist_dir): + """Create a properly mocked Path object that returns builtin_dist_dir for the distributions path.""" + mock_parent_parent_parent = MagicMock() + mock_parent_parent_parent.__truediv__ = ( + lambda self, other: builtin_dist_dir if other == "distributions" else MagicMock() + ) + + mock_path = MagicMock() + mock_path.parent.parent.parent = mock_parent_parent_parent + + return mock_path + + +class TestStackList: + """Test suite for llama stack list command.""" + + def test_builtin_distros_shown_without_running(self, list_stacks_command, mock_distro_dir, tmp_path): + """Test that built-in distributions are shown even before running them.""" + mock_path = create_path_mock(mock_distro_dir) + + # Mock DISTRIBS_BASE_DIR to be a non-existent directory (no custom distributions) + with patch("llama_stack.cli.stack.list_stacks.DISTRIBS_BASE_DIR", tmp_path / "nonexistent"): + with patch("llama_stack.cli.stack.list_stacks.Path") as mock_path_class: + mock_path_class.return_value = mock_path + + distributions = list_stacks_command._get_distribution_dirs() + + # Verify built-in distributions are found + assert len(distributions) > 0, "Should find built-in distributions" + assert all(source_type == "built-in" for _, source_type in distributions.values()), ( + "All should be built-in" + ) + + # Check specific distributions we created + assert "starter" in distributions + assert "nvidia" in distributions + assert "dell" in distributions + + def test_custom_distribution_overrides_builtin(self, list_stacks_command, mock_distro_dir, mock_distribs_base_dir): + """Test that custom distributions override built-in ones with the same name.""" + mock_path = create_path_mock(mock_distro_dir) + + with patch("llama_stack.cli.stack.list_stacks.DISTRIBS_BASE_DIR", mock_distribs_base_dir): + with patch("llama_stack.cli.stack.list_stacks.Path") as mock_path_class: + mock_path_class.return_value = mock_path + + distributions = list_stacks_command._get_distribution_dirs() + + # "starter" should exist and be marked as "custom" (not "built-in") + # because the custom version overrides the built-in one + assert "starter" in distributions + _, source_type = distributions["starter"] + assert source_type == "custom", "Custom distribution should override built-in" + + def test_hidden_directories_ignored(self, list_stacks_command, mock_distro_dir, tmp_path): + """Test that hidden directories (starting with .) are ignored.""" + # Add a hidden directory + hidden_dir = mock_distro_dir / ".hidden" + hidden_dir.mkdir() + (hidden_dir / "build.yaml").write_text("# build") + + # Add a __pycache__ directory + pycache_dir = mock_distro_dir / "__pycache__" + pycache_dir.mkdir() + + mock_path = create_path_mock(mock_distro_dir) + + with patch("llama_stack.cli.stack.list_stacks.DISTRIBS_BASE_DIR", tmp_path / "nonexistent"): + with patch("llama_stack.cli.stack.list_stacks.Path") as mock_path_class: + mock_path_class.return_value = mock_path + + distributions = list_stacks_command._get_distribution_dirs() + + assert ".hidden" not in distributions + assert "__pycache__" not in distributions