diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml index d9d7db9dc1..8cd6d7f63a 100644 --- a/.github/codeowners/areas.yaml +++ b/.github/codeowners/areas.yaml @@ -57,6 +57,7 @@ areas: - tests/test_epd_cli.py - examples/cli/ - tests/test_stack_discovery.py + - tests/test_output_adapter_discovery.py - tests/test_power_diagnostics.py - tests/test_unified_cli.py - tests/test_ngram_cli.py @@ -197,6 +198,8 @@ shared: owners: [replay, sweeper, forward-pass-engine, maintainers] - glob: tests/test_stack_discovery.py owners: [replay, sweeper, maintainers] + - glob: tests/test_output_adapter_discovery.py + owners: [replay, sweeper, forward-pass-engine, maintainers] - glob: python/aisimulate/src/aisimulate/resources.py owners: [replay, sweeper, maintainers] - glob: python/aisimulate/src/aisimulate/resource_scheduler.py diff --git a/CODEOWNERS b/CODEOWNERS index ba969cda3c..9f1f68a856 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -78,6 +78,7 @@ /tests/test_state_cache_config.py @ai-dynamo/aisimulate-replay-codeowners /tests/fixtures/power-contract-v1.json @ai-dynamo/aisimulate-replay-codeowners /tests/test_unified_traffic_runtime.py @ai-dynamo/aisimulate-replay-codeowners +/tests/test_output_adapter_discovery.py @ai-dynamo/aisimulate-replay-codeowners # === mocker (@ai-dynamo/aisimulate-mocker-codeowners) === /crates/core/src/engine/ @ai-dynamo/aisimulate-mocker-codeowners @@ -204,6 +205,7 @@ /tests/fixtures/power-contract-v1.json @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_spline_estimator_config.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /tests/test_unified_traffic_runtime.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain +/tests/test_output_adapter_discovery.py @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /python/aisimulate/src/aisimulate/replay/ @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/access-aisimulate-maintain /.claude/skills/aisimulate-fpe-performance @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain /docs/schemas/power-metrics-v1.schema.json @ai-dynamo/aisimulate-replay-codeowners @ai-dynamo/aisimulate-sweeper-codeowners @ai-dynamo/aisimulate-forward-pass-engine-codeowners @ai-dynamo/access-aisimulate-maintain diff --git a/docs/cli/user-guide.md b/docs/cli/user-guide.md index 08d593f80f..4d30b65f45 100644 --- a/docs/cli/user-guide.md +++ b/docs/cli/user-guide.md @@ -84,6 +84,10 @@ warmup fields and retains both worker pools across the barrier. See | `predict` | Evaluate one concrete deployment under a workload. | `aisimulate predict -c prediction.yaml --output-dir ./prediction-output` | A metrics summary and `prediction-output/prediction.json`. | | `recommend` | Search deployment and load choices for an optimization goal. | `aisimulate recommend -c recommendation.yaml --output-dir ./recommendation-output` | Ranked configurations, `recommendation.json`, and concrete YAML files under `recommendations/`. | +`recommend` also accepts repeatable `--output NAME` options. Each selected output adapter requires +a same-named top-level configuration section, may observe live candidate and round notifications, +and writes additional artifacts after recommendation. + Unless labeled as captured output, metric values in example results are hypothetical and illustrate the output format. Captured detail examples are simulation results, not hardware measurements. @@ -1883,7 +1887,10 @@ after adapter canonicalization and deduplication so it maps one-to-one to the nu Read [Understand your prediction](understand-your-prediction.md) for an annotated report, ITL/TPOT definitions, incomplete-request handling, and SLA interpretation. -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 configures artifact generation and is excluded from recommended prediction +YAML files. Output adapters receive the prepared output directory directly and must preserve +unrelated 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 diff --git a/docs/sweeper/architecture.md b/docs/sweeper/architecture.md index 45bfd73c47..abbb022529 100644 --- a/docs/sweeper/architecture.md +++ b/docs/sweeper/architecture.md @@ -52,6 +52,32 @@ 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. +Recommendation outputs 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. A selected adapter may subscribe to live +candidate and round notifications. After the recommendation completes, AISimulate passes the +adapter's section, the final `SweepResult`, and the prepared `--output-dir` to its writer. 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 required operation and one optional operation: + +```python +adapter.subscribe(config) # optional: RecommendationOutputCallbacks or None +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. AISimulate writes canonical recommendation output before calling +`write()`, and a writer failure makes the command fail. Adapters share that output directory and must +not modify canonical files. `subscribe()` runs in the supervised recommendation worker and may +return `on_candidate` and `on_round` callbacks. These callbacks execute synchronously on the search +path, must return promptly, and abort the recommendation if they fail. The worker and final writer +may use separate adapter instances, so plugins must not depend on shared in-memory state between +callbacks and `write()`. + See [Sweep Configuration Providers](sweep-config-provider.md) for the SDK provider contract. ## Sweep Flow diff --git a/python/aisimulate/src/aisimulate/cli_args.py b/python/aisimulate/src/aisimulate/cli_args.py index 4947d80945..8756fe56c5 100644 --- a/python/aisimulate/src/aisimulate/cli_args.py +++ b/python/aisimulate/src/aisimulate/cli_args.py @@ -5,6 +5,8 @@ from __future__ import annotations import argparse +import importlib.metadata +from collections.abc import Sequence from pathlib import Path from typing import Any @@ -76,6 +78,14 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="pace prediction against the real wall clock instead of virtual time", ) + subparsers.choices["recommend"].add_argument( + "--output", + dest="outputs", + action="append", + default=[], + metavar="NAME", + help="write additional artifacts using an installed output adapter", + ) add_support_parser(subparsers) return parser @@ -93,6 +103,35 @@ def _load_mapping(path: str) -> dict[str, Any]: return value +def _extract_output_configs( + raw: dict[str, Any], outputs: Sequence[str], *, stack: str +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Validate selected output sections before supervision can clear old outputs.""" + + remaining = dict(raw) + configs: dict[str, dict[str, Any]] = {} + if not outputs: + return remaining, configs + + from .config.common import CONFIG_ADAPTER_ENTRY_POINT_GROUP, RECOMMENDATION_CORE_SECTIONS + + config_adapter_names = { + entry.name for entry in importlib.metadata.entry_points().select(group=CONFIG_ADAPTER_ENTRY_POINT_GROUP) + } + for name in dict.fromkeys(outputs): + if not name or "." in name: + raise _CliConfigError(f"invalid --output name {name!r}") + if name in RECOMMENDATION_CORE_SECTIONS or f"{stack}.{name}" in config_adapter_names: + raise _CliConfigError(f"output adapter name {name!r} collides with a recommendation input section") + 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 _apply_overrides(data: dict[str, Any], overrides: list[str], *, command: str) -> None: for assignment in overrides: if "=" not in assignment: diff --git a/python/aisimulate/src/aisimulate/config/common.py b/python/aisimulate/src/aisimulate/config/common.py index 493f40066a..80b8e37e5e 100644 --- a/python/aisimulate/src/aisimulate/config/common.py +++ b/python/aisimulate/src/aisimulate/config/common.py @@ -215,6 +215,7 @@ def load_yaml(path: str | Path) -> dict[str, Any]: return data +CONFIG_ADAPTER_ENTRY_POINT_GROUP = "aisimulate.config_adapters" PREDICTION_CORE_SECTIONS = frozenset({"traffic", "engine", "evaluation", "execution"}) RECOMMENDATION_CORE_SECTIONS = frozenset({*PREDICTION_CORE_SECTIONS, "optimization", "optimizer"}) diff --git a/python/aisimulate/src/aisimulate/config_adapter.py b/python/aisimulate/src/aisimulate/config_adapter.py index b0f560f150..6e3033c4d4 100644 --- a/python/aisimulate/src/aisimulate/config_adapter.py +++ b/python/aisimulate/src/aisimulate/config_adapter.py @@ -11,6 +11,7 @@ from dataclasses import dataclass from typing import Any, Protocol, runtime_checkable +from .config.common import CONFIG_ADAPTER_ENTRY_POINT_GROUP from .sweeper.provider import API_VERSION as SWEEP_PROVIDER_API_VERSION from .sweeper.provider import ( AdapterReplaySpec, @@ -21,7 +22,6 @@ ) CONFIG_ADAPTER_API_VERSION = 3 -CONFIG_ADAPTER_ENTRY_POINT_GROUP = "aisimulate.config_adapters" @dataclass(frozen=True) diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py index 56d60013aa..b18f77de56 100644 --- a/python/aisimulate/src/aisimulate/main.py +++ b/python/aisimulate/src/aisimulate/main.py @@ -14,7 +14,7 @@ from pydantic import ValidationError from .afd_artifacts import write_afd_qualification_artifacts -from .cli_args import _apply_overrides, _CliConfigError, _load_mapping, build_parser +from .cli_args import _apply_overrides, _CliConfigError, _extract_output_configs, _load_mapping, build_parser from .compiler import prediction_to_replay_spec from .config.cli import ( CorePredictionConfig, @@ -40,6 +40,12 @@ write_recommendations, write_requests, ) +from .output_adapter import ( + OutputAdapterExecutionError, + OutputAdapterResolutionError, + resolve_output_adapters, + write_output_adapters, +) from .power import normalize_power_summary from .resources import ( GuardedRunnerFactory, @@ -218,6 +224,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", []), stack=args.stack) + 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) @@ -227,6 +235,7 @@ def _recommend(args: argparse.Namespace, raw: dict[str, Any], factory) -> int: stack=args.stack, runner_factory=factory, providers=adapters, + output_configs=output_configs, show_progress=args.format == "table", ) selected: list[tuple[str, Any, dict[str, Any]]] = [] @@ -268,6 +277,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 3 if getattr(result.counts, "resource_limited", 0) else 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 = [] for index, ((_, candidate, _), path) in enumerate(zip(selected, paths, strict=True), start=1): row = { @@ -329,6 +347,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..0081ab62f2 --- /dev/null +++ b/python/aisimulate/src/aisimulate/output_adapter.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Discovery and invocation for recommendation output adapters.""" + +from __future__ import annotations + +import importlib.metadata +from collections.abc import Callable, Iterable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +from .sweeper.config import Candidate +from .sweeper.result import CandidateRecord, SweepResult + +OUTPUT_ADAPTER_API_VERSION = 1 +OUTPUT_ADAPTER_ENTRY_POINT_GROUP = "aisimulate.output_adapters" + + +@dataclass(frozen=True) +class RecommendationOutputCallbacks: + """Optional live recommendation callbacks requested by an output adapter.""" + + on_candidate: Callable[[CandidateRecord], None] | None = None + on_round: Callable[[int, list[Candidate]], None] | None = None + + +@runtime_checkable +class RecommendationOutputAdapter(Protocol): + """Observe a recommendation and write additional artifacts when it completes.""" + + 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 during execution.""" + + +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") + subscribe = getattr(adapter, "subscribe", None) + if subscribe is not None and not callable(subscribe): + raise OutputAdapterResolutionError( + f"output adapter {requested_name!r} has non-callable optional attribute: subscribe" + ) + 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 _dispatch_callback(name: str, event: str, callback: Callable[..., None], *args: Any) -> None: + try: + callback(*args) + except KeyboardInterrupt: + raise + except Exception as exc: + raise OutputAdapterExecutionError( + f"output adapter {name!r} {event} callback failed: {type(exc).__name__}: {exc}" + ) from exc + + +def resolve_output_callbacks( + configs: Mapping[str, Mapping[str, Any]], + *, + injected: Mapping[str, RecommendationOutputAdapter] | None = None, + entry_points: Iterable[importlib.metadata.EntryPoint] | None = None, +) -> RecommendationOutputCallbacks: + """Resolve live subscriptions from fresh worker-local adapter instances.""" + + adapters = resolve_output_adapters(configs, injected=injected, entry_points=entry_points) + candidate_callbacks: list[tuple[str, Callable[[CandidateRecord], None]]] = [] + round_callbacks: list[tuple[str, Callable[[int, list[Candidate]], None]]] = [] + for name, adapter in adapters.items(): + subscribe = getattr(adapter, "subscribe", None) + if subscribe is None: + continue + try: + callbacks = subscribe(configs[name]) + except KeyboardInterrupt: + raise + except Exception as exc: + raise OutputAdapterExecutionError( + f"output adapter {name!r} subscription failed: {type(exc).__name__}: {exc}" + ) from exc + if callbacks is None: + continue + if not isinstance(callbacks, RecommendationOutputCallbacks): + raise OutputAdapterExecutionError( + f"output adapter {name!r} subscribe() must return RecommendationOutputCallbacks or None" + ) + if callbacks.on_candidate is not None: + candidate_callbacks.append((name, callbacks.on_candidate)) + if callbacks.on_round is not None: + round_callbacks.append((name, callbacks.on_round)) + + def on_candidate(record: CandidateRecord) -> None: + for name, callback in candidate_callbacks: + _dispatch_callback(name, "on_candidate", callback, deepcopy(record)) + + def on_round(round_no: int, candidates: list[Candidate]) -> None: + for name, callback in round_callbacks: + _dispatch_callback(name, "on_round", callback, round_no, deepcopy(candidates)) + + return RecommendationOutputCallbacks( + on_candidate=on_candidate if candidate_callbacks else None, + on_round=on_round if round_callbacks else None, + ) + + +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", + "RecommendationOutputCallbacks", + "resolve_output_adapters", + "resolve_output_callbacks", + "validate_output_adapter", + "write_output_adapters", +] diff --git a/python/aisimulate/src/aisimulate/recommend.py b/python/aisimulate/src/aisimulate/recommend.py index e2a58a7772..8997154b0d 100644 --- a/python/aisimulate/src/aisimulate/recommend.py +++ b/python/aisimulate/src/aisimulate/recommend.py @@ -20,6 +20,7 @@ RecommendationAdapterContext, SimulationConfigAdapter, ) +from .output_adapter import RecommendationOutputAdapter, resolve_output_callbacks from .resources import GuardedRunnerFactory, discover_host, resolve_budget from .sweeper.afd_perfmodel import AFDPerformanceModel from .sweeper.config import SmartSearchConfig @@ -35,6 +36,7 @@ def run_recommendation( stack: str, runner_factory: RunnerFactory, providers: Mapping[str, SimulationConfigAdapter] | None = None, + output_configs: Mapping[str, Mapping[str, Any]] | None = None, afd_performance_model: AFDPerformanceModel | None = None, show_progress: bool = True, output_requirements: ReplayOutputRequirements | None = None, @@ -48,6 +50,7 @@ def run_recommendation( stack=stack, runner_factory=runner_factory, providers=providers, + output_configs=output_configs, afd_performance_model=afd_performance_model, show_progress=show_progress, output_requirements=output_requirements, @@ -64,6 +67,8 @@ def _run_recommendation( stack: str, runner_factory: RunnerFactory, providers: Mapping[str, SimulationConfigAdapter] | None = None, + output_configs: Mapping[str, Mapping[str, Any]] | None = None, + output_adapters: Mapping[str, RecommendationOutputAdapter] | None = None, afd_performance_model: AFDPerformanceModel | None = None, show_progress: bool = True, output_requirements: ReplayOutputRequirements | None = None, @@ -113,7 +118,13 @@ def _run_recommendation( ), afd_performance_model=afd_performance_model, ) - return sweeper.run(smart, top_n=None) + output_callbacks = resolve_output_callbacks(output_configs or {}, injected=output_adapters) + return sweeper.run( + smart, + top_n=None, + on_candidate=output_callbacks.on_candidate, + on_round=output_callbacks.on_round, + ) def recommendation_to_sweeper( diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py index 682f86153a..a4e7f27d45 100644 --- a/python/aisimulate/src/aisimulate/supervision.py +++ b/python/aisimulate/src/aisimulate/supervision.py @@ -336,7 +336,7 @@ def _save_report(output: str, filename: str, report: dict[str, Any], *, overwrit def main(argv: Sequence[str] | None = None) -> int: - from .cli_args import _apply_overrides, _load_mapping, build_parser + from .cli_args import _apply_overrides, _extract_output_configs, _load_mapping, build_parser from .output import prepare_output_directory arguments = list(sys.argv[1:] if argv is None else argv) @@ -367,7 +367,10 @@ def main(argv: Sequence[str] | None = None) -> int: from .config.cli import CorePredictionConfig, CoreRecommendationConfig from .config.common import split_config_sections - core_raw, adapter_raw = split_config_sections(raw, command=args.command) + validation_raw = dict(raw) + if args.command == "recommend": + validation_raw, _ = _extract_output_configs(raw, args.outputs, stack=args.stack) + core_raw, adapter_raw = split_config_sections(validation_raw, command=args.command) config_type = CorePredictionConfig if args.command == "predict" else CoreRecommendationConfig config = config_type.model_validate(core_raw) if config.engine.workers.encoder is not None: diff --git a/tests/test_host_resources.py b/tests/test_host_resources.py index 055811afdd..495aa64c18 100644 --- a/tests/test_host_resources.py +++ b/tests/test_host_resources.py @@ -223,6 +223,7 @@ def capabilities(self): def test_recommendation_applies_slots_without_changing_suggestion_batches(monkeypatch): import aisimulate.recommend as recommendation import aisimulate.sweeper.search as search + from aisimulate.output_adapter import OUTPUT_ADAPTER_API_VERSION, RecommendationOutputCallbacks raw = _config() raw["traffic"]["load"]["concurrency"] = 8 @@ -230,19 +231,44 @@ def test_recommendation_applies_slots_without_changing_suggestion_batches(monkey monkeypatch.setattr(recommendation, "resolve_budget", lambda *a, **kw: {"cpu_limit": 2}) seen = [] + class SubscribedAdapter: + name = "dgd" + api_version = OUTPUT_ADAPTER_API_VERSION + + def subscribe(self, config): + assert config == {"socket": "/tmp/events.sock"} + return RecommendationOutputCallbacks(on_candidate=lambda candidate: None) + + def write(self, config, *, result, output_dir): + del config, result, output_dir + return [] + class CaptureSweeper: def __init__(self, **kwargs): pass - def run(self, smart, *, top_n): - seen.append(smart.sweep) + def run(self, smart, *, top_n, on_candidate=None, on_round=None): + del top_n + seen.append((smart.sweep, on_candidate, on_round)) return "result" monkeypatch.setattr(search, "Sweeper", CaptureSweeper) - assert recommendation._run_recommendation(config, stack="engine", runner_factory=object()) == "result" - assert seen[0].parallel_evals == 2 - assert seen[0].candidates_per_round == 8 - assert seen[0].max_trials == 256 + assert ( + recommendation._run_recommendation( + config, + stack="engine", + runner_factory=object(), + output_configs={"dgd": {"socket": "/tmp/events.sock"}}, + output_adapters={"dgd": SubscribedAdapter()}, + ) + == "result" + ) + sweep, on_candidate, on_round = seen[0] + assert sweep.parallel_evals == 2 + assert sweep.candidates_per_round == 8 + assert sweep.max_trials == 256 + assert on_candidate is not None + assert on_round is None assert config.optimizer.parallelism == 8 diff --git a/tests/test_output_adapter_discovery.py b/tests/test_output_adapter_discovery.py new file mode 100644 index 0000000000..759cc2f67e --- /dev/null +++ b/tests/test_output_adapter_discovery.py @@ -0,0 +1,215 @@ +# 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, + RecommendationOutputCallbacks, + resolve_output_adapters, + resolve_output_callbacks, + write_output_adapters, +) +from aisimulate.sweeper.config import Candidate +from aisimulate.sweeper.result import CandidateProvenance, CandidateRecord, CandidateStatus + + +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" + + +def test_output_adapter_subscribes_to_live_recommendation_callbacks() -> None: + received = [] + + class SubscribedAdapter(_Adapter): + def subscribe(self, config): + assert config == {"channel": "private"} + return RecommendationOutputCallbacks( + on_candidate=lambda record: received.append(("candidate", record)), + on_round=lambda round_no, candidates: received.append(("round", round_no, candidates)), + ) + + callbacks = resolve_output_callbacks( + {"dgd": {"channel": "private"}}, + injected={"dgd": SubscribedAdapter("dgd")}, + entry_points=[], + ) + record = CandidateRecord( + candidate_id="candidate-000001", + status=CandidateStatus.FEASIBLE, + config={"backend": "vllm"}, + used_gpus=1, + score=1.0, + provenance=CandidateProvenance(model="example/model", hardware="h200_sxm"), + ) + candidate = Candidate(config={"backend": "vllm"}, used_gpus=1, score=1.0, metrics={}) + + assert callbacks.on_candidate is not None + callbacks.on_candidate(record) + assert callbacks.on_round is not None + callbacks.on_round(2, [candidate]) + + assert received == [("candidate", record), ("round", 2, [candidate])] + + +def test_output_adapter_callbacks_are_isolated_between_subscribers() -> None: + observed = [] + + class MutatingAdapter(_Adapter): + def subscribe(self, config): + del config + + def mutate(record): + record.config["nested"]["value"] = 2 + + return RecommendationOutputCallbacks(on_candidate=mutate) + + class ObservingAdapter(_Adapter): + def subscribe(self, config): + del config + return RecommendationOutputCallbacks( + on_candidate=lambda record: observed.append(record.config["nested"]["value"]) + ) + + callbacks = resolve_output_callbacks( + {"mutating": {}, "observing": {}}, + injected={ + "mutating": MutatingAdapter("mutating"), + "observing": ObservingAdapter("observing"), + }, + entry_points=[], + ) + record = CandidateRecord( + candidate_id="candidate-000001", + status=CandidateStatus.FEASIBLE, + config={"nested": {"value": 1}}, + used_gpus=1, + score=1.0, + provenance=CandidateProvenance(model="example/model", hardware="h200_sxm"), + ) + + assert callbacks.on_candidate is not None + callbacks.on_candidate(record) + + assert observed == [1] + assert record.config["nested"]["value"] == 1 + + +def test_output_adapter_callback_failure_identifies_adapter_and_event() -> None: + class FailingAdapter(_Adapter): + def subscribe(self, config): + del config + + def fail(record): + del record + raise ValueError("closed") + + return RecommendationOutputCallbacks(on_candidate=fail) + + callbacks = resolve_output_callbacks( + {"dgd": {}}, + injected={"dgd": FailingAdapter("dgd")}, + entry_points=[], + ) + record = CandidateRecord( + candidate_id="candidate-000001", + status=CandidateStatus.FEASIBLE, + config={}, + used_gpus=1, + score=1.0, + provenance=CandidateProvenance(model="example/model", hardware="h200_sxm"), + ) + + assert callbacks.on_candidate is not None + with pytest.raises(OutputAdapterExecutionError, match="dgd.*on_candidate.*ValueError: closed"): + callbacks.on_candidate(record) diff --git a/tests/test_resource_supervision.py b/tests/test_resource_supervision.py index 4dfafc7145..0b9fc9f6f3 100644 --- a/tests/test_resource_supervision.py +++ b/tests/test_resource_supervision.py @@ -147,6 +147,81 @@ def test_invalid_cli_input_preserves_outputs_and_stack_error_order(tmp_path, fai assert "invalid --set" not in result.stderr +@pytest.mark.parametrize( + "name,section,error", + [ + ("traffic", {}, "collides"), + ("placement", {}, "collides"), + ("bad.name", {}, "invalid --output name"), + ("missing", {}, "requires a top-level"), + ("artifact", [], "must be a mapping"), + ], +) +def test_recommend_output_validation_preserves_outputs(tmp_path, name, section, error): + import os + import subprocess + + raw = { + "engine": { + "mode": "aggregated", + "model": "example/model", + "hardware": "h200_sxm", + "context_length": 4096, + "workers": {"aggregated": {}}, + }, + "optimization": { + "target": "throughput", + "constraints": {"max_candidate_gpus": 1}, + }, + } + if name != "missing": + raw[name] = section + config = tmp_path / "config.yaml" + config.write_text(yaml.safe_dump(raw)) + # Metadata discovery must reject this collision without loading its provider. + metadata = tmp_path / "test_config_adapter-1.0.dist-info" + metadata.mkdir() + (metadata / "METADATA").write_text("Metadata-Version: 2.1\nName: test-config-adapter\nVersion: 1.0\n") + (metadata / "entry_points.txt").write_text( + "[aisimulate.config_adapters]\nengine.placement = nonexistent_adapter:factory\n" + ) + output = tmp_path / "output" + (output / "recommendations").mkdir(parents=True) + original = { + "recommendation.json": "old result", + "recommendation.csv": "old candidates", + "recommendations/0001.yaml": "old config", + "resource-runtime.json": "old diagnostics", + "keep.txt": "unrelated file", + } + for filename, contents in original.items(): + (output / filename).write_text(contents) + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(tmp_path), env.get("PYTHONPATH")))) + result = subprocess.run( + [ + sys.executable, + "-m", + "aisimulate", + "recommend", + "--config", + str(config), + "--output", + name, + "--output-dir", + str(output), + "--overwrite", + ], + env=env, + capture_output=True, + text=True, + timeout=15, + ) + assert result.returncode == 2, result.stderr + assert error in result.stderr + assert {str(path.relative_to(output)): path.read_text() for path in output.rglob("*") if path.is_file()} == original + + def _cli_arguments(tmp_path): config = tmp_path / "config.yaml" config.write_text("engine:\n model: example/model\n hardware: h200_sxm\n workers:\n aggregated: {}\n") @@ -519,6 +594,7 @@ def test_supervisor_argument_and_output_setup_do_not_import_runtime(): """ import sys import aisimulate.supervision, aisimulate.cli_args, aisimulate.output +aisimulate.cli_args._extract_output_configs({'artifact': {}}, ['artifact'], stack='engine') runtime = ('aisimulate._runtime', 'aisimulate.sweeper', 'numpy', 'pandas') assert not any(name == root or name.startswith(root + '.') for name in sys.modules for root in runtime) diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py index ea5ddfe8fd..e522b85556 100644 --- a/tests/test_unified_cli.py +++ b/tests/test_unified_cli.py @@ -13,6 +13,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 @@ -222,6 +223,26 @@ 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)] + + +class _MissingArtifactOutputAdapter(_DGDOutputAdapter): + def write(self, config, *, result, output_dir): + self.calls.append((config, result, output_dir)) + return ["missing.yaml"] + + def test_predict_is_the_single_concrete_cli(tmp_path, monkeypatch, capsys) -> None: config_path = tmp_path / "prediction.yaml" config_path.write_text( @@ -948,6 +969,110 @@ def test_recommendation_outputs_each_concrete_prediction_once(tmp_path, monkeypa ) +@pytest.mark.parametrize("stack", ["engine", "dynamo"]) +@pytest.mark.parametrize("missing_artifact", [False, True]) +def test_recommendation_invokes_selected_output_adapter(tmp_path, monkeypatch, capsys, stack, missing_artifact) -> 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 = [] + recommendation_kwargs = {} + adapter = _MissingArtifactOutputAdapter() if missing_artifact else _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=[], + ), + ) + + def run_recommendation(*args, **kwargs): + del args + recommendation_kwargs.update(kwargs) + return result + + monkeypatch.setattr("aisimulate.recommend.run_recommendation", run_recommendation) + + 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", + ] + ) == (1 if missing_artifact else 0) + + captured = capsys.readouterr() + 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 recommendation_kwargs["output_configs"] == {"dgd": output_config} + assert received_result is result + assert received_dir == output + for filename in ("recommendation.json", "recommendation.csv", "recommendations/0001.yaml"): + assert (output / filename).is_file() + assert "dgd" not in yaml.safe_load((output / "recommendations" / "0001.yaml").read_text()) + if missing_artifact: + assert "reported missing artifact" in captured.err + else: + assert (output / "qwen.yaml").read_text() == "kind: DynamoGraphDeployment\n" + assert json.loads(captured.out)[0]["score"] == 1.0 + + def test_overwrite_only_removes_known_outputs(tmp_path) -> None: root = tmp_path / "out" recommendations = root / "recommendations"