diff --git a/environment_servers/legacy_agent/task_data.py b/environment_servers/legacy_agent/task_data.py new file mode 100644 index 00000000000..ec50b9ddcd9 --- /dev/null +++ b/environment_servers/legacy_agent/task_data.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Task fields of a ``legacy_agent`` relay: a single-agent run request for the agent's ``/run``. + +The agent's resources server, or the agent itself when it is self-contained, declares the other fields. +""" + +from nemo_gym.task_data import SingleAgentTaskData as TaskData + + +__all__ = ["TaskData"] diff --git a/environment_servers/nemo_user_sim/task_data.py b/environment_servers/nemo_user_sim/task_data.py new file mode 100644 index 00000000000..ab3549cfdfd --- /dev/null +++ b/environment_servers/nemo_user_sim/task_data.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Task fields that the ``nemo_user_sim`` environment server reads: its resources server's ``TaskData``. + +The environment server runs the user and assistant agents itself and reads the same task fields its +resources server declares, so it shares that schema instead of keeping a copy. +""" + +from resources_servers.nemo_user_sim.task_data import TaskData + + +__all__ = ["TaskData"] diff --git a/environment_servers/single_agent_turn/task_data.py b/environment_servers/single_agent_turn/task_data.py new file mode 100644 index 00000000000..4b96231a2e8 --- /dev/null +++ b/environment_servers/single_agent_turn/task_data.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Task fields that ``single_agent_turn`` reads itself. + +It sends ``responses_create_params`` to its agent and passes every other field to its resources +server, whose ``TaskData`` declares them. ``SingleAgentTurnTaskInput`` is the runtime model. +""" + +from nemo_gym.task_data import SingleAgentTaskData as TaskData + + +__all__ = ["TaskData"] diff --git a/environment_servers/single_agent_turn_legacy/task_data.py b/environment_servers/single_agent_turn_legacy/task_data.py new file mode 100644 index 00000000000..75d4130eca0 --- /dev/null +++ b/environment_servers/single_agent_turn_legacy/task_data.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Task fields that ``single_agent_turn_legacy`` reads itself: the same as ``single_agent_turn``.""" + +from nemo_gym.task_data import SingleAgentTaskData as TaskData + + +__all__ = ["TaskData"] diff --git a/fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx b/fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx index 58be9d87445..fa94151d21f 100644 --- a/fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx +++ b/fern/versions/latest/pages/environment-tutorials/resources-server-apis.mdx @@ -459,6 +459,24 @@ The source JSONL stays unchanged and can still be shared with flat-row recipes. Omit `taskset` to keep the current flat-row format. This option accepts flat input only; already materialized input should be passed directly with `--input`. +Collation validates each source row against two `task_data.py` schemas. The first +belongs to the Environment Server the dataset routes to, under +`environment_servers//task_data.py`, and declares the fields that +server reads itself. The second belongs to the resources server that server binds, +and declares the remaining task fields. Each schema validates only the fields it +declares and the fields neither declares, so a resources schema with +`extra="forbid"` does not need to list the Environment Server's fields. Fields that +neither schema declares are reported against the resources server's schema. + +The built-in single-agent Environment Servers declare `responses_create_params`, +one agent's request to one model, with `SingleAgentTaskData` from +`nemo_gym/task_data.py`. Datasets without `taskset` route by agent, so they are +single-agent run requests and use the same schema. An Environment Server that runs +several agents declares its own fields instead, and its rows need no +`responses_create_params`. Rows that a preparation script already materialized +name their own taskset, so collation checks them only against the resources +server's schema, and their Environment Server validates them at dispatch. + For example, add these keys to a config whose `swebench_pro_hermes` Environment Server binds the `swebench_pro_hermes_resources_server` instance: diff --git a/nemo_gym/global_config.py b/nemo_gym/global_config.py index 96d59270c9f..829c9ebefa8 100644 --- a/nemo_gym/global_config.py +++ b/nemo_gym/global_config.py @@ -1977,6 +1977,36 @@ def taskset_environment_server_name(global_config_dict: DictConfig, taskset: Opt return name if isinstance(name, str) else None +@dataclass(frozen=True) +class TasksetEnvironmentServer: + """The Environment Server a taskset routes to.""" + + name: str + implementation: str + config: DictConfig + + +def taskset_environment_server(global_config_dict: DictConfig, taskset: str) -> TasksetEnvironmentServer: + """Resolve a taskset through ``environment_server_routes`` to its Environment Server instance. + + Raises ``ConfigError`` when the taskset has no route or the route names no Environment Server. + """ + routes = global_config_dict.get(ENVIRONMENT_SERVER_ROUTES_KEY_NAME) + environment_name = routes.get(taskset) if isinstance(routes, DictConfig) else None + if not isinstance(environment_name, str) or not environment_name: + raise ConfigError(f"No Environment Server route for taskset {taskset!r}; set environment_server_routes.") + environment = global_config_dict.get(environment_name) + servers = environment.get(ENVIRONMENT_SERVER_TYPE_KEY_NAME) if isinstance(environment, DictConfig) else None + if not isinstance(servers, DictConfig) or len(servers) != 1: + raise ConfigError(f"Taskset {taskset!r} route {environment_name!r} must name an Environment Server.") + implementation, environment_config = next(iter(servers.items())) + if not isinstance(environment_config, DictConfig): + raise ConfigError(f"Environment Server {environment_name!r} must have a configuration mapping.") + return TasksetEnvironmentServer( + name=environment_name, implementation=str(implementation), config=environment_config + ) + + def resolve_dataset_agent( global_config_dict: DictConfig, declaring_instance_name: str, @@ -2006,17 +2036,8 @@ def resolve_dataset_agent( is_agent = isinstance(block, DictConfig) and "responses_api_agents" in block if taskset is not None: - routes = global_config_dict.get(ENVIRONMENT_SERVER_ROUTES_KEY_NAME) - environment_name = routes.get(taskset) if isinstance(routes, DictConfig) else None - if not isinstance(environment_name, str) or not environment_name: - raise ConfigError(f"No Environment Server route for taskset {taskset!r}; set environment_server_routes.") - environment = global_config_dict.get(environment_name) - servers = environment.get(ENVIRONMENT_SERVER_TYPE_KEY_NAME) if isinstance(environment, DictConfig) else None - if not isinstance(servers, DictConfig) or len(servers) != 1: - raise ConfigError(f"Taskset {taskset!r} route {environment_name!r} must name an Environment Server.") - environment_config = next(iter(servers.values())) - if not isinstance(environment_config, DictConfig): - raise ConfigError(f"Environment Server {environment_name!r} must have a configuration mapping.") + environment = taskset_environment_server(global_config_dict, taskset) + environment_name, environment_config = environment.name, environment.config agent_names = [] for agent_ref in environment_server_agent_refs(environment_config): name = agent_ref.get("name") diff --git a/nemo_gym/task_data.py b/nemo_gym/task_data.py index 17fe9f6aa0d..3bbf86230e9 100644 --- a/nemo_gym/task_data.py +++ b/nemo_gym/task_data.py @@ -8,6 +8,12 @@ A self-contained agent (one that declares datasets but no ``resources_server`` reference) ships the same module under ``responses_api_agents//`` and owns its rows' schema; an agent WITH a reference uses that resources server's schema instead. +An environment server ships the same module under ``environment_servers//`` to +declare the task fields it reads itself. Rows are validated against the environment server they +route to, composed with the bound resources server's schema: each schema validates the fields it +declares plus the fields neither declares. Single-agent environment servers declare +``responses_create_params``, one agent's request to one model, through ``SingleAgentTaskData``. +Rows routed by agent rather than by taskset are single-agent run requests and use it too. It describes the task-owned fields of that server's dataset rows, written FLAT in the planned end-state shape: the fields as they will appear inside the unified ``task_data`` row key after the row-format migration. Framework-owned keys (see ``RESERVED_ROW_KEYS``) are never part of @@ -45,9 +51,9 @@ import sys from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Iterable, List, Mapping, Optional +from typing import Any, Dict, Iterable, List, Mapping, Optional, Union, get_args -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from nemo_gym.episode_types import is_materialized_task_row @@ -56,12 +62,13 @@ TASK_DATA_EXPORT_NAME = "TaskData" # Framework-owned top-level row keys. Everything else in a row is task-owned and is the subject of -# the server's TaskData schema. (The row-format migration will extend this set with reserved keys such as +# the environment server's and resources server's TaskData schemas. ``responses_create_params`` is +# not one of them: it is a task field of single-agent environment servers (``SingleAgentTaskData``). +# (The row-format migration will extend this set with reserved keys such as # ``task_id``/``subset_for_metrics``/``provenance``; until then some servers legitimately use those # names as task fields, so they stay task-owned here.) RESERVED_ROW_KEYS = frozenset( { - "responses_create_params", "agent_ref", "task_source", "_ng_task_index", @@ -70,6 +77,32 @@ ) +class SingleAgentResponsesCreateParams(BaseModel): + """The part of a Responses API request that a dataset row supplies. + + ``input`` may be missing from a source row whose dataset declares a ``prompt_config``, which + fills it during collation. The environment server validates the full request at runtime; this + keeps schemas free of server dependencies. + """ + + model_config = ConfigDict(extra="allow") + + input: Optional[Union[str, List[Any]]] = None + + +class SingleAgentTaskData(BaseModel): + """Task fields that every single-agent environment server reads: one agent's request to one model. + + The bound resources server's ``TaskData`` declares the remaining fields. + """ + + model_config = ConfigDict(extra="allow") + + responses_create_params: SingleAgentResponsesCreateParams = Field( + json_schema_extra={"consumed_by": ["prompt"]}, + ) + + class TaskDataSchemaError(Exception): """A server's ``task_data.py`` exists but does not satisfy the protocol.""" @@ -197,9 +230,8 @@ def summary(self) -> str: return "\n".join(parts) -def legacy_metadata_fields(adapter: TypeAdapter) -> frozenset: - """Schema fields annotated ``legacy_location: verifier_metadata`` (today's wire reads them there).""" - from typing import get_args +def _schema_models(adapter: TypeAdapter) -> List[type[BaseModel]]: + """The models a schema is built from: the model itself, or each member of a union.""" def models_of(tp, out): if isinstance(tp, type) and issubclass(tp, BaseModel): @@ -209,8 +241,24 @@ def models_of(tp, out): models_of(arg, out) return out + return models_of(getattr(adapter, "_type", None), []) + + +def declared_fields(adapter: TypeAdapter) -> frozenset: + """The field names, and their aliases, that any model of a schema declares.""" + names = set() + for model in _schema_models(adapter): + for field_name, info in model.model_fields.items(): + names.add(field_name) + if isinstance(info.alias, str): + names.add(info.alias) + return frozenset(names) + + +def legacy_metadata_fields(adapter: TypeAdapter) -> frozenset: + """Schema fields annotated ``legacy_location: verifier_metadata`` (today's wire reads them there).""" names = set() - for model in models_of(getattr(adapter, "_type", None), []): + for model in _schema_models(adapter): for field_name, info in model.model_fields.items(): extra = info.json_schema_extra if isinstance(extra, dict) and extra.get("legacy_location") == LEGACY_METADATA_KEY: @@ -218,16 +266,69 @@ def models_of(tp, out): return frozenset(names) +SINGLE_AGENT_TASK_DATA = TypeAdapter(SingleAgentTaskData) + + class TaskDataValidator: - """Validates dataset rows against a server's ``TaskData`` schema, accumulating a report.""" + """Validates dataset rows against an environment server's and a resources server's ``TaskData``. + + ``environment_adapter`` is the schema of the environment server the rows route to, such as + ``SINGLE_AGENT_TASK_DATA`` for rows routed by agent. Each schema validates the fields it + declares plus the fields neither declares, so a resources schema with ``extra="forbid"`` does not + see the environment server's fields. Undeclared fields are reported against the resources server's + schema. Either adapter may be ``None`` when its server ships no schema. The environment schema does + not validate materialized rows (see ``validate_row``). + """ - def __init__(self, server_name: str, adapter: TypeAdapter, dataset_fpath: str): + def __init__( + self, + server_name: str, + adapter: Optional[TypeAdapter], + dataset_fpath: str, + *, + environment_adapter: Optional[TypeAdapter] = None, + ): self._adapter = adapter - self._legacy_fields = legacy_metadata_fields(adapter) + self._environment_adapter = environment_adapter + self._fields = declared_fields(adapter) if adapter is not None else frozenset() + self._environment_fields = ( + declared_fields(environment_adapter) if environment_adapter is not None else frozenset() + ) + self._legacy_fields = legacy_metadata_fields(adapter) if adapter is not None else frozenset() self.report = TaskDataValidationReport(server_name=server_name, dataset_fpath=dataset_fpath) + def _validate_fields(self, subject: Dict[str, Any], *, environment_applies: bool) -> tuple[List[str], set]: + """Validate one row's task fields against each schema; return its errors and undeclared fields.""" + errors: List[str] = [] + undeclared: set = set() + # The environment server's fields stay out of the resources schema's subject either way, so + # they are never reported as fields the resources server does not declare. + environment_adapter = self._environment_adapter if environment_applies else None + for adapter, own, other, reports_unknown in ( + (environment_adapter, self._environment_fields, self._fields, False), + (self._adapter, self._fields, self._environment_fields, True), + ): + if adapter is None: + continue + fields = {key: value for key, value in subject.items() if key in own or key not in other} + try: + validated = adapter.validate_python(fields) + except Exception as e: + errors.append(str(e).strip().replace("\n", "; ")) + continue + # Pydantic returns the concrete model (the selected union member for union schemas), and + # with extra="allow" it stores undeclared inputs on __pydantic_extra__ — so unknown-field + # reporting works uniformly for plain models and discriminated unions. + if reports_unknown and isinstance(validated, BaseModel): + undeclared = set(getattr(validated, "__pydantic_extra__", None) or {}) + return errors, undeclared + def validate_row(self, row_index: int, row: Dict[str, Any]) -> None: self.report.rows += 1 + # A materialized row already names its taskset, and the environment server that taskset routes + # to validates its input at dispatch. The environment schema does not validate it here, so such a + # row in a dataset routed by agent is not required to be a single-agent run request. + environment_applies = not is_materialized_task_row(row) # New materialized tasks keep the flat source fields; historical single-agent # tasks use task_input.task_data. Normalize both through the existing schema path. task_input = row.get("task_input") @@ -270,27 +371,31 @@ def validate_row(self, row_index: int, row: Dict[str, Any]) -> None: subject, conflicts = normalize_task_fields(row) for key in conflicts: self.report.conflicting_keys[key] = self.report.conflicting_keys.get(key, 0) + 1 - try: - validated = self._adapter.validate_python(subject) - except Exception as e: + errors, undeclared = self._validate_fields(subject, environment_applies=environment_applies) + if errors: self.report.error_rows += 1 if len(self.report.errors) < TaskDataValidationReport.MAX_RECORDED_ERRORS: - first_line = str(e).strip().replace("\n", "; ") - self.report.errors.append(f"{row_index}: {first_line[:400]}") + self.report.errors.append(f"{row_index}: {'; '.join(error[:400] for error in errors)}") return - # Pydantic returns the concrete model (the selected union member for union schemas), and - # with extra="allow" it stores undeclared inputs on __pydantic_extra__ — so unknown-field - # reporting works uniformly for plain models and discriminated unions. - if isinstance(validated, BaseModel): - for key in getattr(validated, "__pydantic_extra__", None) or {}: - self.report.unknown_keys[key] = self.report.unknown_keys.get(key, 0) + 1 + for key in undeclared: + self.report.unknown_keys[key] = self.report.unknown_keys.get(key, 0) + 1 def validate_jsonl_rows( - server_name: str, adapter: TypeAdapter, dataset_fpath: str, lines: Iterable[str] + server_name: str, + adapter: Optional[TypeAdapter], + dataset_fpath: str, + lines: Iterable[str], + *, + environment_adapter: Optional[TypeAdapter] = None, ) -> TaskDataValidationReport: """Validate an iterable of JSONL lines against one schema; entry point for whole-file validation.""" - validator = TaskDataValidator(server_name=server_name, adapter=adapter, dataset_fpath=dataset_fpath) + validator = TaskDataValidator( + server_name=server_name, + adapter=adapter, + dataset_fpath=dataset_fpath, + environment_adapter=environment_adapter, + ) for i, line in enumerate(lines): line = line.strip() if not line: diff --git a/nemo_gym/train_data_utils.py b/nemo_gym/train_data_utils.py index a84d4a24e67..ab10b01a00c 100644 --- a/nemo_gym/train_data_utils.py +++ b/nemo_gym/train_data_utils.py @@ -25,7 +25,7 @@ from devtools import pprint from omegaconf import DictConfig, OmegaConf -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, TypeAdapter, ValidationError from tqdm.auto import tqdm from nemo_gym import _resolve_under_cwd_or_install @@ -58,12 +58,14 @@ GlobalConfigDictParser, get_global_config_dict, resolve_dataset_agent, + taskset_environment_server, ) from nemo_gym.hf_utils import ( download_hf_dataset_as_jsonl, ) from nemo_gym.prompt import apply_prompt_to_row, load_prompt_config, validate_prompt_compatibility from nemo_gym.task_data import ( + SINGLE_AGENT_TASK_DATA, TaskDataSchemaError, TaskDataValidationReport, TaskDataValidator, @@ -528,7 +530,29 @@ def _prepare_command(global_config: DictConfig, dataset: DatasetConfig) -> str: return "gym eval prepare --config " +class TasksetSchemas(BaseModel): + """Where a taskset's rows find their schemas: its environment server and that server's resources server.""" + + environment_implementation: str + resources_implementation: Optional[str] = None + + @classmethod + def for_taskset(cls, global_config_dict: DictConfig, taskset: str) -> "TasksetSchemas": + environment = taskset_environment_server(global_config_dict, taskset) + resources_ref = environment.config.get("resources_server") + resources_name = resources_ref.get("name") if isinstance(resources_ref, DictConfig) else None + resources = global_config_dict.get(resources_name) if isinstance(resources_name, str) else None + servers = resources.get("resources_servers") if isinstance(resources, DictConfig) else None + return cls( + environment_implementation=environment.implementation, + resources_implementation=next(iter(servers)) if isinstance(servers, DictConfig) and servers else None, + ) + + class TrainDataProcessor(BaseModel): + # Filled while loading configs, which is the only step with the run's environment_server_routes. + _taskset_schemas: Dict[str, TasksetSchemas] = PrivateAttr(default_factory=dict) + def run(self, global_config_dict: DictConfig): # pragma: no cover """ See the README section "How To: Prepare and validate data for PR submission or RL training" @@ -638,6 +662,9 @@ def load_and_validate_server_instance_configs( pin=dataset.agent if isinstance(dataset, BenchmarkDatasetConfig) else None, taskset=dataset.taskset, ) + self._taskset_schemas[dataset.taskset] = TasksetSchemas.for_taskset( + global_config_dict, dataset.taskset + ) inner_config = agent_config.get_inner_run_server_config() inner_config.datasets = in_scope_datasets @@ -1015,37 +1042,65 @@ def _owning_resources_server_impl( return next(iter(other.resources_servers)) return None + @staticmethod + def _load_task_data_schema(impl_key: str, base_folder: str) -> Optional[TypeAdapter]: + """Load ``//task_data.py``, or None when the server ships no usable schema.""" + server_dir = find_server_dir(impl_key, base_folder=base_folder) + if server_dir is None: + return None + try: + return load_task_data_schema(server_dir) + except TaskDataSchemaError as e: + warnings.warn(f"Skipping task_data validation for {impl_key}: {e}", stacklevel=2) + return None + @classmethod def _task_data_validator_for( cls, c: ServerInstanceConfig, d: Union[DatasetConfig, BenchmarkDatasetConfig], server_instance_configs: List[ServerInstanceConfig], + *, + taskset_schemas: Optional["TasksetSchemas"] = None, ) -> Optional[TaskDataValidator]: + """The validator for one dataset: its environment server's schema composed with its task owner's. + + A dataset routed by taskset uses the routed environment server's schema (``taskset_schemas``), + and the schema of the resources server that environment server binds. A dataset routed by agent + is a single-agent run request, so it uses ``SingleAgentTaskData``, and the schema of the + resources server or self-contained agent that owns its data. + """ + # A taskset whose route was not resolved has no known environment server, so it is not + # assumed to be a single-agent run request. + environment_adapter: Optional[TypeAdapter] = SINGLE_AGENT_TASK_DATA if d.taskset is None else None impl_key = cls._owning_resources_server_impl(c, server_instance_configs) base_folder = "resources_servers" - if impl_key is None: + if taskset_schemas is not None: + environment_adapter = cls._load_task_data_schema( + taskset_schemas.environment_implementation, "environment_servers" + ) + if taskset_schemas.resources_implementation is not None: + impl_key = taskset_schemas.resources_implementation + if impl_key is None and c.SERVER_TYPE == "responses_api_agents": # No resources server owns this data. A self-contained agent (one that declares no # resources_server reference at all) may own a schema itself, under # responses_api_agents//task_data.py. An agent whose reference is - # dangling is skipped: its schema home is the (missing) resources server. - if c.SERVER_TYPE != "responses_api_agents": - return None - if getattr(c.get_inner_run_server_config(), "resources_server", None) is not None: - return None - impl_key = next(iter(c.responses_api_agents)) - base_folder = "responses_api_agents" - server_dir = find_server_dir(impl_key, base_folder=base_folder) - if server_dir is None: + # dangling has no task owner: its schema home is the (missing) resources server. + if getattr(c.get_inner_run_server_config(), "resources_server", None) is None: + impl_key = next(iter(c.responses_api_agents)) + base_folder = "responses_api_agents" + adapter = cls._load_task_data_schema(impl_key, base_folder) if impl_key is not None else None + if adapter is None and environment_adapter is None: return None - try: - adapter = load_task_data_schema(server_dir) - except TaskDataSchemaError as e: - warnings.warn(f"Skipping task_data validation for {impl_key}: {e}", stacklevel=2) - return None - if adapter is None: - return None - return TaskDataValidator(server_name=impl_key, adapter=adapter, dataset_fpath=str(d.jsonl_fpath)) + server_name = impl_key if adapter is not None else None + if server_name is None: + server_name = taskset_schemas.environment_implementation if taskset_schemas is not None else "single-agent" + return TaskDataValidator( + server_name=server_name, + adapter=adapter, + dataset_fpath=str(d.jsonl_fpath), + environment_adapter=environment_adapter, + ) @classmethod def _dataset_metrics_hook_for( @@ -1121,7 +1176,12 @@ def _collate_samples_single_type( positional_task_id_rows = 0 validator = None if task_data_validation != "off": - validator = self._task_data_validator_for(c, d, server_instance_configs) + validator = self._task_data_validator_for( + c, + d, + server_instance_configs, + taskset_schemas=self._taskset_schemas.get(d.taskset) if d.taskset is not None else None, + ) with open(prepare_path, "w") as target: for row_index, line in enumerate(self._iter_dataset_lines(d)): row = json.loads(line) diff --git a/resources_servers/terminal_bench_4/tests/test_benchmark_preparation.py b/resources_servers/terminal_bench_4/tests/test_benchmark_preparation.py index d58fdc9d16c..0c8cf2d7967 100644 --- a/resources_servers/terminal_bench_4/tests/test_benchmark_preparation.py +++ b/resources_servers/terminal_bench_4/tests/test_benchmark_preparation.py @@ -8,14 +8,19 @@ from omegaconf import OmegaConf from benchmarks.terminal_bench_4 import prepare as preparation -from nemo_gym.task_data import TaskDataValidator, load_task_data_schema +from nemo_gym.task_data import SINGLE_AGENT_TASK_DATA, TaskDataValidator, load_task_data_schema SERVER_DIR = Path(__file__).resolve().parents[1] def task_validator() -> TaskDataValidator: - return TaskDataValidator("terminal_bench_4", load_task_data_schema(SERVER_DIR), "test.jsonl") + return TaskDataValidator( + "terminal_bench_4", + load_task_data_schema(SERVER_DIR), + "test.jsonl", + environment_adapter=SINGLE_AGENT_TASK_DATA, + ) def test_prepared_names_match_pinned_manifest(tmp_path, monkeypatch): diff --git a/tests/unit_tests/test_benchmarks.py b/tests/unit_tests/test_benchmarks.py index 84118308294..5590ae8a871 100644 --- a/tests/unit_tests/test_benchmarks.py +++ b/tests/unit_tests/test_benchmarks.py @@ -1208,10 +1208,13 @@ def test_agent_declared_taskset_must_route_to_that_agent(self, config): @pytest.fixture def multi_agent_config(self, config): config.user = {"responses_api_agents": {"impl": {"entrypoint": "app.py"}}} - server = config.environment.environment_servers.single_agent_turn + # A conversation server's own implementation: collation selects the task schema by implementation, + # and single_agent_turn's schema requires one agent's responses_create_params. + server = config.environment.environment_servers.pop("single_agent_turn") del server["agent_server"] server.user_agent = {"type": "responses_api_agents", "name": "user"} server.assistant_agent = {"type": "responses_api_agents", "name": "agent"} + config.environment.environment_servers.conversation = server return config def test_multi_agent_route_resolves_without_a_pin(self, multi_agent_config): @@ -1219,7 +1222,7 @@ def test_multi_agent_route_resolves_without_a_pin(self, multi_agent_config): # The server fronts two agents and names no agent_server, so no single agent is reported. assert resolve_dataset_agent(multi_agent_config, "resources", taskset="swe:test") is None - multi_agent_config.environment.environment_servers.single_agent_turn.agent_server = { + multi_agent_config.environment.environment_servers.conversation.agent_server = { "type": "responses_api_agents", "name": "agent", } @@ -1245,7 +1248,7 @@ def test_multi_agent_route_rejects_a_reference_to_a_non_agent(self, multi_agent_ from nemo_gym.config_types import ConfigError from nemo_gym.global_config import resolve_dataset_agent - multi_agent_config.environment.environment_servers.single_agent_turn.user_agent.name = "resources" + multi_agent_config.environment.environment_servers.conversation.user_agent.name = "resources" with pytest.raises(ConfigError, match="references 'resources', which is not an agent"): resolve_dataset_agent(multi_agent_config, "resources", taskset="swe:test") diff --git a/tests/unit_tests/test_task_data.py b/tests/unit_tests/test_task_data.py index 3cb0ddd01e6..c15b9259ed2 100644 --- a/tests/unit_tests/test_task_data.py +++ b/tests/unit_tests/test_task_data.py @@ -18,6 +18,7 @@ ) from nemo_gym.task_data import ( RESERVED_ROW_KEYS, + SINGLE_AGENT_TASK_DATA, TaskDataSchemaError, TaskDataValidator, find_server_dir, @@ -25,13 +26,14 @@ normalize_task_fields, validate_jsonl_rows, ) -from nemo_gym.train_data_utils import TrainDataProcessor +from nemo_gym.train_data_utils import TrainDataProcessor, TrainDataProcessorConfig REPO_ROOT = Path(__file__).resolve().parents[2] SCHEMA_FILES = sorted( list((REPO_ROOT / "resources_servers").glob("*/task_data.py")) + list((REPO_ROOT / "responses_api_agents").glob("*/task_data.py")) + + list((REPO_ROOT / "environment_servers").glob("*/task_data.py")) ) @@ -46,7 +48,8 @@ def test_splices_verifier_metadata_contents(self): fields, conflicts = normalize_task_fields( {"responses_create_params": {}, "verifier_metadata": {"label": "safe", "type": "homonyms"}} ) - assert fields == {"label": "safe", "type": "homonyms"} + # responses_create_params is a single-agent task field, not a framework key. + assert fields == {"responses_create_params": {}, "label": "safe", "type": "homonyms"} assert conflicts == [] def test_equal_duplicate_is_harmless(self): @@ -67,7 +70,7 @@ def test_splices_task_data_contents(self): fields, conflicts = normalize_task_fields( {"responses_create_params": {}, "task_data": {"expected_city": "Tokyo"}} ) - assert fields == {"expected_city": "Tokyo"} + assert fields == {"responses_create_params": {}, "expected_city": "Tokyo"} assert conflicts == [] def test_empty_task_data_container_leaves_no_residue(self): @@ -142,13 +145,18 @@ class _Schema(BaseModel): class TestTaskDataValidator: - def _validator(self): + def _validator(self, environment_adapter=None): from pydantic import TypeAdapter - return TaskDataValidator(server_name="s", adapter=TypeAdapter(_Schema), dataset_fpath="d.jsonl") + return TaskDataValidator( + server_name="s", + adapter=TypeAdapter(_Schema), + dataset_fpath="d.jsonl", + environment_adapter=environment_adapter, + ) def test_clean_rows(self): - v = self._validator() + v = self._validator(SINGLE_AGENT_TASK_DATA) v.validate_row(0, {"responses_create_params": {}, "question": "q", "expected_answer": "a"}) assert v.report.clean and v.report.rows == 1 and v.report.error_rows == 0 @@ -182,7 +190,7 @@ def test_conflicting_duplicate_marks_report_dirty(self): assert v.report.conflicting_keys == {"question": 1} def test_materialized_task_validates_its_task_data(self): - v = self._validator() + v = self._validator(SINGLE_AGENT_TASK_DATA) task = {"taskset": "t", "task_id": "0"} v.validate_row( 0, @@ -210,10 +218,23 @@ def test_generic_materialized_task_validates_flat_fields(self): def test_materialized_verifier_metadata_preserves_schema_normalization(self, nested): fields = {"verifier_metadata": {"question": "q", "expected_answer": "a"}} task_input = {"responses_create_params": {}, **({"task_data": fields} if nested else fields)} - v = self._validator() + v = self._validator(SINGLE_AGENT_TASK_DATA) v.validate_row(0, {"task_id": {"taskset": "t", "task_id": "0"}, "task_input": task_input}) assert v.report.clean + def test_materialized_row_is_not_required_to_be_a_single_agent_run_request(self): + # A row that already names its taskset is validated by its environment server at dispatch, so a + # dataset routed by agent does not apply the single-agent schema to it. + v = self._validator(SINGLE_AGENT_TASK_DATA) + v.validate_row( + 0, {"task_id": {"taskset": "t", "task_id": "0"}, "task_input": {"question": "q", "expected_answer": "a"}} + ) + assert v.report.clean, v.report.summary() + # A plain row in the same dataset is a run request and still needs one. + v.validate_row(1, {"question": "q", "expected_answer": "a"}) + assert v.report.error_rows == 1 + assert "responses_create_params" in v.report.errors[0] + def test_mixed_materialized_task_reports_conflicting_fields(self): v = self._validator() v.validate_row( @@ -286,7 +307,7 @@ def test_dangling_rs_reference_is_skipped(self): class TestShippedSchemas: - """CI enforcement for every committed task_data.py (resources servers and agents).""" + """CI enforcement for every committed task_data.py (resources servers, agents, and environment servers).""" ALLOWED_IMPORT_PREFIXES = ("pydantic", "nemo_gym.task_data") @@ -354,7 +375,9 @@ def test_committed_example_data_validates_clean(self, server): assert server_dir is not None adapter = load_task_data_schema(server_dir) example = server_dir / "data" / "example.jsonl" - report = validate_jsonl_rows(server, adapter, str(example), example.read_text().splitlines()) + report = validate_jsonl_rows( + server, adapter, str(example), example.read_text().splitlines(), environment_adapter=SINGLE_AGENT_TASK_DATA + ) assert report.rows > 0 assert report.clean, report.summary() @@ -401,7 +424,7 @@ def _agent(self, impl, rs_ref): def _dataset(self): from types import SimpleNamespace - return SimpleNamespace(jsonl_fpath="data/example.jsonl") + return SimpleNamespace(jsonl_fpath="data/example.jsonl", taskset=None) def test_agent_without_rs_reference_uses_its_own_schema(self, tmp_path, monkeypatch): agent_dir = tmp_path / "responses_api_agents" / "tau2" @@ -418,13 +441,25 @@ def test_agent_without_rs_reference_uses_its_own_schema(self, tmp_path, monkeypa assert validator is not None assert validator.report.server_name == "tau2" - def test_agent_without_rs_reference_and_no_schema_is_skipped(self, tmp_path, monkeypatch): + @staticmethod + def _assert_validates_only_single_agent_fields(validator): + # Rows routed by agent are single-agent run requests, so responses_create_params is still checked. + # With no task owner schema, other fields are neither validated nor reported as unknown. + validator.validate_row(0, {"responses_create_params": {"input": "q"}, "anything": 1}) + assert validator.report.clean + validator.validate_row(1, {"anything": 1}) + assert validator.report.error_rows == 1 + assert "responses_create_params" in validator.report.errors[0] + + def test_agent_without_rs_reference_and_no_schema_validates_single_agent_fields(self, tmp_path, monkeypatch): (tmp_path / "responses_api_agents" / "tau2").mkdir(parents=True) monkeypatch.chdir(tmp_path) agent = self._agent("tau2", None) - assert TrainDataProcessor._task_data_validator_for(agent, self._dataset(), [agent]) is None + self._assert_validates_only_single_agent_fields( + TrainDataProcessor._task_data_validator_for(agent, self._dataset(), [agent]) + ) - def test_agent_with_dangling_rs_reference_is_skipped(self, tmp_path, monkeypatch): + def test_agent_with_dangling_rs_reference_validates_single_agent_fields(self, tmp_path, monkeypatch): agent_dir = tmp_path / "responses_api_agents" / "tau2" agent_dir.mkdir(parents=True) (agent_dir / "task_data.py").write_text( @@ -432,7 +467,10 @@ def test_agent_with_dangling_rs_reference_is_skipped(self, tmp_path, monkeypatch ) monkeypatch.chdir(tmp_path) agent = self._agent("tau2", {"type": "resources_servers", "name": "missing_rs"}) - assert TrainDataProcessor._task_data_validator_for(agent, self._dataset(), [agent]) is None + # The agent's own schema is not used: its schema home is the missing resources server. + self._assert_validates_only_single_agent_fields( + TrainDataProcessor._task_data_validator_for(agent, self._dataset(), [agent]) + ) def test_agent_with_valid_rs_reference_uses_that_servers_schema(self, tmp_path, monkeypatch): # The benchmark shape: the dataset is declared on the agent, the agent references a @@ -532,6 +570,103 @@ def test_every_resources_server_ships_a_schema(self): "describe its dataset rows (see nemo_gym/task_data.py for the protocol)." ) + def test_every_environment_server_ships_a_schema(self): + missing = sorted( + d.name + for d in (REPO_ROOT / "environment_servers").iterdir() + if d.is_dir() and (d / "app.py").exists() and not (d / "task_data.py").exists() + ) + assert not missing, ( + f"environment servers without a task_data.py schema: {missing}. Every environment server must " + "declare the task fields it reads itself (see nemo_gym/task_data.py for the protocol)." + ) + + def test_single_agent_task_data_matches_the_runtime_input(self): + from nemo_gym.single_agent_turn_types import SingleAgentTurnTaskInput + from nemo_gym.task_data import SingleAgentTaskData + + # The runtime model keeps every other task field in its task_data container. + assert set(SingleAgentTurnTaskInput.model_fields) == set(SingleAgentTaskData.model_fields) | {"task_data"} + + +class TestEnvironmentServerSchemas: + """Collation validates each dataset against the environment server it routes to.""" + + @pytest.fixture + def servers(self, tmp_path, monkeypatch): + (tmp_path / "resources_servers" / "qa").mkdir(parents=True) + (tmp_path / "resources_servers" / "qa" / "task_data.py").write_text( + "from pydantic import BaseModel, ConfigDict\n" + "class TaskData(BaseModel):\n" + " model_config = ConfigDict(extra='forbid')\n" + " expected_answer: str\n" + ) + (tmp_path / "environment_servers" / "dialogue").mkdir(parents=True) + (tmp_path / "environment_servers" / "dialogue" / "task_data.py").write_text( + "from pydantic import BaseModel, ConfigDict\n" + "class TaskData(BaseModel):\n" + " model_config = ConfigDict(extra='allow')\n" + " persona: str\n" + ) + monkeypatch.chdir(tmp_path) + return tmp_path + + def _collate(self, tmp_path, rows, *, environment=None): + source = tmp_path / "source.jsonl" + source.write_text("".join(json.dumps(row) + "\n" for row in rows)) + dataset = {"name": "d", "type": "example", "jsonl_fpath": str(source)} + config = { + "resources": {"resources_servers": {"qa": {"entrypoint": "app.py", "domain": "other"}}}, + "agent": { + "responses_api_agents": { + "simple_agent": { + "entrypoint": "app.py", + "resources_server": {"type": "resources_servers", "name": "resources"}, + } + } + }, + } + if environment is None: + config["resources"]["resources_servers"]["qa"]["datasets"] = [dataset] + else: + config["resources"]["resources_servers"]["qa"]["datasets"] = [dict(dataset, taskset="t")] + config["environment_server_routes"] = {"t": "environment"} + config["environment"] = { + "environment_servers": { + environment: { + "entrypoint": "app.py", + "agent_server": {"type": "responses_api_agents", "name": "agent"}, + "resources_server": {"type": "resources_servers", "name": "resources"}, + } + } + } + global_config = OmegaConf.create(config | {"mode": "example_validation", "output_dirpath": str(tmp_path)}) + processor = TrainDataProcessor() + configs = processor.load_and_validate_server_instance_configs( + TrainDataProcessorConfig.model_validate(global_config), global_config + ) + return processor._collate_samples_single_type("example", configs, task_data_validation="error") + + @pytest.mark.parametrize("environment", [None, "single_agent_turn"]) + def test_single_agent_rows_need_responses_create_params(self, servers, environment): + # The resources schema forbids extra fields but never sees responses_create_params. + self._collate( + servers, [{"responses_create_params": {"input": "q"}, "expected_answer": "a"}], environment=environment + ) + with pytest.raises(ValueError, match=r"SingleAgentTaskData; responses_create_params;\s+Field required"): + self._collate(servers, [{"expected_answer": "a"}], environment=environment) + + def test_environment_server_with_its_own_fields_needs_no_responses_create_params(self, servers): + self._collate(servers, [{"persona": "p", "expected_answer": "a"}], environment="dialogue") + # Both schemas report their errors for one row. + with pytest.raises(ValueError, match=r"persona;\s+Field required.*expected_answer;\s+Field required"): + self._collate(servers, [{"question": "q"}], environment="dialogue") + + def test_environment_server_without_a_schema_validates_the_resources_fields(self, servers): + self._collate(servers, [{"expected_answer": "a"}], environment="unschematized") + with pytest.raises(ValueError, match="expected_answer"): + self._collate(servers, [{"answer": "a"}], environment="unschematized") + class TestRepoDataMatchesSchemas: """The drift gate: every committed dataset row must validate against its server's schema. @@ -618,7 +753,16 @@ def test_all_committed_rows_validate(self, monkeypatch): ) owner = ("resources_servers", rs) if rs else None if owner and str(d["jsonl_fpath"]) in tracked: - server_files.setdefault(owner, set()).add(str(d["jsonl_fpath"])) + # A dataset that declares a taskset routes to an environment server and + # uses its schema; one that does not routes by agent and is a single-agent + # run request. Mirrors collation's choice. + environment = None + if d.get("taskset"): + route = (full.get("environment_server_routes") or {}).get(d["taskset"]) + block = full.get(route) if route else None + servers = block.get("environment_servers") if isinstance(block, dict) else None + environment = next(iter(servers)) if isinstance(servers, dict) else "" + server_files.setdefault(owner, {})[str(d["jsonl_fpath"])] = environment assert len(server_files) > 50, "mapping looks broken: too few servers with committed data" agent_owned = [o for o in server_files if o[0] == "responses_api_agents"] @@ -631,8 +775,14 @@ def test_all_committed_rows_validate(self, monkeypatch): if adapter is None: failures.append(f"{server}: no loadable schema") continue - for f in sorted(files): - report = validate_jsonl_rows(server, adapter, f, open(f).read().splitlines()) + for f, environment in sorted(files.items()): + environment_adapter = SINGLE_AGENT_TASK_DATA + if environment is not None: + environment_dir = find_server_dir(environment, base_folder="environment_servers") + environment_adapter = load_task_data_schema(environment_dir) if environment_dir else None + report = validate_jsonl_rows( + server, adapter, f, open(f).read().splitlines(), environment_adapter=environment_adapter + ) rows_checked += report.rows if not report.clean: failures.append(report.summary())