[iguehring-nvidia]

What does this PR do?

Adds gym eval plan, which prepares one benchmark and writes a plan file: every task once, sorted by task id, with a SHA-256 digest of the exact input a run receives.

Why: Gym has no single command that answers "which tasks does this benchmark have, and with exactly what inputs?" in a stable form. The prepared JSONL depends on preparation details, its row order is not guaranteed, and many rows have no id, so collation falls back to the row position. gym eval plan produces that answer as one file: stable task ids, a digest of each task's exact run input, and a versioned, documented format. The plan stores ids and digests, not the inputs themselves, so it stays small; a run re-prepares the data, and the digests show whether it still matches the plan.

What the command does:

• Selects the benchmark with the existing --benchmark NAME[/FLAVOR] lookup and runs its prepare.py (shared code with gym eval prepare).
• Takes each task's id from task_id, problem_id or instance_id, the same precedence as materialize_task (feat(data): collate declared tasksets #3818); the field list is now defined once, as TASK_ID_FIELDS in nemo_gym/task_materialization.py, and used by materialize_task, rollout collection and the plan. Unlike collation, there is no fallback to the row position: a position-based id names a different task once rows are reordered, so benchmarks without an id field are refused with an actionable error.
• Each task is {task_id: {taskset, task_id}, digest}. The digest is sha256(canonical JSON of task_input), where task_input is the MaterializedTask input with routing and bookkeeping fields (agent_ref, task_source, _ng_*) excluded. The taskset is the dataset's declared taskset, or the benchmark name when it declares none.
• Records the Gym version and commit and the selected config paths (relative to the checkout).
• Output is deterministic (sorted tasks, canonical JSON): same Gym commit and same prepared data give a byte-identical file.
• --task-ids plans a subset; --schema prints the JSON Schema; the format is labelled nemo-gym.eval-plan.v1 and is meant to change only additively within that version.
• Benchmarks with a prompt template (prompt_config) are refused for now; applying the template is a follow-up.

prepare_benchmark is split into _datasets_to_prepare and _prepare_datasets so both commands share them; its behaviour is unchanged, except that the "no preparable dataset" error no longer names gym eval prepare (it suggests --config <config path> or --benchmark <name>). The plan considers only the benchmark datasets among the datasets to prepare.

No issue filed: this is a new, additive command. Related: the Harbor task work (#3673 and follow-ups) produces the same MaterializedTask rows and pins, so plans for Harbor datasets should line up with it. Feedback on the file format is very welcome while this is a draft.

Validation

• pytest tests/unit_tests/test_eval_plan.py: 20 passed (id precedence, ids and digests only (no inputs) in the file, declared taskset used, refusal of missing and duplicate ids, refusal of already-materialized rows as a config error, routing fields excluded from input and digest, subset selection, deterministic output, schema round trip, CLI: file written, subset, prompt-template refusal, several-benchmarks refusal, --out required, --schema).
• pytest tests/unit_tests/test_cli_main.py tests/unit_tests/test_cli_compat.py tests/unit_tests/test_rollout_collection.py tests/unit_tests/test_task_materialization.py tests/unit_tests/test_benchmarks.py: all pass after rebasing onto main, except test_benchmarks.py::test_every_repo_benchmark_appears_in_listing, the DNS gaierror on my machine noted below.
• pytest tests/unit_tests/test_benchmarks.py tests/unit_tests/test_cli_main.py tests/unit_tests/test_cli_compat.py: 335 passed, 1 failed. The failure (test_every_repo_benchmark_appears_in_listing) is a DNS gaierror while resolving four configs on my machine and fails identically on unmodified main.
• pre-commit run --all-files: all hooks pass.
• Real runs:
  • gym eval plan --benchmark swebench/verified/opencode --out plan.json: 500 tasks keyed by instance_id; running it twice gives byte-identical files.
  • --task-ids astropy__astropy-12907,django__django-11099: plan with those 2 tasks.
  • --benchmark terminal_bench_2_1/terminus_2 (fresh preparation, on main with fix(terminal_bench_2_1): give prepared rows a stable task_id and order #4135): 89 tasks, keyed by task folder name.
  • --benchmark aime24: refused because of its prompt template.

Rollouts

N/A: the command only prepares data and writes a file; no server, agent or model code changes.

Compatibility and benchmark impact

None. New command; gym eval prepare behaviour unchanged.

Checklist

• I have read the contributing guidelines.
• The change is focused; unrelated "drive-by" edits are tracked as separate issues/PRs.
• Tests added or updated and pass locally, or N/A for docs-only / non-code changes (so CI unit/server checks pass when applicable).
• Documentation added or updated, or N/A with justification.
• Pre-commit checks pass locally (pre-commit run --all-files) (so CI lint/format checks pass).
• New source files include the required Apache-2.0 SPDX header.
• All commits have DCO sign-off (git commit -s) (so the DCO check passes).

[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.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-pr-4131-259e1873e9b7.docs.buildwithfern.com/nemo/gym

[sephmard]

Looks good overall — nice refactor, solid test coverage, left one small comment on a test gap. Approving.

[iguehring-nvidia]

Needs a review from Tim Ruffles (who needs to join Nemo org before he can be assigned) before merging.

[timruffles]

Exactly what we need, TY!

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/37815277147