diff --git a/src/llama_stack_configuration.py b/src/llama_stack_configuration.py index 1a563bc33..c704800a3 100644 --- a/src/llama_stack_configuration.py +++ b/src/llama_stack_configuration.py @@ -998,6 +998,125 @@ def deep_merge_list_replace( return result +def _replace_or_append_inference_provider( + inference_list: list[Any], entry: dict[str, Any] +) -> None: + """Replace an inference entry with the same provider_id, else append. + + Parameters: + inference_list: Mutable providers.inference list. + entry: New provider entry to install. + """ + provider_id = entry["provider_id"] + for index, existing in enumerate(inference_list): + if isinstance(existing, dict) and existing.get("provider_id") == provider_id: + logger.info( + "Replacing existing inference provider with " + "provider_id=%r; a later high-level entry overwrote it", + provider_id, + ) + inference_list[index] = entry + return + inference_list.append(entry) + + +def _build_inference_entry( + provider: dict[str, Any], emitted_id: str, ls_provider_type: str +) -> tuple[dict[str, Any], list[str]]: + """Build a providers.inference entry from one high-level provider. + + Parameters: + provider: One high-level ``inference.providers`` entry. + emitted_id: The provider_id to emit (explicit id or hyphenated type). + ls_provider_type: Llama Stack provider_type from :data:`PROVIDER_TYPE_MAP`. + + Returns: + tuple[dict[str, Any], list[str]]: The provider entry, and its + ``allowed_models`` (empty list when unset). + """ + entry: dict[str, Any] = { + "provider_id": emitted_id, + "provider_type": ls_provider_type, + } + + provider_config: dict[str, Any] = {} + if provider.get("extra"): + provider_config.update(provider["extra"]) + if provider.get("api_key_env"): + key_field = API_KEY_FIELD_MAP.get(ls_provider_type, "api_key") + provider_config[key_field] = "${env." + provider["api_key_env"] + "}" + allowed_models = provider.get("allowed_models") or [] + if allowed_models: + provider_config["allowed_models"] = allowed_models + if provider_config: + entry["config"] = provider_config + + return entry, allowed_models + + +class _LLMModelRegistrar: # pylint: disable=too-few-public-methods + """Owns ``registered_resources.models`` writes for one synthesis call. + + ``apply_high_level_inference`` needs to register a model per + ``allowed_models`` entry, dedupe against models registered before the + call (baseline, native_override, BYOK, ...), and — when a later + high-level entry reuses a ``provider_id`` — evict the models it + registered for that provider's earlier declaration. Bundling those three + pieces of state here keeps that bookkeeping out of the caller instead of + threading a list, a set, and a dict through free-function parameters. + """ + + def __init__(self, ls_config: dict[str, Any]) -> None: + self._models = ls_config.setdefault("registered_resources", {}).setdefault( + "models", [] + ) + self._known_ids = { + m.get("model_id") for m in self._models if isinstance(m, dict) + } + self._owned_by_provider: dict[str, list[str]] = {} + + def sync(self, provider_id: str, allowed_models: list[str]) -> None: + """Register ``allowed_models`` for ``provider_id``, replacing its prior set. + + A later high-level entry with the same emitted ``provider_id`` fully + replaces the earlier one's provider config (see + ``_replace_or_append_inference_provider``), so any model this + registrar added for it earlier in the same call is stale and must be + evicted first — otherwise a model no longer served by the replaced + provider would linger in ``registered_resources.models``. + + Parameters: + provider_id: provider_id of the inference provider offering the + models. + allowed_models: Model names to register. + """ + stale = set(self._owned_by_provider.pop(provider_id, [])) + if stale: + self._models[:] = [ + m + for m in self._models + if not (isinstance(m, dict) and m.get("model_id") in stale) + ] + self._known_ids.difference_update(stale) + + added = [] + for model_name in allowed_models: + if model_name in self._known_ids: + continue + self._models.append( + { + "model_id": model_name, + "model_type": "llm", + "provider_id": provider_id, + "provider_model_id": model_name, + } + ) + self._known_ids.add(model_name) + added.append(model_name) + if added: + self._owned_by_provider[provider_id] = added + + def apply_high_level_inference( ls_config: dict[str, Any], inference: dict[str, Any] ) -> None: @@ -1014,6 +1133,14 @@ def apply_high_level_inference( appended. Secrets are emitted as ``${env.}`` references, never resolved values (R6). + Each of the provider's ``allowed_models`` is also registered as an LLM + entry in ``registered_resources.models`` (skipping any ``model_id`` already + present), so the model is usable even when the provider endpoint is + unreachable at startup — Llama Stack's auto-discovery otherwise requires a + live connection to list models. Replacing a provider (same emitted id) + also evicts the LLM entries this function registered for the earlier + declaration, so a superseded provider's models don't linger. + Parameters: ls_config: The Llama Stack configuration being synthesized (modified in place). @@ -1029,39 +1156,18 @@ def apply_high_level_inference( providers_section = ls_config.setdefault("providers", {}) inference_list = providers_section.setdefault("inference", []) + model_registrar = _LLMModelRegistrar(ls_config) for provider in providers: provider_type = provider["type"] emitted_id = provider.get("id") or provider_type.replace("_", "-") ls_provider_type = PROVIDER_TYPE_MAP[provider_type] - entry: dict[str, Any] = { - "provider_id": emitted_id, - "provider_type": ls_provider_type, - } + entry, allowed_models = _build_inference_entry( + provider, emitted_id, ls_provider_type + ) - provider_config: dict[str, Any] = {} - if provider.get("extra"): - provider_config.update(provider["extra"]) - if provider.get("api_key_env"): - key_field = API_KEY_FIELD_MAP.get(ls_provider_type, "api_key") - provider_config[key_field] = "${env." + provider["api_key_env"] + "}" - if provider.get("allowed_models"): - provider_config["allowed_models"] = provider["allowed_models"] - if provider_config: - entry["config"] = provider_config - - # Replace a baseline provider with the same id, else append. - for index, existing in enumerate(inference_list): - if isinstance(existing, dict) and existing.get("provider_id") == emitted_id: - logger.info( - "Replacing existing inference provider with " - "provider_id=%r; a later high-level entry overwrote it", - emitted_id, - ) - inference_list[index] = entry - break - else: - inference_list.append(entry) + _replace_or_append_inference_provider(inference_list, entry) + model_registrar.sync(emitted_id, allowed_models) logger.info( "Applied %d high-level inference provider(s) to synthesized config", @@ -1375,7 +1481,7 @@ def generate_configuration( def main() -> None: """CLI entry point.""" parser = ArgumentParser( - description="Enrich Llama Stack config with Lightspeed values", + description="Enrich or synthesize Llama Stack config from Lightspeed values", ) parser.add_argument( "-c", @@ -1387,20 +1493,32 @@ def main() -> None: "-i", "--input", default="run.yaml", - help="Input Llama Stack config (default: run.yaml)", + help="Input Llama Stack config for legacy enrichment mode; ignored " + "with --synthesize (default: run.yaml)", ) parser.add_argument( "-o", "--output", default="run_.yaml", - help="Output enriched config (default: run_.yaml)", + help="Output config file (default: run_.yaml)", + ) + parser.add_argument( + "--synthesize", + action="store_true", + help="Build a complete run.yaml from -c alone instead of enriching " + "an existing run.yaml given by -i", ) args = parser.parse_args() with open(args.config, "r", encoding="utf-8") as f: - config = yaml.safe_load(f) + config = yaml.safe_load(f) or {} - generate_configuration(args.input, args.output, config) + if args.synthesize: + synthesize_to_file( + config, args.output, config_file_dir=str(Path(args.config).parent) + ) + else: + generate_configuration(args.input, args.output, config) if __name__ == "__main__": diff --git a/tests/unit/test_llama_stack_synthesize.py b/tests/unit/test_llama_stack_synthesize.py index 15927ae99..632b19686 100644 --- a/tests/unit/test_llama_stack_synthesize.py +++ b/tests/unit/test_llama_stack_synthesize.py @@ -6,8 +6,11 @@ write-to-file step (persistent path, mode 0600). """ +# pylint: disable=too-many-lines + import os import stat +import sys from pathlib import Path from typing import Any, Optional, get_args @@ -20,6 +23,7 @@ deep_merge_list_replace, ensure_mcp_tool_runtime, load_default_baseline, + main, migrate_config_dumb, synthesize_configuration, synthesize_to_file, @@ -412,6 +416,133 @@ def test_apply_high_level_inference_empty_is_noop() -> None: assert ls_config["providers"]["inference"] == [{"provider_id": "x"}] +def test_apply_high_level_inference_registers_llm_model() -> None: + """An allowed model is registered as an LLM resource pointing at its provider.""" + ls_config: dict[str, Any] = {"providers": {"inference": []}} + inference = { + "providers": [ + { + "type": "openai", + "api_key_env": "OPENAI_API_KEY", + "allowed_models": ["gpt-4o-mini"], + } + ] + } + apply_high_level_inference(ls_config, inference) + models = ls_config["registered_resources"]["models"] + assert models == [ + { + "model_id": "gpt-4o-mini", + "model_type": "llm", + "provider_id": "openai", + "provider_model_id": "gpt-4o-mini", + } + ] + + +def test_apply_high_level_inference_registers_multiple_allowed_models() -> None: + """Every allowed model for a provider is registered, not just the first.""" + ls_config: dict[str, Any] = {"providers": {"inference": []}} + inference = { + "providers": [ + { + "type": "vllm", + "id": "vllm-prod", + "allowed_models": ["model-a", "model-b"], + } + ] + } + apply_high_level_inference(ls_config, inference) + model_ids = {m["model_id"] for m in ls_config["registered_resources"]["models"]} + assert model_ids == {"model-a", "model-b"} + assert all( + m["provider_id"] == "vllm-prod" + for m in ls_config["registered_resources"]["models"] + ) + + +def test_apply_high_level_inference_no_allowed_models_no_registration() -> None: + """A provider without allowed_models registers no LLM model.""" + ls_config: dict[str, Any] = {"providers": {"inference": []}} + inference = {"providers": [{"type": "sentence_transformers"}]} + apply_high_level_inference(ls_config, inference) + assert ls_config["registered_resources"]["models"] == [] + + +def test_apply_high_level_inference_skips_already_registered_model() -> None: + """A model already present in registered_resources.models is not duplicated.""" + ls_config: dict[str, Any] = { + "providers": {"inference": []}, + "registered_resources": { + "models": [ + { + "model_id": "gpt-4o-mini", + "model_type": "llm", + "provider_id": "stale-provider", + "provider_model_id": "gpt-4o-mini", + } + ] + }, + } + inference = {"providers": [{"type": "openai", "allowed_models": ["gpt-4o-mini"]}]} + apply_high_level_inference(ls_config, inference) + models = ls_config["registered_resources"]["models"] + assert len(models) == 1 + assert models[0]["provider_id"] == "stale-provider" + + +def test_apply_high_level_inference_replacing_provider_evicts_its_stale_models() -> ( + None +): + """A later entry with the same provider_id drops the earlier entry's models.""" + ls_config: dict[str, Any] = {"providers": {"inference": []}} + inference = { + "providers": [ + { + "type": "vllm", + "id": "vllm-shared", + "allowed_models": ["model-old"], + }, + { + "type": "vllm", + "id": "vllm-shared", + "allowed_models": ["model-new"], + }, + ] + } + apply_high_level_inference(ls_config, inference) + models = ls_config["registered_resources"]["models"] + model_ids = {m["model_id"] for m in models} + assert model_ids == {"model-new"} + assert all(m["provider_id"] == "vllm-shared" for m in models) + + +def test_apply_high_level_inference_replacing_provider_keeps_baseline_models() -> None: + """Eviction only removes models this call registered, not pre-existing ones.""" + ls_config: dict[str, Any] = { + "providers": {"inference": []}, + "registered_resources": { + "models": [ + { + "model_id": "baseline-model", + "model_type": "llm", + "provider_id": "vllm-shared", + "provider_model_id": "baseline-model", + } + ] + }, + } + inference = { + "providers": [ + {"type": "vllm", "id": "vllm-shared", "allowed_models": ["model-a"]}, + {"type": "vllm", "id": "vllm-shared", "allowed_models": ["model-b"]}, + ] + } + apply_high_level_inference(ls_config, inference) + model_ids = {m["model_id"] for m in ls_config["registered_resources"]["models"]} + assert model_ids == {"baseline-model", "model-b"} + + def test_provider_type_map_covers_every_literal_value() -> None: """Every UnifiedInferenceProvider.type value has a PROVIDER_TYPE_MAP entry.""" literal_values = set( @@ -788,6 +919,103 @@ def test_migrate_config_dumb_rejects_non_mapping_inputs(tmp_path: Path) -> None: migrate_config_dumb(str(empty_run), lcs_path, out_path) +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- + + +def test_main_synthesize_flag_builds_run_yaml( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """--synthesize builds a complete run.yaml from -c alone.""" + config_path = tmp_path / "lightspeed-stack.yaml" + config_path.write_text( + yaml.dump( + { + "llama_stack": { + "config": { + "baseline": "empty", + "native_override": {"version": 2, "apis": ["inference"]}, + } + } + } + ), + encoding="utf-8", + ) + output_path = tmp_path / "run.yaml" + + monkeypatch.setattr( + sys, + "argv", + [ + "llama_stack_configuration.py", + "-c", + str(config_path), + "-o", + str(output_path), + "--synthesize", + ], + ) + main() + + result = yaml.safe_load(output_path.read_text(encoding="utf-8")) + assert result == {"version": 2, "apis": ["inference"]} + + +def test_main_default_uses_legacy_enrichment( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Without --synthesize, the CLI keeps calling legacy generate_configuration.""" + config_path = tmp_path / "lightspeed-stack.yaml" + config_path.write_text(yaml.dump({}), encoding="utf-8") + input_path = tmp_path / "run.yaml" + input_path.write_text(yaml.dump({"version": 2}), encoding="utf-8") + output_path = tmp_path / "run_.yaml" + + monkeypatch.setattr( + sys, + "argv", + [ + "llama_stack_configuration.py", + "-c", + str(config_path), + "-i", + str(input_path), + "-o", + str(output_path), + ], + ) + main() + + result = yaml.safe_load(output_path.read_text(encoding="utf-8")) + assert result["version"] == 2 + + +def test_main_synthesize_flag_handles_empty_config_file( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An empty/comment-only -c file loads as {} instead of crashing on None.""" + config_path = tmp_path / "lightspeed-stack.yaml" + config_path.write_text("# only a comment\n", encoding="utf-8") + output_path = tmp_path / "run.yaml" + + monkeypatch.setattr( + sys, + "argv", + [ + "llama_stack_configuration.py", + "-c", + str(config_path), + "-o", + str(output_path), + "--synthesize", + ], + ) + main() + + assert output_path.exists() + + # --------------------------------------------------------------------------- # reference profiles (LCORE-2346) # ---------------------------------------------------------------------------