feat(benchmarks): explicit agent pin on benchmark datasets; scaffold generates the decoupled layout

nemo_gym/benchmarks.py

@@ -24,13 +24,14 @@
 from pydantic import BaseModel
 
 from nemo_gym import PARENT_DIR
-from nemo_gym.config_types import BenchmarkDatasetConfig
+from nemo_gym.config_types import BenchmarkDatasetConfig, ConfigError
 from nemo_gym.discovery import _parse_no_environment_tolerating_unset_values, discover_components
 from nemo_gym.global_config import (
     POLICY_MODEL_KEY_NAME,
     GlobalConfigDictParser,
     GlobalConfigDictParserConfig,
     get_first_server_config_dict,
+    resolve_dataset_agent,
 )
 
 
@@ -69,11 +70,17 @@ def from_initial_config_dict(
         else:
             global_config_dict = _parse_no_environment_tolerating_unset_values(initial_config_dict)
 
+        # A benchmark dataset may be declared by an agent block (legacy) or a resources server
+        # block (decoupled layout). `resolve_dataset_agent` is the same resolver rollout
+        # dispatch uses, so the listed agent is always the agent that runs.
         datasets: List[BenchmarkDatasetConfig] = []
-        candidate_agent_server_instance_names: List[str] = []
+        declaring_instance_names: List[str] = []
         for server_instance_name in global_config_dict:
             server_config = global_config_dict[server_instance_name]
-            if not isinstance(server_config, (dict, DictConfig)) or "responses_api_agents" not in server_config:
+            if not isinstance(server_config, (dict, DictConfig)):
+                continue
+            is_agent = "responses_api_agents" in server_config
+            if not is_agent and "resources_servers" not in server_config:
                 continue
 
             inner_server_config = get_first_server_config_dict(global_config_dict, server_instance_name)
@@ -83,7 +90,7 @@ def from_initial_config_dict(
                     continue
 
                 datasets.append(BenchmarkDatasetConfig.model_validate(dataset))
-                candidate_agent_server_instance_names.append(server_instance_name)
+                declaring_instance_names.append(str(server_instance_name))
 
         if len(datasets) < 1:
             return
@@ -93,10 +100,15 @@ def from_initial_config_dict(
 
         dataset = datasets[0]
 
+        try:
+            agent_name = resolve_dataset_agent(global_config_dict, declaring_instance_names[0], pin=dataset.agent)
+        except ConfigError as e:
+            raise ConfigError(f"Benchmark config {path}: dataset {dataset.name!r}: {e}") from e
+
         return cls(
             name=dataset.name,
             path=path,
-            agent_name=candidate_agent_server_instance_names[0],
+            agent_name=agent_name,
             num_repeats=dataset.num_repeats,
             dataset=dataset,
         )

nemo_gym/cli/eval.py

@@ -51,6 +51,7 @@
     GlobalConfigDictParserConfig,
     get_first_server_config_dict,
     get_global_config_dict,
+    resolve_dataset_agent,
 )
 
 
@@ -214,11 +215,17 @@ def prepare_benchmark() -> None:
     )
     prepare_benchmark_config = PrepareBenchmarkConfig.model_validate(global_config_dict)
 
+    # A benchmark dataset may be declared by an agent block (legacy) or a resources server
+    # block (decoupled layout). `resolve_dataset_agent` is the same resolver rollout dispatch
+    # uses, so preparation and rollout always agree.
     benchmarks_dict: Dict[str, BenchmarkConfig] = dict()
     inspected_server_instances: List[str] = []
     for server_instance_name in global_config_dict:
         server_config = global_config_dict[server_instance_name]
-        if not isinstance(server_config, (dict, DictConfig)) or "responses_api_agents" not in server_config:
+        if not isinstance(server_config, (dict, DictConfig)):
+            continue
+        is_agent = "responses_api_agents" in server_config
+        if not is_agent and "resources_servers" not in server_config:
             continue
 
         inspected_server_instances.append(server_instance_name)
@@ -243,10 +250,17 @@ def prepare_benchmark() -> None:
 
         dataset = datasets[0]
 
-        benchmarks_dict[server_instance_name] = BenchmarkConfig(
+        try:
+            agent_name = resolve_dataset_agent(global_config_dict, str(server_instance_name), pin=dataset.agent)
+        except ConfigError as e:
+            raise ConfigError(f"Benchmark dataset {dataset.name!r}: {e}") from e
+
+        # Keyed by the declaring instance: two declarations may resolve to the same agent, and
+        # keying by agent would silently drop all but the last.
+        benchmarks_dict[str(server_instance_name)] = BenchmarkConfig(
             name=dataset.name,
             path=Path(""),
-            agent_name=server_instance_name,
+            agent_name=agent_name,
             num_repeats=dataset.num_repeats,
             dataset=dataset,
         )

nemo_gym/config_types.py

@@ -533,6 +533,16 @@ class BenchmarkDatasetConfig(BaseModel):
     prepare_script: Path
     prompt_config: Optional[Path] = None
     num_repeats: int = Field(default=1, ge=1)
+    agent: Optional[str] = Field(
+        default=None,
+        description=(
+            "Agent instance that runs this benchmark (a top-level key of the merged config). "
+            "Only needed when the config is ambiguous: the dataset is declared on a resources "
+            "server that several agents reference. The pin must name one of those agents — rows "
+            "are dispatched along the agent -> resources server edge, so any other value is a "
+            "config error. Unambiguous configs resolve without it."
+        ),
+    )
 
 
 ########################################

nemo_gym/environment/scaffold.py

@@ -277,11 +277,21 @@ def _manifest(composition: _Composition) -> EnvironmentManifest:
 
 
 def _asset_config(composition: _Composition) -> str:
+    # Datasets are declared on the resources server (which defines what the rows mean and how
+    # they are scored), never on the agent — the agent is a run-time choice. The scaffolded
+    # config has exactly one agent referencing the resources server, so benchmark datasets
+    # resolve their harness from that edge; the dataset-level `agent:` key stays reserved for
+    # genuinely ambiguous configs.
+    dataset_dict = composition.dataset.model_dump(mode="json", exclude_none=True)
     agent_config: dict[str, Any] = {
         "resources_server": {"type": "resources_servers", "name": composition.resource_instance},
         "model_server": {"type": "responses_api_models", "name": "policy_model"},
-        "datasets": [composition.dataset.model_dump(mode="json", exclude_none=True)],
     }
+    # When reusing a verifier from another config we do not know its inner implementation key, so
+    # a partial resources-server override could not be merged safely; the dataset stays on the
+    # agent block there (still supported).
+    if composition.reused_verifier is not None:
+        agent_config["datasets"] = [dataset_dict]
     reused_agent = composition.reused_verifier.agent_instance if composition.reused_verifier else None
     if reused_agent:
         agent_entry: dict[str, Any] = {
@@ -296,6 +306,10 @@ def _asset_config(composition: _Composition) -> str:
         "config_paths": [composition.config_reference],
         composition.agent_instance: agent_entry,
     }
+    if composition.reused_verifier is None:
+        config[composition.resource_instance] = {
+            "resources_servers": {composition.resource_implementation: {"datasets": [dataset_dict]}}
+        }
     if composition.rollout_driver:
         config["rollout_collection_driver"] = composition.rollout_driver
 
@@ -490,24 +504,14 @@ def _standalone_resources_server_config(module_name: str) -> str:
     return dedent(
         f"""\
         # Resources server: owns this environment's task verification (verify()) and reward.
+        # Datasets are declared here — the resources server defines what the rows mean and how
+        # they are scored. Which agent runs them is chosen at run time.
         {module_name}_resources_server:          # instance name — how agents/CLI refer to this server
           resources_servers:                    # server type: resources_servers | responses_api_agents | responses_api_models
             {module_name}:                      # implementation directory under resources_servers/
               entrypoint: app.py                # server entry module
               domain: other                     # task domain; change to the closest supported domain
               verified: false                   # set true once the benchmark has been baselined and reviewed
-
-        # Pair the server with the default agent and dataset slots.
-        {module_name}_simple_agent:             # pass this instance as --agent to gym eval run
-          responses_api_agents:
-            simple_agent:
-              entrypoint: app.py
-              resources_server:
-                type: resources_servers
-                name: {module_name}_resources_server
-              model_server:
-                type: responses_api_models
-                name: policy_model
               datasets:
               - name: train
                 type: train
@@ -524,6 +528,18 @@ def _standalone_resources_server_config(module_name: str) -> str:
                 type: example
                 jsonl_fpath: resources_servers/{module_name}/data/example.jsonl
                 num_repeats: 1
+
+        # Pair the server with the default agent.
+        {module_name}_simple_agent:             # pass this instance as --agent to gym eval run
+          responses_api_agents:
+            simple_agent:
+              entrypoint: app.py
+              resources_server:
+                type: resources_servers
+                name: {module_name}_resources_server
+              model_server:
+                type: responses_api_models
+                name: policy_model
         """
     )
 
@@ -551,9 +567,8 @@ def _standalone_data_gitignore() -> str:
         """\
         *train.jsonl
         *validation.jsonl
-        *train_prepare.jsonl
-        *validation_prepare.jsonl
-        *example_prepare.jsonl
+        *_prepare.jsonl
+        *_prepare.*.jsonl
         """
     )
 

nemo_gym/environment/validation.py

@@ -30,7 +30,12 @@
     dump_manifest,
     load_manifest,
 )
-from nemo_gym.global_config import GlobalConfigDictParser, GlobalConfigDictParserConfig
+from nemo_gym.global_config import (
+    GlobalConfigDictParser,
+    GlobalConfigDictParserConfig,
+    dataset_agent_pins,
+    resolve_dataset_agent,
+)
 from nemo_gym.prompt import apply_prompt_to_row, load_prompt_config, validate_prompt_compatibility
 
 
@@ -144,6 +149,20 @@ def _select_dataset_agent(agents: list[Any]) -> Any:
     )
 
 
+def _resolve_dataset_owner_agent(resolved: DictConfig, owner_instance_name: str) -> str:
+    """The agent rollout dispatch routes ``owner_instance_name``'s datasets to (validation-time check)."""
+    pins = dataset_agent_pins(resolved, owner_instance_name)
+    if len(pins) > 1:
+        raise EnvironmentValidationError(
+            f"Datasets on {owner_instance_name!r} pin conflicting agents ({sorted(pins)}); rows route by the "
+            "declaring instance alone, so all pins on one instance must agree."
+        )
+    try:
+        return resolve_dataset_agent(resolved, owner_instance_name, pin=pins[0] if pins else None)
+    except ConfigError as e:
+        raise EnvironmentValidationError(f"Datasets on {owner_instance_name!r}: {e}") from e
+
+
 def _resolve_manifest_composition(config_path: Path) -> ResolvedComposition:
     """Resolve manifest wiring without probing or materializing runtime services."""
     initial = OmegaConf.merge(
@@ -163,17 +182,46 @@ def _resolve_manifest_composition(config_path: Path) -> ResolvedComposition:
     servers = parser.filter_for_server_instance_configs(resolved)
     by_instance = {server.name: server for server in servers}
     agents = [server for server in servers if server.SERVER_TYPE == "responses_api_agents"]
-    selected_agent = _select_dataset_agent(agents)
+
+    # Datasets may live on a resources server (decoupled layout) or on an agent (legacy and
+    # self-contained layouts). The agent is resolved by `resolve_dataset_agent` — the same
+    # resolver rollout dispatch uses — so a config that validates is a config that routes.
+    rs_with_data = [s for s in servers if s.SERVER_TYPE == "resources_servers" and s.datasets]
+    dataset_rs = None
+    if len(rs_with_data) > 1:
+        names = ", ".join(s.name for s in rs_with_data)
+        raise EnvironmentValidationError(
+            f"Workload config must define exactly one dataset-bearing instance; found resources servers: {names}."
+        )
+    if rs_with_data:
+        dataset_rs = rs_with_data[0]
+        agent_name = _resolve_dataset_owner_agent(resolved, dataset_rs.name)
+        selected_agent = by_instance.get(agent_name)
+        if selected_agent is None:
+            raise EnvironmentValidationError(
+                f"Datasets on {dataset_rs.name!r} route to agent {agent_name!r}, which is not defined in the config."
+            )
+    else:
+        selected_agent = _select_dataset_agent(agents)
+        # Rejects an `agent:` pin naming anyone but the declaring agent (dispatch would ignore it).
+        _resolve_dataset_owner_agent(resolved, selected_agent.name)
 
     agent_server = _implementation_name(selected_agent)
     agent_config = selected_agent.get_inner_run_server_config_dict()
     resources_ref = agent_config.get("resources_server") or {}
     model_ref = agent_config.get("model_server") or {}
-    resources_instance = by_instance.get(resources_ref.get("name")) if isinstance(resources_ref, DictConfig) else None
+    resources_instance = (
+        dataset_rs
+        if dataset_rs is not None
+        else by_instance.get(resources_ref.get("name"))
+        if isinstance(resources_ref, DictConfig)
+        else None
+    )
 
     resources_server = _implementation_name(resources_instance) if resources_instance is not None else None
     model_server = model_ref.get("name") if isinstance(model_ref, DictConfig) else None
-    datasets = tuple(_manifest_dataset(dataset) for dataset in (selected_agent.datasets or []))
+    dataset_owner = dataset_rs if dataset_rs is not None else selected_agent
+    datasets = tuple(_manifest_dataset(dataset) for dataset in (dataset_owner.datasets or []))
     grading_mode = None
     if resources_instance is not None:
         grading_mode = resources_instance.get_inner_run_server_config_dict().get("grading_mode")

nemo_gym/global_config.py

@@ -1415,6 +1415,72 @@ def agents_by_resources_server(global_config_dict: DictConfig) -> Dict[str, List
     return result
 
 
+def dataset_agent_pins(global_config_dict: DictConfig, instance_name: str) -> List[str]:
+    """Distinct dataset-level ``agent:`` pins declared by one server instance, in declaration order."""
+    try:
+        inner = get_first_server_config_dict(global_config_dict, instance_name)
+    except Exception:
+        return []
+    pins: List[str] = []
+    for dataset in inner.get("datasets") or []:
+        if not isinstance(dataset, (dict, DictConfig)):
+            continue
+        pin = dataset.get("agent")
+        if isinstance(pin, str) and pin not in pins:
+            pins.append(pin)
+    return pins
+
+
+def resolve_dataset_agent(
+    global_config_dict: DictConfig, declaring_instance_name: str, pin: Optional[str] = None
+) -> str:
+    """Resolve the agent that runs a dataset declared by ``declaring_instance_name``.
+
+    Single source of truth for dataset -> agent routing, shared by benchmark discovery,
+    preparation, manifest validation, and rollout dispatch, so they can never disagree.
+    First hit wins:
+
+    1. ``pin`` (the dataset's ``agent:`` key) — validated: it must name the declaring agent
+       itself, or an agent referencing the declaring resources server. Anything else is a hard
+       error, never a silent re-route (rows carry only the declaring instance name).
+    2. The declaring agent block itself.
+    3. The unique agent referencing the declaring resources server; zero or 2+ is a hard error.
+    """
+    block = global_config_dict.get(declaring_instance_name)
+    is_agent = isinstance(block, DictConfig) and "responses_api_agents" in block
+
+    if is_agent:
+        if pin is not None and pin != declaring_instance_name:
+            raise ConfigError(
+                f"dataset declared by agent {declaring_instance_name!r} pins agent {pin!r}, but rows are "
+                f"dispatched to the declaring agent, so the pin would silently not apply. Drop the `agent:` "
+                "key, or declare the dataset on the resources server and pin there."
+            )
+        return str(declaring_instance_name)
+
+    candidates = agents_by_resources_server(global_config_dict).get(str(declaring_instance_name), [])
+    if pin is not None:
+        if pin not in candidates:
+            raise ConfigError(
+                f"the dataset pins agent {pin!r}, but no agent of that name references resources server "
+                f"{declaring_instance_name!r} (rows are dispatched along the agent -> resources server edge). "
+                f"Agents referencing it: {sorted(candidates) if candidates else 'none'}."
+            )
+        return pin
+    if not candidates:
+        raise ConfigError(
+            f"no agent in the running config references resources server {declaring_instance_name!r}. "
+            "Point an agent's `resources_server` at it, or declare the dataset on the agent."
+        )
+    if len(candidates) > 1:
+        raise ConfigError(
+            f"{len(candidates)} agents reference this resources server ({sorted(candidates)}). "
+            "Pin the harness with an `agent:` key on the dataset declaration, or pass "
+            f"+agent_map={{{declaring_instance_name}: <agent>}} to pick one at run time."
+        )
+    return candidates[0]
+
+
 def find_open_port(
     disallowed_ports: Optional[List[int]] = None,
     max_retries: int = 50,