From 18bba1a9999f788fedcc7ced2ba380a71309c142 Mon Sep 17 00:00:00 2001 From: Severin Klingler Date: Mon, 3 Aug 2026 11:13:59 +0200 Subject: [PATCH 01/59] refactor(experimentalist): a runner and a context, so the loop is just a strategy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_run` held everything: input resolution, all backend I/O, the entity lifecycle, the round loop, and PR publishing. There was no seam a strategy could be swapped at, because the runner and the sole strategy were the same code. Split them. `runner.py` is the composition root — it prepares inputs, runs one strategy, then persists and publishes. It is the only code that holds a Backend. `context.py` is everything a strategy may reach: its datasets, its candidates, and the verbs for measuring and recording them. `deps.py` is gone; config is a constructor argument. Four things move out of the loop and into the host, because they were never the strategy's to decide: - The insight/no-insight branch, with Eval Author now behind a lazy import: a run with no Insight never authors a suite and must not fail to import without the package. - Run resume. The runner re-opens the ExperimentRun; the strategy rebuilds itself from `ctx.candidates()`. A strategy that declares `supports_resume = False` is refused loudly rather than silently restarted — these runs cost hours, so the silent restart is the expensive failure. - The winner's copy-out. Its skip list mixed three owners (backend metadata, strategy documentation, evaluator scaffolding), so it is now composed from three named sets instead of one literal that would strip a third-party strategy's real output. - The Insight-suite report sections, which are a reward channel's epilogue. `report_progress(completed, total, unit, note)` replaces `rounds_completed` end to end. Not every strategy has rounds — DSPy's compile() is one opaque call — and the ones that most need reporting are exactly the ones that cannot produce a fraction. A counter is always producible, so `ExperimentRun` now carries one plus its unit, and a consumer renders a bar only when a total is actually known. Tests get the shared doubles the suite conspicuously lacked: an in-memory backend, a fake evaluator, and a context factory, in tests/experimentalist/. Signed-off-by: Severin Klingler --- .../nemo_experimentalist_plugin/entities.py | 20 +- .../experimentalist/components/loop.py | 1022 ++++------------- .../experimentalist/context.py | 301 +++++ .../experimentalist/deps.py | 71 -- .../experimentalist/experiment_mirror.py | 10 +- .../experimentalist_backend.py | 2 +- .../experimentalist/result.py | 4 +- .../experimentalist/run.py | 26 +- .../experimentalist/runner.py | 441 +++++++ .../tests/experimentalist/conftest.py | 28 + .../tests/experimentalist/doubles.py | 181 +++ ...loop.py => test_dataset_staging_runner.py} | 61 +- .../experimentalist/test_loop_failure.py | 129 --- .../test_loop_insight_suite.py | 136 +-- .../experimentalist/test_loop_reporting.py | 164 +-- .../tests/experimentalist/test_reporting.py | 48 +- .../tests/experimentalist/test_runner.py | 453 ++++++++ .../nemo-experimentalist/tests/test_deps.py | 55 - .../tests/test_experiment_mirror_mapping.py | 12 +- .../tests/test_experimentalist_backend.py | 4 +- .../tests/test_experimentalist_run.py | 116 +- .../tests/test_local_backend_projection.py | 4 +- 22 files changed, 1848 insertions(+), 1440 deletions(-) create mode 100644 plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/context.py delete mode 100644 plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/deps.py create mode 100644 plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/runner.py create mode 100644 plugins/nemo-experimentalist/tests/experimentalist/conftest.py create mode 100644 plugins/nemo-experimentalist/tests/experimentalist/doubles.py rename plugins/nemo-experimentalist/tests/experimentalist/{test_dataset_staging_loop.py => test_dataset_staging_runner.py} (68%) delete mode 100644 plugins/nemo-experimentalist/tests/experimentalist/test_loop_failure.py create mode 100644 plugins/nemo-experimentalist/tests/experimentalist/test_runner.py delete mode 100644 plugins/nemo-experimentalist/tests/test_deps.py diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py index 45c3cb3e26..0bba13206a 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py @@ -377,9 +377,25 @@ class ExperimentRun(NemoEntity, entity_type="experiment_run"): default="running", description="Lifecycle status of this optimization run.", ) - rounds_completed: int = Field( + progress_completed: int = Field( default=0, - description="Number of full optimization rounds completed so far.", + description="Units of work the strategy reports finished so far. Display only.", + ) + progress_total: int | None = Field( + default=None, + description=( + "Units of work expected in total, when the strategy can say. None means it " + "cannot — an opaque strategy has no honest denominator, so consumers show a " + "counter rather than a bar." + ), + ) + progress_unit: str = Field( + default="step", + description="What one unit of progress is: 'round' for the evolutionary loop, 'trial' for a search.", + ) + progress_note: str | None = Field( + default=None, + description="What the strategy is currently doing, for strategies with no meaningful total.", ) winner_agent: str | None = Field( default=None, diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py index 2e111dc889..e80c54756c 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/components/loop.py @@ -16,22 +16,18 @@ import shutil from collections import defaultdict from pathlib import Path -from typing import Any, Literal, cast, get_args +from typing import Any, ClassVar, Literal, cast, get_args -from nemo_eval_author_plugin.eval_author.agent import EvalAuthor from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig from nemo_experimentalist_plugin.entities import ( Candidate, Dataset, EvaluationResult, - ExperimentRun, + RewardRecord, TrialResult, ) from nemo_experimentalist_plugin.experimentalist.components.analyzer import AgentAnalyzer from nemo_experimentalist_plugin.experimentalist.components.coder import Coder, CoderConfig -from nemo_experimentalist_plugin.experimentalist.components.dataset_staging import stage_eval_author_inputs -from nemo_experimentalist_plugin.experimentalist.components.evaluator import Evaluator -from nemo_experimentalist_plugin.experimentalist.components.evaluator.factory import DatasetFactory, EvaluatorFactory from nemo_experimentalist_plugin.experimentalist.components.goal_tree import ( GoalTree, GoalTreeConfig, @@ -47,11 +43,8 @@ candidate_metric_keys, candidate_suite_identity, insight_suite_provenance, - select_insight_promotion_suggestions, stamp_insight_evaluation_result, validate_insight_evaluation_result, - write_insight_comparison_section, - write_insight_promotion_section, ) from nemo_experimentalist_plugin.experimentalist.components.model_config import ( get_fast_model, @@ -72,12 +65,7 @@ from nemo_experimentalist_plugin.experimentalist.components.trace_scorer import ( GroupLeafScorer, ) -from nemo_experimentalist_plugin.experimentalist.deps import ExperimentalistDeps -from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ( - ExperimentalistBackend, -) -from nemo_experimentalist_plugin.experimentalist.reporting import reward_scalar -from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult +from nemo_experimentalist_plugin.experimentalist.context import ExperimentContext from nemo_platform import AsyncNeMoPlatform from nooa import Agent, CodeActStrategy, strategy from nooa.agentdoc import doc, spec @@ -109,16 +97,6 @@ _EXCLUDE_GLOBS = {"*traces*", "*eval-and-optimize_*"} -def _warn_persistence_failure(operation: Literal["archive", "publish"], candidate: str, exc: Exception) -> None: - """Log best-effort persistence failure context.""" - logger.warning( - "[PERSISTENCE] %s failed for candidate %s; continuing: %s", - operation, - candidate, - exc, - ) - - def _ignore_patterns(directory: str, contents: list[str]) -> set[str]: ignored = set() for name in contents: @@ -328,6 +306,10 @@ class EvolutionaryOptimizer(Agent): rounds, mirroring the AAD ``EvolutionaryOptimizer``. """ + #: This loop resumes from its own round-analysis files plus ``ctx.candidates()``, + #: so the runner may re-open an existing run and hand it back. + supports_resume: ClassVar[bool] = True + def __init__( self, working_dir: Path, @@ -367,246 +349,68 @@ def _coder_config(config: EvolutionaryOptimizerConfig) -> CoderConfig: # Public entry point # ------------------------------------------------------------------ - async def run(self, deps: ExperimentalistDeps) -> ExperimentalistResult: + async def run(self, ctx: ExperimentContext) -> Candidate | None: """Run optimization and always close the owned shell session.""" try: - return await self._run(deps) + return await self._run(ctx) finally: await self.shell.close() - async def _run(self, deps: ExperimentalistDeps) -> ExperimentalistResult: + async def _run(self, ctx: ExperimentContext) -> Candidate | None: """Run the Pareto evolutionary optimization loop. Args: - deps: Per-run dependencies (workspace, insight_id, dataset, - backend, optional config override). + ctx: The run's context — its datasets, its candidates, and the verbs for + measuring and recording them. Nothing else is reachable from here. Returns: - An :class:`ExperimentalistResult` committed to the entity store - via ``backend.persist_result()``. + The winning Candidate, or None when no candidate was ever scored. The + runner turns that into the run's terminal result. """ - if deps.backend is None: - raise ValueError("deps.backend must be set before calling run()") - backend = deps.backend - workspace = deps.workspace - config = deps.config if deps.config is not None else self.config - reporter = getattr(deps, "reporter", None) - - # ---- Preflight: fail fast when persistence is enabled but git is missing. - if (config.storage.archive_candidates or config.storage.publish_winner) and shutil.which("git") is None: - raise ValueError( - "Candidate persistence is enabled (storage.archive_candidates/publish_winner) " - "but 'git' is not on PATH, so nothing can be persisted. Install git, or disable " - "storage to run without persistence." - ) - - # ---- Working directory structure --------------------------------- - agents_dir, analysis_dir, results_dir = self._init_structure() - evaluator_factory = EvaluatorFactory() - evaluator = evaluator_factory.build_evaluator( - deps.evaluator_type, - config.evaluator, - experiment_dir=self.working_dir, - ) - dataset_factory = DatasetFactory() - - train_dataset_ref = deps.train_dataset - validation_dataset_ref = deps.validation_dataset - task_template_ref = deps.task_template - insight_eval_dataset: Dataset | None = None - if deps.insight is not None: - if task_template_ref is None: - raise ValueError("Task template is required for insight trace analysis") - if backend.client is None: - raise ValueError("Platform client is required for insight task template loading") - staged_inputs = await stage_eval_author_inputs( - self.working_dir, - train_dataset=train_dataset_ref, - validation_dataset=validation_dataset_ref, - task_template=task_template_ref, - client=backend.client, - workspace=workspace, - ) - train_dataset_ref = staged_inputs.train_dataset - validation_dataset_ref = staged_inputs.validation_dataset - task_template_ref = staged_inputs.task_template - - # ---- Resolve datasets to evaluator-domain objects ----------------- - train_eval_dataset = dataset_factory.build_dataset( - deps.evaluator_type, - train_dataset_ref, - ) - - validation_eval_dataset = dataset_factory.build_dataset( - deps.evaluator_type, - validation_dataset_ref, - ) - - # ---- Resolve insight (Mode 1) vs local agent (Mode 2) ----------- - insight = ( - await backend.get_insight(workspace=workspace, insight_id=str(deps.insight)) - if deps.insight is not None - else None - ) - agent_ref: str | Path | None = deps.agent - if agent_ref is None and insight is not None: - agent_ref = insight.agent - if agent_ref is None: - raise ValueError("Insight or agent is required") - - agent_path = self.working_dir / "eval-and-optimize" / "source-agent" - await backend.get_agent_code( - workspace=workspace, - agent=agent_ref, - dest=agent_path, - clone_depth=config.source.clone_depth, - ) - agent_name = str(agent_ref) - - agent_spec_path: Path | None = None - if deps.agent_spec is not None: - agent_spec_path = await backend.get_agent_spec( - workspace=workspace, - spec=deps.agent_spec, - dest=self.working_dir / "AGENT-SPEC.md", - ) - - if insight is not None: - insight_ref: str = str(deps.insight) - # run the eval_author - if backend.client is None: - raise ValueError("Platform client is required for insight trace loading") - assert task_template_ref is not None - eval_author = EvalAuthor( - experiment_dir=self.working_dir, - config=config.eval_author, - reporter=reporter, - ) - eval_author_result = await eval_author.run( - insight=insight, - agent_path=agent_path, - task_template=dataset_factory.build_task_template(deps.evaluator_type, task_template_ref), - train_dataset=train_eval_dataset, - validation_dataset=validation_eval_dataset, - client=backend.client, - ) - train_eval_dataset = eval_author_result.train_dataset - validation_eval_dataset = eval_author_result.validation_dataset - insight_eval_dataset = eval_author_result.insight_suite - else: - # Mode 2: local agent directory as baseline, no insight required. - insight = None - insight_ref = "" + config = self.config + agents_dir, analysis_dir, _ = self._init_structure() + train_eval_dataset = ctx.datasets["train"] + validation_eval_dataset = ctx.datasets["validation"] + insight_eval_dataset = ctx.datasets.get("insight") + agent_spec_path = ctx.agent_spec # ---- Resume or fresh start --------------------------------------- + # Round analysis files are this strategy's own private state, so detecting the + # last round is its own business; ``ctx.resuming`` only tells it that the runner + # re-opened an existing run. if (round_num := self._detect_last_round()) is not None: logger.info(f"[RESUME] round {round_num}") self._delete_all_artifacts(from_round=round_num) evolution_tree = EvolutionTree.from_dir(agents_dir) candidates: list[Candidate] = list(evolution_tree.survivors(round_num)) - if reporter: - # agent-0 is not re-evaluated on resume; seed the delta baseline - # from its cached validation reward so later deltas are correct. - baseline_node = next( - (n for n in evolution_tree.nodes.values() if n.label == _BASELINE_AGENT_LABEL), - None, - ) - if baseline_node is not None and baseline_node.val_reward: - reporter.seed_baseline(reward_scalar(baseline_node.val_reward)) - run_entity = self._load_run_entity() or await self._create_experiment_run( - workspace=workspace, - backend=backend, - agent_name=agent_name or None, - agent_path=agent_path, - insight_ref=insight_ref or None, - config=config, - ) else: round_num = 0 logger.info("phase=baseline round=0") - if reporter: - reporter.progress(phase="baseline", completed=0, total=config.max_rounds) - run_entity = await self._create_experiment_run( - workspace=workspace, - backend=backend, - agent_name=agent_name or None, - agent_path=agent_path, - insight_ref=insight_ref or None, - config=config, - ) + await ctx.report_progress(completed=0, total=config.max_rounds, unit="round", note="baseline") - try: - # ---- Fetch + build baseline agent (agent-0) -------------- - baseline = await self._create_baseline_agent( - workspace=workspace, - backend=backend, - agents_dir=agents_dir, - agent_name=agent_name, - agent_path=agent_path, - run_id=run_entity.id or "", - config=config, - ) - await self._update_candidate( - baseline, - workspace=workspace, - backend=backend, - run_id=run_entity.id or "", - ) - evolution_tree = EvolutionTree.from_dir(agents_dir) - candidates = list(evolution_tree.survivors(0)) - - # ---- Baseline validation evaluation (round 0) ------------ - validation_candidate_results = await self._evaluate_validation_candidates( - dataset=validation_eval_dataset, - evaluator=evaluator, - candidates=candidates, - ) - validation_result = validation_candidate_results[candidates[0].label] - await backend.persist_evaluation( - workspace=workspace, - result=validation_result, - candidate=candidates[0], - split="validation", - ) - candidates[0].record_reward( - "validation", - metrics=validation_result.aggregate_metrics, - trials=validation_result.trials, - ) - await self._update_candidate( - candidates[0], - workspace=workspace, - backend=backend, - run_id=run_entity.id or "", - ) - if reporter: - reporter.candidate_evaluated( - label=candidates[0].label, - split="validation", - reward=reward_scalar(validation_result.aggregate_metrics), - artifacts=self._results_dir(validation_result.id), - ) - except Exception: - run_entity.status = "failed" - await backend.update_run(workspace=workspace, run=run_entity) - raise + # ---- Fetch + build baseline agent (agent-0) ------------------ + baseline = await self._create_baseline_agent(ctx=ctx, agents_dir=agents_dir, config=config) + await ctx.save_candidate(baseline) + evolution_tree = EvolutionTree.from_dir(agents_dir) + candidates = list(evolution_tree.survivors(0)) - run_id = run_entity.id or "" + # ---- Baseline validation evaluation (round 0) ---------------- + validation_candidate_results = await self._evaluate_validation_candidates( + ctx=ctx, + candidates=candidates, + ) + await ctx.record_reward( + candidates[0], + channel="validation", + result=validation_candidate_results[candidates[0].label], + ) if insight_eval_dataset is not None: - try: - await self._evaluate_and_persist_insight_candidates( - dataset=insight_eval_dataset, - evaluator=evaluator, - candidates=candidates, - workspace=workspace, - backend=backend, - run_id=run_entity.id or "", - ) - except Exception: - run_entity.status = "failed" - await backend.update_run(workspace=workspace, run=run_entity) - raise + await self._evaluate_and_persist_insight_candidates( + ctx=ctx, + dataset=insight_eval_dataset, + candidates=candidates, + ) # ---- Initial goal tree (idempotent) ------------------------------ await self._generate_initial_goal_tree( @@ -619,278 +423,149 @@ async def _run(self, deps: ExperimentalistDeps) -> ExperimentalistResult: phase: Literal["exploration", "exploitation"] = "exploration" if round_num % 2 == 0 else "exploitation" # ---- Pareto optimization loop (shared by fresh start and resume) -- - try: - while True: - prior_analysis = ( - await self._load_round_analysis(analysis_dir=analysis_dir, round_num=round_num - 1) - if round_num > 0 - else None - ) - decision = await self.terminator.run( - round_num=round_num, - evolution_tree=evolution_tree, - prior_analysis=prior_analysis, - config=config, - ) - if decision.stop: - logger.info(f"phase=terminate reason={decision.reason}") - break - - survivors = ( - await self._select_survivors([c.slim() for c in candidates], k=config.max_survivors) - if len(candidates) > 1 - else list(candidates) - ) - survivor_labels = {s.label for s in survivors} - killed = [c for c in candidates if c.label not in survivor_labels] - for candidate in killed: - await self._update_candidate( - candidate, - workspace=workspace, - backend=backend, - run_id=run_id, - updates={"killed_round": round_num}, + while True: + prior_analysis = ( + await self._load_round_analysis(analysis_dir=analysis_dir, round_num=round_num - 1) + if round_num > 0 + else None + ) + decision = await self.terminator.run( + round_num=round_num, + evolution_tree=evolution_tree, + prior_analysis=prior_analysis, + config=config, + ) + if decision.stop: + logger.info(f"phase=terminate reason={decision.reason}") + break + + survivors = ( + await self._select_survivors([c.slim() for c in candidates], k=config.max_survivors) + if len(candidates) > 1 + else list(candidates) + ) + survivor_labels = {s.label for s in survivors} + for candidate in [c for c in candidates if c.label not in survivor_labels]: + await ctx.save_candidate(candidate, updates={"killed_round": round_num}) + + train_candidate_results = await self._evaluate_train_candidates( + ctx=ctx, + survivors=survivors, + round_num=round_num, + max_train_batch_tasks=config.max_train_batch_tasks, + train_batch_seed=config.train_batch_seed, + ) + for survivor in survivors: + if survivor.label in train_candidate_results: + await ctx.record_reward( + survivor, + channel="train", + result=train_candidate_results[survivor.label], ) + analysis = await self._analyze_round( + analysis_dir=analysis_dir, + dataset=train_eval_dataset, + evaluations=train_candidate_results, + survivors=[c.slim() for c in survivors], + round_num=round_num, + config=config, + client=ctx.client, + nmp_workspace=ctx.workspace, + agent_spec_path=agent_spec_path, + ) + await self._update_goal_tree( + analysis_dir=analysis_dir, + round_num=round_num, + analysis=analysis, + dataset=train_eval_dataset, + config=config, + agent_spec_path=agent_spec_path, + ) - train_candidate_results = await self._evaluate_train_candidates( - dataset=train_eval_dataset, - evaluator=evaluator, - survivors=survivors, - round_num=round_num, - max_train_batch_tasks=config.max_train_batch_tasks, - train_batch_seed=config.train_batch_seed, - ) - for survivor in survivors: - if survivor.label in train_candidate_results: - await backend.persist_evaluation( - workspace=workspace, - result=train_candidate_results[survivor.label], - candidate=survivor, - split="train", - ) - survivor.record_reward( - "train", - metrics=train_candidate_results[survivor.label].aggregate_metrics, - trials=train_candidate_results[survivor.label].trials, - ) - await self._update_candidate( - survivor, - workspace=workspace, - backend=backend, - run_id=run_id, - ) - if reporter: - reporter.candidate_evaluated( - label=survivor.label, - split="train", - reward=reward_scalar(train_candidate_results[survivor.label].aggregate_metrics), - artifacts=self._results_dir(train_candidate_results[survivor.label].id), - ) - analysis = await self._analyze_round( - analysis_dir=analysis_dir, - dataset=train_eval_dataset, - evaluations=train_candidate_results, - survivors=[c.slim() for c in survivors], - round_num=round_num, - config=config, - client=backend.client, - nmp_workspace=workspace, - agent_spec_path=agent_spec_path, - ) - await self._update_goal_tree( - analysis_dir=analysis_dir, - round_num=round_num, - analysis=analysis, - dataset=train_eval_dataset, - config=config, - agent_spec_path=agent_spec_path, - ) - - improvements = await self._propose_improvements( - workspace=workspace, - backend=backend, - analysis=analysis, - evolution_tree=evolution_tree, - round_num=round_num, - phase=phase, - config=config, + improvements = await self._propose_improvements( + analysis=analysis, + evolution_tree=evolution_tree, + round_num=round_num, + phase=phase, + config=config, + ) + new_candidates = [ + self._create_agent( + agents_dir=agents_dir, + improvement=imp, + round_num=round_num + 1, + run_id=ctx.run_id, ) - new_candidates = [ - self._create_agent( - agents_dir=agents_dir, - improvement=imp, - round_num=round_num + 1, - run_id=run_entity.id or "", - ) - for imp in improvements - ] - # Persist metadata.json before Coder runs so snapshot can read it. - for candidate in new_candidates: - await self._update_candidate( - candidate, - workspace=workspace, - backend=backend, - run_id=run_entity.id or "", - ) - new_candidates = await self._implement_candidates( - workspace=workspace, - backend=backend, - dataset=train_eval_dataset, - evaluator=evaluator, + for imp in improvements + ] + # Persist metadata.json before Coder runs so snapshot can read it. + for candidate in new_candidates: + await ctx.save_candidate(candidate) + new_candidates = await self._implement_candidates( + ctx=ctx, + dataset=train_eval_dataset, + candidates=new_candidates, + config=config, + ) + for candidate in new_candidates: + await ctx.save_candidate(candidate) + if insight_eval_dataset is not None: + await self._evaluate_and_persist_insight_candidates( + ctx=ctx, + dataset=insight_eval_dataset, candidates=new_candidates, - config=config, ) - for candidate in new_candidates: - await self._update_candidate( + for c in new_candidates: + evolution_tree.add(c) + + candidates = survivors + new_candidates + round_num += 1 + phase = "exploration" if round_num % 2 == 0 else "exploitation" + await ctx.report_progress( + completed=round_num, + total=config.max_rounds, + unit="round", + note="evaluating candidates", + ) + # Announce candidates before the (batched) validation eval, so the narration + # reports work beginning, not completed. Mirror + # _evaluate_validation_candidates' own filter so we only announce candidates + # that will actually be evaluated (cached survivors already carry one). + pending_validation = [c for c in candidates if "validation" not in c.rewards] + for i, candidate in enumerate(pending_validation, start=1): + ctx.note(f"{candidate.label} ({i}/{len(pending_validation)}): {candidate.optimization}") + + validation_candidate_results = await self._evaluate_validation_candidates( + ctx=ctx, + candidates=candidates, + ) + for candidate in candidates: + if candidate.label in validation_candidate_results: + await ctx.record_reward( candidate, - workspace=workspace, - backend=backend, - run_id=run_id, - ) - if insight_eval_dataset is not None: - await self._evaluate_and_persist_insight_candidates( - dataset=insight_eval_dataset, - evaluator=evaluator, - candidates=new_candidates, - workspace=workspace, - backend=backend, - run_id=run_id, + channel="validation", + result=validation_candidate_results[candidate.label], ) - for c in new_candidates: - evolution_tree.add(c) - - candidates = survivors + new_candidates - round_num += 1 - phase = "exploration" if round_num % 2 == 0 else "exploitation" - if reporter: - reporter.progress( - phase="evaluating candidates", - completed=round_num, - total=config.max_rounds, - ) - # Announce candidates before the (batched) validation eval, - # so the narration reports work beginning, not completed. - # Mirror _evaluate_validation_candidates' own filter so we - # only announce candidates that will actually be evaluated - # (cached survivors already carry a validation_reward). - pending_validation = [c for c in candidates if "validation" not in c.rewards] - for i, candidate in enumerate(pending_validation, start=1): - reporter.candidate_started( - label=candidate.label, - optimization=candidate.optimization, - i=i, - n=len(pending_validation), - ) - - validation_candidate_results = await self._evaluate_validation_candidates( + if not config.disable_trajectory_scoring: + trajectory_results = await self._reward_trajectories( + ctx=ctx, dataset=validation_eval_dataset, - evaluator=evaluator, candidates=candidates, + config=config, ) for candidate in candidates: - if candidate.label in validation_candidate_results: - await backend.persist_evaluation( - workspace=workspace, - result=validation_candidate_results[candidate.label], - candidate=candidate, - split="validation", - ) - candidate.record_reward( - "validation", - metrics=validation_candidate_results[candidate.label].aggregate_metrics, - trials=validation_candidate_results[candidate.label].trials, - ) - await self._update_candidate( + if candidate.label in trajectory_results: + candidate.trajectory_detail = trajectory_results[candidate.label]["details"] + await ctx.record_reward( candidate, - workspace=workspace, - backend=backend, - run_id=run_id, + channel="validation-trajectory", + result=RewardRecord(metrics=trajectory_results[candidate.label]["reward"]), ) - if reporter: - reporter.candidate_evaluated( - label=candidate.label, - split="validation", - reward=reward_scalar(validation_candidate_results[candidate.label].aggregate_metrics), - artifacts=self._results_dir(validation_candidate_results[candidate.label].id), - ) - if not config.disable_trajectory_scoring: - trajectory_results = await self._reward_trajectories( - workspace=workspace, - backend=backend, - dataset=validation_eval_dataset, - candidates=candidates, - config=config, - client=backend.client, - ) - for candidate in candidates: - if candidate.label in trajectory_results: - candidate.record_reward( - "validation-trajectory", - metrics=trajectory_results[candidate.label]["reward"], - ) - await self._update_candidate( - candidate, - workspace=workspace, - backend=backend, - run_id=run_id, - updates={ - "trajectory_detail": trajectory_results[candidate.label]["details"], - }, - ) - - if config.storage.archive_candidates: - for candidate in new_candidates: - try: - await backend.archive_candidate(workspace=workspace, candidate=candidate) - except Exception as exc: # noqa: BLE001 - archival must never fail the run - _warn_persistence_failure("archive", candidate.label, exc) - - run_entity.rounds_completed = round_num - await backend.update_run(workspace=workspace, run=run_entity) - - except Exception: - run_entity.status = "failed" - await backend.update_run(workspace=workspace, run=run_entity) - raise - - # ---- Finalize ---------------------------------------------------- - winner_entity = await self._finalize( - workspace=workspace, - backend=backend, - agents_dir=agents_dir, - run_entity=run_entity, - evolution_tree=evolution_tree, - agent_name=agent_name, - insight_dataset=insight_eval_dataset, - ) - - baseline_entity = next( - (node.candidate for node in evolution_tree.nodes.values() if node.round == 0), - None, - ) - result = ExperimentalistResult( - summary=self._render_summary( - rounds_completed=round_num, - baseline=baseline_entity, - winner=winner_entity, - ), - run_id=run_id, - rounds_completed=round_num, - winner=winner_entity, - ) - # Persist the terminal result - await backend.persist_result(workspace=workspace, result=result) + for candidate in new_candidates: + await ctx.archive_candidate(candidate) - # Publish the winner as a draft PR/MR - if config.storage.publish_winner and winner_entity is not None and winner_entity.round != 0: - try: - url = await backend.publish_candidate(workspace=workspace, candidate=winner_entity) - if url: - logger.info(f"[TERMINATOR] opened draft PR/MR for winner {winner_entity.label}: {url}") - except Exception as exc: # noqa: BLE001 - publishing must never fail the run - _warn_persistence_failure("publish", winner_entity.label, exc) - return result + return await self._finalize(evolution_tree=evolution_tree) @strategy(CodeActStrategy(config=CodeActConfig(max_iterations=100, cell_timeout=3600.0))) async def select_diverse_survivors(self, ranked: list[Candidate], k: int) -> list[Candidate]: # pyright: ignore[reportReturnType] @@ -973,7 +648,7 @@ async def merge_analysis( Trial Analysis; Complementary Failures; Failure Patterns; Root Causes; Mechanical/Infrastructure Errors). - If at least one agent has a non-empty `insight_rewards` entry, the round analysis must name + If at least one agent has a non-empty `insight_reward`, the round analysis must name every available Insight Suite dimension and show its values in the separate Insight Suite Reward table. Never blend those metrics into train/validation rewards or imply that they affected ranking. These metrics may steer this analysis, the goal tree, and @@ -1034,19 +709,6 @@ def _init_structure(self) -> tuple[Path, Path, Path]: results_dir.mkdir(parents=True, exist_ok=True) return agents_dir, analysis_dir, results_dir - def _load_run_entity(self) -> ExperimentRun | None: - """Read run.json from the workspace; return None if absent or unparseable.""" - run_path = self.working_dir / "eval-and-optimize" / "run.json" - if not run_path.exists(): - return None - try: - data = json.loads(run_path.read_text()) - # ExperimentRun._restore_id_from_json validator handles id restoration - return ExperimentRun.model_validate(data) - except Exception as exc: # noqa: BLE001 - logger.warning(f"[RESUME] Could not parse run.json: {exc}") - return None - def _detect_last_round(self) -> int | None: """Return the last completed round by scanning analysis files, or None if starting fresh.""" analysis_dir = self.working_dir / "eval-and-optimize" / "analysis" @@ -1127,38 +789,28 @@ def _delete_all_artifacts(self, from_round: int) -> None: async def _create_baseline_agent( self, *, - workspace: str, - backend: ExperimentalistBackend, + ctx: ExperimentContext, agents_dir: Path, - agent_name: str | None, - agent_path: Path | None, - run_id: str, config: EvolutionaryOptimizerConfig, ) -> Candidate: - """Materialize the source agent into ``agents_dir/agent-0`` and return the baseline candidate. + """Fork the agent under test into ``agents_dir/agent-0`` and return the baseline. - An explicit ``--agent`` (``agent_path``, a local dir or a git clone) takes - precedence and is copied directly; otherwise the code is fetched by the - insight's agent name via ``backend.get_agent_code``. Skips the copy when the - directory already exists (resume case). + The runner has already materialized ``ctx.agent_dir``, whatever the source was; + this only copies it in. Skips the copy when the directory already exists (resume). """ baseline_dir = agents_dir / _BASELINE_AGENT_LABEL if not baseline_dir.exists(): - if agent_path is not None: - shutil.copytree(agent_path, baseline_dir, ignore=_ignore_patterns) - elif agent_name: - await backend.get_agent_code(workspace=workspace, agent=agent_name, dest=baseline_dir) + shutil.copytree(ctx.agent_dir, baseline_dir, ignore=_ignore_patterns) await self._generate_architecture_doc(agent_dir=baseline_dir, config=config) - candidate = Candidate( + return Candidate( name=_BASELINE_AGENT_LABEL, label=_BASELINE_AGENT_LABEL, - workspace=workspace, - run_id=run_id, + workspace=ctx.workspace, + run_id=ctx.run_id, ancestor=None, round=0, optimization="baseline", ) - return candidate def _create_agent( self, @@ -1218,33 +870,6 @@ async def _load_round_analysis( path = analysis_dir / f"round-{round_num}.md" return path.read_text() if path.exists() else None - async def _update_candidate( - self, - candidate: Candidate, - *, - workspace: str, - backend: ExperimentalistBackend, - run_id: str, - updates: dict[str, Any] | None = None, - ) -> None: - """Sync candidates to the entity store. - - Fills in ``workspace`` and ``run_id`` from the call-site context - (which is authoritative) before persisting. On first persist the - backend assigns a store id (``_id``); on subsequent calls it updates - the existing record. - """ - candidate.workspace = workspace - candidate.run_id = run_id - if updates is not None: - for key, value in updates.items(): - setattr(candidate, key, value) - if candidate.id: - await backend.update_candidate(workspace=workspace, candidate=candidate) - else: - result = await backend.create_candidate(workspace=workspace, candidate=candidate) - candidate._id = result._id # type: ignore[attr-defined] - def _goal_tree_path(self, round_num: int) -> Path: return self.working_dir / "eval-and-optimize" / "analysis" / f"round-{round_num}-goal.json" @@ -1280,25 +905,6 @@ def _format_evolution_table(self, evolution_tree: EvolutionTree) -> str: return "(none yet — this is the first round)" return evolution_tree.to_markdown_table() - def _copy_best_to_workspace(self, agent_id: str) -> None: - src = self.working_dir / "eval-and-optimize" / "agents" / agent_id - skip_names = { - "metadata.json", - "harbor_wrapper.py", - "dind_environment.py", - "architecture.md", - } - for entry in src.iterdir(): - if entry.name in skip_names: - continue - dst = self.working_dir / entry.name - if entry.is_dir(): - if dst.exists(): - shutil.rmtree(dst) - shutil.copytree(entry, dst) - else: - shutil.copy2(entry, dst) - def _snapshot_metadata(self, candidate_name: str) -> str | None: """Read and return the metadata.json content for a candidate, or None if absent.""" path = self.working_dir / "eval-and-optimize" / "agents" / candidate_name / "metadata.json" @@ -1319,28 +925,18 @@ def _restore_metadata(self, candidate_name: str, content: str | None) -> None: async def _evaluate_agent( self, + ctx: ExperimentContext, candidate: Candidate, - dataset: Dataset, - evaluator: Evaluator, + split: str, task_ids: list[str] | None = None, minimum_attempts: int | None = None, ) -> tuple[Candidate, EvaluationResult]: - """Run evaluator for one candidate and return the candidate/result pair.""" - eval_dataset = dataset.subset(task_ids) if task_ids is not None else dataset - # Force a unique job name per candidate so concurrent candidates don't - # collide on the same results directory when the user sets a fixed job_name. - options_dict = evaluator.options.model_dump() - options_dict["job_name"] = f"{candidate.label}-{eval_dataset.id}" - if minimum_attempts is not None: - configured_attempts = options_dict.get("n_attempts") - if not isinstance(configured_attempts, int): - raise ValueError("Insight evaluator options must define integer n_attempts") - options_dict["n_attempts"] = max(configured_attempts, minimum_attempts) - per_candidate_options = type(evaluator.options).model_validate(options_dict) - result = await evaluator.run( - agent=self.working_dir / "eval-and-optimize" / "agents" / candidate.label, - dataset=eval_dataset, - options=per_candidate_options, + """Evaluate one candidate and return the candidate/result pair, for ``gather``.""" + result = await ctx.evaluate( + candidate, + split=split, + task_ids=task_ids, + minimum_attempts=minimum_attempts, ) return (candidate, result) @@ -1348,27 +944,6 @@ async def _evaluate_agent( # Private step methods — implementations # ------------------------------------------------------------------ - async def _create_experiment_run( - self, - *, - workspace: str, - backend: ExperimentalistBackend, - agent_name: str | None, - agent_path: Path | None, - insight_ref: str | None, - config: EvolutionaryOptimizerConfig, - ) -> ExperimentRun: - """Create an ExperimentRun entity; return it with its store-assigned id.""" - run = ExperimentRun( - workspace=workspace, - agent=agent_name or str(agent_path or ""), - insight=insight_ref, - config_snapshot=config.model_dump(mode="json"), - status="running", - rounds_completed=0, - ) - return await backend.create_run(workspace=workspace, run=run) - async def _generate_architecture_doc( self, *, @@ -1392,41 +967,26 @@ async def _generate_architecture_doc( async def _evaluate_validation_candidates( self, *, - dataset: Dataset, - evaluator: Evaluator, + ctx: ExperimentContext, candidates: list[Candidate], ) -> dict[str, EvaluationResult]: """Evaluate candidates on the validation split; skip any that already have a reward.""" pending = [c for c in candidates if "validation" not in c.rewards] if not pending: return {} - if pending: - splits = frozenset({"validation"}) - restore_heldout_splits(self.working_dir, splits=splits) - try: - candidate_results = await asyncio.gather( - *[ - self._evaluate_agent( - c, - dataset, - evaluator, - ) - for c in pending - ] - ) - finally: - ensure_heldout_hidden(self.working_dir, splits=splits) - return { - candidate_result[0].label: candidate_result[1] - for candidate_result in candidate_results - if candidate_result is not None - } + splits = frozenset({"validation"}) + restore_heldout_splits(self.working_dir, splits=splits) + try: + candidate_results = await asyncio.gather(*[self._evaluate_agent(ctx, c, "validation") for c in pending]) + finally: + ensure_heldout_hidden(self.working_dir, splits=splits) + return {candidate.label: result for candidate, result in candidate_results} async def _evaluate_insight_candidates( self, *, + ctx: ExperimentContext, dataset: Dataset, - evaluator: Evaluator, candidates: list[Candidate], ) -> dict[str, EvaluationResult]: """Evaluate candidates that do not yet have metrics for this Insight suite.""" @@ -1444,33 +1004,22 @@ async def _evaluate_insight_candidates( or not candidate_metric_keys(candidate) ] evaluated = await asyncio.gather( - *[ - self._evaluate_agent( - candidate, - dataset, - evaluator, - minimum_attempts=2, - ) - for candidate in pending - ] + *[self._evaluate_agent(ctx, candidate, "insight", minimum_attempts=2) for candidate in pending] ) return {candidate.label: result for candidate, result in evaluated} async def _evaluate_and_persist_insight_candidates( self, *, + ctx: ExperimentContext, dataset: Dataset, - evaluator: Evaluator, candidates: list[Candidate], - workspace: str, - backend: ExperimentalistBackend, - run_id: str, ) -> None: """Evaluate and persist Insight-suite metrics for the supplied candidates.""" provenance = insight_suite_provenance(dataset) results = await self._evaluate_insight_candidates( + ctx=ctx, dataset=dataset, - evaluator=evaluator, candidates=candidates, ) dataset_metric_keys = dataset.metadata.get("insight_metric_keys") @@ -1501,24 +1050,12 @@ async def _evaluate_and_persist_insight_candidates( if expected_metric_keys is None: expected_metric_keys = metric_keys result = stamp_insight_evaluation_result(result, provenance) - await backend.persist_evaluation( - workspace=workspace, + await ctx.record_reward( + candidate, + channel="insight", result=result, - candidate=candidate, - split="insight", - ) - candidate.record_reward( - "insight", - metrics=result.aggregate_metrics, - trials=result.trials, metadata={"suite_identity": provenance.identity, "metric_keys": list(metric_keys)}, ) - await self._update_candidate( - candidate, - workspace=workspace, - backend=backend, - run_id=run_id, - ) if expected_metric_keys is not None: dataset.metadata["insight_metric_keys"] = list(expected_metric_keys) @@ -1560,8 +1097,7 @@ async def _select_survivors( async def _evaluate_train_candidates( self, *, - dataset: Dataset, - evaluator: Evaluator, + ctx: ExperimentContext, max_train_batch_tasks: int | None, train_batch_seed: int, survivors: list[Candidate], @@ -1575,13 +1111,14 @@ async def _evaluate_train_candidates( changes per round, so cached rewards are not comparable within a round: every survivor is re-evaluated on the same freshly sampled batch. """ + dataset = ctx.datasets["train"] if max_train_batch_tasks is None: # Full-dataset mode: the eval set is identical every round, so evaluate only # survivors that lack a reward and reuse each survivor's cached train reward. pending = [s for s in survivors if "train" not in s.rewards] evaluated: list[tuple[Candidate, EvaluationResult]] = [] if pending: - evaluated = await asyncio.gather(*[self._evaluate_agent(c, dataset, evaluator) for c in pending]) + evaluated = await asyncio.gather(*[self._evaluate_agent(ctx, c, "train") for c in pending]) results = {candidate.label: result for candidate, result in evaluated} for survivor in survivors: if survivor.label in results: @@ -1606,9 +1143,7 @@ async def _evaluate_train_candidates( rng = random.Random(train_batch_seed + round_num) batch_size = min(max_train_batch_tasks, len(all_task_ids)) task_ids = sorted(rng.sample(all_task_ids, batch_size)) - evaluated = await asyncio.gather( - *[self._evaluate_agent(c, dataset, evaluator, task_ids=task_ids) for c in survivors] - ) + evaluated = await asyncio.gather(*[self._evaluate_agent(ctx, c, "train", task_ids=task_ids) for c in survivors]) return {candidate.label: result for candidate, result in evaluated} async def _analyze_round( @@ -1691,8 +1226,6 @@ async def _update_goal_tree( async def _propose_improvements( self, *, - workspace: str, - backend: ExperimentalistBackend, analysis: str, evolution_tree: EvolutionTree, round_num: int, @@ -1717,10 +1250,8 @@ async def _propose_improvements( async def _implement_candidates( self, *, - workspace: str, - backend: ExperimentalistBackend, + ctx: ExperimentContext, dataset: Dataset, - evaluator: Evaluator, candidates: list[Candidate], config: EvolutionaryOptimizerConfig, ) -> list[Candidate]: @@ -1736,7 +1267,7 @@ async def _implement_candidates( ).run( c, dataset, - evaluator, + ctx.evaluation, source_path=config.source.source_path, entrypoint=config.source.entrypoint, ) @@ -1747,50 +1278,24 @@ async def _implement_candidates( finally: for c in candidates: self._restore_metadata(c.name, snapshots[c.name]) - # `return_exceptions=True` above means a build failure arrives as a value, not a - # raise, so the reason is only ever seen if we log it here. Keep the exception - # itself: "Impl failed: agent-1" with no cause is not diagnosable after the fact, - # and a killed candidate is the one thing a run cannot reproduce cheaply. - # Cancellation is not a build failure and must not be swallowed as one: - # `CancelledError` derives from BaseException, so the `Exception` filter below - # would let a cancelled candidate through as if it had built, and it would go on - # to be evaluated and ranked. Re-raise so the whole round unwinds instead. - for result in results: - if isinstance(result, asyncio.CancelledError): - raise result - failures = {c.name: r for c, r in zip(candidates, results, strict=True) if isinstance(r, Exception)} + failed = {c.name for c, r in zip(candidates, results, strict=True) if isinstance(r, Exception)} # Persist a killed marker on failed candidates so a later resume via # EvolutionTree.from_dir does not resurrect them as active survivors # (a node is a survivor exactly when killed_round is None). for candidate in candidates: - if candidate.name not in failures: + if candidate.name not in failed: continue - error = failures[candidate.name] - logger.warning( - "Impl failed: %s — %s: %s", - candidate.name, - type(error).__name__, - error, - exc_info=error, - ) - await self._update_candidate( - candidate, - workspace=workspace, - backend=backend, - run_id=candidate.run_id, - updates={"killed_round": candidate.round}, - ) - return [c for c in candidates if c.name not in failures] + logger.warning(f"Impl failed: {candidate.name}") + await ctx.save_candidate(candidate, updates={"killed_round": candidate.round}) + return [c for c in candidates if c.name not in failed] async def _reward_trajectories( self, *, - workspace: str, - backend: ExperimentalistBackend, + ctx: ExperimentContext, dataset: Dataset, candidates: list[Candidate], config: EvolutionaryOptimizerConfig, - client: AsyncNeMoPlatform | None = None, ) -> dict[str, dict[str, Any]]: """Score candidates against the goal tree; return trajectory results keyed by candidate label.""" tree_path = self._latest_goal_tree_path() @@ -1852,7 +1357,7 @@ async def _reward_trajectories( keys = [(node.id, task_id) for node in nodes for task_id in traces_by_task] logger.info(f"[TRAJ] Starting {len(keys)} GRA scoring tasks...") - scorer = GroupLeafScorer(workspace=self.working_dir, client=client, nmp_workspace=workspace) + scorer = GroupLeafScorer(workspace=self.working_dir, client=ctx.client, nmp_workspace=ctx.workspace) scoring_results = await asyncio.gather( *[scorer.run(node, traces_by_task[task_id], dataset) for node in nodes for task_id in traces_by_task] ) @@ -1887,99 +1392,24 @@ async def _reward_trajectories( return trajectory_results - async def _finalize( - self, - *, - workspace: str, - backend: ExperimentalistBackend, - agents_dir: Path, - run_entity: ExperimentRun, - evolution_tree: EvolutionTree, - agent_name: str, - insight_dataset: Dataset | None, - ) -> Candidate | None: - """Select the winner, copy to workspace root, write final report.""" + async def _finalize(self, *, evolution_tree: EvolutionTree) -> Candidate | None: + """Pick the winner and write this strategy's own report; return the winner. + + Everything host-owned that used to happen here — restoring the held-out splits, + copying the winner into the workspace, closing out the run entity, and the + Insight-suite report sections — belongs to the runner now. + """ # Only survivors that actually have a validation reward are eligible winners. scored = [n for n in evolution_tree.nodes.values() if n.is_survivor and n.val_reward] front = pareto_front(scored, lambda n: n.val_reward) if scored else [] best_id = front[0].label if front else None - - restore_heldout_splits(self.working_dir) - if best_id is None: logger.warning("[FINAL] no candidates to finalize") - run_entity.status = "completed" - await backend.update_run(workspace=workspace, run=run_entity) return None evolution_tree.mark_best(best_id) - self._copy_best_to_workspace(best_id) - - winner = evolution_tree.nodes[best_id].candidate - baseline = next( - (node.candidate for node in evolution_tree.nodes.values() if node.round == 0), - None, - ) - report_path = self.working_dir / "eval-and-optimize" / "OPTIMIZATION.md" - final_report_failed = False try: await self.write_final_report(best_id) - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 - the runner falls back to a compact summary logger.warning(f"[FINAL] Failed to write final report: {exc}") - final_report_failed = True - if not report_path.exists() or not report_path.read_text().strip(): - final_report_failed = True - if final_report_failed: - summary = self._render_summary( - rounds_completed=run_entity.rounds_completed, - baseline=baseline, - winner=winner, - ) - report_path.write_text(f"# Optimization Report\n\n## Compact Run Summary\n\n{summary}\n") - - if insight_dataset is not None: - try: - provenance = insight_suite_provenance(insight_dataset) - if baseline is not None: - write_insight_comparison_section( - report_path, - baseline, - winner, - provenance, - ) - suggestions = select_insight_promotion_suggestions( - insight_dataset, - [node.candidate for node in evolution_tree.nodes.values()], - winner=winner, - ) - write_insight_promotion_section( - report_path, - suggestions, - ) - except ValueError as exc: - logger.warning(f"[FINAL] Skipping Insight Suite report sections: {exc}") - - run_entity.status = "completed" - run_entity.winner_agent = best_id - await backend.update_run(workspace=workspace, run=run_entity) - - return winner - - def _render_summary( - self, - rounds_completed: int, - baseline: Candidate | None, - winner: Candidate | None, - ) -> str: - """Render a human-readable summary of the run outcome.""" - winner_str = winner.name if winner else "none" - details: list[str] = [] - if winner: - if winner.reward("validation").metrics: - details.append(f"validation_reward={winner.reward('validation').metrics}") - if baseline is not None and baseline.reward("insight").metrics and winner.reward("insight").metrics: - details.append( - f"insight_suite=(baseline={baseline.reward('insight').metrics}, winner={winner.reward('insight').metrics})" - ) - suffix = f", {', '.join(details)}" if details else "" - return f"Optimization complete: {rounds_completed} round(s) completed, winner={winner_str}{suffix}" + return evolution_tree.nodes[best_id].candidate diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/context.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/context.py new file mode 100644 index 0000000000..5df0f3a729 --- /dev/null +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/context.py @@ -0,0 +1,301 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The plugin-facing seam: everything a strategy is allowed to reach. + +A strategy receives one :class:`ExperimentContext` and nothing else. The context holds +the :class:`~nemo_experimentalist_plugin.experimentalist.experimentalist_backend.ExperimentalistBackend` +privately, so no component ever sees ``create_run``, ``publish_candidate``, or the +platform client; the runner that built the context is the only code that holds a +backend. + +The two keyword arguments below are deliberately different words. ``evaluate(split=…)`` +names a *dataset split* to run against; ``record_reward(channel=…)`` names a *reward +channel* to store under. Usually a run on the validation split lands in the +``validation`` channel, but trajectory scoring produces a second channel from the same +split, which is why one is not spelled with the other's name. +""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from nemo_experimentalist_plugin.entities import ( + Candidate, + Dataset, + DataValue, + EvaluationResult, + ExperimentRun, + RewardRecord, +) +from nemo_experimentalist_plugin.experimentalist.components.evaluator import Evaluator +from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ( + ExperimentalistBackend, +) +from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter, reward_scalar +from nemo_platform import AsyncNeMoPlatform + +logger = logging.getLogger(__name__) + +#: The split every run has, and what ``evaluate()`` measures unless told otherwise. +#: It is also the channel a selector ranks on by default. +PRIMARY_SPLIT = "validation" + + +class ExperimentContext: + """One run's view of the platform, handed to exactly one strategy. + + Args: + backend: Data-access backend. Held privately; never handed to a component. + workspace: NeMo Platform workspace this run belongs to. + run: The ``ExperimentRun`` entity the runner created (or re-opened on resume). + root: Working directory for the run's artifacts. + agent_dir: The agent under test, materialized by the runner. A strategy forks + candidates from here; it must not write into it. + agent_spec: Optional markdown description of the agent under test. + datasets: Evaluator-domain datasets keyed by split. ``validation`` is always + present; ``train`` and ``insight`` are present when the run has them. + evaluator: Evaluation component the run was configured with. + resuming: True when the runner re-opened an existing run, so the strategy + should rebuild its state from :meth:`candidates` instead of starting over. + reporter: Optional human narration sink. Best-effort and never load-bearing. + """ + + def __init__( + self, + *, + backend: ExperimentalistBackend, + workspace: str, + run: ExperimentRun, + root: Path, + agent_dir: Path, + agent_spec: Path | None = None, + datasets: Mapping[str, Dataset], + evaluator: Evaluator, + resuming: bool = False, + reporter: RunReporter | None = None, + ) -> None: + if PRIMARY_SPLIT not in datasets: + raise ValueError(f"ExperimentContext requires a {PRIMARY_SPLIT!r} dataset; got {sorted(datasets)}") + self._backend = backend + self._run = run + self._evaluator = evaluator + self._reporter = reporter + self.workspace = workspace + self.root = root + self.agent_dir = agent_dir + self.agent_spec = agent_spec + self.datasets: Mapping[str, Dataset] = dict(datasets) + self.resuming = resuming + + # -- Inputs -------------------------------------------------------------- + + @property + def dataset(self) -> Dataset: + """The run's primary dataset — what ``evaluate()`` measures by default.""" + return self.datasets[PRIMARY_SPLIT] + + @property + def run_id(self) -> str: + """Durable id of this run, and the key every Candidate is grouped under.""" + return self._run.id or "" + + @property + def evaluation(self) -> Evaluator: + """The evaluation component this run was configured with. + + Exposed so a composite strategy can hand it to a component it owns — the Coder + runs smoke evals of its own. Once the registry lands this becomes + ``ctx.component("evaluation", …)`` and this property goes away. + """ + return self._evaluator + + @property + def client(self) -> AsyncNeMoPlatform | None: + """Platform client, for reading ``intake://`` traces. Transitional. + + This is the last piece of backend that still reaches a component: the trace + readers resolve ``intake://`` trial traces themselves, so they need a client. + A ``ctx`` verb that loads a trace by reference would close it; until that + exists, a strategy whose components read traces needs this, and a strategy + that does not should ignore it. + """ + return self._backend.client + + # -- Candidates ---------------------------------------------------------- + + async def candidates(self) -> list[Candidate]: + """Every Candidate committed to this run, in store order. + + This is what makes resume possible for a strategy the host did not write: the + population is persisted, so a strategy rebuilds it rather than checkpointing it. + """ + return await self._backend.list_candidates(workspace=self.workspace, run_id=self.run_id) + + def candidate_dir(self, candidate: Candidate) -> Path: + """Directory holding *candidate*'s code. + + Still derived from the label, because a Candidate is still a directory. The + candidate contract replaces this with ``candidate.artifact``. + """ + return self.root / "eval-and-optimize" / "agents" / candidate.label + + async def save_candidate(self, candidate: Candidate, *, updates: dict[str, Any] | None = None) -> Candidate: + """Create or update *candidate* in the entity store. + + Fills in ``workspace`` and ``run_id`` from the run (which is authoritative) + before persisting. On first save the backend assigns a store id; later saves + update the existing record. + """ + candidate.workspace = self.workspace + candidate.run_id = self.run_id + for key, value in (updates or {}).items(): + setattr(candidate, key, value) + if candidate.id: + return await self._backend.update_candidate(workspace=self.workspace, candidate=candidate) + stored = await self._backend.create_candidate(workspace=self.workspace, candidate=candidate) + candidate._id = stored._id # type: ignore[attr-defined] + return candidate + + async def archive_candidate(self, candidate: Candidate) -> None: + """Persist *candidate*'s code to durable storage, if the run archives at all. + + Best-effort in both directions: a backend that cannot archive returns nothing, + and a failure is logged rather than raised — archival must never fail a run. + """ + if not self._backend.storage.archive_candidates: + return + try: + await self._backend.archive_candidate(workspace=self.workspace, candidate=candidate) + except Exception as exc: # noqa: BLE001 - archival must never fail the run + logger.warning("[PERSISTENCE] archive failed for candidate %s; continuing: %s", candidate.label, exc) + + # -- Measurement --------------------------------------------------------- + + async def record_reward( + self, + candidate: Candidate, + *, + channel: str, + result: EvaluationResult | RewardRecord, + metadata: dict[str, DataValue] | None = None, + ) -> None: + """Store one measurement of *candidate* on *channel*, and persist the candidate. + + An ``EvaluationResult`` is the outcome of running the candidate, so its traces + are persisted before the record is stored; a ``RewardRecord`` is a measurement + the strategy computed itself (trajectory scoring, a self-scoring strategy) and + is stored as given. Either way the channel is an open key — adding one costs no + entity change. + """ + if isinstance(result, EvaluationResult): + await self._backend.persist_evaluation( + workspace=self.workspace, + result=result, + candidate=candidate, + split=channel, + ) + record = RewardRecord( + metrics={k: float(v) for k, v in result.aggregate_metrics.items()}, + trials=list(result.trials), + metadata=dict(metadata or {}), + ) + else: + record = result if metadata is None else result.model_copy(update={"metadata": dict(metadata)}) + candidate.set_reward( + channel, + metrics=record.metrics, + summary=record.summary, + trials=record.trials, + metadata=record.metadata, + ) + await self.save_candidate(candidate) + + async def evaluate( + self, + candidate: Candidate, + *, + split: str = PRIMARY_SPLIT, + task_ids: Sequence[str] | None = None, + minimum_attempts: int | None = None, + ) -> EvaluationResult: + """Run the configured evaluation component over *candidate*'s artifact. + + Optional by design: a strategy that scores itself skips this and calls + :meth:`record_reward` directly. The association between the result and the + candidate is owned here, not by the evaluation component. + + Args: + candidate: Whose artifact to evaluate. + split: Which of :attr:`datasets` to run against. + task_ids: Restrict the run to these task ids. + minimum_attempts: Raise the evaluator's attempt count to at least this. + + Raises: + KeyError: if *split* is not one of the run's datasets. + """ + dataset = self.datasets[split] + if task_ids is not None: + dataset = dataset.subset(list(task_ids)) + # Force a unique job name per candidate so concurrent candidates never collide on + # one results directory when a fixed job_name is configured. job_name/n_attempts + # are Harbor's vocabulary reaching a generic call site — an evaluator leak the + # registered evaluation component closes. + options = self._evaluator.options.model_dump() + options["job_name"] = f"{candidate.label}-{dataset.id}" + if minimum_attempts is not None: + configured = options.get("n_attempts") + if not isinstance(configured, int): + raise ValueError("Evaluator options must define integer n_attempts to raise the attempt floor") + options["n_attempts"] = max(configured, minimum_attempts) + result = await self._evaluator.run( + agent=self.candidate_dir(candidate), + dataset=dataset, + options=type(self._evaluator.options).model_validate(options), + ) + if self._reporter is not None: + self._reporter.candidate_evaluated( + label=candidate.label, + split=split, + reward=reward_scalar(result.aggregate_metrics), + artifacts=self.root / "eval-and-optimize" / "results" / result.id, + ) + return result + + # -- Progress ------------------------------------------------------------ + + async def report_progress( + self, + *, + completed: int, + total: int | None = None, + unit: str = "step", + note: str | None = None, + ) -> None: + """Report how far the strategy has got, as a counter rather than a fraction. + + A counter is always producible; a fraction usually is not — an opaque strategy + cannot say how many trials it will run, and even a round-based one stops early + on convergence. When *total* is known a consumer may render a bar; when it is + not, ``note`` is where a strategy says what it is doing instead. + """ + self._run.progress_completed = completed + self._run.progress_total = total + self._run.progress_unit = unit + self._run.progress_note = note + await self._backend.update_run(workspace=self.workspace, run=self._run) + if self._reporter is not None: + self._reporter.progress(phase=note or unit, completed=completed, total=total, unit=unit) + + def note(self, message: str) -> None: + """Say what is happening, for a human watching the run. + + Narration only: it touches no entity and never raises, so a strategy may call it + as freely as it likes. Use :meth:`report_progress` for anything a consumer + should be able to read back off the run. + """ + if self._reporter is not None: + self._reporter.note(message) diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/deps.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/deps.py deleted file mode 100644 index f2547335a5..0000000000 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/deps.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Run-time dependencies injected into every Experimentalist tool call. - -A single :class:`ExperimentalistDeps` instance carries the resolved run -configuration through the Experimentalist loop. Keeping it in its own module -avoids an import cycle between the agent definition and the components it -coordinates. -""" - -from __future__ import annotations - -from pathlib import Path - -from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig -from nemo_experimentalist_plugin.entities import DatasetRef -from nemo_experimentalist_plugin.experimentalist.components.evaluator.base import EvaluatorType -from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ( - ExperimentalistBackend, -) -from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter -from pydantic import BaseModel, model_validator - - -class ExperimentalistDeps(BaseModel): - """Per-run configuration injected into every Experimentalist tool. - - At least one of ``insight`` or ``agent`` must be set: - - Mode 1 (``insight``): agent code defaults to the NMP insight entity's agent. - When ``agent`` is also set, it overrides the insight's agent. - - Mode 2 (``agent``): a local agent directory is used as the baseline. - - Attributes: - workspace: NMP workspace the Experimentalist operates in. - insight: Entity id of the Insight being processed (Mode 1). - agent: Local path to the baseline agent directory (Mode 2), or an - override for the agent referenced by ``insight``. - dataset: Local directory path OR fileset ID for the evaluation dataset. - A :class:`~pathlib.Path` means a local directory; a plain string - means a fileset ID that the backend will resolve at evaluation time. - Defaults to None and must be set before ``run()``. - backend: Shared data-access backend used by every tool. The CLI - owns the backend's client lifecycle — tools must not close it. - config: Optional per-run override of the EvolutionaryOptimizerConfig. - When None the optimizer uses its own default config. - """ - - model_config = {"arbitrary_types_allowed": True} - - workspace: str = "default" - insight: Path | str | None = None - agent: Path | str | None = None - train_dataset: DatasetRef - validation_dataset: DatasetRef - task_template: DatasetRef | None = None - evaluator_type: EvaluatorType = "harbor" - agent_spec: str | None = None - backend: ExperimentalistBackend | None = None - reporter: RunReporter | None = None - config: EvolutionaryOptimizerConfig | None = None - - @model_validator(mode="after") - def _require_insight_or_agent(self) -> ExperimentalistDeps: - has_insight = self.insight is not None - has_agent = self.agent is not None - if not has_insight and not has_agent: - raise ValueError("One of 'insight' (Mode 1) or 'agent' (Mode 2) must be set.") - if has_insight and self.task_template is None: - raise ValueError("'task_template' is required when 'insight' is set.") - return self diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experiment_mirror.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experiment_mirror.py index 659710e1ae..da73847b34 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experiment_mirror.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experiment_mirror.py @@ -70,15 +70,21 @@ def group_metadata(run: ExperimentRun) -> dict[str, str]: """The ExperimentRun fields with no first-class ExperimentGroup home (spec §4.1). Platform ``metadata`` is ``dict[str, str]``: every value must be a string. Non-string - fields are serialized (``config_snapshot`` as JSON, ``rounds_completed`` via ``str``); + fields are serialized (``config_snapshot`` as JSON, the progress counter via ``str``); ``winner_candidate`` is omitted until a winner exists rather than sent as ``None``. + + Progress is reported as a counter with its unit, not a fraction, because the + strategies that most need reporting are exactly the ones that cannot compute one. """ md = { "agent": run.agent, "config_snapshot": json.dumps(run.config_snapshot, sort_keys=True), "status": run.status, - "rounds_completed": str(run.rounds_completed), + "progress_completed": str(run.progress_completed), + "progress_unit": run.progress_unit, } + if run.progress_total is not None: + md["progress_total"] = str(run.progress_total) if run.winner_agent is not None: md["winner_candidate"] = run.winner_agent return md diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py index 3afe17bbd7..edf2a367b6 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/experimentalist_backend.py @@ -735,7 +735,7 @@ async def persist_result(self, *, workspace: str, result: ExperimentalistResult) if run_path.exists(): run: ExperimentRun = _load_entity(ExperimentRun, run_path) run.status = "completed" - run.rounds_completed = result.rounds_completed + run.progress_completed = result.progress_completed run.summary = result.summary # so publish_candidate's _compose_pr_body reads the real summary if result.winner is not None: run.winner_agent = result.winner.id diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/result.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/result.py index 0d4a4038f1..7c7b349605 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/result.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/result.py @@ -33,9 +33,9 @@ class ExperimentalistResult(BaseModel): min_length=1, description="ExperimentRun entity id updated by persist_result.", ) - rounds_completed: int = Field( + progress_completed: int = Field( ge=0, - description="Number of full optimization rounds that ran.", + description="Units of work the strategy reported finished — rounds, trials, whatever it counts.", ) winner: Candidate | None = Field( default=None, diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py index b122a5363a..9fba250431 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/run.py @@ -11,11 +11,11 @@ from nemo_experimentalist_plugin.entities import DatasetRef from nemo_experimentalist_plugin.experimentalist.agent import build_experimentalist_agent from nemo_experimentalist_plugin.experimentalist.components.loop import EvolutionaryOptimizerConfig -from nemo_experimentalist_plugin.experimentalist.deps import ExperimentalistDeps from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ( make_experimentalist_backend, ) from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter, Verbosity +from nemo_experimentalist_plugin.experimentalist.runner import ExperimentRunner from nemo_platform import AsyncNeMoPlatform logger = logging.getLogger(__name__) @@ -42,7 +42,7 @@ def build_run_reporter( async def run_experimentalist( *, - agent: str | None = None, + agent: str | Path | None = None, agent_spec: str | None = None, insight: Path | str | None, train_dataset: DatasetRef, @@ -58,7 +58,7 @@ async def run_experimentalist( Args: agent: Optional baseline agent for Mode 2, or an override for the agent - referenced by ``insight``. A local directory or a git ``url@ref``; a git + referenced by ``insight``. A local directory path or a git ``url@ref``; a git source is fetched by the backend and enables opening a draft PR/MR for the winner against that ref. agent_spec: Optional URI of a markdown file describing the agent under test. @@ -101,24 +101,24 @@ async def run_experimentalist( experiments_output=str(experiment_dir), storage=config.storage, ) - deps = ExperimentalistDeps( + result = await ExperimentRunner( + backend=backend, + strategy=build_experimentalist_agent( + working_dir=experiment_dir, + config=config, + framework_skills_dirs=framework_skills_dirs, + ), + config=config, workspace=workspace, + root=experiment_dir, agent=agent, agent_spec=agent_spec, insight=insight, train_dataset=train_dataset, validation_dataset=validation_dataset, task_template=task_template, - backend=backend, reporter=reporter, - config=config, - ) - experimentalist = build_experimentalist_agent( - working_dir=experiment_dir, - config=config, - framework_skills_dirs=framework_skills_dirs, - ) - result = await experimentalist.run(deps) + ).run() winner = result.winner reporter.run_finished( winner=winner.label if winner is not None else None, diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/runner.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/runner.py new file mode 100644 index 0000000000..b623872aa7 --- /dev/null +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/experimentalist/runner.py @@ -0,0 +1,441 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The experiment runner — the composition root. + +It does three things: **prepare inputs → run one strategy → persist and publish.** +Inputs are an agent directory, an optional ``Insight``, and the eval datasets; the +strategy is whatever the run was configured with; persisting and publishing are the +host's, not the strategy's. + +The runner is the only code that holds an +:class:`~nemo_experimentalist_plugin.experimentalist.experimentalist_backend.ExperimentalistBackend`. +Everything a strategy is allowed to reach goes through the +:class:`~nemo_experimentalist_plugin.experimentalist.context.ExperimentContext` it builds. +""" + +from __future__ import annotations + +import logging +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol + +from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig +from nemo_experimentalist_plugin.entities import ( + Candidate, + Dataset, + DatasetRef, + ExperimentRun, +) +from nemo_experimentalist_plugin.experimentalist.components.dataset_staging import stage_eval_author_inputs +from nemo_experimentalist_plugin.experimentalist.components.evaluator.base import EvaluatorType +from nemo_experimentalist_plugin.experimentalist.components.evaluator.factory import ( + DatasetFactory, + EvaluatorFactory, +) +from nemo_experimentalist_plugin.experimentalist.components.holdout_utils import restore_heldout_splits +from nemo_experimentalist_plugin.experimentalist.components.insight_promotion import ( + insight_suite_provenance, + select_insight_promotion_suggestions, + write_insight_comparison_section, + write_insight_promotion_section, +) +from nemo_experimentalist_plugin.experimentalist.context import ExperimentContext +from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ( + ExperimentalistBackend, +) +from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter, reward_scalar +from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult + +logger = logging.getLogger(__name__) + +#: Host-owned names the winner's directory must not carry into the user's workspace. +#: Composed from three owners, because the skip list has three: the backend's entity +#: metadata, the strategy's generated documentation, and the evaluator's harness +#: scaffolding. A third-party strategy's real output must survive this. +_BACKEND_ARTIFACTS = frozenset({"metadata.json"}) +_STRATEGY_ARTIFACTS = frozenset({"architecture.md"}) +_EVALUATOR_ARTIFACTS = frozenset({"harbor_wrapper.py", "dind_environment.py"}) + + +class Strategy(Protocol): + """What the runner requires of a strategy. + + ``supports_resume`` is a class attribute rather than an optional ``resume()`` so the + runner can check it without instantiating anything, and so a future + ``strategies list`` gets it for free. There is one entry point, not two: a strategy + reads ``ctx.resuming`` and those that do not care ignore it. + """ + + supports_resume: bool + + async def run(self, ctx: ExperimentContext) -> Candidate | None: + """Optimize the agent under test and return the winning Candidate, or None.""" + ... + + +@dataclass(frozen=True) +class PreparedInputs: + """Everything the strategy runs against, after the runner has materialized it.""" + + agent_dir: Path + agent_name: str + agent_spec: Path | None + insight_ref: str | None + datasets: dict[str, Dataset] + + +class ExperimentRunner: + """Wires one run up and calls it in order. + + Args: + backend: The run's data-access backend. Not shared with the strategy. + strategy: The optimization strategy to run. + config: Resolved run configuration. + workspace: NeMo Platform workspace. + root: Working directory for run artifacts. + agent: Baseline agent — a local directory, a git ``url@ref``, or None to take + the agent the Insight names. + agent_spec: Optional URI of a markdown description of the agent under test. + insight: Optional Insight id or local Insight file. + train_dataset: Dataset reference the strategy develops against. + validation_dataset: Dataset reference the winner is selected on. + task_template: Evaluator-specific task template; required with an Insight. + evaluator_type: Which evaluation implementation to build. + reporter: Optional human narration sink. + """ + + def __init__( + self, + *, + backend: ExperimentalistBackend, + strategy: Strategy, + config: EvolutionaryOptimizerConfig, + workspace: str, + root: Path, + agent: str | Path | None, + agent_spec: str | None = None, + insight: str | Path | None = None, + train_dataset: DatasetRef, + validation_dataset: DatasetRef, + task_template: DatasetRef | None = None, + evaluator_type: EvaluatorType = "harbor", + reporter: RunReporter | None = None, + ) -> None: + if agent is None and insight is None: + raise ValueError("One of 'insight' or 'agent' must be set.") + if insight is not None and task_template is None: + raise ValueError("'task_template' is required when 'insight' is set.") + self._backend = backend + self._strategy = strategy + self._config = config + self._workspace = workspace + self._root = root.resolve() + self._agent = agent + self._agent_spec = agent_spec + self._insight = insight + self._train_dataset = train_dataset + self._validation_dataset = validation_dataset + self._task_template = task_template + self._evaluator_type: EvaluatorType = evaluator_type + self._reporter = reporter + self._eo = self._root / "eval-and-optimize" + + async def run(self) -> ExperimentalistResult: + """Prepare inputs, run the strategy, then persist and publish the outcome.""" + self._preflight() + for subdir in ("agents", "analysis", "results"): + (self._eo / subdir).mkdir(parents=True, exist_ok=True) + + evaluator = EvaluatorFactory().build_evaluator( + self._evaluator_type, self._config.evaluator, experiment_dir=self._root + ) + inputs = await self._prepare_inputs() + run, resuming = await self._open_run(inputs) + ctx = ExperimentContext( + backend=self._backend, + workspace=self._workspace, + run=run, + root=self._root, + agent_dir=inputs.agent_dir, + agent_spec=inputs.agent_spec, + datasets=inputs.datasets, + evaluator=evaluator, + resuming=resuming, + reporter=self._reporter, + ) + if resuming: + await self._seed_narration_baseline(ctx) + + try: + winner = await self._strategy.run(ctx) + except Exception: + run.status = "failed" + await self._backend.update_run(workspace=self._workspace, run=run) + raise + + return await self._finalize(ctx, run, inputs, winner) + + # -- Prepare ------------------------------------------------------------- + + def _preflight(self) -> None: + """Fail fast when persistence is enabled but the tool it needs is missing.""" + storage = self._config.storage + if (storage.archive_candidates or storage.publish_winner) and shutil.which("git") is None: + raise ValueError( + "Candidate persistence is enabled (storage.archive_candidates/publish_winner) " + "but 'git' is not on PATH, so nothing can be persisted. Install git, or disable " + "storage to run without persistence." + ) + + async def _prepare_inputs(self) -> PreparedInputs: + """Materialize the agent, the datasets, and — with an Insight — its eval suite.""" + dataset_factory = DatasetFactory() + train_ref, validation_ref, template_ref = ( + self._train_dataset, + self._validation_dataset, + self._task_template, + ) + + insight = None + if self._insight is not None: + insight = await self._backend.get_insight(workspace=self._workspace, insight_id=str(self._insight)) + if self._backend.client is None: + raise ValueError("Platform client is required for insight task template loading") + assert template_ref is not None + staged = await stage_eval_author_inputs( + self._root, + train_dataset=train_ref, + validation_dataset=validation_ref, + task_template=template_ref, + client=self._backend.client, + workspace=self._workspace, + ) + train_ref, validation_ref, template_ref = ( + staged.train_dataset, + staged.validation_dataset, + staged.task_template, + ) + + datasets: dict[str, Dataset] = { + "train": dataset_factory.build_dataset(self._evaluator_type, train_ref), + "validation": dataset_factory.build_dataset(self._evaluator_type, validation_ref), + } + + agent_ref = self._agent if self._agent is not None else (insight.agent if insight is not None else None) + if agent_ref is None: + raise ValueError("Insight or agent is required") + agent_dir = self._eo / "source-agent" + await self._backend.get_agent_code( + workspace=self._workspace, + agent=agent_ref, + dest=agent_dir, + clone_depth=self._config.source.clone_depth, + ) + + agent_spec: Path | None = None + if self._agent_spec is not None: + agent_spec = await self._backend.get_agent_spec( + workspace=self._workspace, + spec=self._agent_spec, + dest=self._root / "AGENT-SPEC.md", + ) + + if insight is not None: + assert template_ref is not None + assert self._backend.client is not None + # Lazy: a run without an Insight never authors an eval suite, so it must not + # fail to import when the Eval Author package is absent. + from nemo_eval_author_plugin.eval_author.agent import EvalAuthor # noqa: PLC0415 + + authored = await EvalAuthor(experiment_dir=self._root, config=self._config.eval_author).run( + insight=insight, + agent_path=agent_dir, + task_template=dataset_factory.build_task_template(self._evaluator_type, template_ref), + train_dataset=datasets["train"], + validation_dataset=datasets["validation"], + client=self._backend.client, + ) + datasets["train"] = authored.train_dataset + datasets["validation"] = authored.validation_dataset + if authored.insight_suite is not None: + datasets["insight"] = authored.insight_suite + + return PreparedInputs( + agent_dir=agent_dir, + agent_name=str(agent_ref), + agent_spec=agent_spec, + insight_ref=str(self._insight) if self._insight is not None else None, + datasets=datasets, + ) + + async def _open_run(self, inputs: PreparedInputs) -> tuple[ExperimentRun, bool]: + """Re-open this run if one already exists here, else create it. + + Resume is strategy-independent at this level: the runner restores the + ``ExperimentRun`` and its inputs, and the strategy rebuilds its own state from + ``ctx.candidates()``. A strategy that cannot do that is refused loudly rather + than silently restarted — these runs cost hours, so the silent restart is the + expensive failure. + """ + existing = self._load_run() + if existing is not None: + if not self._strategy.supports_resume: + raise ValueError( + f"{type(self._strategy).__name__} does not support resume, but " + f"{self._eo / 'run.json'} already holds run {existing.id!r}. Point --experiment-dir " + "at a fresh directory, or delete that run to start over." + ) + logger.info("[RESUME] re-opening run %s", existing.id) + existing.status = "running" + await self._backend.update_run(workspace=self._workspace, run=existing) + return existing, True + + run = ExperimentRun( + workspace=self._workspace, + agent=inputs.agent_name, + insight=inputs.insight_ref, + config_snapshot=self._config.model_dump(mode="json"), + status="running", + ) + return await self._backend.create_run(workspace=self._workspace, run=run), False + + async def _seed_narration_baseline(self, ctx: ExperimentContext) -> None: + """Point the reporter's delta at the baseline it will not see re-evaluated. + + On resume the baseline is not scored again, so without this the first newly + evaluated candidate becomes the delta reference and every later delta is + measured against the wrong thing. Narration only; never fails a run. + """ + if self._reporter is None: + return + try: + baseline = next((c for c in await ctx.candidates() if c.ancestor is None), None) + except Exception as exc: # noqa: BLE001 - narration must never fail the run + logger.debug("[RESUME] could not seed the narration baseline: %s", exc) + return + if baseline is not None and baseline.reward("validation").metrics: + self._reporter.seed_baseline(reward_scalar(baseline.reward("validation").metrics)) + + def _load_run(self) -> ExperimentRun | None: + """Read this directory's ``run.json``, or None when it is absent or unreadable.""" + run_path = self._eo / "run.json" + if not run_path.exists(): + return None + try: + # ExperimentRun's wrap validator restores the computed id. + return ExperimentRun.model_validate_json(run_path.read_text()) + except Exception as exc: # noqa: BLE001 + logger.warning("[RESUME] Could not parse run.json: %s", exc) + return None + + # -- Persist and publish ------------------------------------------------- + + async def _finalize( + self, + ctx: ExperimentContext, + run: ExperimentRun, + inputs: PreparedInputs, + winner: Candidate | None, + ) -> ExperimentalistResult: + """Close the run out: report, workspace copy, terminal result, draft PR.""" + restore_heldout_splits(self._root) + report_path = self._eo / "OPTIMIZATION.md" + candidates = await ctx.candidates() + baseline = next((c for c in candidates if c.ancestor is None), None) + summary = _render_summary(run.progress_completed, run.progress_unit, baseline, winner) + + # The strategy writes the real report. When it produced nothing usable, leave a + # compact one behind so the Insight sections below have a document to append to. + if not report_path.exists() or not report_path.read_text().strip(): + report_path.write_text(f"# Optimization Report\n\n## Compact Run Summary\n\n{summary}\n") + + if winner is not None: + self._copy_winner_to_workspace(ctx.candidate_dir(winner)) + self._write_insight_sections(report_path, inputs, candidates, baseline, winner) + + result = ExperimentalistResult( + summary=summary, + run_id=run.id or "", + progress_completed=run.progress_completed, + winner=winner, + ) + run.status = "completed" + run.winner_agent = winner.id if winner is not None else None + await self._backend.update_run(workspace=self._workspace, run=run) + await self._backend.persist_result(workspace=self._workspace, result=result) + + if self._config.storage.publish_winner and winner is not None and winner.ancestor is not None: + try: + url = await self._backend.publish_candidate(workspace=self._workspace, candidate=winner) + if url: + logger.info("[PUBLISH] opened draft PR/MR for winner %s: %s", winner.label, url) + except Exception as exc: # noqa: BLE001 - publishing must never fail the run + logger.warning("[PERSISTENCE] publish failed for candidate %s; continuing: %s", winner.label, exc) + return result + + def _copy_winner_to_workspace(self, winner_dir: Path) -> None: + """Copy the winner's artifact to the workspace root, minus every owner's scaffolding.""" + skip = _BACKEND_ARTIFACTS | _STRATEGY_ARTIFACTS | _EVALUATOR_ARTIFACTS + if not winner_dir.is_dir(): + logger.warning("[FINAL] winner artifact %s is not a directory; skipping workspace copy", winner_dir) + return + for entry in winner_dir.iterdir(): + if entry.name in skip: + continue + dst = self._root / entry.name + if entry.is_dir(): + if dst.exists(): + shutil.rmtree(dst) + shutil.copytree(entry, dst) + else: + shutil.copy2(entry, dst) + + def _write_insight_sections( + self, + report_path: Path, + inputs: PreparedInputs, + candidates: list[Candidate], + baseline: Candidate | None, + winner: Candidate, + ) -> None: + """Append the Insight-suite comparison and promotion sections to the report. + + The insight suite is a reward channel the selector ignores, so it never affects + the winner — it only reports whether the scenarios authored for the motivating + Insight improved, and which of them are worth promoting. + """ + insight_dataset = inputs.datasets.get("insight") + if insight_dataset is None or not report_path.exists(): + return + try: + provenance = insight_suite_provenance(insight_dataset) + if baseline is not None: + write_insight_comparison_section(report_path, baseline, winner, provenance) + write_insight_promotion_section( + report_path, + select_insight_promotion_suggestions(insight_dataset, candidates, winner=winner), + ) + except ValueError as exc: + logger.warning("[FINAL] Skipping Insight Suite report sections: %s", exc) + + +def _render_summary(completed: int, unit: str, baseline: Candidate | None, winner: Candidate | None) -> str: + """One-line outcome summary, used when the strategy's own report is missing.""" + details: list[str] = [] + if winner is not None: + if winner.reward("validation").metrics: + details.append(f"validation_reward={winner.reward('validation').metrics}") + if baseline is not None and baseline.reward("insight").metrics and winner.reward("insight").metrics: + details.append( + f"insight_suite=(baseline={baseline.reward('insight').metrics}, " + f"winner={winner.reward('insight').metrics})" + ) + suffix = f", {', '.join(details)}" if details else "" + winner_str = winner.label if winner is not None else "none" + return f"Optimization complete: {completed} {unit}(s) completed, winner={winner_str}{suffix}" + + +__all__ = ["ExperimentRunner", "PreparedInputs", "Strategy"] diff --git a/plugins/nemo-experimentalist/tests/experimentalist/conftest.py b/plugins/nemo-experimentalist/tests/experimentalist/conftest.py new file mode 100644 index 0000000000..a04b1b8557 --- /dev/null +++ b/plugins/nemo-experimentalist/tests/experimentalist/conftest.py @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixture wrappers around the doubles in ``doubles.py``.""" + +from pathlib import Path + +import pytest +from doubles import FakeBackend, FakeEvaluator, make_context +from nemo_experimentalist_plugin.experimentalist.context import ExperimentContext + + +@pytest.fixture +def fake_backend() -> FakeBackend: + """An in-memory backend that records every entity it is handed.""" + return FakeBackend() + + +@pytest.fixture +def fake_evaluator() -> FakeEvaluator: + """An evaluator that scores everything 0.5 and records the options each run got.""" + return FakeEvaluator() + + +@pytest.fixture +def experiment_context(tmp_path: Path, fake_backend: FakeBackend, fake_evaluator: FakeEvaluator) -> ExperimentContext: + """A context over ``tmp_path`` wired to the fake backend and evaluator.""" + return make_context(root=tmp_path, backend=fake_backend, evaluator=fake_evaluator) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/doubles.py b/plugins/nemo-experimentalist/tests/experimentalist/doubles.py new file mode 100644 index 0000000000..31a28e2494 --- /dev/null +++ b/plugins/nemo-experimentalist/tests/experimentalist/doubles.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared doubles for the runner, the context, and anything a strategy touches. + +Before the runner existed every test that wanted to drive the loop had to assemble a +backend, an evaluator and a bag of ``deps`` by hand. These three build the same things +once: an in-memory backend that records what it was asked to persist, an evaluator that +returns a fixed result, and a context wired from both. + +``conftest.py`` re-exports these as fixtures; import the classes directly when a test +needs to build more than one. +""" + +from collections.abc import Mapping, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +from nemo_experimentalist_plugin.config import CandidateStorageConfig +from nemo_experimentalist_plugin.entities import ( + Candidate, + Dataset, + EvaluationResult, + ExperimentRun, + TrialResult, +) +from nemo_experimentalist_plugin.experimentalist.components.evaluator.base import Evaluator, EvaluatorConfig +from nemo_experimentalist_plugin.experimentalist.components.repository import AgentSource +from nemo_experimentalist_plugin.experimentalist.context import ExperimentContext +from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import ExperimentalistBackend +from nemo_experimentalist_plugin.experimentalist.reporting import RunReporter +from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult +from nemo_insights_plugin.entities import Insight +from nemo_platform import AsyncNeMoPlatform + + +def fake_client() -> AsyncNeMoPlatform: + """A stand-in platform client for paths that only check whether one is present.""" + return cast(AsyncNeMoPlatform, SimpleNamespace()) + + +class RecordedEvaluation(dict): + """One ``persist_evaluation`` call, as keyword arguments.""" + + +class FakeBackend(ExperimentalistBackend): + """In-memory backend that keeps everything it was asked to persist. + + Candidates are keyed by label the way the local backend keys them, so a test can + assert on ``backend.candidates["agent-1"]`` without reading any files. + """ + + def __init__( + self, + *, + client: AsyncNeMoPlatform | None = None, + storage: CandidateStorageConfig | None = None, + insight: Insight | None = None, + ) -> None: + super().__init__(client, None, storage) + self.candidates: dict[str, Candidate] = {} + self.runs: list[ExperimentRun] = [] + self.evaluations: list[RecordedEvaluation] = [] + self.results: list[ExperimentalistResult] = [] + self.archived: list[str] = [] + self.published: list[str] = [] + self.agent_code_calls: list[tuple[str | Path, Path]] = [] + self._insight = insight + + async def get_insight(self, *, workspace: str, insight_id: str) -> Insight: + if self._insight is None: + raise ValueError(f"FakeBackend has no insight to return for {insight_id!r}") + return self._insight + + async def create_run(self, *, workspace: str, run: ExperimentRun) -> ExperimentRun: + run._id = run.id or f"run-{len(self.runs) + 1}" # type: ignore[attr-defined] + self.runs.append(run) + return run + + async def update_run(self, *, workspace: str, run: ExperimentRun) -> ExperimentRun: + return run + + async def create_candidate(self, *, workspace: str, candidate: Candidate) -> Candidate: + candidate._id = candidate.id or candidate.label # type: ignore[attr-defined] + self.candidates[candidate.label] = candidate + return candidate + + async def update_candidate(self, *, workspace: str, candidate: Candidate) -> Candidate: + self.candidates[candidate.label] = candidate + return candidate + + async def get_candidate(self, *, workspace: str, candidate_id: str) -> Candidate: + return self.candidates[candidate_id] + + async def list_candidates(self, *, workspace: str, run_id: str) -> list[Candidate]: + return [c for c in self.candidates.values() if c.run_id == run_id] + + async def persist_result(self, *, workspace: str, result: ExperimentalistResult) -> None: + self.results.append(result) + + async def persist_evaluation( + self, *, workspace: str, result: EvaluationResult, candidate: Candidate, split: str + ) -> None: + self.evaluations.append( + RecordedEvaluation(workspace=workspace, result=result, candidate=candidate, split=split) + ) + + async def get_experiment_id(self, *, workspace: str, candidate: Candidate, split: str) -> str: + return f"exp-{candidate.label}-{split}" + + async def get_agent_code( + self, *, workspace: str, agent: str | Path, dest: Path, clone_depth: int | None = None + ) -> AgentSource | None: + self.agent_code_calls.append((agent, dest)) + dest.mkdir(parents=True, exist_ok=True) + return None + + async def archive_candidate(self, *, workspace: str, candidate: Candidate) -> str | None: + self.archived.append(candidate.label) + return f"archived://{candidate.label}" + + async def publish_candidate(self, *, workspace: str, candidate: Candidate) -> str | None: + self.published.append(candidate.label) + return f"https://example.invalid/pr/{candidate.label}" + + async def get_agent_spec(self, *, workspace: str, spec: str, dest: Path) -> Path: + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text(f"# spec from {spec}\n") + return dest + + +class FakeEvaluator(Evaluator): + """Returns a fixed reward and records the options each run was given.""" + + evaluator_type = "harbor" + + def __init__(self, *, reward: float = 0.5, options: EvaluatorConfig | None = None) -> None: + super().__init__(options or EvaluatorConfig(), None) + self.reward = reward + self.runs: list[dict[str, object]] = [] + + async def _run(self, agent: Path, dataset: Dataset, options: EvaluatorConfig) -> Sequence[TrialResult]: + self.runs.append({"agent": agent, "dataset": dataset, "options": options}) + return [ + TrialResult(id=f"{agent.name}-{task.id}", task_id=task.id, status="completed") + for task in dataset.list_tasks() + ] + + async def aggregate_results(self, results: Sequence[TrialResult]) -> dict[str, float | int]: + return {"reward": self.reward} + + +def make_context( + *, + root: Path, + backend: ExperimentalistBackend | None = None, + evaluator: Evaluator | None = None, + datasets: Mapping[str, Dataset] | None = None, + run: ExperimentRun | None = None, + workspace: str = "default", + resuming: bool = False, + reporter: RunReporter | None = None, +) -> ExperimentContext: + """Build a context over *root*, defaulting every collaborator to a double.""" + run = run or ExperimentRun(workspace=workspace, agent="agent-under-test") + if not run.id: + run._id = "run-1" # type: ignore[attr-defined] + agent_dir = root / "eval-and-optimize" / "source-agent" + agent_dir.mkdir(parents=True, exist_ok=True) + return ExperimentContext( + backend=backend or FakeBackend(), + workspace=workspace, + run=run, + root=root, + agent_dir=agent_dir, + datasets=datasets or {"train": Dataset(id="train"), "validation": Dataset(id="validation")}, + evaluator=evaluator or FakeEvaluator(), + resuming=resuming, + reporter=reporter, + ) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_loop.py b/plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_runner.py similarity index 68% rename from plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_loop.py rename to plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_runner.py index d2207f3f48..1b744aaf9e 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_loop.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_dataset_staging_runner.py @@ -1,17 +1,27 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Loop integration: Experimentalist stages Eval Author inputs before calling Eval Author.""" +"""The runner stages Eval Author's inputs before calling it, so the source is untouched.""" from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock import pytest +from doubles import FakeBackend, fake_client from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig from nemo_experimentalist_plugin.entities import DatasetRef -from nemo_experimentalist_plugin.experimentalist.components import loop as loop_module -from nemo_experimentalist_plugin.experimentalist.components.loop import EvolutionaryOptimizer +from nemo_experimentalist_plugin.experimentalist import runner as runner_module +from nemo_experimentalist_plugin.experimentalist.runner import ExperimentRunner +from nemo_insights_plugin.entities import Insight + + +class _UnreachedStrategy: + """Never runs: Eval Author raises while the runner is still preparing inputs.""" + + supports_resume = True + + async def run(self, ctx: object) -> None: + raise AssertionError("the strategy must not start when input preparation fails") def _write_tree(root: Path, content: str) -> None: @@ -62,45 +72,34 @@ async def run( raise RuntimeError("stop after eval_author") monkeypatch.setattr( - loop_module, + runner_module, "EvaluatorFactory", lambda: SimpleNamespace(build_evaluator=lambda *args, **kwargs: object()), ) - monkeypatch.setattr(loop_module, "DatasetFactory", RecordingDatasetFactory) - monkeypatch.setattr(loop_module, "EvalAuthor", MutatingEvalAuthor) - monkeypatch.setattr( - EvolutionaryOptimizer, - "_init_structure", - lambda self: (experiment / "agents", experiment / "analysis", experiment / "results"), - ) + monkeypatch.setattr(runner_module, "DatasetFactory", RecordingDatasetFactory) + monkeypatch.setattr("nemo_eval_author_plugin.eval_author.agent.EvalAuthor", MutatingEvalAuthor) - backend = SimpleNamespace( - client=object(), - get_insight=AsyncMock(return_value=SimpleNamespace(agent=str(tmp_path / "agent"))), - get_agent_code=AsyncMock(), - ) - config = EvolutionaryOptimizerConfig() - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = experiment - optimizer.config = config - optimizer.shell = SimpleNamespace(close=AsyncMock()) - deps = SimpleNamespace( - backend=backend, + agent_dir = tmp_path / "agent" + agent_dir.mkdir() + runner = ExperimentRunner( + backend=FakeBackend( + client=fake_client(), + insight=Insight(workspace="default", agent=str(agent_dir), title="t", description="d"), + ), + strategy=_UnreachedStrategy(), + config=EvolutionaryOptimizerConfig(), workspace="default", - config=config, - evaluator_type="harbor", + root=experiment, + agent=agent_dir, + insight="insight-1", train_dataset=DatasetRef(uri=str(train)), validation_dataset=DatasetRef(uri=str(validation)), task_template=DatasetRef(uri=str(template)), - insight="insight-1", - agent=tmp_path / "agent", - agent_spec=None, ) with pytest.raises(RuntimeError, match="stop after eval_author"): - await optimizer.run(deps) + await runner.run() - optimizer.shell.close.assert_awaited_once() assert captured == { "train": experiment / "dataset" / "train", "validation": experiment / "dataset" / "validation", diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_failure.py b/plugins/nemo-experimentalist/tests/experimentalist/test_loop_failure.py deleted file mode 100644 index cefb0e3ae3..0000000000 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_failure.py +++ /dev/null @@ -1,129 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import asyncio -import json -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest -from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig -from nemo_experimentalist_plugin.entities import Candidate, ExperimentRun -from nemo_experimentalist_plugin.experimentalist.components import loop as loop_module -from nemo_experimentalist_plugin.experimentalist.components.loop import EvolutionaryOptimizer -from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import LocalExperimentalistBackend - - -@pytest.mark.asyncio -@pytest.mark.parametrize("failure_method", ["create", "evaluate"]) -async def test_baseline_failure_marks_run_failed(monkeypatch, tmp_path, failure_method): - candidate = SimpleNamespace(label="agent-0") - evolution_tree = SimpleNamespace(survivors=lambda round_num: [candidate]) - backend = LocalExperimentalistBackend(path=tmp_path) - run_entity = await backend.create_run( - workspace="default", - run=ExperimentRun( - workspace="default", - agent="agent", - config_snapshot={}, - status="running", - rounds_completed=0, - ), - ) - monkeypatch.setattr(backend, "get_agent_code", AsyncMock()) - config = EvolutionaryOptimizerConfig() - - monkeypatch.setattr( - loop_module, - "EvaluatorFactory", - lambda: SimpleNamespace(build_evaluator=lambda *args, **kwargs: object()), - ) - monkeypatch.setattr( - loop_module, - "DatasetFactory", - lambda: SimpleNamespace(build_dataset=lambda *args, **kwargs: object()), - ) - monkeypatch.setattr(loop_module.EvolutionTree, "from_dir", lambda path: evolution_tree) - monkeypatch.setattr(EvolutionaryOptimizer, "_init_structure", lambda self: (tmp_path, tmp_path, tmp_path)) - monkeypatch.setattr(EvolutionaryOptimizer, "_detect_last_round", lambda self: None) - monkeypatch.setattr( - EvolutionaryOptimizer, - "_create_experiment_run", - AsyncMock(return_value=run_entity), - ) - baseline_failure = ValueError("baseline failed") - create_baseline = ( - AsyncMock(side_effect=baseline_failure) if failure_method == "create" else AsyncMock(return_value=candidate) - ) - monkeypatch.setattr(EvolutionaryOptimizer, "_create_baseline_agent", create_baseline) - monkeypatch.setattr(EvolutionaryOptimizer, "_update_candidate", AsyncMock()) - evaluate_validation = ( - AsyncMock(side_effect=baseline_failure) - if failure_method == "evaluate" - else AsyncMock(return_value={"agent-0": object()}) - ) - monkeypatch.setattr(EvolutionaryOptimizer, "_evaluate_validation_candidates", evaluate_validation) - - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = tmp_path - optimizer.config = config - optimizer.shell = SimpleNamespace(close=AsyncMock()) - deps = SimpleNamespace( - backend=backend, - workspace="default", - config=config, - evaluator_type="harbor", - train_dataset=object(), - validation_dataset=object(), - insight=None, - agent=tmp_path / "agent", - agent_spec=None, - task_template=None, - ) - - with pytest.raises(ValueError, match="baseline failed"): - await optimizer.run(deps) - - optimizer.shell.close.assert_awaited_once() - assert run_entity.status == "failed" - saved = json.loads((tmp_path / "eval-and-optimize" / "run.json").read_text(encoding="utf-8")) - assert saved["status"] == "failed" - - -@pytest.mark.asyncio -async def test_a_cancelled_coder_does_not_pass_as_a_built_candidate(monkeypatch, tmp_path): - """Cancellation must unwind the round, not be filed as an ordinary build failure. - - ``asyncio.gather(..., return_exceptions=True)`` hands back ``CancelledError``, which - derives from ``BaseException``. An ``isinstance(r, Exception)`` filter therefore sees - neither a failure nor a success, and the candidate would flow on to evaluation and - ranking as though its source had been written. - """ - - class _CancellingCoder: - def __init__(self, **kwargs): - pass - - async def run(self, *args, **kwargs): - raise asyncio.CancelledError - - monkeypatch.setattr(loop_module, "Coder", _CancellingCoder) - monkeypatch.setattr(EvolutionaryOptimizer, "_snapshot_metadata", lambda self, name: None) - monkeypatch.setattr(EvolutionaryOptimizer, "_restore_metadata", lambda self, name, snap: None) - monkeypatch.setattr(EvolutionaryOptimizer, "_coder_config", lambda self, config: None) - monkeypatch.setattr(EvolutionaryOptimizer, "_update_candidate", AsyncMock()) - - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = tmp_path - optimizer._framework_skills_dirs = [] - candidate = Candidate(run_id="run-1", label="agent-1", round=1, optimization="x") - - with pytest.raises(asyncio.CancelledError): - await optimizer._implement_candidates( - workspace="default", - backend=AsyncMock(), - dataset=object(), - evaluator=object(), - candidates=[candidate], - config=EvolutionaryOptimizerConfig(), - ) diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_insight_suite.py b/plugins/nemo-experimentalist/tests/experimentalist/test_loop_insight_suite.py index 75768c4aa4..13798b403f 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_insight_suite.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_loop_insight_suite.py @@ -3,15 +3,14 @@ from pathlib import Path from types import SimpleNamespace -from typing import Any from unittest.mock import AsyncMock import pytest +from doubles import FakeBackend, FakeEvaluator, make_context from nemo_experimentalist_plugin.config import EvolutionaryOptimizerConfig from nemo_experimentalist_plugin.entities import ( Candidate, Dataset, - DatasetRef, DataValue, EvaluationResult, MetricResult, @@ -69,37 +68,13 @@ async def test_insight_run_evaluates_and_persists_baseline_and_new_candidate_met monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - train_dataset = Dataset(id="train") - validation_dataset = Dataset(id="validation") + """The insight suite is scored once for the baseline and once per new candidate.""" insight_dataset = Dataset( id="insight-suite", source=ResourceRef(uri="file:///experiment/eval-and-optimize/eval_author/insight-1/insight-suite"), tasks=[Task(id="insight-task")], metadata=_suite_metadata(), ) - datasets = { - "train": train_dataset, - "validation": validation_dataset, - } - - class RecordingDatasetFactory: - def build_dataset(self, evaluator_type: str, ref: DatasetRef) -> Dataset: - return datasets[ref.uri] - - def build_task_template(self, evaluator_type: str, ref: DatasetRef) -> Task: - return Task(id="template", uri=ref.uri) - - class ReturningEvalAuthor: - def __init__(self, **kwargs: object) -> None: - pass - - async def run(self, **kwargs: Any) -> SimpleNamespace: - return SimpleNamespace( - train_dataset=kwargs["train_dataset"], - validation_dataset=kwargs["validation_dataset"], - insight_suite=insight_dataset, - ) - baseline = Candidate(run_id="run-1", label="agent-0", round=0, optimization="baseline") new_candidate = Candidate( run_id="run-1", @@ -117,8 +92,8 @@ async def run(self, **kwargs: Any) -> SimpleNamespace: async def evaluate_insight_candidates( self: EvolutionaryOptimizer, *, + ctx: object, dataset: Dataset, - evaluator: object, candidates: list[Candidate], ) -> dict[str, EvaluationResult]: insight_evaluations.append((dataset, candidates)) @@ -139,24 +114,6 @@ async def evaluate_validation_candidates( if "validation" not in candidate.rewards } - async def update_candidate( - self: EvolutionaryOptimizer, - candidate: Candidate, - *, - updates: dict[str, object] | None = None, - **kwargs: object, - ) -> None: - for key, value in (updates or {}).items(): - setattr(candidate, key, value) - - run_entity = SimpleNamespace(id="run-1", status="running", rounds_completed=0) - backend = SimpleNamespace( - client=object(), - get_insight=AsyncMock(return_value=SimpleNamespace(agent="agent-source")), - get_agent_code=AsyncMock(), - persist_evaluation=AsyncMock(), - update_run=AsyncMock(), - ) evolution_tree = SimpleNamespace(survivors=lambda round_num: [baseline], add=lambda candidate: None) class StopAfterOneRoundTerminator: @@ -168,23 +125,9 @@ async def run(self, **kwargs: object) -> SimpleNamespace: raise _StopAfterOneRound return SimpleNamespace(stop=False, reason="continue") - monkeypatch.setattr(loop_module, "DatasetFactory", RecordingDatasetFactory) - monkeypatch.setattr( - loop_module, - "EvaluatorFactory", - lambda: SimpleNamespace(build_evaluator=lambda *args, **kwargs: object()), - ) - monkeypatch.setattr(loop_module, "EvalAuthor", ReturningEvalAuthor) - monkeypatch.setattr( - loop_module, - "stage_eval_author_inputs", - AsyncMock(side_effect=lambda _, **refs: SimpleNamespace(**refs)), - ) monkeypatch.setattr(loop_module.EvolutionTree, "from_dir", lambda path: evolution_tree) monkeypatch.setattr(EvolutionaryOptimizer, "_detect_last_round", lambda self: None) - monkeypatch.setattr(EvolutionaryOptimizer, "_create_experiment_run", AsyncMock(return_value=run_entity)) monkeypatch.setattr(EvolutionaryOptimizer, "_create_baseline_agent", AsyncMock(return_value=baseline)) - monkeypatch.setattr(EvolutionaryOptimizer, "_update_candidate", update_candidate) monkeypatch.setattr( EvolutionaryOptimizer, "_evaluate_validation_candidates", @@ -194,7 +137,6 @@ async def run(self, **kwargs: object) -> SimpleNamespace: EvolutionaryOptimizer, "_evaluate_insight_candidates", evaluate_insight_candidates, - raising=False, ) monkeypatch.setattr(EvolutionaryOptimizer, "_generate_initial_goal_tree", AsyncMock()) monkeypatch.setattr( @@ -216,27 +158,24 @@ async def run(self, **kwargs: object) -> SimpleNamespace: AsyncMock(side_effect=lambda **kwargs: kwargs["candidates"]), ) - config = EvolutionaryOptimizerConfig(disable_trajectory_scoring=True) optimizer = object.__new__(EvolutionaryOptimizer) optimizer.working_dir = tmp_path - optimizer.config = config + optimizer.config = EvolutionaryOptimizerConfig(disable_trajectory_scoring=True) optimizer.shell = SimpleNamespace(close=AsyncMock()) optimizer.terminator = StopAfterOneRoundTerminator() - deps = SimpleNamespace( + backend = FakeBackend() + ctx = make_context( + root=tmp_path, backend=backend, - workspace="default", - config=config, - evaluator_type="harbor", - train_dataset=DatasetRef(uri="train"), - validation_dataset=DatasetRef(uri="validation"), - task_template=DatasetRef(uri="template"), - insight="insight-1", - agent=None, - agent_spec=None, + datasets={ + "train": Dataset(id="train"), + "validation": Dataset(id="validation"), + "insight": insight_dataset, + }, ) with pytest.raises(_StopAfterOneRound): - await optimizer.run(deps) + await optimizer.run(ctx) assert insight_evaluations == [ (insight_dataset, [baseline]), @@ -247,9 +186,7 @@ async def run(self, **kwargs: object) -> SimpleNamespace: assert baseline.rewards["insight"].metadata["suite_identity"] == f"sha256:{'a' * 64}" assert new_candidate.rewards["insight"].metadata["suite_identity"] == f"sha256:{'a' * 64}" assert baseline.rewards["insight"].metadata["metric_keys"] == ["uses_required_tool"] - insight_persistence = [ - call.kwargs for call in backend.persist_evaluation.await_args_list if call.kwargs["split"] == "insight" - ] + insight_persistence = [call for call in backend.evaluations if call["split"] == "insight"] assert [call["candidate"] for call in insight_persistence] == [baseline, new_candidate] assert [call["result"].id for call in insight_persistence] == [ insight_results["agent-0"].id, @@ -268,6 +205,7 @@ async def run(self, **kwargs: object) -> SimpleNamespace: @pytest.mark.asyncio async def test_insight_evaluation_skips_cached_candidates_and_empty_suites( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: cached = Candidate( run_id="run-1", @@ -293,25 +231,26 @@ async def test_insight_evaluation_skips_cached_candidates_and_empty_suites( monkeypatch.setattr(EvolutionaryOptimizer, "_evaluate_agent", evaluate_agent) optimizer = object.__new__(EvolutionaryOptimizer) + ctx = make_context(root=tmp_path) evaluated = await optimizer._evaluate_insight_candidates( + ctx=ctx, dataset=Dataset( id="insight-suite", source=ResourceRef(uri="file:///experiment/eval-and-optimize/eval_author/insight-1/insight-suite"), tasks=[Task(id="insight-task")], metadata=_suite_metadata(), ), - evaluator=object(), # type: ignore[arg-type] candidates=[cached, pending], ) assert evaluated == {"agent-1": result} assert evaluate_agent.await_args is not None - assert evaluate_agent.await_args.args[0] is pending + assert evaluate_agent.await_args.args[1] is pending assert evaluate_agent.await_args.kwargs["minimum_attempts"] == 2 empty = await optimizer._evaluate_insight_candidates( + ctx=ctx, dataset=Dataset(id="empty-insight-suite"), - evaluator=object(), # type: ignore[arg-type] candidates=[pending], ) assert empty == {} @@ -321,6 +260,7 @@ async def test_insight_evaluation_skips_cached_candidates_and_empty_suites( @pytest.mark.asyncio async def test_insight_evaluation_reuses_only_matching_suite_identity( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: cached = Candidate( run_id="run-1", @@ -353,17 +293,18 @@ async def test_insight_evaluation_reuses_only_matching_suite_identity( metadata=_suite_metadata("e"), ) + ctx = make_context(root=tmp_path) assert ( await optimizer._evaluate_insight_candidates( + ctx=ctx, dataset=matching, - evaluator=object(), # type: ignore[arg-type] candidates=[cached], ) == {} ) assert await optimizer._evaluate_insight_candidates( + ctx=ctx, dataset=changed, - evaluator=object(), # type: ignore[arg-type] candidates=[cached], ) == {"agent-0": result} evaluate_agent.assert_awaited_once() @@ -372,6 +313,7 @@ async def test_insight_evaluation_reuses_only_matching_suite_identity( @pytest.mark.asyncio async def test_cached_insight_metric_keys_are_order_independent( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: identity = f"sha256:{'a' * 64}" candidates = [ @@ -419,12 +361,9 @@ async def test_cached_insight_metric_keys_are_order_independent( ) await optimizer._evaluate_and_persist_insight_candidates( + ctx=make_context(root=tmp_path), dataset=dataset, - evaluator=object(), # type: ignore[arg-type] candidates=candidates, - workspace="default", - backend=SimpleNamespace(), - run_id="run-1", ) assert dataset.metadata["insight_metric_keys"] == ["reward", "uses_required_tool"] @@ -442,28 +381,23 @@ async def test_insight_evaluation_uses_at_least_two_attempts_without_changing_ot ) received_attempts: list[int] = [] - class RecordingEvaluator: - options = HarborEvaluatorConfig(n_attempts=1) - - async def run(self, **kwargs: object) -> EvaluationResult: + class RecordingEvaluator(FakeEvaluator): + async def run(self, **kwargs: object) -> EvaluationResult: # type: ignore[override] options = kwargs["options"] assert isinstance(options, HarborEvaluatorConfig) received_attempts.append(options.n_attempts) return _insight_result(candidate.label, 0.5) - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = tmp_path - evaluator = RecordingEvaluator() - dataset = Dataset(id="insight-suite") - - await optimizer._evaluate_agent(candidate, dataset, evaluator) # type: ignore[arg-type] - await optimizer._evaluate_agent( - candidate, - dataset, - evaluator, # type: ignore[arg-type] - minimum_attempts=2, + evaluator = RecordingEvaluator(options=HarborEvaluatorConfig(n_attempts=1)) + ctx = make_context( + root=tmp_path, + evaluator=evaluator, + datasets={"validation": Dataset(id="validation"), "insight": Dataset(id="insight-suite")}, ) + await ctx.evaluate(candidate, split="validation") + await ctx.evaluate(candidate, split="insight", minimum_attempts=2) + assert received_attempts == [1, 2] assert evaluator.options.n_attempts == 1 diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_reporting.py b/plugins/nemo-experimentalist/tests/experimentalist/test_loop_reporting.py index a019d4943a..3900bf5b57 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_loop_reporting.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_loop_reporting.py @@ -3,8 +3,6 @@ import math from pathlib import Path -from types import SimpleNamespace -from unittest.mock import AsyncMock import pytest from nemo_experimentalist_plugin.entities import ( @@ -28,7 +26,7 @@ write_insight_promotion_section, ) from nemo_experimentalist_plugin.experimentalist.components.loop import AnalysisSkill, EvolutionaryOptimizer -from nemo_experimentalist_plugin.experimentalist.components.models import EvolutionTree +from nemo_experimentalist_plugin.experimentalist.runner import _render_summary as render_summary _SUITE_IDENTITY = f"sha256:{'a' * 64}" _SUITE_PATH = Path("/experiment/eval-and-optimize/eval_author/insight-1/insight-suite") @@ -113,9 +111,7 @@ def test_terminal_summary_includes_baseline_and_winner_insight_metrics() -> None "validation": RewardRecord(metrics={"reward": 0.75}), }, ) - optimizer = object.__new__(EvolutionaryOptimizer) - - summary = optimizer._render_summary(rounds_completed=1, baseline=baseline, winner=winner) + summary = render_summary(1, "round", baseline, winner) assert "validation_reward={'reward': 0.75}" in summary assert "insight_suite=(baseline={'uses_required_tool': 0.0}" in summary @@ -125,9 +121,7 @@ def test_terminal_summary_includes_baseline_and_winner_insight_metrics() -> None def test_terminal_summary_omits_insight_comparison_when_unavailable() -> None: baseline = _candidate("agent-0", round_num=0) winner = _candidate("agent-1", round_num=1, rewards={"validation": RewardRecord(metrics={"reward": 0.75})}) - optimizer = object.__new__(EvolutionaryOptimizer) - - summary = optimizer._render_summary(rounds_completed=1, baseline=baseline, winner=winner) + summary = render_summary(1, "round", baseline, winner) assert "insight_suite" not in summary @@ -162,7 +156,7 @@ def test_insight_promotion_suggestions_are_stable_discriminative_and_diverse( Task(id="task-flat", uri=(tmp_path / "task-flat").as_uri()), ] baseline = _candidate("agent-0", round_num=0) - baseline.record_reward( + baseline.set_reward( "insight", trials=[ _insight_trial("task-a", 0.0, attempt=1), @@ -178,10 +172,10 @@ def test_insight_promotion_suggestions_are_stable_discriminative_and_diverse( ], ) winner = _candidate("agent-1", round_num=1) - winner.record_reward( + winner.set_reward( "insight", metadata={**winner.rewards["insight"].metadata, "metric_keys": ["uses_required_tool", "reward"]} ) - winner.record_reward( + winner.set_reward( "insight", trials=[ _insight_trial("task-a", 1.0, attempt=1), @@ -227,11 +221,9 @@ def test_task_evidence_excludes_candidates_from_other_suites(tmp_path: Path) -> baseline = _candidate("agent-0", round_num=0) winner = _candidate("agent-1", round_num=1) stale = _candidate("agent-stale", round_num=1) - stale.record_reward( - "insight", metadata={**stale.rewards["insight"].metadata, "suite_identity": f"sha256:{'e' * 64}"} - ) + stale.set_reward("insight", metadata={**stale.rewards["insight"].metadata, "suite_identity": f"sha256:{'e' * 64}"}) for candidate, score in ((baseline, 0.0), (winner, 1.0), (stale, 0.5)): - candidate.record_reward( + candidate.set_reward( "insight", trials=[ _insight_trial(task.id, score, attempt=1), @@ -340,7 +332,7 @@ def test_invalid_runtime_metrics_cannot_be_promotion_evidence( task = Task(id="task-a") baseline = _candidate("agent-0", round_num=0) winner = _candidate("agent-1", round_num=1) - baseline.record_reward( + baseline.set_reward( "insight", trials=[ _insight_trial("task-a", 0.0, attempt=1), @@ -350,7 +342,7 @@ def test_invalid_runtime_metrics_cannot_be_promotion_evidence( invalid_trial = _insight_trial("task-a", invalid_score, attempt=1) if missing_key: invalid_trial.metrics.pop("uses_required_tool") - winner.record_reward( + winner.set_reward( "insight", trials=[ invalid_trial, @@ -372,8 +364,8 @@ def test_one_attempt_failed_and_incomplete_evidence_do_not_qualify_as_stable() - task = Task(id="task-a") baseline = _candidate("agent-0", round_num=0) winner = _candidate("agent-1", round_num=1) - baseline.record_reward("insight", trials=[_insight_trial("task-a", 0.0)]) - winner.record_reward("insight", trials=[_insight_trial("task-a", 1.0)]) + baseline.set_reward("insight", trials=[_insight_trial("task-a", 0.0)]) + winner.set_reward("insight", trials=[_insight_trial("task-a", 1.0)]) assert ( select_insight_promotion_suggestions( @@ -395,7 +387,7 @@ def test_one_attempt_failed_and_incomplete_evidence_do_not_qualify_as_stable() - == [] ) - winner.record_reward("insight", trials=[]) + winner.set_reward("insight", trials=[]) assert ( select_insight_promotion_suggestions( _insight_dataset([task]), @@ -423,14 +415,14 @@ def test_promotion_requires_baseline_to_winner_improvement( baseline = _candidate("agent-0", round_num=0) winner = _candidate("agent-1", round_num=1) candidates = [baseline, winner] - baseline.record_reward( + baseline.set_reward( "insight", trials=[ _insight_trial("task-a", baseline_score, attempt=1), _insight_trial("task-a", baseline_score, attempt=2), ], ) - winner.record_reward( + winner.set_reward( "insight", trials=[ _insight_trial("task-a", winner_score, attempt=1), @@ -439,7 +431,7 @@ def test_promotion_requires_baseline_to_winner_improvement( ) if bad_score is not None: bad = _candidate("agent-bad", round_num=1) - bad.record_reward( + bad.set_reward( "insight", trials=[ _insight_trial("task-a", bad_score, attempt=1), @@ -481,127 +473,3 @@ def test_deterministic_insight_comparison_section_uses_local_suite_identity( assert str(_SUITE_PATH) in report assert _SUITE_IDENTITY in report assert "| `uses_required_tool` | 0.000 | 1.000 | +1.000 |" in report - - -@pytest.mark.asyncio -async def test_final_report_failure_preserves_compact_summary_and_deterministic_sections( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - baseline = _candidate( - "agent-0", - round_num=0, - rewards={ - "insight": RewardRecord(metrics={"reward": 0.5, "uses_required_tool": 0.0}), - "validation": RewardRecord(metrics={"reward": 0.5}), - }, - ) - winner = _candidate( - "agent-1", - round_num=1, - rewards={ - "insight": RewardRecord(metrics={"reward": 0.75, "uses_required_tool": 1.0}), - "validation": RewardRecord(metrics={"reward": 0.75}), - }, - ) - tree = EvolutionTree() - tree.add(baseline) - tree.add(winner) - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = tmp_path - (tmp_path / "eval-and-optimize").mkdir() - monkeypatch.setattr(optimizer, "_copy_best_to_workspace", lambda best_id: None) - original_report_writer = EvolutionaryOptimizer.write_final_report - type.__setattr__( - EvolutionaryOptimizer, - "write_final_report", - AsyncMock(side_effect=RuntimeError("LLM report failed")), - ) - run = SimpleNamespace(status="running", winner_agent=None, rounds_completed=1) - backend = SimpleNamespace(update_run=AsyncMock()) - - try: - finalized = await optimizer._finalize( - workspace="default", - backend=backend, - agents_dir=tmp_path / "eval-and-optimize" / "agents", - run_entity=run, - evolution_tree=tree, - agent_name="agent", - insight_dataset=_insight_dataset([Task(id="task-a")]), - ) - finally: - type.__setattr__( - EvolutionaryOptimizer, - "write_final_report", - original_report_writer, - ) - - report = (tmp_path / "eval-and-optimize" / "OPTIMIZATION.md").read_text() - assert finalized is winner - assert "## Compact Run Summary" in report - assert "Optimization complete: 1 round(s) completed" in report - assert "## Deterministic Insight Suite Comparison" in report - assert "## Insight Suite Promotion Suggestions" in report - - -@pytest.mark.asyncio -async def test_insight_report_mismatch_does_not_fail_completed_run( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, -) -> None: - baseline = _candidate( - "agent-0", - round_num=0, - rewards={ - "insight": RewardRecord(metrics={"reward": 0.5, "uses_required_tool": 0.0}), - "validation": RewardRecord(metrics={"reward": 0.5}), - }, - ) - winner = _candidate( - "agent-1", - round_num=1, - rewards={ - "insight": RewardRecord(metrics={"reward": 0.75, "uses_required_tool": 1.0}), - "validation": RewardRecord(metrics={"reward": 0.75}), - }, - ) - winner.record_reward( - "insight", metadata={**winner.rewards["insight"].metadata, "suite_identity": f"sha256:{'e' * 64}"} - ) - tree = EvolutionTree() - tree.add(baseline) - tree.add(winner) - optimizer = object.__new__(EvolutionaryOptimizer) - optimizer.working_dir = tmp_path - (tmp_path / "eval-and-optimize").mkdir() - monkeypatch.setattr(optimizer, "_copy_best_to_workspace", lambda best_id: None) - original_report_writer = EvolutionaryOptimizer.write_final_report - type.__setattr__(EvolutionaryOptimizer, "write_final_report", AsyncMock()) - run = SimpleNamespace(status="running", winner_agent=None, rounds_completed=1) - backend = SimpleNamespace(update_run=AsyncMock()) - - try: - with caplog.at_level("WARNING"): - finalized = await optimizer._finalize( - workspace="default", - backend=backend, - agents_dir=tmp_path / "eval-and-optimize" / "agents", - run_entity=run, - evolution_tree=tree, - agent_name="agent", - insight_dataset=_insight_dataset([Task(id="task-a")]), - ) - finally: - type.__setattr__( - EvolutionaryOptimizer, - "write_final_report", - original_report_writer, - ) - - assert finalized is winner - assert run.status == "completed" - assert run.winner_agent == winner.label - backend.update_run.assert_awaited_once() - assert "Skipping Insight Suite report sections" in caplog.text diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_reporting.py b/plugins/nemo-experimentalist/tests/experimentalist/test_reporting.py index 1c5b69fe77..52afb1d50e 100644 --- a/plugins/nemo-experimentalist/tests/experimentalist/test_reporting.py +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_reporting.py @@ -8,6 +8,8 @@ import io from pathlib import Path +import pytest +from doubles import make_context from nemo_experimentalist_plugin.experimentalist.reporting import ( RunReporter, Verbosity, @@ -199,17 +201,37 @@ def test_build_run_reporter_emits_header_and_is_reusable() -> None: assert "baseline" in sink.getvalue() -def test_deps_accepts_a_reporter() -> None: - from nemo_experimentalist_plugin.entities import DatasetRef - from nemo_experimentalist_plugin.experimentalist.deps import ExperimentalistDeps +@pytest.mark.asyncio +async def test_report_progress_narrates_and_records_the_counter(tmp_path: Path) -> None: + """One verb feeds both the human narration and the run entity.""" + sink = io.StringIO() + ctx = make_context(root=tmp_path, reporter=RunReporter(sink=sink)) - reporter = RunReporter(sink=io.StringIO()) - deps = ExperimentalistDeps( - workspace="w", - reporter=reporter, - insight=Path("/test/insight"), - train_dataset=DatasetRef(uri="test-train"), - validation_dataset=DatasetRef(uri="test-val"), - task_template=DatasetRef(uri="test-task"), - ) - assert deps.reporter is reporter + await ctx.report_progress(completed=7, total=15, unit="round", note="evaluating candidates") + + assert "round 7/≤15" in sink.getvalue() + assert "evaluating candidates" in sink.getvalue() + assert (ctx._run.progress_completed, ctx._run.progress_total, ctx._run.progress_unit) == (7, 15, "round") + + +@pytest.mark.asyncio +async def test_report_progress_without_a_total_is_still_honest(tmp_path: Path) -> None: + """An opaque strategy has no denominator, so it reports a bare counter.""" + sink = io.StringIO() + ctx = make_context(root=tmp_path, reporter=RunReporter(sink=sink)) + + await ctx.report_progress(completed=340, unit="trial", note="bootstrapping demos") + + assert "trial 340" in sink.getvalue() + assert "≤" not in sink.getvalue() + assert ctx._run.progress_total is None + + +def test_note_narrates_without_touching_the_run(tmp_path: Path) -> None: + sink = io.StringIO() + ctx = make_context(root=tmp_path, reporter=RunReporter(sink=sink)) + + ctx.note("agent-3 (1/2): add a retrieval step") + + assert "agent-3 (1/2): add a retrieval step" in sink.getvalue() + assert ctx._run.progress_completed == 0 diff --git a/plugins/nemo-experimentalist/tests/experimentalist/test_runner.py b/plugins/nemo-experimentalist/tests/experimentalist/test_runner.py new file mode 100644 index 0000000000..b462c0efef --- /dev/null +++ b/plugins/nemo-experimentalist/tests/experimentalist/test_runner.py @@ -0,0 +1,453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The runner's own contract: prepare inputs, run one strategy, persist and publish.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar + +import pytest +from doubles import FakeBackend, FakeEvaluator, fake_client +from nemo_experimentalist_plugin.config import CandidateStorageConfig, EvolutionaryOptimizerConfig +from nemo_experimentalist_plugin.entities import ( + Candidate, + Dataset, + DatasetRef, + ExperimentRun, + ResourceRef, + RewardRecord, + Task, +) +from nemo_experimentalist_plugin.experimentalist import runner as runner_module +from nemo_experimentalist_plugin.experimentalist.context import ExperimentContext +from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import LocalExperimentalistBackend +from nemo_experimentalist_plugin.experimentalist.runner import ExperimentRunner +from nemo_insights_plugin.entities import Insight + + +class RecordingStrategy: + """Keeps the context it was handed, then returns (or raises) what it was told to.""" + + supports_resume: ClassVar[bool] = True + + def __init__(self, winner: Candidate | None = None, error: Exception | None = None) -> None: + self.winner = winner + self.error = error + self.ctx: ExperimentContext | None = None + + async def run(self, ctx: ExperimentContext) -> Candidate | None: + self.ctx = ctx + if self.error is not None: + raise self.error + if self.winner is not None: + await ctx.save_candidate(self.winner) + return self.winner + + +class NonResumableStrategy(RecordingStrategy): + supports_resume: ClassVar[bool] = False + + +def _stub_factories(monkeypatch: pytest.MonkeyPatch) -> None: + """Replace evaluator/dataset construction, which would otherwise need Harbor.""" + monkeypatch.setattr( + runner_module, + "EvaluatorFactory", + lambda: type("F", (), {"build_evaluator": staticmethod(lambda *a, **k: FakeEvaluator())})(), + ) + monkeypatch.setattr( + runner_module, + "DatasetFactory", + lambda: type( + "F", + (), + { + "build_dataset": staticmethod(lambda _type, ref: Dataset(id=ref.uri)), + "build_task_template": staticmethod(lambda _type, ref: Task(id="template", uri=ref.uri)), + }, + )(), + ) + + +def _make_runner( + tmp_path: Path, + *, + strategy: RecordingStrategy, + monkeypatch: pytest.MonkeyPatch, + backend: FakeBackend | None = None, + config: EvolutionaryOptimizerConfig | None = None, +) -> tuple[ExperimentRunner, FakeBackend]: + backend = backend or FakeBackend() + agent_dir = tmp_path / "agent" + agent_dir.mkdir(exist_ok=True) + _stub_factories(monkeypatch) + return ( + ExperimentRunner( + backend=backend, + strategy=strategy, + config=config or EvolutionaryOptimizerConfig(), + workspace="default", + root=tmp_path / "experiment", + agent=agent_dir, + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + ), + backend, + ) + + +@pytest.mark.asyncio +async def test_strategy_failure_marks_the_run_failed(monkeypatch, tmp_path) -> None: + strategy = RecordingStrategy(error=ValueError("baseline failed")) + runner, _ = _make_runner(tmp_path, strategy=strategy, monkeypatch=monkeypatch) + + with pytest.raises(ValueError, match="baseline failed"): + await runner.run() + + assert strategy.ctx is not None + assert strategy.ctx._run.status == "failed" + + +@pytest.mark.asyncio +async def test_a_completed_run_persists_its_result_and_winner(monkeypatch, tmp_path) -> None: + winner = Candidate( + run_id="run-1", + label="agent-1", + ancestor="agent-0", + round=1, + optimization="add a tool", + rewards={"validation": RewardRecord(metrics={"reward": 0.75})}, + ) + (tmp_path / "experiment" / "eval-and-optimize" / "agents" / "agent-1").mkdir(parents=True) + runner, backend = _make_runner(tmp_path, strategy=RecordingStrategy(winner=winner), monkeypatch=monkeypatch) + + result = await runner.run() + + assert result.winner is winner + assert "winner=agent-1" in result.summary + assert backend.results == [result] + # publish_winner is off by default, so nothing was published. + assert backend.published == [] + + +@pytest.mark.asyncio +async def test_the_winner_is_published_when_storage_asks_for_it(monkeypatch, tmp_path) -> None: + winner = Candidate(run_id="run-1", label="agent-1", ancestor="agent-0", round=1, optimization="add a tool") + (tmp_path / "experiment" / "eval-and-optimize" / "agents" / "agent-1").mkdir(parents=True) + config = EvolutionaryOptimizerConfig(storage=CandidateStorageConfig(publish_winner=True)) + runner, backend = _make_runner( + tmp_path, strategy=RecordingStrategy(winner=winner), config=config, monkeypatch=monkeypatch + ) + + await runner.run() + + assert backend.published == ["agent-1"] + + +@pytest.mark.asyncio +async def test_the_baseline_is_never_published(monkeypatch, tmp_path) -> None: + """``ancestor is None`` means the baseline — there is nothing to open a PR against.""" + baseline = Candidate(run_id="run-1", label="agent-0", ancestor=None, round=0, optimization="baseline") + (tmp_path / "experiment" / "eval-and-optimize" / "agents" / "agent-0").mkdir(parents=True) + config = EvolutionaryOptimizerConfig(storage=CandidateStorageConfig(publish_winner=True)) + runner, backend = _make_runner( + tmp_path, strategy=RecordingStrategy(winner=baseline), config=config, monkeypatch=monkeypatch + ) + + await runner.run() + + assert backend.published == [] + + +@pytest.mark.asyncio +async def test_the_winner_is_copied_out_without_any_owners_scaffolding(monkeypatch, tmp_path) -> None: + winner = Candidate(run_id="run-1", label="agent-1", ancestor="agent-0", round=1, optimization="add a tool") + winner_dir = tmp_path / "experiment" / "eval-and-optimize" / "agents" / "agent-1" + winner_dir.mkdir(parents=True) + (winner_dir / "main.py").write_text("print('hello')\n") + (winner_dir / "metadata.json").write_text("{}") # the backend's + (winner_dir / "architecture.md").write_text("# arch") # the strategy's + (winner_dir / "harbor_wrapper.py").write_text("# harness") # the evaluator's + runner, _ = _make_runner(tmp_path, strategy=RecordingStrategy(winner=winner), monkeypatch=monkeypatch) + + await runner.run() + + root = tmp_path / "experiment" + assert (root / "main.py").read_text() == "print('hello')\n" + for scaffolding in ("metadata.json", "architecture.md", "harbor_wrapper.py"): + assert not (root / scaffolding).exists(), scaffolding + + +@pytest.mark.asyncio +async def test_a_run_with_no_winner_still_completes(monkeypatch, tmp_path) -> None: + runner, backend = _make_runner(tmp_path, strategy=RecordingStrategy(winner=None), monkeypatch=monkeypatch) + + result = await runner.run() + + assert result.winner is None + assert "winner=none" in result.summary + assert backend.results == [result] + + +@pytest.mark.asyncio +async def test_a_missing_strategy_report_falls_back_to_the_compact_summary(monkeypatch, tmp_path) -> None: + runner, _ = _make_runner(tmp_path, strategy=RecordingStrategy(winner=None), monkeypatch=monkeypatch) + + await runner.run() + + report = (tmp_path / "experiment" / "eval-and-optimize" / "OPTIMIZATION.md").read_text() + assert "Compact Run Summary" in report + assert "Optimization complete" in report + + +@pytest.mark.asyncio +async def test_an_existing_report_is_left_alone(monkeypatch, tmp_path) -> None: + eo = tmp_path / "experiment" / "eval-and-optimize" + eo.mkdir(parents=True) + (eo / "OPTIMIZATION.md").write_text("# The strategy's own report\n") + runner, _ = _make_runner(tmp_path, strategy=RecordingStrategy(winner=None), monkeypatch=monkeypatch) + + await runner.run() + + assert (eo / "OPTIMIZATION.md").read_text() == "# The strategy's own report\n" + + +@pytest.mark.asyncio +async def test_resume_refuses_loudly_for_a_strategy_that_cannot(monkeypatch, tmp_path) -> None: + """These runs cost hours, so a silent restart is the expensive failure.""" + eo = tmp_path / "experiment" / "eval-and-optimize" + eo.mkdir(parents=True) + (eo / "run.json").write_text( + json.dumps( + {"id": "run-existing", "name": "run-existing", "workspace": "default", "agent": "a", "status": "running"} + ) + ) + runner, _ = _make_runner(tmp_path, strategy=NonResumableStrategy(), monkeypatch=monkeypatch) + + with pytest.raises(ValueError, match="does not support resume"): + await runner.run() + + +@pytest.mark.asyncio +async def test_resume_reopens_the_existing_run_rather_than_creating_one(monkeypatch, tmp_path) -> None: + eo = tmp_path / "experiment" / "eval-and-optimize" + eo.mkdir(parents=True) + (eo / "run.json").write_text( + json.dumps( + {"id": "run-existing", "name": "run-existing", "workspace": "default", "agent": "a", "status": "failed"} + ) + ) + strategy = RecordingStrategy() + runner, backend = _make_runner(tmp_path, strategy=strategy, monkeypatch=monkeypatch) + + await runner.run() + + assert strategy.ctx is not None + assert strategy.ctx.run_id == "run-existing" + assert strategy.ctx.resuming is True + assert backend.runs == [] # nothing was created + + +@pytest.mark.asyncio +async def test_persistence_without_git_fails_before_any_work(monkeypatch, tmp_path) -> None: + monkeypatch.setattr(runner_module.shutil, "which", lambda _: None) + config = EvolutionaryOptimizerConfig(storage=CandidateStorageConfig(archive_candidates=True)) + strategy = RecordingStrategy() + runner, _ = _make_runner(tmp_path, strategy=strategy, config=config, monkeypatch=monkeypatch) + + with pytest.raises(ValueError, match="'git' is not on PATH"): + await runner.run() + + assert strategy.ctx is None + + +@pytest.mark.asyncio +async def test_a_failed_run_is_visible_on_disk(monkeypatch, tmp_path) -> None: + """The local backend writes run.json, so a failed run survives the process.""" + _stub_factories(monkeypatch) + agent_dir = tmp_path / "agent" + agent_dir.mkdir() + runner = ExperimentRunner( + backend=LocalExperimentalistBackend(path=tmp_path / "experiment"), + strategy=RecordingStrategy(error=ValueError("baseline failed")), + config=EvolutionaryOptimizerConfig(), + workspace="default", + root=tmp_path / "experiment", + agent=agent_dir, + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + ) + + with pytest.raises(ValueError, match="baseline failed"): + await runner.run() + + saved = json.loads((tmp_path / "experiment" / "eval-and-optimize" / "run.json").read_text(encoding="utf-8")) + assert saved["status"] == "failed" + + +def _insight_dataset(identity_char: str = "a") -> Dataset: + suite = Path("/experiment/eval-and-optimize/eval_author/insight-1/insight-suite") + task = Task(id="task-a", uri=(suite / "task-a").as_uri()) + return Dataset( + id="insight", + source=ResourceRef(uri=suite.as_uri()), + tasks=[task], + metadata={ + "insight_suite_identity": f"sha256:{identity_char * 64}", + "insight_suite_scorer_identity": f"sha256:{'b' * 64}", + "insight_suite_task_hashes": { + "task-a": {"content_hash": f"sha256:{'c' * 64}", "verifier_hash": f"sha256:{'d' * 64}"} + }, + }, + ) + + +def _insight_candidate(label: str, *, round_num: int, insight: float, validation: float) -> Candidate: + return Candidate( + run_id="run-1", + label=label, + ancestor=None if round_num == 0 else "agent-0", + round=round_num, + optimization="baseline" if round_num == 0 else "improve required tool use", + rewards={ + "insight": RewardRecord( + metrics={"reward": validation, "uses_required_tool": insight}, + metadata={ + "suite_identity": f"sha256:{'a' * 64}", + "metric_keys": ["reward", "uses_required_tool"], + }, + ), + "validation": RewardRecord(metrics={"reward": validation}), + }, + ) + + +async def _run_with_insight_suite( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *, + dataset: Dataset, + candidates: list[Candidate], +) -> None: + """Drive a run whose Eval Author authored *dataset*, with *candidates* already scored.""" + backend = FakeBackend( + client=fake_client(), # an Insight run needs a platform client to load traces + insight=Insight( + workspace="default", + agent="agent-source", + title="tool use", + description="the agent skips the required tool", + ), + ) + for candidate in candidates: + await backend.create_candidate(workspace="default", candidate=candidate) + winner = candidates[-1] + (tmp_path / "experiment" / "eval-and-optimize" / "agents" / winner.label).mkdir(parents=True) + + _stub_factories(monkeypatch) + monkeypatch.setattr( + runner_module, + "stage_eval_author_inputs", + _async(lambda _root, **refs: SimpleNamespace(**refs)), + ) + monkeypatch.setattr( + "nemo_eval_author_plugin.eval_author.agent.EvalAuthor", + lambda **_: SimpleNamespace( + run=_async( + lambda **kwargs: SimpleNamespace( + train_dataset=kwargs["train_dataset"], + validation_dataset=kwargs["validation_dataset"], + insight_suite=dataset, + ) + ) + ), + ) + agent_dir = tmp_path / "agent" + agent_dir.mkdir(exist_ok=True) + await ExperimentRunner( + backend=backend, + strategy=RecordingStrategy(winner=winner), + config=EvolutionaryOptimizerConfig(), + workspace="default", + root=tmp_path / "experiment", + agent=agent_dir, + insight="insight-1", + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + task_template=DatasetRef(uri="template"), + ).run() + + +def _async(fn): + async def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + return wrapper + + +@pytest.mark.asyncio +async def test_a_fallback_report_still_carries_the_insight_suite_sections(monkeypatch, tmp_path) -> None: + """The strategy wrote no report, but the Insight sections are the runner's to add.""" + baseline = _insight_candidate("agent-0", round_num=0, insight=0.0, validation=0.5) + winner = _insight_candidate("agent-1", round_num=1, insight=1.0, validation=0.75) + + await _run_with_insight_suite(tmp_path, monkeypatch, dataset=_insight_dataset(), candidates=[baseline, winner]) + + report = (tmp_path / "experiment" / "eval-and-optimize" / "OPTIMIZATION.md").read_text() + assert "## Compact Run Summary" in report + # The strategy never reported progress, so the counter is at its default unit. + assert "Optimization complete: 0 step(s) completed" in report + assert "## Deterministic Insight Suite Comparison" in report + assert "## Insight Suite Promotion Suggestions" in report + + +@pytest.mark.asyncio +async def test_an_insight_suite_mismatch_does_not_fail_a_completed_run(monkeypatch, tmp_path, caplog) -> None: + baseline = _insight_candidate("agent-0", round_num=0, insight=0.0, validation=0.5) + winner = _insight_candidate("agent-1", round_num=1, insight=1.0, validation=0.75) + # The winner was scored against a different suite than the run ended up with. + winner.set_reward( + "insight", metadata={**winner.rewards["insight"].metadata, "suite_identity": "sha256:" + "e" * 64} + ) + + with caplog.at_level("WARNING"): + await _run_with_insight_suite(tmp_path, monkeypatch, dataset=_insight_dataset(), candidates=[baseline, winner]) + + assert "Skipping Insight Suite report sections" in caplog.text + + +def test_a_run_needs_either_an_agent_or_an_insight(tmp_path) -> None: + with pytest.raises(ValueError, match="One of 'insight' or 'agent'"): + ExperimentRunner( + backend=FakeBackend(), + strategy=RecordingStrategy(), + config=EvolutionaryOptimizerConfig(), + workspace="default", + root=tmp_path, + agent=None, + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + ) + + +def test_an_insight_run_needs_a_task_template(tmp_path) -> None: + with pytest.raises(ValueError, match="'task_template' is required"): + ExperimentRunner( + backend=FakeBackend(), + strategy=RecordingStrategy(), + config=EvolutionaryOptimizerConfig(), + workspace="default", + root=tmp_path, + agent=None, + insight="insight-1", + train_dataset=DatasetRef(uri="train"), + validation_dataset=DatasetRef(uri="validation"), + ) + + +def test_the_run_entity_carries_a_progress_counter_not_a_round_count() -> None: + """Not every strategy has rounds, so the entity counts units and names the unit.""" + run = ExperimentRun(workspace="default", agent="a") + assert (run.progress_completed, run.progress_total, run.progress_unit) == (0, None, "step") diff --git a/plugins/nemo-experimentalist/tests/test_deps.py b/plugins/nemo-experimentalist/tests/test_deps.py deleted file mode 100644 index 21bf8bc69d..0000000000 --- a/plugins/nemo-experimentalist/tests/test_deps.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""ExperimentalistDeps validation: insight and/or agent may be set (or combined).""" - -from pathlib import Path - -import pytest -from nemo_experimentalist_plugin.entities import DatasetRef -from nemo_experimentalist_plugin.experimentalist.deps import ExperimentalistDeps - - -def _datasets(tmp_path: Path) -> dict: - # DatasetRef is a lazy URI handle resolved at evaluation time, so the paths - # need not exist here. - return { - "train_dataset": DatasetRef(uri=str(tmp_path / "train")), - "validation_dataset": DatasetRef(uri=str(tmp_path / "val")), - } - - -def _task_template(tmp_path: Path) -> DatasetRef: - return DatasetRef(uri=str(tmp_path / "template")) - - -def test_agent_only_ok(tmp_path: Path) -> None: - ExperimentalistDeps(agent="ssh://git@h/g/r.git@main", **_datasets(tmp_path)) - - -def test_insight_only_ok(tmp_path: Path) -> None: - ExperimentalistDeps(insight="ins-1", task_template=_task_template(tmp_path), **_datasets(tmp_path)) - - -def test_insight_and_agent_combined_ok(tmp_path: Path) -> None: - # The Mode-1 PR workflow: an insight guides optimization while a git agent - # supplies the code + PR target. Both set together is now valid. - deps = ExperimentalistDeps( - insight="ins-1", - agent="ssh://git@h/g/r.git@main", - task_template=_task_template(tmp_path), - **_datasets(tmp_path), - ) - assert deps.insight == "ins-1" - assert deps.agent == "ssh://git@h/g/r.git@main" - - -def test_insight_without_task_template_raises(tmp_path: Path) -> None: - # Mode 1 needs a task template to fill from production traces. - with pytest.raises(ValueError, match="task_template"): - ExperimentalistDeps(insight="ins-1", **_datasets(tmp_path)) - - -def test_neither_raises(tmp_path: Path) -> None: - with pytest.raises(ValueError, match="must be set"): - ExperimentalistDeps(**_datasets(tmp_path)) diff --git a/plugins/nemo-experimentalist/tests/test_experiment_mirror_mapping.py b/plugins/nemo-experimentalist/tests/test_experiment_mirror_mapping.py index 7add67fe5e..19e49db719 100644 --- a/plugins/nemo-experimentalist/tests/test_experiment_mirror_mapping.py +++ b/plugins/nemo-experimentalist/tests/test_experiment_mirror_mapping.py @@ -32,17 +32,21 @@ def test_group_metadata_carries_run_fields(): insight=None, config_snapshot={"k": 1}, status="running", - rounds_completed=2, + progress_completed=2, + progress_total=5, + progress_unit="round", winner_agent="agent-3", ) md = m.group_metadata(run) # Platform metadata is dict[str, str]: config_snapshot is JSON-serialized and the - # round counter is stringified, so the create/update body passes server validation. + # progress counter is stringified, so the create/update body passes server validation. assert md == { "agent": "a", "config_snapshot": '{"k": 1}', "status": "running", - "rounds_completed": "2", + "progress_completed": "2", + "progress_total": "5", + "progress_unit": "round", "winner_candidate": "agent-3", } @@ -54,7 +58,7 @@ def test_group_metadata_omits_winner_until_present(): insight=None, config_snapshot={"k": 1}, status="running", - rounds_completed=0, + progress_completed=0, winner_agent=None, ) md = m.group_metadata(run) diff --git a/plugins/nemo-experimentalist/tests/test_experimentalist_backend.py b/plugins/nemo-experimentalist/tests/test_experimentalist_backend.py index 5928578d8b..5abe02e8c6 100644 --- a/plugins/nemo-experimentalist/tests/test_experimentalist_backend.py +++ b/plugins/nemo-experimentalist/tests/test_experimentalist_backend.py @@ -334,7 +334,7 @@ async def test_persist_result_writes_run_summary(tmp_path: Path) -> None: run._id = "run-1" # type: ignore[attr-defined] (backend._eo / "run.json").write_text(run.model_dump_json(indent=2)) - result = ExperimentalistResult(summary="the real run summary", run_id="run-1", rounds_completed=2, winner=None) + result = ExperimentalistResult(summary="the real run summary", run_id="run-1", progress_completed=2, winner=None) await backend.persist_result(workspace="w", result=result) saved = json.loads((backend._eo / "run.json").read_text()) @@ -349,7 +349,7 @@ async def test_persist_result_preserves_generated_optimization_report(tmp_path: report_path = backend._eo / "OPTIMIZATION.md" report_path.write_text("# Full optimization report\n\nInsight Suite Metrics") - result = ExperimentalistResult(summary="compact run summary", run_id="run-1", rounds_completed=2, winner=None) + result = ExperimentalistResult(summary="compact run summary", run_id="run-1", progress_completed=2, winner=None) await backend.persist_result(workspace="w", result=result) assert report_path.read_text() == "# Full optimization report\n\nInsight Suite Metrics" diff --git a/plugins/nemo-experimentalist/tests/test_experimentalist_run.py b/plugins/nemo-experimentalist/tests/test_experimentalist_run.py index 2b88ee807d..96db854057 100644 --- a/plugins/nemo-experimentalist/tests/test_experimentalist_run.py +++ b/plugins/nemo-experimentalist/tests/test_experimentalist_run.py @@ -1,17 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import logging -from dataclasses import dataclass +"""``run_experimentalist`` wires the CLI's inputs into one :class:`ExperimentRunner`. + +What the runner then does with them is covered by the runner's own tests; these check +the hand-off, and that the caller keeps ownership of its platform client. +""" + +from dataclasses import dataclass, field from pathlib import Path -from typing import cast +from typing import Any, cast import pytest from nemo_experimentalist_plugin.entities import DatasetRef from nemo_experimentalist_plugin.experimentalist import run as experimentalist_run -from nemo_experimentalist_plugin.experimentalist.components import loop as loop_module from nemo_experimentalist_plugin.experimentalist.components.loop import EvolutionaryOptimizerConfig -from nemo_experimentalist_plugin.experimentalist.deps import ExperimentalistDeps from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import LocalExperimentalistBackend from nemo_experimentalist_plugin.experimentalist.result import ExperimentalistResult from nemo_platform import AsyncNeMoPlatform @@ -46,24 +49,17 @@ class AgentFactoryCall: @dataclass -class FakeExperimentalist: - deps: ExperimentalistDeps | None = None +class RecordingRunner: + """Stands in for the real runner and keeps the kwargs it was constructed with.""" - async def run(self, deps: ExperimentalistDeps) -> ExperimentalistResult: - self.deps = deps - return ExperimentalistResult(summary="optimization complete", run_id="run-1", rounds_completed=1) + calls: list[dict[str, Any]] = field(default_factory=list) + def __call__(self, **kwargs: Any) -> "RecordingRunner": + self.calls.append(kwargs) + return self -def test_persistence_warning_includes_exception_message( - caplog: pytest.LogCaptureFixture, -) -> None: - caplog.set_level(logging.WARNING, logger=loop_module.__name__) - - loop_module._warn_persistence_failure("archive", "agent-2", RuntimeError("push rejected by remote")) - - assert "archive" in caplog.text - assert "agent-2" in caplog.text - assert "push rejected by remote" in caplog.text + async def run(self) -> ExperimentalistResult: + return ExperimentalistResult(summary="optimization complete", run_id="run-1", progress_completed=1) def _make_run_paths(tmp_path: Path) -> ExperimentRunPaths: @@ -89,7 +85,8 @@ async def test_run_experimentalist_builds_and_runs_complete_local_contract( client = ClosingClient() backend = LocalExperimentalistBackend(path=tmp_path / "backend") optimizer_config = EvolutionaryOptimizerConfig(max_rounds=2) - experimentalist = FakeExperimentalist() + strategy = object() + runner = RecordingRunner() backend_calls: list[BackendFactoryCall] = [] agent_calls: list[AgentFactoryCall] = [] litellm_calls: list[bool] = [] @@ -100,23 +97,19 @@ def make_backend( experiments_output: str, storage: object = None, ) -> LocalExperimentalistBackend: - backend_calls.append( - BackendFactoryCall( - client=client, - experiments_output=experiments_output, - ) - ) + backend_calls.append(BackendFactoryCall(client=client, experiments_output=experiments_output)) return backend def build_agent( *, working_dir: Path, config: EvolutionaryOptimizerConfig, framework_skills_dirs: list[Path] | None - ) -> FakeExperimentalist: + ) -> object: assert framework_skills_dirs is None agent_calls.append(AgentFactoryCall(working_dir=working_dir, config=config)) - return experimentalist + return strategy monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", make_backend) monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", build_agent) + monkeypatch.setattr(experimentalist_run, "ExperimentRunner", runner) monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: litellm_calls.append(True)) train_dataset = DatasetRef(uri=str(paths.train)) @@ -135,25 +128,23 @@ def build_agent( assert summary == "optimization complete" assert paths.experiment.is_dir() - assert backend_calls == [ - BackendFactoryCall( - client=client, - experiments_output=str(paths.experiment.resolve()), - ) - ] + assert backend_calls == [BackendFactoryCall(client=client, experiments_output=str(paths.experiment.resolve()))] assert agent_calls == [AgentFactoryCall(working_dir=paths.experiment.resolve(), config=optimizer_config)] assert litellm_calls == [True] assert not client.closed - assert experimentalist.deps is not None - assert experimentalist.deps.workspace == "workspace-a" - # ``agent`` is forwarded verbatim (it may be a git url@ref); the loop resolves it. - assert experimentalist.deps.agent == paths.agent - assert experimentalist.deps.insight is None - assert experimentalist.deps.train_dataset == train_dataset - assert experimentalist.deps.validation_dataset == validation_dataset - assert experimentalist.deps.backend is backend - assert experimentalist.deps.config is optimizer_config - assert experimentalist.deps.agent_spec is None + + (call,) = runner.calls + assert call["backend"] is backend + assert call["strategy"] is strategy + assert call["config"] is optimizer_config + assert call["workspace"] == "workspace-a" + assert call["root"] == paths.experiment.resolve() + # ``agent`` is forwarded verbatim (it may be a git url@ref); the runner resolves it. + assert call["agent"] == paths.agent + assert call["insight"] is None + assert call["train_dataset"] == train_dataset + assert call["validation_dataset"] == validation_dataset + assert call["agent_spec"] is None @pytest.mark.asyncio @@ -162,20 +153,11 @@ async def test_run_experimentalist_forwards_platform_insight_id_verbatim( tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) - client = ClosingClient() - backend = LocalExperimentalistBackend(path=tmp_path / "backend") - experimentalist = FakeExperimentalist() + runner = RecordingRunner() - monkeypatch.setattr( - experimentalist_run, - "make_experimentalist_backend", - lambda **_: backend, - ) - monkeypatch.setattr( - experimentalist_run, - "build_experimentalist_agent", - lambda **_: experimentalist, - ) + monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", lambda **_: object()) + monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", lambda **_: object()) + monkeypatch.setattr(experimentalist_run, "ExperimentRunner", runner) monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: None) await experimentalist_run.run_experimentalist( @@ -185,26 +167,25 @@ async def test_run_experimentalist_forwards_platform_insight_id_verbatim( task_template=DatasetRef(uri=str(paths.train)), experiment_dir=paths.experiment, workspace="workspace-a", - client=cast(AsyncNeMoPlatform, client), + client=cast(AsyncNeMoPlatform, ClosingClient()), config=EvolutionaryOptimizerConfig(), ) - assert experimentalist.deps is not None # A str id is not resolved to a Path — it flows through untouched to the backend. - assert experimentalist.deps.insight == "insight-remote-123" + assert runner.calls[0]["insight"] == "insight-remote-123" @pytest.mark.asyncio -async def test_run_experimentalist_forwards_agent_spec_uri_to_deps( +async def test_run_experimentalist_forwards_agent_spec_uri( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: paths = _make_run_paths(tmp_path) - backend = LocalExperimentalistBackend(path=tmp_path / "backend") - experimentalist = FakeExperimentalist() + runner = RecordingRunner() - monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", lambda **_: backend) - monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", lambda **_: experimentalist) + monkeypatch.setattr(experimentalist_run, "make_experimentalist_backend", lambda **_: object()) + monkeypatch.setattr(experimentalist_run, "build_experimentalist_agent", lambda **_: object()) + monkeypatch.setattr(experimentalist_run, "ExperimentRunner", runner) monkeypatch.setattr(experimentalist_run, "_enable_litellm_drop_params", lambda: None) spec_uri = "/path/to/AGENT-SPEC.md" @@ -220,8 +201,7 @@ async def test_run_experimentalist_forwards_agent_spec_uri_to_deps( config=EvolutionaryOptimizerConfig(), ) - assert experimentalist.deps is not None - assert experimentalist.deps.agent_spec == spec_uri + assert runner.calls[0]["agent_spec"] == spec_uri @pytest.mark.asyncio diff --git a/plugins/nemo-experimentalist/tests/test_local_backend_projection.py b/plugins/nemo-experimentalist/tests/test_local_backend_projection.py index f815b6407d..9785ae7b55 100644 --- a/plugins/nemo-experimentalist/tests/test_local_backend_projection.py +++ b/plugins/nemo-experimentalist/tests/test_local_backend_projection.py @@ -20,7 +20,7 @@ def _run() -> ExperimentRun: - return ExperimentRun(workspace="default", agent="a", config_snapshot={}, status="running", rounds_completed=0) + return ExperimentRun(workspace="default", agent="a", config_snapshot={}, status="running", progress_completed=0) async def test_no_projection_without_client(tmp_path: Path) -> None: @@ -56,7 +56,7 @@ async def test_persist_result_projects_finalize(tmp_path: Path) -> None: be._mirrors["default"] = AsyncMock() await be.persist_result( workspace="default", - result=ExperimentalistResult(summary="done", run_id="run-1", rounds_completed=1, winner=None), + result=ExperimentalistResult(summary="done", run_id="run-1", progress_completed=1, winner=None), ) be._mirrors["default"].finalize.assert_awaited_once() From 2d0262cba8c1045c927427efdb14aaa8b064988f Mon Sep 17 00:00:00 2001 From: Severin Klingler Date: Mon, 3 Aug 2026 11:29:29 +0200 Subject: [PATCH 02/59] refactor(experimentalist): a Candidate is metadata and an artifact reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A Candidate was four jobs in one object: a proposed change, an identity, a set of measurements, and the optimized artifact. The artifact was always a directory, the identity was derived from that directory's name, and storage of the first three was implemented by walking the fourth — `list_candidates` globbed `agents/*/metadata.json`. A strategy that does not produce one directory per candidate could not store, list, or resume anything. Separate them. `Proposal` carries the build request; the runner stores identity and measurements at `eval-and-optimize/candidates/.json`; and the finished work is *addressed* through `artifact: ResourceRef` rather than contained. That is the same pattern `persist_evaluation` already uses for traces, applied to the one entity that lacked it. The lifecycle inverts with it. A candidate used to be created and then mutated as the Coder filled its directory in, which is why a failed build needed a killed marker so resume would not resurrect it. Now `ctx.fork()` reserves and populates a directory, the Builder writes it, and `ctx.commit_candidate()` validates the result and creates the Candidate — so `artifact` can be required, a failed build leaves no record at all, and the killed marker for it goes away. The fields that went with the old shape: - `round` → `ancestor is None` for "this is the baseline", which `entities.py` already documented as the same thing, plus a strategy-supplied `generation` for grouping. DSPy leaves it 0; our loop sets the round. - `optimization`/`optimization_type`/`task_ids` → `description` plus the embedded Proposal, whose payload is owned by the Proposer/Builder pair. The Coder gets its own `BuildRequest` view and never sees the entity. - `ancestor` is a candidate id, not a directory name, so `benchmarks/run.py` reads the winner's location from its artifact instead of building `agents/{label}`. Candidate validation rejects a record whose `ancestor` or `description` disagrees with the Proposal they were derived from — two accounts of one candidate's origin must not be able to drift. Signed-off-by: Severin Klingler --- .../nemo-experimentalist/benchmarks/run.py | 22 +- plugins/nemo-experimentalist/pyproject.toml | 2 +- .../nemo_experimentalist_plugin/entities.py | 131 +++++--- .../experimentalist/components/coder.py | 61 +++- .../components/insight_promotion.py | 2 +- .../experimentalist/components/loop.py | 294 ++++++------------ .../experimentalist/components/models.py | 90 +++--- .../experimentalist/components/proposer.py | 69 +++- .../experimentalist/components/terminator.py | 6 +- .../experimentalist/components/tools.py | 37 ++- .../experimentalist/context.py | 150 ++++++++- .../experimentalist/experiment_mirror.py | 14 +- .../experimentalist_backend.py | 86 +++-- .../experimentalist/run.py | 2 +- .../tests/{experimentalist => }/doubles.py | 50 +++ .../test_candidate_contract.py | 184 +++++++++++ .../test_evolution_tree_rendering.py | 11 +- .../test_loop_insight_suite.py | 134 +++++--- .../experimentalist/test_loop_reporting.py | 16 +- .../tests/experimentalist/test_runner.py | 27 +- .../tests/experimentalist/test_terminator.py | 2 +- .../tests/test_experiment_mirror.py | 25 +- .../tests/test_experiment_mirror_mapping.py | 24 +- .../tests/test_experimentalist_backend.py | 3 +- .../tests/test_local_backend_projection.py | 18 +- 25 files changed, 970 insertions(+), 490 deletions(-) rename plugins/nemo-experimentalist/tests/{experimentalist => }/doubles.py (81%) create mode 100644 plugins/nemo-experimentalist/tests/experimentalist/test_candidate_contract.py diff --git a/plugins/nemo-experimentalist/benchmarks/run.py b/plugins/nemo-experimentalist/benchmarks/run.py index 19c4331ad2..6d053ef97f 100644 --- a/plugins/nemo-experimentalist/benchmarks/run.py +++ b/plugins/nemo-experimentalist/benchmarks/run.py @@ -27,6 +27,7 @@ HarborEvaluator, HarborEvaluatorConfig, ) +from nemo_experimentalist_plugin.experimentalist.experimentalist_backend import load_candidate from nemo_experimentalist_plugin.resolve import resolve_dataset from pydantic import BaseModel, Field, model_validator @@ -329,7 +330,7 @@ def summarize_experimentalist_jobs(results_dir: Path) -> dict[str, Any]: if not isinstance(payload, dict) or not isinstance(payload.get("stats"), dict): continue stats = payload["stats"] - job = { + job: dict[str, Any] = { "name": result_path.parent.name, "trials": int(payload.get("n_total_trials") or 0), "completed_trials": int(stats.get("n_completed_trials") or 0), @@ -350,7 +351,7 @@ def summarize_experimentalist_jobs(results_dir: Path) -> dict[str, Any]: "cache_tokens", "output_tokens", ): - totals[key] += job[key] + totals[key] = int(totals[key]) + int(job[key]) if job["cost_usd"] is not None: costs.append(float(job["cost_usd"])) return {**totals, "cost_usd": sum(costs) if costs else None, "job_results": jobs} @@ -395,9 +396,11 @@ async def run_benchmark(args: argparse.Namespace) -> Path: package_client = PackageDatasetClient() metadata = await package_client.get_dataset_metadata(suite.dataset.requested_reference) - canonical_task_ids = {task.name for task in metadata.task_ids} + # Harbor's task ids are a union of source-specific types; every variant names the + # task the same way, but only some spell it ``name``. + canonical_task_ids = {str(getattr(task, "name", task)) for task in metadata.task_ids} canonical_task_ids = validate_canonical_suite( - suite, canonical_task_ids=canonical_task_ids, resolved_ref=metadata.version + suite, canonical_task_ids=canonical_task_ids, resolved_ref=metadata.version or "" ) # Resolved up front: the optimizer only needs these after the baseline evaluation, # which is hours of image builds to discover a typo'd skill name. @@ -462,10 +465,15 @@ async def run_benchmark(args: argparse.Namespace) -> Path: framework_skills_dirs=framework_skills_dirs, ) run_document = json.loads((experimentalist_dir / "eval-and-optimize" / "run.json").read_text(encoding="utf-8")) - winner_label = run_document.get("winner_agent") - if not isinstance(winner_label, str) or not winner_label: + winner_id = run_document.get("winner_agent") + if not isinstance(winner_id, str) or not winner_id: raise RuntimeError("Experimentalist completed without a selected winner") - winner_dir = experimentalist_dir / "eval-and-optimize" / "agents" / winner_label + # ``winner_agent`` is a candidate id, not a directory name — the winner's location + # comes from its own artifact reference. + winner_record = experimentalist_dir / "eval-and-optimize" / "candidates" / f"{winner_id}.json" + winner_candidate = load_candidate(winner_record) + winner_label = winner_candidate.label + winner_dir = local_path_from_uri(winner_candidate.artifact.uri, context="Winner artifact") winner = await _evaluate_heldout( label="winner", agent_dir=winner_dir, diff --git a/plugins/nemo-experimentalist/pyproject.toml b/plugins/nemo-experimentalist/pyproject.toml index 842f9ff161..35f7fda7f1 100644 --- a/plugins/nemo-experimentalist/pyproject.toml +++ b/plugins/nemo-experimentalist/pyproject.toml @@ -33,5 +33,5 @@ packages = ["src/nemo_experimentalist_plugin"] [tool.pytest.ini_options] asyncio_mode = "auto" -pythonpath = ["src"] +pythonpath = ["src", "tests"] testpaths = ["tests"] diff --git a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py index 0bba13206a..33d524a793 100644 --- a/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py +++ b/plugins/nemo-experimentalist/src/nemo_experimentalist_plugin/entities.py @@ -444,22 +444,55 @@ class RewardRecord(BaseModel): metadata: dict[str, DataValue] = Field(default_factory=dict, description="Provenance for this measurement.") +class Proposal(BaseModel): + """A request to build one candidate — the Proposer → Builder contract. + + A transient component message, not a separately persisted entity and not an + unfinished Candidate: a proposal describes work to perform, and a Candidate is the + durable result after a Builder completes it. Failed proposals produce no Candidate + and are not retained; a successful one is embedded in ``Candidate.generated_from``. + + ``kind`` is an opaque compatibility discriminator, not a global enumeration: it + routes the proposal to a Builder that declares it can accept it. ``payload`` is + owned by that Proposer/Builder pair — Layer A stores and transports it and + interprets neither. + """ + + ancestor: str | None = Field( + default=None, + description="Parent Candidate id to build from. None means the baseline.", + ) + description: str = Field( + min_length=1, + description="Human-readable explanation of the proposed variant.", + ) + kind: str = Field( + min_length=1, + description="Builder compatibility discriminator, e.g. 'code-change' or 'parameters'.", + ) + payload: dict[str, DataValue] = Field( + default_factory=dict, + description="Build instructions, validated by the component-owned schema for this kind.", + ) + + class Candidate(NemoEntity, entity_type="candidate"): """A candidate agent version produced during an Experimentalist run. - Lives in the entity store so candidates are queryable, resumable, and - survive the local working directory being deleted. ``run_id`` groups all - candidates for a single ExperimentRun. + Metadata and measurements live in the entity store; the completed work is + *addressed*, not contained. ``artifact`` points at the resource that defines this + candidate and, when external evaluation is used, is directly consumable by the + run's evaluation component. Its format belongs to the components that produce and + consume that candidate kind — the host only stores, transports, archives and + publishes the reference. - Like every entity in this plugin, a Candidate's durable identity is its - store-assigned ``id`` (``name`` is left for the store to auto-slug). - ``label`` is the run-scoped handle ("agent-0", "agent-1", ...) used for the - working directory, evolution-tree key, and ``ancestor`` references; it is - unique within a run, not globally. + A Candidate is only ever created once its artifact exists and validates, so + ``artifact`` is required and no durable record points at partial work. Incomplete + work is a runner-owned path, not a Candidate. - A candidate's completed artifact is not stored here yet: it is still the - ``agents/