refactor: consolidate dynamic provider config parsing (#4985)

Author: mattf

src/llama_stack/cli/stack/_list_deps.py

@@ -12,10 +12,9 @@
 from termcolor import cprint
 
 from llama_stack.core.build import get_provider_dependencies
-from llama_stack.core.datatypes import Provider, StackConfig
-from llama_stack.core.distribution import get_provider_registry
+from llama_stack.core.datatypes import StackConfig
+from llama_stack.core.stack import run_config_from_dynamic_config_spec
 from llama_stack.log import get_logger
-from llama_stack_api import Api
 
 TEMPLATES_PATH = Path(__file__).parent.parent.parent / "templates"
 
@@ -91,39 +90,11 @@ def run_stack_list_deps_command(args: argparse.Namespace) -> None:
                     )
                     sys.exit(1)
     elif args.providers:
-        provider_list: dict[str, list[Provider]] = dict()
-        for api_provider in args.providers.split(","):
-            if "=" not in api_provider:
-                cprint(
-                    "Could not parse `--providers`. Please ensure the list is in the format api1=provider1,api2=provider2",
-                    color="red",
-                    file=sys.stderr,
-                )
-                sys.exit(1)
-            api, provider_type = api_provider.split("=")
-            providers_for_api = get_provider_registry().get(Api(api), None)
-            if providers_for_api is None:
-                cprint(
-                    f"{api} is not a valid API.",
-                    color="red",
-                    file=sys.stderr,
-                )
-                sys.exit(1)
-            if provider_type in providers_for_api:
-                provider = Provider(
-                    provider_type=provider_type,
-                    provider_id=provider_type.split("::")[1],
-                    module=None,
-                )
-                provider_list.setdefault(api, []).append(provider)
-            else:
-                cprint(
-                    f"{provider_type} is not a valid provider for the {api} API.",
-                    color="red",
-                    file=sys.stderr,
-                )
-                sys.exit(1)
-        config = StackConfig(providers=provider_list, distro_name="providers-run")
+        try:
+            config = run_config_from_dynamic_config_spec(args.providers)
+        except ValueError as e:
+            cprint(str(e), color="red", file=sys.stderr)
+            sys.exit(1)
 
     normal_deps, special_deps, external_provider_dependencies = get_provider_dependencies(config)
     normal_deps += SERVER_DEPENDENCIES

src/llama_stack/cli/stack/run.py

@@ -16,21 +16,10 @@
 from termcolor import cprint
 
 from llama_stack.cli.subcommand import Subcommand
-from llama_stack.core.datatypes import Api, Provider, StackConfig
-from llama_stack.core.distribution import get_provider_registry
-from llama_stack.core.stack import cast_distro_name_to_string, replace_env_vars
-from llama_stack.core.storage.datatypes import (
-    InferenceStoreReference,
-    KVStoreReference,
-    ServerStoresConfig,
-    SqliteKVStoreConfig,
-    SqliteSqlStoreConfig,
-    SqlStoreReference,
-    StorageConfig,
-)
+from llama_stack.core.datatypes import StackConfig
+from llama_stack.core.stack import cast_distro_name_to_string, replace_env_vars, run_config_from_dynamic_config_spec
 from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
 from llama_stack.core.utils.config_resolution import resolve_config_or_distro
-from llama_stack.core.utils.dynamic import instantiate_class_type
 from llama_stack.log import LoggingConfig, get_logger
 
 REPO_ROOT = Path(__file__).parent.parent.parent.parent
@@ -92,50 +81,20 @@ def _run_stack_run_cmd(self, args: argparse.Namespace) -> None:
             except ValueError as e:
                 self.parser.error(str(e))
         elif args.providers:
-            provider_list: dict[str, list[Provider]] = dict()
-            for api_provider in args.providers.split(","):
-                if "=" not in api_provider:
-                    cprint(
-                        "Could not parse `--providers`. Please ensure the list is in the format api1=provider1,api2=provider2",
-                        color="red",
-                        file=sys.stderr,
-                    )
-                    sys.exit(1)
-                api, provider_type = api_provider.split("=")
-                providers_for_api = get_provider_registry().get(Api(api), None)
-                if providers_for_api is None:
-                    cprint(
-                        f"{api} is not a valid API.",
-                        color="red",
-                        file=sys.stderr,
-                    )
-                    sys.exit(1)
-                if provider_type in providers_for_api:
-                    config_type = instantiate_class_type(providers_for_api[provider_type].config_class)
-                    if config_type is not None and hasattr(config_type, "sample_run_config"):
-                        config = config_type.sample_run_config(__distro_dir__="~/.llama/distributions/providers-run")
-                    else:
-                        config = {}
-                    provider = Provider(
-                        provider_type=provider_type,
-                        config=config,
-                        provider_id=provider_type.split("::")[1],
-                    )
-                    provider_list.setdefault(api, []).append(provider)
-                else:
-                    cprint(
-                        f"{provider} is not a valid provider for the {api} API.",
-                        color="red",
-                        file=sys.stderr,
-                    )
-                    sys.exit(1)
-            run_config = self._generate_run_config_from_providers(providers=provider_list)
+            distro_dir = DISTRIBS_BASE_DIR / "providers-run"
+            os.makedirs(distro_dir, exist_ok=True)
+            try:
+                run_config = run_config_from_dynamic_config_spec(
+                    dynamic_config_spec=args.providers,
+                    distro_dir=distro_dir,
+                    distro_name="providers-run",
+                )
+            except ValueError as e:
+                cprint(str(e), color="red", file=sys.stderr)
+                sys.exit(1)
             config_dict = run_config.model_dump(mode="json")
 
-            # Write config to disk in providers-run directory
-            distro_dir = DISTRIBS_BASE_DIR / "providers-run"
             config_file = distro_dir / "config.yaml"
-
             logger.info(f"Writing generated config to: {config_file}")
             with open(config_file, "w") as f:
                 yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False)
@@ -261,44 +220,3 @@ def _start_ui_development_server(self, stack_server_port: int):
             )
         except Exception as e:
             logger.error(f"Failed to start UI development server in {ui_dir}: {e}")
-
-    def _generate_run_config_from_providers(self, providers: dict[str, list[Provider]]):
-        apis = list(providers.keys())
-        distro_dir = DISTRIBS_BASE_DIR / "providers-run"
-        # need somewhere to put the storage.
-        os.makedirs(distro_dir, exist_ok=True)
-        storage = StorageConfig(
-            backends={
-                "kv_default": SqliteKVStoreConfig(
-                    db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/kvstore.db",
-                ),
-                "sql_default": SqliteSqlStoreConfig(
-                    db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/sql_store.db",
-                ),
-            },
-            stores=ServerStoresConfig(
-                metadata=KVStoreReference(
-                    backend="kv_default",
-                    namespace="registry",
-                ),
-                inference=InferenceStoreReference(
-                    backend="sql_default",
-                    table_name="inference_store",
-                ),
-                conversations=SqlStoreReference(
-                    backend="sql_default",
-                    table_name="openai_conversations",
-                ),
-                prompts=KVStoreReference(
-                    backend="kv_default",
-                    namespace="prompts",
-                ),
-            ),
-        )
-
-        return StackConfig(
-            distro_name="providers-run",
-            apis=apis,
-            providers=providers,
-            storage=storage,
-        )

src/llama_stack/core/stack.py

@@ -10,6 +10,7 @@
 import os
 import re
 import tempfile
+from pathlib import Path
 from typing import Any, get_type_hints
 
 import yaml
@@ -728,26 +729,38 @@ def get_stack_run_config_from_distro(distro: str) -> StackConfig:
     return StackConfig(**replace_env_vars(run_config))
 
 
-def run_config_from_adhoc_config_spec(
-    adhoc_config_spec: str, provider_registry: ProviderRegistry | None = None
+def run_config_from_dynamic_config_spec(
+    dynamic_config_spec: str,
+    provider_registry: ProviderRegistry | None = None,
+    distro_dir: Path | None = None,
+    distro_name: str = "dynamic-distro",
 ) -> StackConfig:
     """
-    Create an adhoc distribution from a list of API providers.
+    Create a dynamic distribution from a list of API providers.
 
     The list should be of the form "api=provider", e.g. "inference=fireworks". If you have
     multiple pairs, separate them with commas or semicolons, e.g. "inference=fireworks,safety=llama-guard,agents=meta-reference"
     """
 
-    api_providers = adhoc_config_spec.replace(";", ",").split(",")
-    provider_registry = provider_registry or get_provider_registry()
+    api_providers = dynamic_config_spec.replace(";", ",").split(",")
+    provider_registry = get_provider_registry() if provider_registry is None else provider_registry
 
-    distro_dir = tempfile.mkdtemp()
+    distro_dir = distro_dir or Path(tempfile.mkdtemp())
     provider_configs_by_api = {}
     for api_provider in api_providers:
+        if "=" not in api_provider:
+            raise ValueError(
+                f"Failed to parse provider spec '{api_provider}'. Expected format: api=provider (e.g. inference=fireworks)"
+            )
         api_str, provider = api_provider.split("=")
-        api = Api(api_str)
+        try:
+            api = Api(api_str)
+        except ValueError:
+            raise ValueError(f"Failed to parse provider spec: '{api_str}' is not a valid API") from None
 
-        providers_by_type = provider_registry[api]
+        providers_by_type = provider_registry.get(api)
+        if providers_by_type is None:
+            raise ValueError(f"Failed to find providers for API '{api_str}'")
         provider_spec = providers_by_type.get(provider)
         if not provider_spec:
             provider_spec = providers_by_type.get(f"inline::{provider}")
@@ -761,23 +774,27 @@ def run_config_from_adhoc_config_spec(
 
         # call method "sample_run_config" on the provider spec config class
         provider_config_type = instantiate_class_type(provider_spec.config_class)
-        provider_config = replace_env_vars(provider_config_type.sample_run_config(__distro_dir__=distro_dir))
+        provider_config = replace_env_vars(provider_config_type.sample_run_config(__distro_dir__=str(distro_dir)))
 
-        provider_configs_by_api[api_str] = [
+        provider_configs_by_api.setdefault(api_str, []).append(
             Provider(
                 provider_id=provider_spec.provider_type.split("::")[-1],
                 provider_type=provider_spec.provider_type,
                 config=provider_config,
             )
-        ]
+        )
     config = StackConfig(
-        distro_name="distro-test",
+        distro_name=distro_name,
         apis=list(provider_configs_by_api.keys()),
         providers=provider_configs_by_api,
         storage=StorageConfig(
             backends={
-                "kv_default": SqliteKVStoreConfig(db_path=f"{distro_dir}/kvstore.db"),
-                "sql_default": SqliteSqlStoreConfig(db_path=f"{distro_dir}/sql_store.db"),
+                "kv_default": SqliteKVStoreConfig(
+                    db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/kvstore.db",
+                ),
+                "sql_default": SqliteSqlStoreConfig(
+                    db_path=f"${{env.SQLITE_STORE_DIR:={distro_dir}}}/sql_store.db",
+                ),
             },
             stores=ServerStoresConfig(
                 metadata=KVStoreReference(backend="kv_default", namespace="registry"),

tests/integration/conftest.py

@@ -15,7 +15,7 @@
 import pytest
 from dotenv import load_dotenv
 
-from llama_stack.core.stack import run_config_from_adhoc_config_spec
+from llama_stack.core.stack import run_config_from_dynamic_config_spec
 from llama_stack.log import get_logger
 from llama_stack.testing.api_recorder import patch_httpx_for_test_id
 
@@ -148,7 +148,7 @@ def pytest_configure(config):
     if getattr(config.option, "embedding_model", None) is None:
         stack_config = config.getoption("--stack-config", default=None)
         if stack_config and "=" in stack_config:
-            run_config = run_config_from_adhoc_config_spec(stack_config)
+            run_config = run_config_from_dynamic_config_spec(stack_config)
             inference_providers = run_config.providers.get("inference", [])
             if any("sentence-transformers" in p.provider_type for p in inference_providers):
                 config.option.embedding_model = "sentence-transformers/nomic-ai/nomic-embed-text-v1.5"
@@ -162,7 +162,7 @@ def pytest_addoption(parser):
             a 'pointer' to the stack. this can be either be:
             (a) a template name like `starter`, or
             (b) a path to a config.yaml file, or
-            (c) an adhoc config spec, e.g. `inference=fireworks,safety=llama-guard,agents=meta-reference`, or
+            (c) a dynamic config spec, e.g. `inference=fireworks,safety=llama-guard,agents=meta-reference`, or
             (d) a server config like `server:ci-tests`, or
             (e) a docker config like `docker:ci-tests` (builds and runs container)
             """
@@ -256,7 +256,7 @@ def pytest_generate_tests(metafunc):
         config_str = metafunc.config.getoption("--stack-config", default=None) or os.environ.get("LLAMA_STACK_CONFIG")
         providers = None
         if config_str and "=" in config_str:
-            run_config = run_config_from_adhoc_config_spec(config_str)
+            run_config = run_config_from_dynamic_config_spec(config_str)
             providers = [p.provider_id for p in run_config.providers.get("vector_io", [])]
         if providers is None:
             inference_mode = os.environ.get("LLAMA_STACK_TEST_INFERENCE_MODE")

tests/integration/fixtures/common.py

@@ -27,7 +27,7 @@
 
 from llama_stack.core.datatypes import QualifiedModel, VectorStoresConfig
 from llama_stack.core.library_client import LlamaStackAsLibraryClient
-from llama_stack.core.stack import run_config_from_adhoc_config_spec
+from llama_stack.core.stack import run_config_from_dynamic_config_spec
 from llama_stack.core.utils.config_resolution import resolve_config_or_distro
 from llama_stack.env import get_env_or_fail
 
@@ -321,7 +321,7 @@ def instantiate_llama_stack_client(session):
         pass
 
     if "=" in config:
-        run_config = run_config_from_adhoc_config_spec(config)
+        run_config = run_config_from_dynamic_config_spec(config)
 
         # --stack-config bypasses template so need this to set default embedding model
         if "vector_io" in config and "inference" in config:

tests/unit/cli/test_stack_config.py

@@ -229,26 +229,13 @@ def test_parse_and_maybe_upgrade_config_preserves_custom_external_providers_dir(
 
 
 def test_generate_run_config_from_providers():
-    """Test that _generate_run_config_from_providers creates a valid config"""
-    import argparse
-
-    from llama_stack.cli.stack.run import StackRun
-    from llama_stack.core.datatypes import Provider
+    """Test that run_config_from_dynamic_config_spec creates a valid config for the providers-run distro"""
+    from llama_stack.core.stack import run_config_from_dynamic_config_spec
 
-    parser = argparse.ArgumentParser()
-    subparsers = parser.add_subparsers()
-    stack_run = StackRun(subparsers)
-
-    providers = {
-        "inference": [
-            Provider(
-                provider_type="inline::meta-reference",
-                provider_id="meta-reference",
-            )
-        ]
-    }
-
-    config = stack_run._generate_run_config_from_providers(providers=providers)
+    config = run_config_from_dynamic_config_spec(
+        "inference=remote::openai",
+        distro_name="providers-run",
+    )
     config_dict = config.model_dump(mode="json")
 
     # Verify basic structure