[ananthsub]

What does this PR do?

Validates each dataset row against the schema of the environment server it routes to, composed with the bound resources server's schema.

responses_create_params was a framework-reserved row key, so collation never validated it, and each dataset was checked only against its resources server's TaskData. An environment server that runs several agents had no way to declare the fields it reads, and a single-agent row without a request passed validation.

Where this fits

Datasets reach an environment server in one of two ways today: run request rows route by their agent, and datasets that declare taskset: are materialized during collation and route by environment_server_routes. The goal is one path, where every dataset declares a taskset and routes to an environment server:

flowchart LR
    P["Preparation script<br/>task rows with a stable task_id"] --> C["Collation<br/>validates, computes metrics,<br/>materializes"]
    Y["Dataset YAML<br/>file + taskset"] --> C
    S["task_data.py schemas<br/>environment server + resources server"] -.-> C
    C --> R["Rollout collection<br/>routes by taskset"]
    T["environment_server_routes"] --> R
    R --> E["Environment server<br/>single-agent, multi-agent,<br/>or legacy_agent relay"]

Loading

The UserSim stack (#3800 and #3802) has merged. UserSim's datasets already declare taskset: and commit plain task rows, so they use the one path. Each step is additive, and datasets without taskset: keep working until the last one:

flowchart TB
    S4053["PR 4053 (merged)<br/>tasksets route to servers<br/>with several agents"]
    S4060["PR 4060 (merged)<br/>metrics for<br/>materialized rows"]
    S3800["PR 3800 (merged)<br/>UserSim runtime"]
    S3802["PR 3802 (merged)<br/>UserSim validation tasks"]
    S4096["PR 4096<br/>stable, unique<br/>task IDs"]
    S4097["PR 4097<br/>metrics for<br/>every row"]
    S4098["PR 4098<br/>schema per<br/>environment server"]
    S4099["PR 4099<br/>legacy_agent relays<br/>serve tasksets"]
    SREF["Remove the bridge<br/>for already-materialized<br/>rows"]
    SNEXT["Later<br/>configs declare taskset:,<br/>then agent routing retires"]
    S4053 --> S4060 --> S3800 --> S3802 --> S4096 --> S4097 --> S4098 --> S4099 --> SREF --> SNEXT
    classDef done fill:#F1EFE8,stroke:#5F5E5A,color:#444441
    classDef prereq fill:#FAEEDA,stroke:#854F0B,color:#633806
    classDef this fill:#EEEDFE,stroke:#534AB7,color:#26215C,stroke-width:2px
    classDef next fill:#E1F5EE,stroke:#0F6E56,color:#085041
    class S4053,S4060,S3800,S3802 done
    class S4098 this
    class S4096,S4097,S4099,SREF,SNEXT next

Loading

Gray: merged. Purple: this PR. Green: other steps.

Base branch. This stack is based on main. Review only this PR's own diff.

Changes

• responses_create_params is a task field of single-agent environment servers, declared once by SingleAgentTaskData in nemo_gym/task_data.py. It is no longer in RESERVED_ROW_KEYS. Its input is optional, because a dataset's prompt_config fills it during collation.
• Environment servers ship task_data.py under environment_servers/<implementation>/. single_agent_turn, single_agent_turn_legacy, and legacy_agent export SingleAgentTaskData.
• TaskDataValidator takes an optional environment_adapter. Each schema validates the fields it declares plus the fields neither declares. A resources schema with extra="forbid" never sees the environment server's fields, so none of the existing schemas needs an edit. Errors from both schemas are reported for one row, and undeclared fields are reported against the resources server's schema.
• Collation picks the schemas by route:
  • A dataset with taskset: uses the routed environment server's schema and the schema of the resources server it binds. TasksetSchemas records them while configs load, the only step that has the run's routes.
  • A dataset without taskset: routes by agent, so it is a single-agent run request. It uses SingleAgentTaskData and the schema of the resources server or self-contained agent that owns its data, as before.
  • A taskset whose route was not resolved is not assumed to be single-agent.
• taskset_environment_server in global_config.py resolves a taskset's route, shared by resolve_dataset_agent and collation.
• Rows that a preparation script already materialized name their own taskset, and their environment server validates them at dispatch. The environment schema does not validate them during collation, so they are not mistaken for single-agent run requests. Its fields are still kept out of the resources schema's undeclared fields.
• environment_servers/nemo_user_sim/task_data.py shares UserSim's resources TaskData, since its environment server reads the same task fields.
• Schema tests now cover environment_servers/, require every environment server to ship a schema, and check that SingleAgentTaskData matches the runtime SingleAgentTurnTaskInput.
• resources-server-apis.mdx describes the two schemas.

Out of scope, for follow-up PRs: environment manifest validation, the Hugging Face upload check, and the environment scaffold still assume single-agent rows. Manifests and the scaffold only describe agent-routed environments today.

No issue is linked. This follows the dataset routing design for environment servers.

Validation

• tests/unit_tests/test_task_data.py::TestEnvironmentServerSchemas: through load_and_validate_server_instance_configs and collation, single-agent rows need responses_create_params whether they route by agent or by taskset, a resources schema with extra="forbid" accepts them, an environment server with its own schema needs no responses_create_params and both schemas report errors, and an environment server without a schema still validates the resources fields.
• tests/unit_tests/test_task_data.py: a materialized row in a dataset routed by agent needs no responses_create_params, while a plain row in the same dataset still does. This test fails without the change.
• tests/unit_tests/test_task_data.py: the committed-data sweep uses the routed environment server's schema for datasets that declare taskset:, as collation does, and the single-agent schema otherwise.
• gym dataset collate +mode=example_validation on resources_servers/nemo_user_sim/configs/nemo_user_sim.yaml: UserSim's committed task rows materialize with their own task IDs, schema validation is clean, and the committed metrics sidecar is unchanged.
• resources_servers/nemo_user_sim/tests and environment_servers/nemo_user_sim/tests fail only where they need the usersim package, identically on main; gym env test --resources-server nemo_user_sim installs it.
• pytest tests/unit_tests: the tests that fail here fail identically on main; they need the sandbox, telemetry, or Slurm extras.
• pytest environment_servers resources_servers/terminal_bench_4/tests/test_benchmark_preparation.py: all pass.
• pre-commit run on the changed files: all pass.

Rollouts

N/A. The change affects offline validation during collation and writes no different rows. The collation check described in the last PR of this stack ran on real repository configs with this change.

Compatibility and benchmark impact

• Every committed dataset validates clean.
• Datasets routed by agent whose owner ships no task_data.py are now checked for responses_create_params, which they need to run.
• Callers of TaskDataValidator or validate_jsonl_rows that validate single-agent rows pass environment_adapter=SINGLE_AGENT_TASK_DATA; without it, responses_create_params is reported as an undeclared field.
• No benchmark impact.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.