From 3e16906e7253657e65d1fafe897595cf4e16d20b Mon Sep 17 00:00:00 2001
From: "Dr. Stefan Schimanski" <stefan.schimanski@gmail.com>
Date: Fri, 4 Sep 2026 17:03:21 +0200
Subject: [PATCH 1/9] feat(cli): prototype recommendation output adapters

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 docs/cli/user-guide.md                        |   8 +-
 docs/sweeper/architecture.md                  |  20 +++
 python/aisimulate/src/aisimulate/cli_args.py  |   8 +
 python/aisimulate/src/aisimulate/main.py      |  37 +++++
 .../src/aisimulate/output_adapter.py          | 143 ++++++++++++++++++
 tests/test_output_adapter_discovery.py        | 104 +++++++++++++
 tests/test_unified_cli.py                     | 140 +++++++++++++++++
 7 files changed, 459 insertions(+), 1 deletion(-)
 create mode 100644 python/aisimulate/src/aisimulate/output_adapter.py
 create mode 100644 tests/test_output_adapter_discovery.py

diff --git a/docs/cli/user-guide.md b/docs/cli/user-guide.md
index 08d593f80f..61fb613be6 100644
--- a/docs/cli/user-guide.md
+++ b/docs/cli/user-guide.md
@@ -84,6 +84,9 @@ 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 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 +1886,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..49e215d67f 100644
--- a/docs/sweeper/architecture.md
+++ b/docs/sweeper/architecture.md
@@ -52,6 +52,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.
+
 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..a9b4331c4e 100644
--- a/python/aisimulate/src/aisimulate/cli_args.py
+++ b/python/aisimulate/src/aisimulate/cli_args.py
@@ -76,6 +76,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
 
diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py
index 817679d156..878d1cb64e 100644
--- a/python/aisimulate/src/aisimulate/main.py
+++ b/python/aisimulate/src/aisimulate/main.py
@@ -39,6 +39,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,
@@ -90,6 +96,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)
@@ -210,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)
@@ -260,6 +287,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 = {
@@ -321,6 +357,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 '<none>'}"
+            )
+        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 ea5ddfe8fd..32907f0737 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,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(
@@ -948,6 +963,131 @@ def test_recommendation_outputs_each_concrete_prediction_once(tmp_path, monkeypa
     )
 
 
+@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"

From f023ec5451687fbf7f9c1f51e12a7ac652d12607 Mon Sep 17 00:00:00 2001
From: devivasudevan <49675305+devivasudevan@users.noreply.github.com>
Date: Wed, 23 Sep 2026 15:49:13 +0200
Subject: [PATCH 2/9] feat(sweeper): add on_candidate callback to Sweeper.run

Sweeper.run() currently exposes on_round for round-level progress, but
nothing fires per candidate outcome. DEP ai-dynamo/dynamo#15073 needs
this granularity for its search_resolved event (per-candidate
materialization outcome, not just per-round), and no such hook exists
today for any caller to build on.

on_candidate is invoked once per recorded CandidateRecord (feasible,
infeasible, unsupported, timed-out, or failed), at the same call site
_record() already uses for the existing resource_aware checkpoint()
call. Purely additive: defaults to None, zero behavior change for
every existing caller.

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 .../src/aisimulate/sweeper/search.py          |  9 +++++
 tests/sweeper/test_search.py                  | 34 +++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py
index e726bd3068..9e60b76e39 100644
--- a/python/aisimulate/src/aisimulate/sweeper/search.py
+++ b/python/aisimulate/src/aisimulate/sweeper/search.py
@@ -1113,6 +1113,7 @@ def run(
         top_n: int | None = 5,
         candidate_retention: CandidateRetention | str = CandidateRetention.ALL,
         on_round: Callable[[int, list[Candidate]], None] | None = None,
+        on_candidate: Callable[[CandidateRecord], None] | None = None,
     ) -> SweepResult:
         """Run the sweep and return the canonical schema-versioned result.
 
@@ -1126,6 +1127,12 @@ def run(
         unsupported, timed-out, and failed candidate. ``"feasible"`` retains only
         feasible rows and ``"views"`` retains only the scalar top-N or Pareto front;
         run-wide counts always describe the complete run.
+
+        ``on_candidate``, when provided, is invoked once per recorded candidate
+        outcome (feasible, infeasible, unsupported, timed-out, or failed) with the
+        same ``CandidateRecord`` appended to the run's candidate ledger, in
+        evaluation order. It complements ``on_round``, which only reports the
+        cumulative feasible-candidate list at round boundaries.
         """
         if top_n is not None and top_n < 1:
             raise ValueError(f"top_n must be positive or None, got {top_n}")
@@ -1462,6 +1469,8 @@ def _record(
                     ),
                 )
                 candidate_records.append(record)
+                if on_candidate is not None:
+                    on_candidate(record)
                 if resource_aware:
                     from ..supervision import checkpoint
 
diff --git a/tests/sweeper/test_search.py b/tests/sweeper/test_search.py
index 501dcbb5c4..43e723909f 100644
--- a/tests/sweeper/test_search.py
+++ b/tests/sweeper/test_search.py
@@ -21,6 +21,7 @@
     ReplaySpec,
     RunnerCapabilities,
 )
+from aisimulate.sweeper.result import CandidateRecord, CandidateStatus
 from aisimulate.sweeper.sampler import Suggestion
 from aisimulate.sweeper.search import Sweeper
 from aisimulate.sweeper.search_space import BranchSpace
@@ -1307,3 +1308,36 @@ def artifact(sample, spec):
         assert payload == {
             k: v for k, v in identity.items() if k not in {"enable_eplb", "moe_backend", "wideep_num_slots"}
         }
+
+
+def test_on_candidate_fires_once_per_feasible_outcome_in_evaluation_order(monkeypatch):
+    branch = _branch(_pc())  # 8 GPUs
+    _stub(monkeypatch, branch)
+    factory = _FakeRunnerFactory()
+    seen: list[CandidateRecord] = []
+
+    Sweeper(
+        runner_factory=factory,
+        sampler_factory=_FakeSampler,
+        show_progress=False,
+    ).run(_config(), top_n=None, on_candidate=seen.append)
+
+    assert [record.status for record in seen] == [CandidateStatus.FEASIBLE] * 3
+    # Evaluation order (increasing max_num_seqs), not the best-first ranked order.
+    assert [record.score for record in seen] == [256.0, 512.0, 768.0]
+
+
+def test_on_candidate_reports_infeasible_outcomes_too(monkeypatch):
+    branch = _branch(_pc(tp=16, replicas=4))  # 64 GPUs, exceeds a 32 GPU budget
+    _stub(monkeypatch, branch)
+    seen: list[CandidateRecord] = []
+
+    Sweeper(
+        runner_factory=_FakeRunnerFactory(),
+        sampler_factory=_FakeSampler,
+        show_progress=False,
+    ).run(_config(gpu_budget=32), top_n=None, on_candidate=seen.append)
+
+    assert len(seen) == 33
+    assert all(record.status == CandidateStatus.INFEASIBLE for record in seen)
+    assert all("over gpu_budget" in (record.reason or "") for record in seen)

From 606fa47d94e1ea23d8fba3aa9421999f6ccfa0fe Mon Sep 17 00:00:00 2001
From: devivasudevan <49675305+devivasudevan@users.noreply.github.com>
Date: Tue, 29 Sep 2026 18:14:02 +0200
Subject: [PATCH 3/9] fix(sweeper): detach on_candidate records and close
 wave-timeout race

on_candidate now receives a deep copy of each CandidateRecord instead
of the live ledger entry, so callback mutations can't affect the
SweepResult that gets built from the same ledger. The docstring is
updated to document all six outcomes on_candidate can report and to
spell out the detach guarantee.

evaluate_waves() now drains a second, non-blocking wait() immediately
after resuming from each yield, before checking the deadline. Without
it, a sibling future that completed while the caller's on_candidate
callback was running (which can take arbitrarily long, since the
generator is suspended at the yield for its duration) was misreported
as timed out instead of as its real outcome. Adds a regression test
that reproduces the race deterministically via a fake wait()/clock.

Addresses review comments from coderabbitai and tedzhouhk on #319.

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 .../src/aisimulate/resource_scheduler.py      |  35 +++-
 .../src/aisimulate/sweeper/search.py          |  14 +-
 tests/test_resource_scheduler.py              | 156 ++++++++++++++++--
 3 files changed, 181 insertions(+), 24 deletions(-)

diff --git a/python/aisimulate/src/aisimulate/resource_scheduler.py b/python/aisimulate/src/aisimulate/resource_scheduler.py
index 51f1e408f6..8021f57ece 100644
--- a/python/aisimulate/src/aisimulate/resource_scheduler.py
+++ b/python/aisimulate/src/aisimulate/resource_scheduler.py
@@ -76,6 +76,19 @@ def evaluate_waves(specs, *, factory, initializer, evaluate, workers: int, timeo
             initialized = False
             initialization_timeout = getattr(getattr(factory, "policy", None), "initialization_timeout_seconds", 60.0)
             deadline = started + timeout if timeout else None
+
+            def _drain(done_futures):
+                nonlocal interrupted
+                for future in done_futures:
+                    index = active[future]
+                    try:
+                        result = future.result()
+                    except ResourceLimitError as exc:
+                        interrupted = InterruptedEvaluation(str(exc), True, exc.plan)
+                        continue
+                    del active[future]
+                    yield index, result
+
             while active:
                 if not initialized:
                     worker_pids = {str(pid) for pid in (getattr(pool, "_processes", None) or {})}
@@ -87,15 +100,19 @@ def evaluate_waves(specs, *, factory, initializer, evaluate, workers: int, timeo
                         interrupted = InterruptedEvaluation("worker initialization timed out", False, {})
                         break
                 done, _ = wait(active, timeout=0.05, return_when=FIRST_COMPLETED)
-                for future in done:
-                    index = active[future]
-                    try:
-                        result = future.result()
-                    except ResourceLimitError as exc:
-                        interrupted = InterruptedEvaluation(str(exc), True, exc.plan)
-                        continue
-                    del active[future]
-                    yield index, result
+                yield from _drain(done)
+                if interrupted:
+                    break
+                if not active:
+                    break
+                # Each yield above can suspend this generator for an arbitrary
+                # duration -- the caller's on_candidate callback runs while
+                # paused there. Drain anything that finished during that pause
+                # before judging the deadline below, so a sibling that
+                # completed while the callback was running is reported as its
+                # real outcome instead of being misclassified as timed out.
+                still_done, _ = wait(active, timeout=0, return_when=FIRST_COMPLETED)
+                yield from _drain(still_done)
                 if interrupted:
                     break
                 if not active:
diff --git a/python/aisimulate/src/aisimulate/sweeper/search.py b/python/aisimulate/src/aisimulate/sweeper/search.py
index 9e60b76e39..3ad5c4578b 100644
--- a/python/aisimulate/src/aisimulate/sweeper/search.py
+++ b/python/aisimulate/src/aisimulate/sweeper/search.py
@@ -1129,10 +1129,12 @@ def run(
         run-wide counts always describe the complete run.
 
         ``on_candidate``, when provided, is invoked once per recorded candidate
-        outcome (feasible, infeasible, unsupported, timed-out, or failed) with the
-        same ``CandidateRecord`` appended to the run's candidate ledger, in
-        evaluation order. It complements ``on_round``, which only reports the
-        cumulative feasible-candidate list at round boundaries.
+        outcome (feasible, infeasible, unsupported, resource-limited, timed-out, or
+        failed) with a deep copy of the ``CandidateRecord`` appended to the run's
+        candidate ledger, in evaluation order. The copy is detached: mutating it
+        has no effect on the ledger used to build ``SweepResult``. It complements
+        ``on_round``, which only reports the cumulative feasible-candidate list at
+        round boundaries.
         """
         if top_n is not None and top_n < 1:
             raise ValueError(f"top_n must be positive or None, got {top_n}")
@@ -1470,7 +1472,9 @@ def _record(
                 )
                 candidate_records.append(record)
                 if on_candidate is not None:
-                    on_candidate(record)
+                    # Detach: on_candidate must not be able to mutate the ledger
+                    # entry that SweepResult is built from.
+                    on_candidate(record.model_copy(deep=True))
                 if resource_aware:
                     from ..supervision import checkpoint
 
diff --git a/tests/test_resource_scheduler.py b/tests/test_resource_scheduler.py
index d8a4be33af..00168e9136 100644
--- a/tests/test_resource_scheduler.py
+++ b/tests/test_resource_scheduler.py
@@ -10,7 +10,13 @@
 
 from aisimulate.config.common import ResourceConfig
 from aisimulate.resource_scheduler import InterruptedEvaluation, evaluate_waves
-from aisimulate.resources import GB, GuardedRunnerFactory, HostResources, ResourceEstimate, ResourceLimitError
+from aisimulate.resources import (
+    GB,
+    GuardedRunnerFactory,
+    HostResources,
+    ResourceEstimate,
+    ResourceLimitError,
+)
 
 
 @dataclass
@@ -41,7 +47,16 @@ def _evaluate(spec):
 def test_admission_splits_waves_and_skips_only_oversized_candidates():
     factory = _BudgetFactory()
     specs = [{"id": i, "cost": cost} for i, cost in enumerate([1, 1, 3, 1])]
-    results = dict(evaluate_waves(specs, factory=factory, initializer=_init, evaluate=_evaluate, workers=4, timeout=10))
+    results = dict(
+        evaluate_waves(
+            specs,
+            factory=factory,
+            initializer=_init,
+            evaluate=_evaluate,
+            workers=4,
+            timeout=10,
+        )
+    )
     assert results[0] == {"id": 0}
     assert results[1] == {"id": 1}
     assert isinstance(results[2], InterruptedEvaluation)
@@ -50,11 +65,15 @@ def test_admission_splits_waves_and_skips_only_oversized_candidates():
     assert factory.admitted == [[0, 1], [3]]
 
 
-def test_checkpoint_failure_stops_before_workers_without_reclassifying_candidates(monkeypatch):
+def test_checkpoint_failure_stops_before_workers_without_reclassifying_candidates(
+    monkeypatch,
+):
     from aisimulate import resource_scheduler as scheduler
 
     factory = _BudgetFactory()
-    failure = ResourceLimitError("execution evidence exceeds the bounded checkpoint budget")
+    failure = ResourceLimitError(
+        "execution evidence exceeds the bounded checkpoint budget"
+    )
 
     def fail_checkpoint(event, plan):
         raise failure
@@ -66,7 +85,16 @@ def unexpected_pool(**kwargs):
     monkeypatch.setattr(scheduler, "ProcessPoolExecutor", unexpected_pool)
     specs = [{"id": i, "cost": 1} for i in range(2)]
     with pytest.raises(ResourceLimitError) as caught:
-        next(evaluate_waves(specs, factory=factory, initializer=_init, evaluate=_evaluate, workers=2, timeout=10))
+        next(
+            evaluate_waves(
+                specs,
+                factory=factory,
+                initializer=_init,
+                evaluate=_evaluate,
+                workers=2,
+                timeout=10,
+            )
+        )
     assert caught.value is failure
     assert factory.admitted == [[0, 1]]
 
@@ -74,7 +102,16 @@ def unexpected_pool(**kwargs):
 def test_resource_retries_are_bounded_and_preserve_completed_work():
     factory = _BudgetFactory()
     specs = [{"id": 0, "cost": 1}, {"id": 1, "cost": 1, "refuse": True}]
-    results = list(evaluate_waves(specs, factory=factory, initializer=_init, evaluate=_evaluate, workers=2, timeout=10))
+    results = list(
+        evaluate_waves(
+            specs,
+            factory=factory,
+            initializer=_init,
+            evaluate=_evaluate,
+            workers=2,
+            timeout=10,
+        )
+    )
     assert dict(results)[0] == {"id": 0}
     assert dict(results)[1].resource_limited
     assert len(results) == 2
@@ -90,7 +127,9 @@ def estimate_host_resources(self, workload, *, concurrency=None):
 def test_whole_wave_reserves_sum_against_one_live_snapshot(monkeypatch):
     from aisimulate import resources
 
-    monkeypatch.setattr(resources, "discover_host", lambda: HostResources(16 * GB, 9 * GB, 8))
+    monkeypatch.setattr(
+        resources, "discover_host", lambda: HostResources(16 * GB, 9 * GB, 8)
+    )
     factory = GuardedRunnerFactory(_EstimateFactory(), "custom", ResourceConfig())
     spec = SimpleNamespace(workload={}, concurrency=1)
     assert factory.admit_wave([spec])["status"] == "admitted"
@@ -98,7 +137,9 @@ def test_whole_wave_reserves_sum_against_one_live_snapshot(monkeypatch):
         factory.admit_wave([spec, spec])
     assert caught.value.plan["required_bytes"] == 12 * GB
     assert caught.value.plan["available_bytes"] < 8 * GB
-    monkeypatch.setattr(resources, "discover_host", lambda: HostResources(16 * GB, 6 * GB, 8))
+    monkeypatch.setattr(
+        resources, "discover_host", lambda: HostResources(16 * GB, 6 * GB, 8)
+    )
     with pytest.raises(ResourceLimitError):
         factory.admit_wave([spec])
 
@@ -139,7 +180,14 @@ def stop(pool):
     monkeypatch.setattr(scheduler, "close_pool", stop)
     specs = [{"id": i, "cost": 1} for i in range(2)]
     results = list(
-        evaluate_waves(specs, factory=Factory(), initializer=_init, evaluate=_evaluate, workers=2, timeout=2)
+        evaluate_waves(
+            specs,
+            factory=Factory(),
+            initializer=_init,
+            evaluate=_evaluate,
+            workers=2,
+            timeout=2,
+        )
     )
     assert results == [(0, {"id": 0}), (1, {"id": 1})]
     assert waves == [[0, 1], [1]]
@@ -204,5 +252,93 @@ def advance(futures, **kwargs):
     monkeypatch.setattr(scheduler, "wait", advance)
     monkeypatch.setattr(scheduler.time, "monotonic", lambda: elapsed[0])
     specs = [{"id": i, "cost": 1} for i in range(2)]
-    results = dict(evaluate_waves(specs, factory=factory, initializer=_init, evaluate=_evaluate, workers=2, timeout=2))
+    results = dict(
+        evaluate_waves(
+            specs,
+            factory=factory,
+            initializer=_init,
+            evaluate=_evaluate,
+            workers=2,
+            timeout=2,
+        )
+    )
     assert results == {0: {"id": 0}, 1: {"id": 1}}
+
+
+def test_sibling_finished_during_a_yield_pause_is_not_reported_timed_out(monkeypatch):
+    """Regression test: evaluate_waves() suspends at each ``yield index, result``
+    while the caller (Sweeper._record's on_candidate callback) runs. If a
+    sibling future completes during that pause, it must still be reported as
+    its real outcome, not misclassified as timed out just because the
+    deadline was crossed while the callback for an earlier candidate was
+    running. Uses a controlled fake ``wait()``/clock, not a real sleeping
+    callback, so the regression is deterministic rather than timing-sensitive.
+    """
+    from concurrent.futures import Future
+    from pathlib import Path
+
+    from aisimulate import resource_scheduler as scheduler
+
+    future0 = Future()
+    future0.set_result({"id": 0})
+    future1 = Future()
+    future1.set_result({"id": 1})
+
+    class Pool:
+        def __init__(self, **kwargs):
+            self._processes = {123: object()}
+            Path(kwargs["initargs"][2], "123").touch()
+
+        def submit(self, evaluate, spec):
+            return future0 if spec["id"] == 0 else future1
+
+    elapsed = [0.0]
+    calls = [0]
+
+    def fake_wait(futures, *, timeout, return_when):
+        calls[0] += 1
+        if calls[0] == 1:
+            # Initial poll: only candidate 0 has finished so far.
+            return {future0}, {future1}
+        if calls[0] == 2:
+            # The drain called immediately after resuming from the yield for
+            # candidate 0 -- this is the regression target. Candidate 1
+            # finished while the (simulated) on_candidate callback for
+            # candidate 0 was running, represented below by advancing the
+            # fake clock past the deadline before this call happens.
+            return {future1}, set()
+        raise AssertionError(
+            "no further wait() calls expected once both futures are drained"
+        )
+
+    factory = _BudgetFactory()
+    monkeypatch.setattr(scheduler, "ProcessPoolExecutor", Pool)
+    monkeypatch.setattr(scheduler, "close_pool", lambda pool: None)
+    monkeypatch.setattr(scheduler, "terminate_pool", lambda pool: None)
+    monkeypatch.setattr(scheduler, "wait", fake_wait)
+    monkeypatch.setattr(scheduler.time, "monotonic", lambda: elapsed[0])
+    specs = [{"id": 0, "cost": 1}, {"id": 1, "cost": 1}]
+    gen = evaluate_waves(
+        specs,
+        factory=factory,
+        initializer=_init,
+        evaluate=_evaluate,
+        workers=2,
+        timeout=1.0,
+    )
+
+    assert next(gen) == (0, {"id": 0})
+
+    # Simulate on_candidate for candidate 0 running long enough to cross the
+    # deadline -- candidate 1 has already finished by the time it returns.
+    elapsed[0] = 1.5
+
+    assert next(gen) == (1, {"id": 1}), (
+        "candidate 1 finished before the callback for candidate 0 returned and "
+        "must be reported as feasible, not swept into the deadline-exceeded "
+        "InterruptedEvaluation just because the clock advanced while that "
+        "callback was running"
+    )
+
+    with pytest.raises(StopIteration):
+        next(gen)

From e4e940e58d926135df0440c2035dacab8d54a280 Mon Sep 17 00:00:00 2001
From: devivasudevan <49675305+devivasudevan@users.noreply.github.com>
Date: Tue, 29 Sep 2026 19:50:55 +0200
Subject: [PATCH 4/9] fix(sweeper): repeat non-blocking drain until empty, not
 once

The previous fix drained one extra time after the first wait() to catch
a sibling that finished while the caller's on_candidate callback was
running. With three or more active futures that's not enough: a second
sibling can finish while the caller is still handling the first
extra-drained yield. Loop the non-blocking poll-and-drain until a poll
comes back empty, rather than doing it once, before checking pressure
or the deadline.

Also fixes the existing regression test's fake wait()/Future timing,
which asserted a callback-pause timeline while pre-resolving both
futures up front and contradicting Future.done() semantics in the fake
wait(). fake_wait now checks .done() directly, and each future's result
is set at the point in time the test claims it completes. Adds a
three-candidate regression case for the multi-sibling drain gap above.

Addresses coderabbitai's follow-up review on #319.

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 .../src/aisimulate/resource_scheduler.py      |  20 +++-
 tests/test_resource_scheduler.py              | 110 ++++++++++++++----
 2 files changed, 101 insertions(+), 29 deletions(-)

diff --git a/python/aisimulate/src/aisimulate/resource_scheduler.py b/python/aisimulate/src/aisimulate/resource_scheduler.py
index 8021f57ece..cc0f28c2f3 100644
--- a/python/aisimulate/src/aisimulate/resource_scheduler.py
+++ b/python/aisimulate/src/aisimulate/resource_scheduler.py
@@ -107,12 +107,20 @@ def _drain(done_futures):
                     break
                 # Each yield above can suspend this generator for an arbitrary
                 # duration -- the caller's on_candidate callback runs while
-                # paused there. Drain anything that finished during that pause
-                # before judging the deadline below, so a sibling that
-                # completed while the callback was running is reported as its
-                # real outcome instead of being misclassified as timed out.
-                still_done, _ = wait(active, timeout=0, return_when=FIRST_COMPLETED)
-                yield from _drain(still_done)
+                # paused there. A sibling can finish during ANY such pause,
+                # not just the one right after the first wait -- with three or
+                # more active futures, a later sibling can complete while the
+                # caller is still handling an earlier yield from this same
+                # non-blocking drain. So keep polling and draining until a
+                # poll comes back empty, rather than doing it once, before
+                # judging pressure or the deadline below.
+                while True:
+                    still_done, _ = wait(active, timeout=0, return_when=FIRST_COMPLETED)
+                    if not still_done:
+                        break
+                    yield from _drain(still_done)
+                    if interrupted or not active:
+                        break
                 if interrupted:
                     break
                 if not active:
diff --git a/tests/test_resource_scheduler.py b/tests/test_resource_scheduler.py
index 00168e9136..87202a02eb 100644
--- a/tests/test_resource_scheduler.py
+++ b/tests/test_resource_scheduler.py
@@ -271,8 +271,10 @@ def test_sibling_finished_during_a_yield_pause_is_not_reported_timed_out(monkeyp
     sibling future completes during that pause, it must still be reported as
     its real outcome, not misclassified as timed out just because the
     deadline was crossed while the callback for an earlier candidate was
-    running. Uses a controlled fake ``wait()``/clock, not a real sleeping
-    callback, so the regression is deterministic rather than timing-sensitive.
+    running. ``fake_wait`` mirrors real ``wait()`` semantics by checking
+    ``Future.done()`` directly rather than hardcoding which call returns
+    what, so the fakes can't quietly assert a timeline a real ``Future``
+    couldn't produce.
     """
     from concurrent.futures import Future
     from pathlib import Path
@@ -280,9 +282,7 @@ def test_sibling_finished_during_a_yield_pause_is_not_reported_timed_out(monkeyp
     from aisimulate import resource_scheduler as scheduler
 
     future0 = Future()
-    future0.set_result({"id": 0})
     future1 = Future()
-    future1.set_result({"id": 1})
 
     class Pool:
         def __init__(self, **kwargs):
@@ -292,25 +292,11 @@ def __init__(self, **kwargs):
         def submit(self, evaluate, spec):
             return future0 if spec["id"] == 0 else future1
 
-    elapsed = [0.0]
-    calls = [0]
-
     def fake_wait(futures, *, timeout, return_when):
-        calls[0] += 1
-        if calls[0] == 1:
-            # Initial poll: only candidate 0 has finished so far.
-            return {future0}, {future1}
-        if calls[0] == 2:
-            # The drain called immediately after resuming from the yield for
-            # candidate 0 -- this is the regression target. Candidate 1
-            # finished while the (simulated) on_candidate callback for
-            # candidate 0 was running, represented below by advancing the
-            # fake clock past the deadline before this call happens.
-            return {future1}, set()
-        raise AssertionError(
-            "no further wait() calls expected once both futures are drained"
-        )
+        done = {future for future in futures if future.done()}
+        return done, set(futures) - done
 
+    elapsed = [0.0]
     factory = _BudgetFactory()
     monkeypatch.setattr(scheduler, "ProcessPoolExecutor", Pool)
     monkeypatch.setattr(scheduler, "close_pool", lambda pool: None)
@@ -327,10 +313,13 @@ def fake_wait(futures, *, timeout, return_when):
         timeout=1.0,
     )
 
+    future0.set_result({"id": 0})
     assert next(gen) == (0, {"id": 0})
 
-    # Simulate on_candidate for candidate 0 running long enough to cross the
-    # deadline -- candidate 1 has already finished by the time it returns.
+    # Simulate candidate 1 finishing, and the deadline being crossed, while
+    # the (simulated) on_candidate callback for candidate 0 is running --
+    # i.e. during the pause at the yield above.
+    future1.set_result({"id": 1})
     elapsed[0] = 1.5
 
     assert next(gen) == (1, {"id": 1}), (
@@ -342,3 +331,78 @@ def fake_wait(futures, *, timeout, return_when):
 
     with pytest.raises(StopIteration):
         next(gen)
+
+
+def test_multiple_siblings_finishing_across_successive_yield_pauses_are_drained(
+    monkeypatch,
+):
+    """Follow-up regression: a single extra non-blocking drain after the first
+    wait only catches one sibling completing during one pause. With three or
+    more active futures, a second sibling can finish while the caller is
+    still handling the *first* extra-drained yield, so the drain has to
+    repeat until a poll comes back empty rather than run once and move on to
+    checking pressure/the deadline. Candidate 1's completion here is
+    discovered during the pause after candidate 0's yield, and candidate 2's
+    is discovered during the pause after candidate 1's yield -- both before
+    evaluate_waves ever reaches the deadline check, even though the clock has
+    already crossed it by the time candidate 2 is drained.
+    """
+    from concurrent.futures import Future
+    from pathlib import Path
+
+    from aisimulate import resource_scheduler as scheduler
+
+    future0, future1, future2 = Future(), Future(), Future()
+
+    class Pool:
+        def __init__(self, **kwargs):
+            self._processes = {123: object()}
+            Path(kwargs["initargs"][2], "123").touch()
+
+        def submit(self, evaluate, spec):
+            return {0: future0, 1: future1, 2: future2}[spec["id"]]
+
+    def fake_wait(futures, *, timeout, return_when):
+        done = {future for future in futures if future.done()}
+        return done, set(futures) - done
+
+    elapsed = [0.0]
+    factory = _BudgetFactory(capacity=3)
+    monkeypatch.setattr(scheduler, "ProcessPoolExecutor", Pool)
+    monkeypatch.setattr(scheduler, "close_pool", lambda pool: None)
+    monkeypatch.setattr(scheduler, "terminate_pool", lambda pool: None)
+    monkeypatch.setattr(scheduler, "wait", fake_wait)
+    monkeypatch.setattr(scheduler.time, "monotonic", lambda: elapsed[0])
+    specs = [{"id": i, "cost": 1} for i in range(3)]
+    gen = evaluate_waves(
+        specs,
+        factory=factory,
+        initializer=_init,
+        evaluate=_evaluate,
+        workers=3,
+        timeout=1.0,
+    )
+
+    future0.set_result({"id": 0})
+    assert next(gen) == (0, {"id": 0})
+
+    # Candidate 1 finishes during the pause for candidate 0's yield.
+    future1.set_result({"id": 1})
+    assert next(gen) == (1, {"id": 1})
+
+    # Candidate 2 finishes during the pause for candidate 1's yield -- the
+    # *second* extra pause -- with the deadline already crossed by the time
+    # evaluate_waves resumes. Draining it here, without the test making an
+    # extra explicit call, is exactly what the repeat-until-empty loop is for.
+    future2.set_result({"id": 2})
+    elapsed[0] = 1.5
+
+    assert next(gen) == (2, {"id": 2}), (
+        "candidate 2 finished during the pause for candidate 1's yield and "
+        "must be drained in the same non-blocking pass, not reported as "
+        "timed out because the deadline was already crossed by the time "
+        "evaluate_waves got back to checking it"
+    )
+
+    with pytest.raises(StopIteration):
+        next(gen)

From ece264bda6a8d07d7d2116f4a2c36b7a68bd39ff Mon Sep 17 00:00:00 2001
From: "Dr. Stefan Schimanski" <stefan.schimanski@gmail.com>
Date: Wed, 30 Sep 2026 14:25:21 +0200
Subject: [PATCH 5/9] feat(cli): expose live recommendation callbacks

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 .github/codeowners/areas.yaml                 |   3 +
 CODEOWNERS                                    |   2 +
 docs/cli/user-guide.md                        |   3 +-
 docs/sweeper/architecture.md                  |  23 ++--
 python/aisimulate/src/aisimulate/main.py      |   1 +
 .../src/aisimulate/output_adapter.py          |  87 +++++++++++++-
 python/aisimulate/src/aisimulate/recommend.py |  13 +-
 tests/test_host_resources.py                  |  38 +++++-
 tests/test_output_adapter_discovery.py        | 111 ++++++++++++++++++
 tests/test_unified_cli.py                     |  12 +-
 10 files changed, 267 insertions(+), 26 deletions(-)

diff --git a/.github/codeowners/areas.yaml b/.github/codeowners/areas.yaml
index 6c3b33bb01..8c86328580 100644
--- a/.github/codeowners/areas.yaml
+++ b/.github/codeowners/areas.yaml
@@ -52,6 +52,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
@@ -184,6 +185,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 274892cd0a..a864b9d3a8 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -73,6 +73,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
@@ -194,6 +195,7 @@
 /python/aisimulate/.github/codeowners/                   @ai-dynamo/aisimulate-infra-codeowners @ai-dynamo/access-aisimulate-maintain
 /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_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 61fb613be6..4d30b65f45 100644
--- a/docs/cli/user-guide.md
+++ b/docs/cli/user-guide.md
@@ -85,7 +85,8 @@ warmup fields and retains both worker pools across the barrier. See
 | `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 and writes additional artifacts after recommendation.
+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.
diff --git a/docs/sweeper/architecture.md b/docs/sweeper/architecture.md
index 49e215d67f..09bbbc8978 100644
--- a/docs/sweeper/architecture.md
+++ b/docs/sweeper/architecture.md
@@ -52,25 +52,30 @@ 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
+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, 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.
+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 operation:
+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. 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.
+directory and exist after the call. A final `write()` failure leaves canonical recommendation output
+intact and makes the command fail. `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.
 
diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py
index 878d1cb64e..c2d9ec1b7d 100644
--- a/python/aisimulate/src/aisimulate/main.py
+++ b/python/aisimulate/src/aisimulate/main.py
@@ -246,6 +246,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]]] = []
diff --git a/python/aisimulate/src/aisimulate/output_adapter.py b/python/aisimulate/src/aisimulate/output_adapter.py
index 9e21692d38..0081ab62f2 100644
--- a/python/aisimulate/src/aisimulate/output_adapter.py
+++ b/python/aisimulate/src/aisimulate/output_adapter.py
@@ -1,24 +1,35 @@
 # 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."""
+"""Discovery and invocation for recommendation output adapters."""
 
 from __future__ import annotations
 
 import importlib.metadata
-from collections.abc import Iterable, Mapping, Sequence
+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.result import SweepResult
+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):
-    """Write additional artifacts for a completed recommendation."""
+    """Observe a recommendation and write additional artifacts when it completes."""
 
     name: str
     api_version: int
@@ -38,7 +49,7 @@ class OutputAdapterResolutionError(RuntimeError):
 
 
 class OutputAdapterExecutionError(RuntimeError):
-    """A selected recommendation output adapter failed to write its artifacts."""
+    """A selected recommendation output adapter failed during execution."""
 
 
 def validate_output_adapter(adapter: Any, *, requested_name: str) -> RecommendationOutputAdapter:
@@ -56,6 +67,11 @@ def validate_output_adapter(adapter: Any, *, requested_name: str) -> Recommendat
         )
     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
 
 
@@ -100,6 +116,65 @@ def resolve_output_adapters(
     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]],
@@ -137,7 +212,9 @@ def write_output_adapters(
     "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 ff902df5c1..1761467381 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,
 ) -> SweepResult:
@@ -47,6 +49,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,
     )
@@ -62,6 +65,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,
 ) -> SweepResult:
@@ -109,7 +114,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/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
index 0016914296..759cc2f67e 100644
--- a/tests/test_output_adapter_discovery.py
+++ b/tests/test_output_adapter_discovery.py
@@ -12,9 +12,13 @@
     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:
@@ -102,3 +106,110 @@ def test_output_adapter_writes_into_supplied_directory(tmp_path) -> None:
 
     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_unified_cli.py b/tests/test_unified_cli.py
index 32907f0737..36b778be1b 100644
--- a/tests/test_unified_cli.py
+++ b/tests/test_unified_cli.py
@@ -1000,6 +1000,7 @@ def test_recommendation_invokes_selected_output_adapter(tmp_path, monkeypatch, c
     )
     result = _RecommendationResult([candidate])
     selected_stack = []
+    recommendation_kwargs = {}
     adapter = _DGDOutputAdapter()
 
     def resolve_stack(name):
@@ -1016,10 +1017,12 @@ def resolve_stack(name):
             entry_points=[],
         ),
     )
-    monkeypatch.setattr(
-        "aisimulate.recommend.run_recommendation",
-        lambda *args, **kwargs: result,
-    )
+    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 (
@@ -1051,6 +1054,7 @@ def resolve_stack(name):
         "renderer": "aic",
         "format": "manifest",
     }
+    assert recommendation_kwargs["output_configs"] == {"dgd": output_config}
     assert received_result is result
     assert received_dir == output
     assert (output / "qwen.yaml").read_text() == "kind: DynamoGraphDeployment\n"

From 3479fe7bfb09a1e6a147206fd245a2715838740e Mon Sep 17 00:00:00 2001
From: "Dr. Stefan Schimanski" <stefan.schimanski@gmail.com>
Date: Wed, 30 Sep 2026 14:32:29 +0200
Subject: [PATCH 6/9] style: format changed test files

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 tests/test_resource_scheduler.py | 12 +++---------
 tests/test_unified_cli.py        |  1 +
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/tests/test_resource_scheduler.py b/tests/test_resource_scheduler.py
index 87202a02eb..cabf9219d3 100644
--- a/tests/test_resource_scheduler.py
+++ b/tests/test_resource_scheduler.py
@@ -71,9 +71,7 @@ def test_checkpoint_failure_stops_before_workers_without_reclassifying_candidate
     from aisimulate import resource_scheduler as scheduler
 
     factory = _BudgetFactory()
-    failure = ResourceLimitError(
-        "execution evidence exceeds the bounded checkpoint budget"
-    )
+    failure = ResourceLimitError("execution evidence exceeds the bounded checkpoint budget")
 
     def fail_checkpoint(event, plan):
         raise failure
@@ -127,9 +125,7 @@ def estimate_host_resources(self, workload, *, concurrency=None):
 def test_whole_wave_reserves_sum_against_one_live_snapshot(monkeypatch):
     from aisimulate import resources
 
-    monkeypatch.setattr(
-        resources, "discover_host", lambda: HostResources(16 * GB, 9 * GB, 8)
-    )
+    monkeypatch.setattr(resources, "discover_host", lambda: HostResources(16 * GB, 9 * GB, 8))
     factory = GuardedRunnerFactory(_EstimateFactory(), "custom", ResourceConfig())
     spec = SimpleNamespace(workload={}, concurrency=1)
     assert factory.admit_wave([spec])["status"] == "admitted"
@@ -137,9 +133,7 @@ def test_whole_wave_reserves_sum_against_one_live_snapshot(monkeypatch):
         factory.admit_wave([spec, spec])
     assert caught.value.plan["required_bytes"] == 12 * GB
     assert caught.value.plan["available_bytes"] < 8 * GB
-    monkeypatch.setattr(
-        resources, "discover_host", lambda: HostResources(16 * GB, 6 * GB, 8)
-    )
+    monkeypatch.setattr(resources, "discover_host", lambda: HostResources(16 * GB, 6 * GB, 8))
     with pytest.raises(ResourceLimitError):
         factory.admit_wave([spec])
 
diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py
index 36b778be1b..6819868559 100644
--- a/tests/test_unified_cli.py
+++ b/tests/test_unified_cli.py
@@ -1017,6 +1017,7 @@ def resolve_stack(name):
             entry_points=[],
         ),
     )
+
     def run_recommendation(*args, **kwargs):
         del args
         recommendation_kwargs.update(kwargs)

From 26c99f4bff0ec93b7f395dca8bb1ed9e33ef09ba Mon Sep 17 00:00:00 2001
From: "Dr. Stefan Schimanski" <stefan.schimanski@gmail.com>
Date: Wed, 30 Sep 2026 18:27:27 +0200
Subject: [PATCH 7/9] fix(cli): validate output adapter boundaries

Signed-off-by: Dr. Stefan Schimanski <stefan.schimanski@gmail.com>
---
 docs/sweeper/architecture.md                  |  13 +-
 python/aisimulate/src/aisimulate/main.py      |  13 +-
 .../aisimulate/src/aisimulate/supervision.py  |   6 +-
 tests/test_unified_cli.py                     | 115 ++++++++++++++++++
 4 files changed, 137 insertions(+), 10 deletions(-)

diff --git a/docs/sweeper/architecture.md b/docs/sweeper/architecture.md
index 09bbbc8978..abbb022529 100644
--- a/docs/sweeper/architecture.md
+++ b/docs/sweeper/architecture.md
@@ -70,12 +70,13 @@ 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 final `write()` failure leaves canonical recommendation output
-intact and makes the command fail. `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()`.
+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.
 
diff --git a/python/aisimulate/src/aisimulate/main.py b/python/aisimulate/src/aisimulate/main.py
index c2d9ec1b7d..05b750d0d1 100644
--- a/python/aisimulate/src/aisimulate/main.py
+++ b/python/aisimulate/src/aisimulate/main.py
@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import argparse
+import importlib.metadata
 import json
 import sys
 from collections.abc import Sequence
@@ -21,8 +22,9 @@
     CoreRecommendationConfig,
     prediction_mapping,
 )
-from .config.common import split_config_sections
+from .config.common import RECOMMENDATION_CORE_SECTIONS, split_config_sections
 from .config_adapter import (
+    CONFIG_ADAPTER_ENTRY_POINT_GROUP,
     ConfigAdapterResolutionError,
     PredictionAdapterContext,
     SimulationConfigAdapter,
@@ -97,15 +99,20 @@ def _compile_prediction_adapters(
 
 
 def _extract_output_configs(
-    raw: dict[str, Any], outputs: Sequence[str]
+    raw: dict[str, Any], outputs: Sequence[str], *, stack: 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]] = {}
+    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")
@@ -235,7 +242,7 @@ 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", []))
+    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)
diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py
index 682f86153a..2414544029 100644
--- a/python/aisimulate/src/aisimulate/supervision.py
+++ b/python/aisimulate/src/aisimulate/supervision.py
@@ -367,7 +367,11 @@ 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":
+                for name in dict.fromkeys(args.outputs):
+                    validation_raw.pop(name, None)
+            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_unified_cli.py b/tests/test_unified_cli.py
index 6819868559..10c2bc6437 100644
--- a/tests/test_unified_cli.py
+++ b/tests/test_unified_cli.py
@@ -237,6 +237,12 @@ def write(self, config, *, result, output_dir):
         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(
@@ -1063,6 +1069,115 @@ def run_recommendation(*args, **kwargs):
     assert json.loads(capsys.readouterr().out)[0]["score"] == 1.0
 
 
+def test_output_adapter_failure_preserves_canonical_recommendation_files(tmp_path, monkeypatch, capsys) -> 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": "qwen"},
+            }
+        )
+    )
+    result = _RecommendationResult(
+        [
+            Candidate(
+                config={"backend": "vllm"},
+                used_gpus=1,
+                score=1.0,
+                metrics={},
+                prediction_config=prediction,
+            )
+        ]
+    )
+    adapter = _MissingArtifactOutputAdapter()
+    monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner()))
+    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",
+                "--config",
+                str(config_path),
+                "--output",
+                "dgd",
+                "--output-dir",
+                str(output),
+            ]
+        )
+        == 1
+    )
+
+    assert (output / "recommendation.json").is_file()
+    assert (output / "recommendation.csv").is_file()
+    assert (output / "recommendations" / "0001.yaml").is_file()
+    assert "reported missing artifact" in capsys.readouterr().err
+
+
+@pytest.mark.parametrize(
+    ("name", "config_adapter_names"),
+    [
+        ("traffic", []),
+        ("placement", ["engine.placement"]),
+    ],
+)
+def test_output_adapter_name_must_not_collide_with_input_section(
+    tmp_path, monkeypatch, capsys, name, config_adapter_names
+) -> 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": {}},
+                },
+                name: {},
+            }
+        )
+    )
+    entry_points = SimpleNamespace(
+        select=lambda **kwargs: [SimpleNamespace(name=entry_name) for entry_name in config_adapter_names]
+    )
+    monkeypatch.setattr(cli.importlib.metadata, "entry_points", lambda: entry_points)
+    monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner()))
+
+    with pytest.raises(SystemExit, match="2"):
+        cli.main(
+            [
+                "recommend",
+                "--config",
+                str(config_path),
+                "--output",
+                name,
+            ]
+        )
+
+    assert f"output adapter name '{name}' collides" in capsys.readouterr().err
+
+
 def test_selected_output_requires_matching_configuration_section(tmp_path, monkeypatch, capsys) -> None:
     config_path = tmp_path / "recommend.yaml"
     config_path.write_text(

From 1f9cfd66454799d950c2f7731635673b81e5c207 Mon Sep 17 00:00:00 2001
From: hongkuanz <hongkuanz@nvidia.com>
Date: Wed, 30 Sep 2026 14:07:41 -0700
Subject: [PATCH 8/9] fix: validate output sections before overwrite cleanup

Share lightweight output-section validation between the supervisor and CLI worker so invalid names and configuration are rejected before existing results are removed. Cover core and installed-adapter collisions, invalid sections, and the runtime import boundary.

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
---
 python/aisimulate/src/aisimulate/cli_args.py  | 31 +++++++
 .../src/aisimulate/config/common.py           |  1 +
 .../src/aisimulate/config_adapter.py          |  2 +-
 python/aisimulate/src/aisimulate/main.py      | 30 +------
 .../aisimulate/src/aisimulate/supervision.py  |  5 +-
 tests/test_resource_supervision.py            | 88 +++++++++++++++++++
 tests/test_unified_cli.py                     |  2 +-
 7 files changed, 126 insertions(+), 33 deletions(-)

diff --git a/python/aisimulate/src/aisimulate/cli_args.py b/python/aisimulate/src/aisimulate/cli_args.py
index a9b4331c4e..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
 
@@ -101,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 1d9fa8df11..7817d87847 100644
--- a/python/aisimulate/src/aisimulate/config/common.py
+++ b/python/aisimulate/src/aisimulate/config/common.py
@@ -210,6 +210,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 05b750d0d1..21391b1256 100644
--- a/python/aisimulate/src/aisimulate/main.py
+++ b/python/aisimulate/src/aisimulate/main.py
@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import argparse
-import importlib.metadata
 import json
 import sys
 from collections.abc import Sequence
@@ -15,16 +14,15 @@
 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,
     CoreRecommendationConfig,
     prediction_mapping,
 )
-from .config.common import RECOMMENDATION_CORE_SECTIONS, split_config_sections
+from .config.common import split_config_sections
 from .config_adapter import (
-    CONFIG_ADAPTER_ENTRY_POINT_GROUP,
     ConfigAdapterResolutionError,
     PredictionAdapterContext,
     SimulationConfigAdapter,
@@ -98,30 +96,6 @@ def _compile_prediction_adapters(
     return compiled
 
 
-def _extract_output_configs(
-    raw: dict[str, Any], outputs: Sequence[str], *, stack: 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]] = {}
-    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 _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)
diff --git a/python/aisimulate/src/aisimulate/supervision.py b/python/aisimulate/src/aisimulate/supervision.py
index 2414544029..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)
@@ -369,8 +369,7 @@ def main(argv: Sequence[str] | None = None) -> int:
 
             validation_raw = dict(raw)
             if args.command == "recommend":
-                for name in dict.fromkeys(args.outputs):
-                    validation_raw.pop(name, None)
+                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)
diff --git a/tests/test_resource_supervision.py b/tests/test_resource_supervision.py
index 9dea4e6c6f..8b0c3ca6c4 100644
--- a/tests/test_resource_supervision.py
+++ b/tests/test_resource_supervision.py
@@ -147,6 +147,93 @@ 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",
+    [
+        *[
+            (name, {}, "collides")
+            for name in (
+                "engine",
+                "traffic",
+                "evaluation",
+                "execution",
+                "optimization",
+                "optimizer",
+            )
+        ],
+        ("placement", {}, "collides"),
+        ("bad.name", {}, "invalid --output name"),
+        ("", {}, "invalid --output name"),
+        ("missing", {}, "requires a top-level"),
+        ("artifact", None, "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" and name not in raw:
+        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")
@@ -518,6 +605,7 @@ def test_supervisor_argument_and_output_setup_do_not_import_runtime():
             "-c",
             (
                 "import sys; import aisimulate.supervision, aisimulate.cli_args, aisimulate.output; "
+                "aisimulate.cli_args._extract_output_configs({'artifact': {}}, ['artifact'], stack='engine'); "
                 "assert not any(name in sys.modules for name in "
                 "('aisimulate._runtime', 'aisimulate.sweeper', 'numpy', 'pandas', 'aisimulate_core'))"
             ),
diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py
index 10c2bc6437..78c8325b0b 100644
--- a/tests/test_unified_cli.py
+++ b/tests/test_unified_cli.py
@@ -1161,7 +1161,7 @@ def test_output_adapter_name_must_not_collide_with_input_section(
     entry_points = SimpleNamespace(
         select=lambda **kwargs: [SimpleNamespace(name=entry_name) for entry_name in config_adapter_names]
     )
-    monkeypatch.setattr(cli.importlib.metadata, "entry_points", lambda: entry_points)
+    monkeypatch.setattr("importlib.metadata.entry_points", lambda: entry_points)
     monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner()))
 
     with pytest.raises(SystemExit, match="2"):

From e1e0ebde65b4a58061b6f0dd1ab684884177f596 Mon Sep 17 00:00:00 2001
From: hongkuanz <hongkuanz@nvidia.com>
Date: Wed, 30 Sep 2026 14:19:34 -0700
Subject: [PATCH 9/9] test: streamline output adapter regression coverage

Share CLI setup between successful and failed artifact writers, remove validation cases already covered by supervised subprocess tests, and retain representative invalid-input cases that prove existing outputs survive.

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
---
 tests/test_resource_supervision.py |  16 +--
 tests/test_unified_cli.py          | 191 +++++------------------------
 2 files changed, 30 insertions(+), 177 deletions(-)

diff --git a/tests/test_resource_supervision.py b/tests/test_resource_supervision.py
index 8b0c3ca6c4..27d5224838 100644
--- a/tests/test_resource_supervision.py
+++ b/tests/test_resource_supervision.py
@@ -150,22 +150,10 @@ def test_invalid_cli_input_preserves_outputs_and_stack_error_order(tmp_path, fai
 @pytest.mark.parametrize(
     "name,section,error",
     [
-        *[
-            (name, {}, "collides")
-            for name in (
-                "engine",
-                "traffic",
-                "evaluation",
-                "execution",
-                "optimization",
-                "optimizer",
-            )
-        ],
+        ("traffic", {}, "collides"),
         ("placement", {}, "collides"),
         ("bad.name", {}, "invalid --output name"),
-        ("", {}, "invalid --output name"),
         ("missing", {}, "requires a top-level"),
-        ("artifact", None, "requires a top-level"),
         ("artifact", [], "must be a mapping"),
     ],
 )
@@ -186,7 +174,7 @@ def test_recommend_output_validation_preserves_outputs(tmp_path, name, section,
             "constraints": {"max_candidate_gpus": 1},
         },
     }
-    if name != "missing" and name not in raw:
+    if name != "missing":
         raw[name] = section
     config = tmp_path / "config.yaml"
     config.write_text(yaml.safe_dump(raw))
diff --git a/tests/test_unified_cli.py b/tests/test_unified_cli.py
index 78c8325b0b..e522b85556 100644
--- a/tests/test_unified_cli.py
+++ b/tests/test_unified_cli.py
@@ -970,7 +970,8 @@ def test_recommendation_outputs_each_concrete_prediction_once(tmp_path, monkeypa
 
 
 @pytest.mark.parametrize("stack", ["engine", "dynamo"])
-def test_recommendation_invokes_selected_output_adapter(tmp_path, monkeypatch, capsys, stack) -> None:
+@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",
@@ -1007,7 +1008,7 @@ def test_recommendation_invokes_selected_output_adapter(tmp_path, monkeypatch, c
     result = _RecommendationResult([candidate])
     selected_stack = []
     recommendation_kwargs = {}
-    adapter = _DGDOutputAdapter()
+    adapter = _MissingArtifactOutputAdapter() if missing_artifact else _DGDOutputAdapter()
 
     def resolve_stack(name):
         selected_stack.append(name)
@@ -1032,27 +1033,25 @@ def run_recommendation(*args, **kwargs):
     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",
-            ]
-        )
-        == 0
-    )
+    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]
@@ -1064,148 +1063,14 @@ def run_recommendation(*args, **kwargs):
     assert recommendation_kwargs["output_configs"] == {"dgd": output_config}
     assert received_result is result
     assert received_dir == output
-    assert (output / "qwen.yaml").read_text() == "kind: DynamoGraphDeployment\n"
+    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())
-    assert json.loads(capsys.readouterr().out)[0]["score"] == 1.0
-
-
-def test_output_adapter_failure_preserves_canonical_recommendation_files(tmp_path, monkeypatch, capsys) -> 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": "qwen"},
-            }
-        )
-    )
-    result = _RecommendationResult(
-        [
-            Candidate(
-                config={"backend": "vllm"},
-                used_gpus=1,
-                score=1.0,
-                metrics={},
-                prediction_config=prediction,
-            )
-        ]
-    )
-    adapter = _MissingArtifactOutputAdapter()
-    monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner()))
-    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",
-                "--config",
-                str(config_path),
-                "--output",
-                "dgd",
-                "--output-dir",
-                str(output),
-            ]
-        )
-        == 1
-    )
-
-    assert (output / "recommendation.json").is_file()
-    assert (output / "recommendation.csv").is_file()
-    assert (output / "recommendations" / "0001.yaml").is_file()
-    assert "reported missing artifact" in capsys.readouterr().err
-
-
-@pytest.mark.parametrize(
-    ("name", "config_adapter_names"),
-    [
-        ("traffic", []),
-        ("placement", ["engine.placement"]),
-    ],
-)
-def test_output_adapter_name_must_not_collide_with_input_section(
-    tmp_path, monkeypatch, capsys, name, config_adapter_names
-) -> 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": {}},
-                },
-                name: {},
-            }
-        )
-    )
-    entry_points = SimpleNamespace(
-        select=lambda **kwargs: [SimpleNamespace(name=entry_name) for entry_name in config_adapter_names]
-    )
-    monkeypatch.setattr("importlib.metadata.entry_points", lambda: entry_points)
-    monkeypatch.setattr(cli, "resolve_runner_factory", lambda stack: _Factory(_Runner()))
-
-    with pytest.raises(SystemExit, match="2"):
-        cli.main(
-            [
-                "recommend",
-                "--config",
-                str(config_path),
-                "--output",
-                name,
-            ]
-        )
-
-    assert f"output adapter name '{name}' collides" in capsys.readouterr().err
-
-
-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
+    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: