[fix][sft] Make nested pretokenized shard order deterministic

[zhang-bofan]

What does this PR do?

Give recursively discovered pretokenized SFT shards a deterministic order, so resuming over the same shard tree preserves the dataset index-to-sample mapping.

Why it happens

_collect_data_files sorts filenames but leaves subdirectory traversal in filesystem order. The loader concatenates shards in that order, while StatefulSequentialSampler checkpoints only a position. Different directory enumeration can therefore repeat or omit samples after resume.

In a local reproduction, identical JSONL shard trees on /tmp and /dev/shm loaded in orders b,c,a and c,b,a. After consuming b, uninterrupted iteration returned c,a, but restoring the dataloader checkpoint against the second tree returned b,a.

Solution

Sort the visible dirnames in place before os.walk recurses. This retains the existing hidden-directory filtering and per-directory filename sorting.

Test plan and validation

The regression test varies directory enumeration at two nesting levels and checks real StatefulDataLoader save/restore with StatefulSequentialSampler, parametrized over Parquet, JSONL, and Arrow.

• Before the fix: all 3 new cases failed at the resumed-sample assertion.
• Current local validation: all 105 tests in test_checkpoint_loading.py, test_sft_pretokenized.py, and test_sft_dataloader.py passed.
• The separate real-filesystem reproducer now loads a,b,c on both filesystems and resumes with b,c on both.
• pre-commit run --all-files --config .pre-commit-config.yaml: Ruff, Black, and Gitleaks passed; the patch was staged for the Gitleaks check.

Local validation used Python 3.12.14, PyTorch 2.9.0, datasets 5.0.1, and torchdata 0.11.0 in a separate locked environment. The test command was the following, with checkout and validation-environment paths abbreviated:

PYTHONPATH="$SKYRL_CHECKOUT:$SKYRL_CHECKOUT/skyrl-gym" CUDA_VISIBLE_DEVICES='' \
  uv run --isolated --locked --project "$SKYRL_CPU_ENV" \
  pytest --noconftest tests/backends/skyrl_train/test_checkpoint_loading.py \
  tests/train/test_sft_pretokenized.py tests/train/test_sft_dataloader.py -q

--noconftest skips the session fixture that starts and stops Ray; these tests exercise the loader and sampler locally. Local runs did not exercise the repository's full dependency combination or GPU/E2E training.

CPU test fixture

The checkpoint-loading test constructs SkyRLTrainBackend with object.__new__, so it must initialize _inference_engines_initialized=False as the real constructor does. #2065 added the call that reads this flag during checkpoint loading, leaving the test setup incomplete. This PR includes the one-line fixture repair; the optimizer/scheduler assertions remain intact.

The initial PR run and the same main baseline failed the same two checkpoint cases. Locally, the checkpoint module changed from 2 failed / 1 passed to 3 passed; the combined relevant suite above passes all 105 tests.

GitHub validation for 9845a83: the training CPU job now passes all three stages: 1705 passed / 16 skipped without vLLM, 206 passed / 1 skipped with vLLM, and 72 passed / 22 skipped for Tinker.

The general CPU job also passed (25 and 208 tests), including the engine benchmark step. The remaining failing checks are Anyscale authentication before GPU tests and Vercel deployment authorization.

Compatibility

Nested shard directories gain a deterministic traversal order. This can differ from the order used by earlier versions, so the change does not preserve an old unsorted order when resuming a checkpoint created before the fix.

AI assistance: Codex prepared this change and ran the local reproducer and checks.

---

Note

Medium Risk
Changes dataset sample order for nested shard layouts versus older runs, so checkpoints from before the fix may not align; resume correctness for new runs is the intended win.

Overview
Fixes non-deterministic nested pretokenized shard ordering by sorting subdirectory names during os.walk in _collect_data_files, alongside the existing per-directory filename sort. That keeps the concatenated dataset order stable across filesystems so SFT resume with StatefulDataLoader / StatefulSequentialSampler does not repeat or skip samples when the shard tree is reloaded.

Adds test_nested_shard_order_is_stable_on_resume (Parquet, JSONL, Arrow) that flips simulated directory enumeration order and asserts checkpoint restore continues the same sample sequence. Unrelated test tweak: sets _inference_engines_initialized = False on a mocked SkyRLTrainBackend in checkpoint-loading tests.

^{Reviewed by Cursor Bugbot for commit 9845a83. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request ensures that data files are collected in a deterministic order by sorting the directory names in os.walk within _collect_data_files. It also adds a corresponding unit test, test_nested_shard_order_is_stable_on_resume, to verify that the shard order remains stable and consistent across different filesystem directory orderings when resuming. There are no review comments, so I have no feedback to provide.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with the deterministic traversal fix correctly covered by resume-focused regression tests.

_{Reviews (1) · Last reviewed commit: "fix(sft): stabilize nested pretokenized ..."}