[qiaochuz-nv]

What does this PR do?

python environments/instruction_following/prepare.py and python environments/reasoning_gym/prepare.py, the
README steps that download each environment's training data, now write the published file. Before this
change, both crashed inside Arrow before writing anything.

Fixes NVBug 6805516. Closes #4510.

Cause

Each Hugging Face repo publishes one JSONL file, and both scripts loaded it with datasets.load_dataset(...)
only to json.dumps every row back out. They have done this since they were added (#1403 and #1378).
load_dataset reads JSONL with Arrow's JSON reader, which rejects both files, though the data itself is valid:

• nvidia/Nemotron-RL-instruction_following: kwargs holds one entry per instruction, and 29,268 of
  46,391 rows have a null entry there. Arrow corrupts the list offsets for those rows
  ([Python][JSON] Leading null list elements produce invalid list offsets during type inference apache/arrow#51305), so the read fails with ArrowIndexError.
• nvidia/Nemotron-RL-ReasoningGym-v1: 12 metadata keys change JSON type between rows (for example,
  solution is a number in one task and an array in another), so Arrow cannot infer one column type and
  the read fails with ArrowInvalid.

Fix

Download the file with hf_hub_download and copy it to data/train.jsonl, the same shape as
environments/simple_strands_math/prepare.py and environments/calendar/prepare.py. The script copies
rather than renames, so the shared Hugging Face cache keeps its file (the point of #3128). The Wrote N rows to <path> line is kept, with N counted from the written file. The prepared file is byte-identical to
the published one, so the null entries and mixed metadata types reach the verifiers as published, which
is also how the committed data/example.jsonl rows look.

I kept the scripts unpinned (no revision=), as they were before. The other 14 environments/*/prepare.py
scripts that call load_dataset read Parquet datasets, where load_dataset is the right tool, so they
are unchanged.

Before the fix

Gym 882073970 (main), fresh HF_HOME:

python environments/instruction_following/prepare.py
python environments/reasoning_gym/prepare.py

pyarrow.lib.ArrowIndexError: array slice would exceed array length
datasets.exceptions.DatasetGenerationError: An error occurred while generating the dataset

pyarrow.lib.ArrowInvalid: JSON parse error: Column(/metadata/puzzle) changed from object to array in row 1
datasets.exceptions.DatasetGenerationError: An error occurred while generating the dataset

After the fix

Gym 72e7b4ccb (this branch), fresh HF_HOME, same commands:

Wrote 46391 rows to .../environments/instruction_following/data/train.jsonl
Wrote 15000 rows to .../environments/reasoning_gym/data/train.jsonl

The written files have the same SHA-256 as the published files (c4b03c58… and ebc81871…), and the
cached copies are unchanged.

Validation

The prepare logs above and the new unit test were rerun at 72e7b4ccb. The longer runs below used the
same change one rebase earlier (628012ace on main 494528232). The rebase onto 882073970 touched none
of these files, pyproject.toml or uv.lock.

• New tests/unit_tests/test_jsonl_environment_prepare.py (2 cases). Each fixture holds rows Arrow cannot
  load and checks that the copy is byte for byte, that the cache file is kept, the hf_hub_download
  arguments, and the exact Wrote line. Both cases pass on this branch. On main they fail earlier, at
  ImportError: pip install datasets, because the test stubs huggingface_hub. The Arrow failure itself
  is shown by the logs above.
• pytest tests/unit_tests -n 8 (macOS): 11,127 passed, 40 failed, 2 errors; on main (these two scripts
  swapped back) 11,125 passed, 40 failed, 2 errors. The failures are sandbox, telemetry and launcher tests
  that fail on macOS. The two lists differ by one test each
  (test_harbor_package_store.py::test_short_http_reply_is_a_package_store_error,
  test_docker_compose.py::test_a_working_directory_the_image_lacks_is_created_before_the_first_exec[False]).
  Both pass and fail intermittently when rerun alone on this branch.
• pre-commit run --all-files: passed.
• QA testcase reasoning_gym (6 cases): main 5/6 (download fails on the error above); this branch 6/6.
  The download case now also records the HF revision before running and checks that every prepared row
  equals the published data/train.jsonl at that revision (15,000 rows at 4c25c3b6).
• QA testcase instruction_following (H100, Qwen3-0.6B on vLLM): main 0/2 (both cases stop at the
  prepare error above); this branch 2/2. The prepared file's SHA-256 equals the independently downloaded
  file at revision bcdec57a; the 5 README example rollouts match independently computed instruction
  flags (rewards 0, 0, 0, 1, 1); 24 direct /verify controls passed. Because prepare had always crashed
  first, this run was the first to reach the case's later checks, and two of them expected the verify
  response to echo the row's id. It never has (InstructionFollowingVerifyResponse does not carry it),
  so the QA case is corrected on the QA side, not here.
• Rollout evidence (H100, Qwen3-0.6B on vLLM, gym env start --environment <env> +
  gym eval run --no-serve): rows taken from each freshly prepared data/train.jsonl, chosen to be the
  shapes Arrow rejected. instruction_following: 3 rows with a null kwargs entry, 3 rollouts, 0 failures,
  rewards 1, 1, 0. reasoning_gym: 4 rows whose metadata.solution is null, a list, a float and a string,
  4 rollouts, 0 failures, rewards 1, 0, 0, 0.01. None masked.

Compatibility

data/train.jsonl now holds the published rows unchanged. Before, nothing was written for these dataset
revisions. No change to configs or READMEs.

🤖 Generated with Claude Code

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

[qiaochuz-nv]

/ok to test 72e7b4c