Repository navigation
fix(orchestrator): drop restored work that points at deleted by-handle payload files - #4044
faresobeid wants to merge 34 commits into
Conversation
…th the dataset cursor
…t/payload-resume-filter # Conflicts: # src/prime_rl/orchestrator/train_sink.py
…e payload files Payload files are deleted once the trainer has read them, which can be after the checkpoint was taken. On resume, drop queued traces and open-group episodes whose files are gone instead of crashing the trainer on a missing file.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
…nk has their last result
…nk has their last result
…t/payload-resume-filter
… the main loop starts
…-replay (#3970) into feat/payload-resume-filter
# Conflicts: # src/prime_rl/orchestrator/orchestrator.py
# Conflicts: # tests/unit/orchestrator/test_train_sink.py
# Conflicts: # src/prime_rl/orchestrator/orchestrator.py
…out, outside the curriculum
…t/payload-resume-filter
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 3 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit bcc48f5. Configure here.
| rollouts = self.group_size_for(env_name) - len(fresh) | ||
| requests.append( | ||
| TaskRequest(env_name=env_name, task=task, step=self.progress.step, rollouts=rollouts, group_id=group_id) | ||
| ) |
There was a problem hiding this comment.
Zero-rollout resume requests still dispatch
Medium Severity
load_state_dict always enqueues a TaskRequest, including when rollouts is 0 or negative, and never finalizes a group that already has every member. try_schedule always dispatches one episode for a newly created group, so that request still runs. The extra episode lands in an already-full pending_groups list and process_group then scores an oversized group.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit bcc48f5. Configure here.
| float(sum(len(episode.traces) for episode in group) - sum(len(episode.traces) for episode in fresh)), | ||
| ) | ||
| if fresh: | ||
| group = fresh |
There was a problem hiding this comment.
Stale filter undercounts group budget
Low Severity
n_owed is computed after group is replaced with fresh, so a resumed group that lost some members to max_off_policy_steps reports a smaller episode budget than arrived. Zero-output tracking and _record_group then under-count those members.
Reviewed by Cursor Bugbot for commit bcc48f5. Configure here.
| open_groups = self.dispatcher.open_train_groups() | ||
| open_groups |= {request.group_id: (request.env_name, request.task) for request in self.train_source.resumed} | ||
| return build_state(self.progress, self.train_source, self.train_sink.state_dict(open_groups)) | ||
|
|
There was a problem hiding this comment.
Sweep misses restored queued traces
Medium Severity
oldest_live_payload_version is documented to follow queued work, but it only sees episode_by_trace and pending_groups. Restored queued traces have no episodes, so the orphan payload sweep can treat their v<k> directories as dead and delete files the next trainer read still needs.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit bcc48f5. Configure here.
Stacked on #3902 (by-handle payloads) and includes #3970 (checkpoint queued samples and open train groups); review only the commits not in those PRs.
#3970 saves queued samples and the finished episodes of open groups in the orchestrator checkpoint. With by-handle payloads those samples point into payload files, and #3902 deletes a file once the trainer has read its batch. If that happens after the checkpoint was taken, a resume points at deleted files and the trainer crashes (
FileNotFoundErrorin_pread_file).What
TrainSink.load_state_dictdrops restored queued traces and open-group episodes that point at a missing payload file, and logs how many queued traces were dropped. Dropped open-group episodes are regenerated like any other missing group member. Cost: a resumed by-handle run loses the queued work whose files are gone (typically 1-3 batches). Inline runs are unaffected.Validation
test_resume_drops_queued_traces_with_missing_payload_files.pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs*.py tests/unit/inference tests/unit/transports tests/unit/train/rl: 479 passed; ruff clean;uv lock --checkclean. Not run on GPU.Before merge
Merge #3970 and #3902 first; then this PR's diff is just the resume filter.
Note
Medium Risk
Changes checkpoint format and resume paths for the train pipeline; incorrect restore could lose queued work or crash the trainer, but missing payloads and staleness are explicitly filtered and tested.
Overview
Orchestrator checkpoints now include queued training samples and in-flight rollout groups, saved atomically with progress and curriculum state via
TrainSink.state_dict/ asyncCheckpointManager.save. On resume, restored queued traces ship first, open groups continue with the same task/group id, andTrainSource.resumedschedules missing rollouts before new tasks.Resume safety for by-handle payloads:
TrainSink.load_state_dictdrops queued traces and saved open-group episodes whose payload files no longer exist (trainer may have read and deleted them after the checkpoint), logs how many queued traces were skipped, and regenerates dropped group members like any other missing rollout. Stale work is still filtered bymax_off_policy_stepsat restore and when mixed-version resumed groups are scored; batch assembly prefers oldest policy version first.Supporting wiring: dispatcher
emitted_train_groups+forget_consumed_group,snapshot_stateon ship/checkpoint, and training.md resume semantics. New unit tests cover missing payload files and off-policy scoring on resumed groups.Reviewed by Cursor Bugbot for commit 1d222a7. Bugbot is set up for automated code reviews on this repo. Configure here.