From 61aa3777fc657c2156efd29390abd344de32f44d Mon Sep 17 00:00:00 2001 From: "Dr. Stefan Schimanski" Date: Fri, 4 Sep 2026 17:03:21 +0200 Subject: [PATCH] feat(cli): prototype recommendation output adapters Signed-off-by: Dr. Stefan Schimanski --- docs/cli/design.md | 37 ++++- python/aisimulate/src/aisimulate/main.py | 45 ++++++ .../src/aisimulate/output_adapter.py | 143 ++++++++++++++++++ tests/test_output_adapter_discovery.py | 104 +++++++++++++ tests/test_unified_cli.py | 140 +++++++++++++++++ 5 files changed, 467 insertions(+), 2 deletions(-) create mode 100644 python/aisimulate/src/aisimulate/output_adapter.py create mode 100644 tests/test_output_adapter_discovery.py diff --git a/docs/cli/design.md b/docs/cli/design.md index b9afd5e78c..6090794e39 100644 --- a/docs/cli/design.md +++ b/docs/cli/design.md @@ -23,7 +23,7 @@ The design covers: The design does not cover: - Python APIs, `ReplaySpec`, runner/factory implementation internals, or adapter internals beyond the - stack-discovery contract. + stack- and output-discovery contracts. - How the `engine` and `dynamo` stacks execute a prediction. - AIConfigurator (AIC), Planner, router, optimizer, worker-pool, caching, or timeout internals. - Compatibility shims, migration code, or implementation sequencing. @@ -74,6 +74,12 @@ The `predict` verb is intentional: one pinned configuration predicts serving beh |---|---|---:|---| | `--capture-per-request` | flag | `false` | Write per-request prediction records to `requests.jsonl`. | +`recommend` also accepts: + +| Option | Type | Default | Meaning | +|---|---|---:|---| +| `--output NAME` | repeatable string | None | Invoke an installed post-recommendation output adapter. A matching top-level configuration section is required. | + The CLI deliberately does not expose field-specific flags such as `--request-per-second` or `--num-workers`. YAML is the authoritative semantic configuration surface. @@ -97,6 +103,26 @@ The config-adapter ABI has three operations: compile a concrete prediction secti recommendation section into a search plan, and materialize one candidate. The legacy Sweeper provider ABI remains a separate SDK compatibility surface. +Post-recommendation artifacts are discovered independently through +`aisimulate.output_adapters`. An output adapter name is an accepted `--output` value and identifies +the same-named top-level configuration section. AISimulate removes explicitly selected output +sections before validating simulation configuration, then passes each resolved section, the final +`SweepResult`, and the prepared `--output-dir` to its adapter after writing canonical recommendation +files. The adapter returns the relative paths it wrote. Output adapters do not affect simulation, +ranking, or recommended prediction configurations. + +The output-adapter ABI has one operation: + +```python +adapter.write(config, result=result, output_dir=output_dir) +``` + +Adapters are loaded only when selected. Names must be unique, implementations must declare the +supported output-adapter API version, and reported paths must be relative to the supplied output +directory and exist after the call. A plugin failure leaves canonical recommendation output intact +and makes the command fail. Version 1 invokes adapters only after final Candidate selection; it does +not expose per-round incumbent callbacks. + If a requested optional stack is not installed, the CLI exits with code `2` before loading the configuration and reports an actionable error: @@ -1092,7 +1118,9 @@ after adapter canonicalization and deduplication so it maps one-to-one to the nu ## Outputs -Output controls are CLI-only. They never appear in an input or recommended YAML file. +Core output controls are CLI-only. A selected output adapter owns its same-named top-level input +section; that section is configuration for artifact generation and is excluded from recommended +prediction YAML files. Recommendation output uses the schema-versioned `SweepResult` contract documented in [`docs/sweeper/results.md`](../sweeper/results.md). It preserves run metadata, a candidate-attempt @@ -1115,6 +1143,7 @@ Replay metrics use unit-bearing names such as `*_tok_s`, `*_ms`, `*_w`, and `*_j ```text / ├── recommendation.json +├── # only with --output NAME └── recommendations/ ├── 0001.yaml ├── 0002.yaml @@ -1148,6 +1177,10 @@ Specifically, overwrite may replace `prediction.json`, `recommendation.json`, `r numbered `recommendations/NNNN.yaml` files. Other files, including non-numbered files inside `recommendations/`, are preserved. +Output adapters receive this prepared directory directly; AISimulate does not create an additional +adapter-specific subdirectory. An adapter owns replacement and cleanup of the relative paths it +reports and must preserve unrelated files. + ### Standard Output `--format table` prints a concise human-readable summary. `--format json` prints the same summary as diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index f6c14ee15a..fc6cbf436a 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -37,6 +37,12 @@ write_recommendations, write_requests, ) +from .output_adapter import ( + OutputAdapterExecutionError, + OutputAdapterResolutionError, + resolve_output_adapters, + write_output_adapters, +) from .stack import StackResolutionError, resolve_runner_factory from .sweeper.provider import AdapterReplaySpec from .sweeper.replay import ReplayOutputRequirements @@ -71,6 +77,14 @@ def build_parser() -> argparse.ArgumentParser: child.add_argument("--overwrite", action="store_true") child.add_argument("--format", choices=("table", "json"), default="table") subparsers.choices["predict"].add_argument("--capture-per-request", action="store_true") + subparsers.choices["recommend"].add_argument( + "--output", + dest="outputs", + action="append", + default=[], + metavar="NAME", + help="write additional artifacts using an installed output adapter", + ) return parser @@ -146,6 +160,25 @@ def _compile_prediction_adapters( return compiled +def _extract_output_configs( + raw: dict[str, Any], outputs: Sequence[str] +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Remove explicitly selected output sections from recommendation input.""" + + remaining = dict(raw) + configs: dict[str, dict[str, Any]] = {} + for name in dict.fromkeys(outputs): + if not name or "." in name: + raise _CliConfigError(f"invalid --output name {name!r}") + value = remaining.pop(name, None) + if value is None: + raise _CliConfigError(f"--output {name!r} requires a top-level {name!r} configuration section") + if not isinstance(value, dict): + raise _CliConfigError(f"output section {name!r} must be a mapping") + configs[name] = value + return remaining, configs + + def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: core_raw, adapter_raw = split_config_sections(raw, command="predict") config = CorePredictionConfig.model_validate(core_raw) @@ -202,6 +235,8 @@ def _predict(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: def _recommend(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: from .recommend import run_recommendation + raw, output_configs = _extract_output_configs(raw, getattr(args, "outputs", [])) + output_adapters = resolve_output_adapters(output_configs) core_raw, adapter_raw = split_config_sections(raw, command="recommend") config = CoreRecommendationConfig.model_validate(core_raw) adapters = _resolve_section_adapters(adapter_raw, args.stack) @@ -251,6 +286,15 @@ def _recommend(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: sys.stderr.write(f"no feasible candidate found; saved full result to: {result_path}\n") return 1 paths = write_recommendations(root, [config for _, _, config in selected]) + try: + write_output_adapters( + output_adapters, + output_configs, + result=result, + output_dir=root, + ) + except OutputAdapterExecutionError as exc: + raise _CliExecutionError(str(exc)) from exc rows = [ { "rank": index, @@ -285,6 +329,7 @@ def main(argv: Sequence[str] | None = None) -> int: except ( _CliConfigError, ConfigAdapterResolutionError, + OutputAdapterResolutionError, ValidationError, ValueError, ) as exc: diff --git a/python/aisimulate/src/aisimulate/output_adapter.py b/python/aisimulate/src/aisimulate/output_adapter.py new file mode 100644 index 0000000000..9e21692d38 --- /dev/null +++ b/python/aisimulate/src/aisimulate/output_adapter.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Discovery and invocation for post-recommendation output adapters.""" + +from __future__ import annotations + +import importlib.metadata +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +from .sweeper.result import SweepResult + +OUTPUT_ADAPTER_API_VERSION = 1 +OUTPUT_ADAPTER_ENTRY_POINT_GROUP = "aisimulate.output_adapters" + + +@runtime_checkable +class RecommendationOutputAdapter(Protocol): + """Write additional artifacts for a completed recommendation.""" + + name: str + api_version: int + + def write( + self, + config: Mapping[str, Any], + *, + result: SweepResult, + output_dir: Path, + ) -> Sequence[str | Path]: + """Write artifacts and return their paths relative to ``output_dir``.""" + + +class OutputAdapterResolutionError(RuntimeError): + """A selected recommendation output adapter could not be resolved.""" + + +class OutputAdapterExecutionError(RuntimeError): + """A selected recommendation output adapter failed to write its artifacts.""" + + +def validate_output_adapter(adapter: Any, *, requested_name: str) -> RecommendationOutputAdapter: + """Validate one resolved output adapter.""" + + if getattr(adapter, "name", None) != requested_name: + raise OutputAdapterResolutionError( + f"output adapter {requested_name!r} returned name {getattr(adapter, 'name', None)!r}" + ) + api_version = getattr(adapter, "api_version", None) + if type(api_version) is not int or api_version != OUTPUT_ADAPTER_API_VERSION: + raise OutputAdapterResolutionError( + f"output adapter {requested_name!r} uses API version {api_version!r}; " + f"AISimulate requires {OUTPUT_ADAPTER_API_VERSION}" + ) + if not callable(getattr(adapter, "write", None)): + raise OutputAdapterResolutionError(f"output adapter {requested_name!r} is missing callable: write") + return adapter + + +def resolve_output_adapters( + names: Iterable[str], + *, + injected: Mapping[str, RecommendationOutputAdapter] | None = None, + entry_points: Iterable[importlib.metadata.EntryPoint] | None = None, +) -> dict[str, RecommendationOutputAdapter]: + """Resolve selected output adapters without importing unrelated packages.""" + + requested = list(dict.fromkeys(names)) + injected = injected or {} + installed = ( + list(entry_points) + if entry_points is not None + else list(importlib.metadata.entry_points().select(group=OUTPUT_ADAPTER_ENTRY_POINT_GROUP)) + ) + resolved: dict[str, RecommendationOutputAdapter] = {} + for name in requested: + if name in injected: + resolved[name] = validate_output_adapter(injected[name], requested_name=name) + continue + matches = [entry for entry in installed if entry.name == name] + if not matches: + available = sorted(set(injected) | {entry.name for entry in installed}) + raise OutputAdapterResolutionError( + f"output adapter {name!r} is unavailable; installed adapters: " + f"{', '.join(available) if available else ''}" + ) + if len(matches) > 1: + raise OutputAdapterResolutionError(f"output adapter {name!r} has multiple installed providers") + entry = matches[0] + try: + constructor = entry.load() + adapter = constructor() if callable(constructor) else constructor + except Exception as exc: + raise OutputAdapterResolutionError( + f"failed to load output adapter {name!r} from {entry.value!r}: {type(exc).__name__}: {exc}" + ) from exc + resolved[name] = validate_output_adapter(adapter, requested_name=name) + return resolved + + +def write_output_adapters( + adapters: Mapping[str, RecommendationOutputAdapter], + configs: Mapping[str, Mapping[str, Any]], + *, + result: SweepResult, + output_dir: Path, +) -> dict[str, tuple[Path, ...]]: + """Invoke selected adapters and validate their reported artifact paths.""" + + artifacts: dict[str, tuple[Path, ...]] = {} + for name, adapter in adapters.items(): + try: + reported = adapter.write(configs[name], result=result, output_dir=output_dir) + if isinstance(reported, (str, bytes)): + raise TypeError("must return a sequence of relative paths") + paths: list[Path] = [] + for raw_path in reported: + path = Path(raw_path) + if path.is_absolute() or not path.parts or ".." in path.parts: + raise ValueError(f"returned path outside the output directory: {raw_path!r}") + if not (output_dir / path).exists(): + raise ValueError(f"reported missing artifact: {path}") + paths.append(path) + except KeyboardInterrupt: + raise + except Exception as exc: + raise OutputAdapterExecutionError(f"output adapter {name!r} failed: {type(exc).__name__}: {exc}") from exc + artifacts[name] = tuple(paths) + return artifacts + + +__all__ = [ + "OUTPUT_ADAPTER_API_VERSION", + "OUTPUT_ADAPTER_ENTRY_POINT_GROUP", + "OutputAdapterExecutionError", + "OutputAdapterResolutionError", + "RecommendationOutputAdapter", + "resolve_output_adapters", + "validate_output_adapter", + "write_output_adapters", +] diff --git a/tests/test_output_adapter_discovery.py b/tests/test_output_adapter_discovery.py new file mode 100644 index 0000000000..0016914296 --- /dev/null +++ b/tests/test_output_adapter_discovery.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from aisimulate.output_adapter import ( + OUTPUT_ADAPTER_API_VERSION, + OutputAdapterExecutionError, + OutputAdapterResolutionError, + resolve_output_adapters, + write_output_adapters, +) + + +class _Adapter: + api_version = OUTPUT_ADAPTER_API_VERSION + + def __init__(self, name: str) -> None: + self.name = name + + def write(self, config, *, result, output_dir): + del config, result + path = Path("artifact.txt") + (output_dir / path).write_text("artifact\n") + return [path] + + +@dataclass +class _EntryPoint: + name: str + value: str + provider: object + loads: int = 0 + + def load(self): + self.loads += 1 + if isinstance(self.provider, Exception): + raise self.provider + return self.provider + + +def test_only_selected_output_adapter_is_loaded() -> None: + selected = _EntryPoint("dgd", "example:create_dgd", lambda: _Adapter("dgd")) + unused = _EntryPoint("chart", "example:create_chart", lambda: _Adapter("chart")) + + resolved = resolve_output_adapters(["dgd"], entry_points=[selected, unused]) + + assert list(resolved) == ["dgd"] + assert selected.loads == 1 + assert unused.loads == 0 + + +def test_missing_output_adapter_lists_available_names() -> None: + available = _EntryPoint("chart", "example:create_chart", lambda: _Adapter("chart")) + + with pytest.raises(OutputAdapterResolutionError, match="dgd.*installed adapters: chart"): + resolve_output_adapters(["dgd"], entry_points=[available]) + + assert available.loads == 0 + + +@pytest.mark.parametrize( + ("adapter", "message"), + [ + (_Adapter("wrong"), "returned name"), + (type("OldAdapter", (_Adapter,), {"api_version": 0})("dgd"), "API version"), + (type("IncompleteAdapter", (), {"name": "dgd", "api_version": OUTPUT_ADAPTER_API_VERSION})(), "write"), + ], +) +def test_invalid_output_adapter_abi_is_rejected(adapter, message) -> None: + with pytest.raises(OutputAdapterResolutionError, match=message): + resolve_output_adapters(["dgd"], injected={"dgd": adapter}, entry_points=[]) + + +def test_reported_artifacts_must_be_relative_existing_paths(tmp_path) -> None: + class InvalidPathAdapter(_Adapter): + def write(self, config, *, result, output_dir): + del config, result, output_dir + return ["../outside"] + + with pytest.raises(OutputAdapterExecutionError, match="outside the output directory"): + write_output_adapters( + {"dgd": InvalidPathAdapter("dgd")}, + {"dgd": {}}, + result=object(), + output_dir=tmp_path, + ) + + +def test_output_adapter_writes_into_supplied_directory(tmp_path) -> None: + artifacts = write_output_adapters( + {"dgd": _Adapter("dgd")}, + {"dgd": {"name": "example"}}, + result=object(), + output_dir=tmp_path, + ) + + assert artifacts == {"dgd": (Path("artifact.txt"),)} + assert (tmp_path / "artifact.txt").read_text() == "artifact\n" diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py index 8d56adcf6c..2a581a9d7b 100644 --- a/tests/test_unified_cli.py +++ b/tests/test_unified_cli.py @@ -12,6 +12,7 @@ import aisimulate.main as cli from aisimulate.output import prepare_output_directory +from aisimulate.output_adapter import OUTPUT_ADAPTER_API_VERSION, resolve_output_adapters from aisimulate.sweeper.config import Candidate from aisimulate.sweeper.provider import AdapterReplaySpec, AdapterSearchPlan from aisimulate.sweeper.replay import ReplayReport, RunnerCapabilities @@ -135,6 +136,20 @@ def materialize_candidate(self, plan, selection, context): return AdapterReplaySpec(config=dict(plan.state)) +class _DGDOutputAdapter: + name = "dgd" + api_version = OUTPUT_ADAPTER_API_VERSION + + def __init__(self) -> None: + self.calls = [] + + def write(self, config, *, result, output_dir): + self.calls.append((config, result, output_dir)) + path = output_dir / f"{config['name']}.yaml" + path.write_text("kind: DynamoGraphDeployment\n") + return [path.relative_to(output_dir)] + + def test_predict_is_the_single_concrete_cli(tmp_path, monkeypatch, capsys) -> None: config_path = tmp_path / "prediction.yaml" config_path.write_text( @@ -510,6 +525,131 @@ def test_recommendation_outputs_each_concrete_prediction_once( ) +@pytest.mark.parametrize("stack", ["engine", "dynamo"]) +def test_recommendation_invokes_selected_output_adapter(tmp_path, monkeypatch, capsys, stack) -> None: + prediction = { + "engine": { + "mode": "aggregated", + "model": "example/model", + "hardware": "h200_sxm", + "context_length": 4096, + "workers": {"aggregated": {}}, + } + } + config_path = tmp_path / "recommend.yaml" + config_path.write_text( + yaml.safe_dump( + { + **prediction, + "optimization": { + "target": "throughput", + "constraints": {"max_candidate_gpus": 1}, + }, + "dgd": { + "name": "original", + "renderer": "aic", + "format": "manifest", + }, + } + ) + ) + candidate = Candidate( + config={"backend": "vllm"}, + used_gpus=1, + score=1.0, + metrics={}, + prediction_config=prediction, + ) + result = _RecommendationResult([candidate]) + selected_stack = [] + adapter = _DGDOutputAdapter() + + def resolve_stack(name): + selected_stack.append(name) + return _Factory(_Runner()) + + monkeypatch.setattr(cli, "resolve_runner_factory", resolve_stack) + monkeypatch.setattr( + cli, + "resolve_output_adapters", + lambda names: resolve_output_adapters( + names, + injected={"dgd": adapter}, + entry_points=[], + ), + ) + monkeypatch.setattr( + "aisimulate.recommend.run_recommendation", + lambda *args, **kwargs: result, + ) + + output = tmp_path / "out" + assert ( + cli.main( + [ + "recommend", + "--stack", + stack, + "--config", + str(config_path), + "--set", + "dgd.name=qwen", + "--output", + "dgd", + "--output-dir", + str(output), + "--format", + "json", + ] + ) + == 0 + ) + + assert selected_stack == [stack] + assert len(adapter.calls) == 1 + output_config, received_result, received_dir = adapter.calls[0] + assert output_config == { + "name": "qwen", + "renderer": "aic", + "format": "manifest", + } + assert received_result is result + assert received_dir == output + assert (output / "qwen.yaml").read_text() == "kind: DynamoGraphDeployment\n" + assert "dgd" not in yaml.safe_load((output / "recommendations" / "0001.yaml").read_text()) + assert json.loads(capsys.readouterr().out)[0]["score"] == 1.0 + + +def test_selected_output_requires_matching_configuration_section(tmp_path, monkeypatch, capsys) -> None: + config_path = tmp_path / "recommend.yaml" + config_path.write_text( + yaml.safe_dump( + { + "engine": { + "model": "example/model", + "hardware": "h200_sxm", + "context_length": 4096, + "workers": {"aggregated": {}}, + } + } + ) + ) + monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner())) + + with pytest.raises(SystemExit, match="2"): + cli.main( + [ + "recommend", + "--config", + str(config_path), + "--output", + "dgd", + ] + ) + + assert "requires a top-level 'dgd' configuration section" in capsys.readouterr().err + + def test_overwrite_only_removes_known_outputs(tmp_path) -> None: root = tmp_path / "out" recommendations = root / "recommendations"