Skip to content

fix: cancel pending rollout tasks after a failure - #2350

Open
yuyz-cyber wants to merge 1 commit into
THUDM:mainfrom
yuyz-cyber:fix/cancel-failed-rollout-siblings
Open

fix: cancel pending rollout tasks after a failure#2350
yuyz-cyber wants to merge 1 commit into
THUDM:mainfrom
yuyz-cyber:fix/cancel-failed-rollout-siblings

Conversation

@yuyz-cyber

Copy link
Copy Markdown

Problem

Training and evaluation launch rollout tasks concurrently at three levels: samples within a rollout group, rollout groups within a training batch, and samples within an evaluation batch.

When one task raises an exception, the exception is propagated but the unfinished sibling tasks are not cancelled. They can continue generation or reward computation after the rollout has failed, wasting resources and potentially leaving unobserved task exceptions in the event loop.

Fix

At each concurrency boundary, cancel unfinished sibling tasks and wait for their cleanup before re-raising the original exception. The training path also resets its shared generation state, and both training and evaluation close their progress bars before returning the failure. Successful rollout behavior is unchanged.

Test plan

  • Added a regression test where one generation task fails while a sibling is still running.
  • Verified that the sibling is cancelled and the original exception is preserved.
  • tests/plugin_contracts/test_plugin_generate_contracts.py: 5 passed.
  • tests/test_streaming_rollout.py: 26 passed.
  • Ruff, Black, and isort checks passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant