refactor(eval-author): simplify dataset augmentation

plugins/nemo-eval-author/README.md

@@ -33,7 +33,7 @@ from nemo_eval_author_plugin.eval_author.run import run_eval_author
 # Still borrowed from Experimentalist, and on the way out. Treat these as Eval Author's
 # own types once they move; do not build new code on the Experimentalist paths.
 from nemo_experimentalist_plugin.entities import Dataset, DatasetRef
-from nemo_experimentalist_plugin.experimentalist.components.dataset_staging import stage_task_template
+from nemo_experimentalist_plugin.experimentalist.components.dataset_staging import stage_eval_author_inputs
 from nemo_experimentalist_plugin.experimentalist.components.trace_analyzer import TraceAnalyzer
 from nemo_experimentalist_plugin.experimentalist.components.trace_explorer import TraceExplorer
 ```

plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/README.md

@@ -55,18 +55,16 @@ Each Eval Author invocation fills a fresh candidate suite from the current templ
 and traces. The complete suite is Harbor-validated locally, promoted to the
 experiment-local working copy with backup-and-restore failure handling, and
 analyzed for the Insight's root cause. Eval Author then adds normalized
-Insight-specific verifier metric keys to every materialized task while
-preserving the template's existing task metrics. The metric authoring step is
-scoped to the Insight suite; the user's train and validation datasets remain
-unchanged. The authored verifiers must pass static Harbor validation before the
-local suite is returned to the optimization loop.
+Insight-specific verifier metric keys to every task in the staged train,
+validation, and generated Insight datasets while preserving existing task
+metrics. All three datasets must pass static Harbor validation before they are
+returned.
 
 After authoring and validation, Eval Author hashes every task file and verifier
 file and persists deterministic suite and scorer identities in the local suite's
-manifest. The returned dataset continues to point at the single experiment-local
-suite. Candidate Insight results persist the suite identity, and resume reuses
-those results only when the identity still matches; changed task or verifier
-content is re-evaluated.
+manifest. The returned Insight dataset points at the experiment-local suite.
+Eval Author does not split, merge, or evaluate that suite. Experimentalist
+currently leaves it unconsumed for downstream integration by its owning team.
 
 Task-template inputs may be local paths, `file://` URIs, or NeMo Platform
 `fileset://<workspace>/<fileset>` references. Fileset-backed templates are
@@ -76,12 +74,6 @@ them. The staged template is refreshed on every invocation rather than reused.
 The returned Python contract is documented in the
 [Eval Author Python Reference](REFERENCE.md#evalauthorresult).
 
-Insight metrics remain adaptive development feedback. They may steer round
-analysis, goal-tree updates, and proposals, but validation remains the direct
-Pareto and winner-selection criterion. Promotion suggestions require complete
-repeated baseline-to-winner improvement evidence, remain advisory, and never
-mutate the canonical validation dataset.
-
 ## Intended Invocation
 
 Until a CLI or platform job is wired, Python callers can invoke the runner

plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/REFERENCE.md

@@ -11,13 +11,13 @@ SPDX-License-Identifier: Apache-2.0
 
 | Field | Type | Description |
 | --- | --- | --- |
-| `train_dataset` | `Dataset` | Training dataset supplied to the run. Eval Author does not mutate it. |
-| `validation_dataset` | `Dataset` | Validation dataset supplied to the run. Eval Author does not mutate it. |
-| `insight_suite` | `Dataset \| None` | Finalized experiment-local Insight dataset for immediate evaluation by the optimization loop. |
+| `train_dataset` | `Dataset` | Staged training dataset with the authored metrics and verifier changes. |
+| `validation_dataset` | `Dataset` | Staged validation dataset with the authored metrics and verifier changes. |
+| `insight_suite` | `Dataset \| None` | Finalized experiment-local task set generated from the Insight's trace references. |
 | `insight_suite_identity` | `str \| None` | SHA-256 identity of the finalized Insight task and verifier content. |
+| `metric_keys` | `tuple[str, ...]` | Metric keys authored across all three datasets. |
 | `summary` | `str` | Eval Author's analysis summary. |
 
 When an Insight suite is materialized successfully, `insight_suite` and
-`insight_suite_identity` are both populated. Callers can persist the identity
-with candidate results and reuse those results only while the suite identity
-continues to match.
+`insight_suite_identity` are both populated. Eval Author does not split, merge,
+or evaluate the returned suite; its caller owns that downstream integration.

plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/agent.py

@@ -12,14 +12,20 @@
 from pathlib import Path
 from typing import Any
 
-from nemo_eval_author_plugin.eval_author.materialization import InsightSuite
-from nemo_eval_author_plugin.eval_author.models import EvalAuthorConfig, EvalAuthorResult
+from nemo_eval_author_plugin.eval_author.materialization import InsightSuite, validate_metric_contracts
+from nemo_eval_author_plugin.eval_author.models import EvalAuthorConfig, EvalAuthorResult, MetricAuthoringResult
 from nemo_eval_author_plugin.model_config import (
     bridge_author_env_to_experimentalist,
     get_fast_model,
     get_smart_model,
 )
-from nemo_experimentalist_plugin.entities import Dataset, DatasetValidationError, ResourceRef, Task, TrialResult
+from nemo_experimentalist_plugin.entities import (
+    Dataset,
+    DatasetValidationError,
+    ResourceRef,
+    Task,
+    TrialResult,
+)
 from nemo_experimentalist_plugin.experimentalist.components import cache
 from nemo_experimentalist_plugin.experimentalist.components.tools import GuardedShellTools
 from nemo_experimentalist_plugin.experimentalist.components.trace_analyzer import (
@@ -115,15 +121,19 @@ async def author_insight_metrics(
         self,
         insight: Insight,
         diagnostics: list[tuple[str, Diagnostic]],
+        train_dataset: Dataset,
+        validation_dataset: Dataset,
         insight_suite: Dataset,
         runner_conventions: str,
         validation_feedback: str | None = None,
-    ) -> str:
+    ) -> MetricAuthoringResult:
         """Author verifier metrics for the materialized tasks that capture the insight.
 
         Args:
             insight: The insight whose failure mode the tasks should detect.
             diagnostics: Per-trace ``(trace_ref, Diagnostic)`` pairs for concrete evidence.
+            train_dataset: Staged training tasks to augment for optimization feedback.
+            validation_dataset: Staged validation tasks to augment for scoring.
             insight_suite: The materialized tasks recreated from the Insight's production traces.
             runner_conventions: Summary of how this dataset's runner works (from ``discover_runner``).
                 Use this as the authoritative reference for what artifacts exist at
@@ -135,33 +145,34 @@ async def author_insight_metrics(
         Refer to ``self.context["dataset_documentation"]`` for the dataset-specific API
         and metric authoring conventions (file layout, how to add/remove/modify a metric).
 
-        **Scope: new grades on the materialized Insight tasks**
+        **Scope: every task in all three datasets**
 
-        Add at least one new Insight-specific metric key to every task in ``insight_suite``.
-        Use the same new metric key set and shared scoring semantics across the entire
-        suite. Preserve every existing verifier metric, including the task's ordinary
-        ``reward`` or ``score``; append the Insight signal instead of replacing the
-        task's original notion of success.
+        Add at least one new Insight-specific metric key to every task in
+        ``train_dataset``, ``validation_dataset``, and ``insight_suite``. Use the same
+        new metric key set and scoring semantics everywhere. Preserve every existing
+        verifier metric, including ordinary ``reward`` or ``score`` values.
 
-        Only edit verifier files in the materialized Insight suite. Do not modify the
-        user's train or validation datasets, and do not change task instructions,
-        environments, solutions, or other agent-visible inputs. This work adds new
-        grades to the new rows; it does not add new agent output to old benchmark rows.
+        Task-specific verifier edits are allowed and expected when existing verifier
+        layouts differ. Only edit verifier files. Do not modify task instructions,
+        environments, solutions, or other agent-visible inputs.
 
         Name each new metric after the root-cause behavior, not a trace id or surface
         symptom. Measure the current Harbor run from runtime artifacts such as OTLP
         traces or agent outputs. Do not hard-code scores for the production traces that
         motivated the Insight.
 
+        In every task's configured verifier directory, write ``metric-contract.json``
+        containing exactly ``{"metric_keys": ["key_one", "key_two"]}``. The list must
+        contain the same newly authored keys for every task and must exactly match the
+        ``metric_keys`` returned in ``MetricAuthoringResult``.
+
         **Validate while authoring**
 
-        After every verifier edit, call ``await insight_suite.validate()``. This performs
-        evaluator-specific static checks without launching trials or executing verifier
-        code. If it raises ``DatasetValidationError``, use its task, path, and source
-        location diagnostics to repair the files, then call it again. Do not return until
-        the suite passes validation. If ``validation_feedback`` is provided, the caller's
-        mandatory validation found errors in the previous attempt; fix every reported
-        failure and revalidate the suite.
+        After verifier edits, call ``await train_dataset.validate()``,
+        ``await validation_dataset.validate()``, and ``await insight_suite.validate()``.
+        These perform evaluator-specific static checks without launching trials. If one
+        raises ``DatasetValidationError``, repair its diagnostics and revalidate all three.
+        If ``validation_feedback`` is provided, fix every reported failure before returning.
 
         **Metric quality**
 
@@ -187,9 +198,9 @@ async def author_insight_metrics(
         objects for Y, so X is missing from its context.  Measure whether the agent retrieves
         all required objects, not merely whether X appears in the final answer.
 
-        Return a concise summary naming the new metric key(s), what they measure, and
-        which runtime evidence they score. The caller retains the materialized suite and
-        the user's unchanged train and validation datasets.
+        Return ``MetricAuthoringResult`` with the unique, non-empty ``metric_keys`` added
+        by the verifier edits plus a concise summary. Include at least one
+        Insight-specific key beyond generic ``reward`` or ``score``.
         """  # noqa: D413
         ...
 
@@ -282,8 +293,8 @@ async def _run(
             insight: The Insight to investigate with relevant traces.
             agent_path: Agent root, relative to ``experiment_dir`` or absolute.
             task_template: Parsed evaluator task containing explicit placeholders.
-            train_dataset: The train dataset, returned unchanged.
-            validation_dataset: The validation dataset, returned unchanged.
+            train_dataset: Staged training dataset to augment in place.
+            validation_dataset: Staged validation dataset to augment in place.
             client: Existing NeMo Platform client used for Intake requests.
         """
         reporter = self._reporter
@@ -391,18 +402,42 @@ async def _run(
         self.context["dataset_documentation"] = doc(type(materialized_dataset), inline_depth=1)
         if reporter is not None:
             reporter.progress(phase="eval author · discovering runner")
-        runner_conventions = await self.discover_runner(materialized_dataset)
+        runner_conventions = await self.discover_runner(train_dataset)
         if reporter is not None:
             reporter.progress(phase="eval author · authoring metrics")
-        summary = await self.author_insight_metrics(
-            insight,
-            diagnostics,
-            materialized_dataset,
-            runner_conventions,
-        )
+
+        validation_feedback: str | None = None
         for repair_attempt in range(self._config.max_validation_repair_attempts + 1):
+            authoring_result = await self.author_insight_metrics(
+                insight,
+                diagnostics,
+                train_dataset,
+                validation_dataset,
+                materialized_dataset,
+                runner_conventions,
+                validation_feedback=validation_feedback,
+            )
             try:
-                await materialized_dataset.validate()
+                validation_errors: list[str] = []
+                for label, dataset in (
+                    ("train", train_dataset),
+                    ("validation", validation_dataset),
+                    ("insight", materialized_dataset),
+                ):
+                    try:
+                        await dataset.validate()
+                    except DatasetValidationError as exc:
+                        validation_errors.append(f"{label} dataset:\n{exc}")
+                if validation_errors:
+                    raise DatasetValidationError("\n".join(validation_errors))
+                validate_metric_contracts(
+                    {
+                        "train": train_dataset,
+                        "validation": validation_dataset,
+                        "insight": materialized_dataset,
+                    },
+                    metric_keys=authoring_result.metric_keys,
+                )
             except DatasetValidationError as exc:
                 if repair_attempt >= self._config.max_validation_repair_attempts:
                     raise
@@ -419,24 +454,20 @@ async def _run(
                         total=self._config.max_validation_repair_attempts,
                         unit="attempt",
                     )
-                summary = await self.author_insight_metrics(
-                    insight,
-                    diagnostics,
-                    materialized_dataset,
-                    runner_conventions,
-                    validation_feedback=str(exc),
-                )
-            else:
-                finalized_suite = insight_suite.finalize()
-                if reporter is not None:
-                    reporter.progress(phase="eval author · complete")
-                return EvalAuthorResult(
-                    train_dataset=train_dataset,
-                    validation_dataset=validation_dataset,
-                    insight_suite=finalized_suite.dataset,
-                    insight_suite_identity=finalized_suite.identity,
-                    summary=summary,
-                )
+                validation_feedback = str(exc)
+                continue
+
+            finalized_suite = insight_suite.finalize()
+            if reporter is not None:
+                reporter.progress(phase="eval author · complete")
+            return EvalAuthorResult(
+                train_dataset=train_dataset,
+                validation_dataset=validation_dataset,
+                insight_suite=finalized_suite.dataset,
+                insight_suite_identity=finalized_suite.identity,
+                metric_keys=authoring_result.metric_keys,
+                summary=authoring_result.summary,
+            )
 
         raise AssertionError("unreachable")
 

plugins/nemo-eval-author/src/nemo_eval_author_plugin/eval_author/materialization.py

@@ -18,12 +18,13 @@
 
 import tomlkit
 from harbor.models.task.task import Task as HarborTask
-from nemo_experimentalist_plugin.entities import Task, local_path_from_uri
+from nemo_experimentalist_plugin.entities import Dataset, DatasetValidationError, Task, local_path_from_uri
 from nemo_experimentalist_plugin.experimentalist.components.evaluator.harbor import HarborDataset
 
 _MANIFEST_SCHEMA_VERSION = 3
 _CONTENT_HASH_SCHEMA_VERSION = 1
 _METRIC_CONTRACT_VERSION = 1
+_METRIC_CONTRACT_FILENAME = "metric-contract.json"
 _SLUG_RE = re.compile(r"[^a-z0-9]+")
 
 
@@ -67,6 +68,61 @@ def _verifier_dir(task_dir: Path) -> Path:
     raise ValueError(f"Materialized task has no verifier directory: {task_dir}")
 
 
+def validate_metric_contracts(
+    datasets: dict[str, Dataset],
+    *,
+    metric_keys: tuple[str, ...],
+) -> None:
+    """Require every task to declare the same newly authored metric keys."""
+    expected = set(metric_keys)
+    failures: list[str] = []
+    for dataset_name, dataset in datasets.items():
+        for task in dataset.list_tasks():
+            if not task.uri:
+                failures.append(f"{dataset_name}/{task.id}: task URI is missing")
+                continue
+            try:
+                task_dir = local_path_from_uri(
+                    task.uri,
+                    context=f"{dataset_name} dataset task {task.id!r}",
+                ).resolve()
+                verifier_dir = _verifier_dir(task_dir)
+            except (FileNotFoundError, ValueError, tomllib.TOMLDecodeError) as exc:
+                failures.append(f"{dataset_name}/{task.id}: {exc}")
+                continue
+
+            contract_path = verifier_dir / _METRIC_CONTRACT_FILENAME
+            if not contract_path.is_file():
+                failures.append(f"{dataset_name}/{task.id}: missing {contract_path}")
+                continue
+            try:
+                payload = json.loads(contract_path.read_text(encoding="utf-8"))
+            except (json.JSONDecodeError, OSError) as exc:
+                failures.append(f"{dataset_name}/{task.id}: invalid JSON in {contract_path}: {exc}")
+                continue
+            if not isinstance(payload, dict) or set(payload) != {"metric_keys"}:
+                failures.append(f"{dataset_name}/{task.id}: {contract_path} must contain only a metric_keys list")
+                continue
+            raw_keys = payload["metric_keys"]
+            if (
+                not isinstance(raw_keys, list)
+                or any(not isinstance(key, str) or not key.strip() for key in raw_keys)
+                or len(raw_keys) != len(set(raw_keys))
+            ):
+                failures.append(
+                    f"{dataset_name}/{task.id}: {contract_path} metric_keys must be unique non-empty strings"
+                )
+                continue
+            actual = set(raw_keys)
+            if actual != expected:
+                failures.append(
+                    f"{dataset_name}/{task.id}: {contract_path} declares {sorted(actual)}, expected {sorted(expected)}"
+                )
+    if failures:
+        details = "\n".join(f"- {failure}" for failure in failures)
+        raise DatasetValidationError(f"Authored metric contract validation failed:\n{details}")
+
+
 def _content_provenance(suite_dir: Path, manifest: dict[str, object]) -> tuple[list[dict[str, object]], str, str]:
     raw_tasks = manifest.get("tasks")
     if not isinstance(raw_tasks, list):