[ananthsub]

What does this PR do?

Checks during collation that the task IDs of a declared taskset are stable and unique.

This is the first PR in a stack that makes task data schemas general across environment servers.

A TaskId names one task in a run, and results and reward profiles group by it. #3818 takes the ID from task_id, problem_id, or instance_id, and otherwise falls back to the row's position in the collation. Positional IDs change when datasets are added or reordered, and nothing stopped two rows from sharing an ID.

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 S4096 this
    class S4097,S4098,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

• Collation warns once for each dataset whose taskset rows fall back to positional IDs, and names the fix: a stable task_id from the preparation script.
• Collation reports task IDs that name two source rows in one taskset, including one file declared twice under the same taskset. Repeats of one row (num_repeats) are not duplicates. The report follows task_data_validation: it fails collation in error mode, the default for +mode=example_validation, and is printed in warn mode.
• TASK_ID_FIELDS and source_task_id in task_materialization.py name the ID rule that materialize_task and collation share.
• The source ID field stays in task_input. Many resources servers declare task_id as a task field, and some verifiers read it, so removing it would change what they receive.
• resources-server-apis.mdx documents both checks.

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

Validation

• tests/unit_tests/test_task_materialization.py: one positional-ID warning per dataset, no warning for stable IDs, repeats that are not duplicates, duplicate IDs rejected in error mode for a repeated task_id, a task_id and problem_id with the same value, and one file declared twice, and the report printed in warn mode and skipped in off mode.
• pytest tests/unit_tests/test_benchmarks.py tests/unit_tests/test_global_config.py tests/unit_tests/test_rollout_collection.py tests/unit_tests/test_task_data.py tests/unit_tests/test_task_materialization.py tests/unit_tests/test_train_data_utils.py tests/unit_tests/test_environment_validation.py tests/unit_tests/test_cli_eval.py: all pass on this branch.
• gym dataset collate +mode=example_validation on resources_servers/nemo_user_sim/configs/nemo_user_sim.yaml, the only committed config that declares a taskset: no positional-ID warnings and no duplicate task IDs.
• pre-commit run on the changed files: all pass.

Rollouts

N/A. The change adds collation warnings and a duplicate-ID check, and changes no rows that collation writes. UserSim's example dataset, the only committed dataset that declares a taskset, collates without either.

Compatibility and benchmark impact

• Datasets without taskset: are unaffected.
• A taskset with duplicate task IDs now fails example validation and prints a report during training preparation.
• 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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[iguehring-nvidia]

nit: In line 1089, we are checking if task validation is on before printing the warning with if task_data_validation != "off". Here, we are not. Just pointing this out for consistency. I would argue that we should just always warn and don't even have the option to turn it off? Makes the code simpler and I see little value in that configurability.

[iguehring-nvidia]

currently we would accept an empty string. Is that on purpose?

[iguehring-nvidia]

nit: comment necessary / helpful?

The merge-base changed after approval.