fix: never fast-forward an infinite env's generator on resume - #3253
Draft
parkerpettit wants to merge 1 commit into
Draft
fix: never fast-forward an infinite env's generator on resume#3253parkerpettit wants to merge 1 commit into
parkerpettit wants to merge 1 commit into
Conversation
TrainSource.load_state_dict restored an infinite env's data position by calling next() on the generator cursor-many times. That is only correct for a pure deterministic generator. An infinite taskset backed by live external state (a task server, a curriculum stream) does a real draw on every pull — with a cursor in the thousands, resume leaks thousands of claimed-but-never-run assignments and grinds for hours while the trainer sits in the startup weight broadcast until its timeout kills the job. Restore the cursor as a counter only. Finite envs keep the exact epoch-shuffle position replay.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
TrainSource.load_state_dictrestores an infinite env's data position by callingnext()on the generator cursor-many times. That replay is only correct for a pure deterministic generator — an assumption nowhere declared. An infinite taskset backed by live external state (a task server, a curriculum stream) does a real draw on every pull.Hit on live runs: resuming a run whose infinite env had drawn thousands of episodes made the orchestrator grind at 60% CPU for what projected to hours, leaking hundreds of claimed-but-never-run generation assignments into the external task server, while the trainer sat blocked in the startup weight broadcast until its 3600s timeout killed the job. No log line is emitted during the replay, so the run looks silently wedged (field workaround:
--orchestrator.ckpt.skip-progress, which also drops the RNG/cursor state).Fix
Restore the infinite env's cursor as a counter only — never by consuming the generator. A fresh iterator is not the old one; for stateful streams the position is owned by the stream itself. Finite envs keep the exact epoch-shuffle position replay, which is pure.
Trade-off: a taskset whose generator IS a pure seeded function no longer replays to its old position and will restart its sequence after resume (possible repeats). If that matters, the cleaner contract is an explicit replayability attribute on the taskset — happy to do that instead if preferred.
Tests
tests/unit/orchestrator/test_train_source.py: resume of an infinite env consumes zero generator pulls and continues from the restored counter; finite-env resume still replays the exact epoch-shuffle position. Verified on ar-cluster (2 passed).🤖 Prepared with Claude Code (draft per Parker)