Skip to content

[fix][train] Shut down DataLoader workers on exit - #1998

Open
bvolpato wants to merge 2 commits into
NovaSky-AI:mainfrom
bvolpato:bvolpato/shutdown-dataloader-workers
Open

[fix][train] Shut down DataLoader workers on exit#1998
bvolpato wants to merge 2 commits into
NovaSky-AI:mainfrom
bvolpato:bvolpato/shutdown-dataloader-workers

Conversation

@bvolpato

@bvolpato bvolpato commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Summary

StatefulDataLoader owns worker processes through its iterator. RL and SFT entrypoints only cleaned up dataloaders on successful runs, so training exceptions could leave train or evaluation workers alive.

Changes

  • Add a guarded, idempotent shutdown helper for torchdata 0.11.
  • Shut down RL train and evaluation dataloaders on success and failure.
  • Move SFT shutdown into the entrypoint's finally block and close every train and evaluation dataloader.
  • Log cleanup failures without masking the original training exception.

Normal Python unwinding is covered. Abrupt termination such as SIGKILL remains outside scope.

Testing

  • Focused lifecycle tests (123 passed)
  • Full prior SkyRL Train CPU suite (1467 passed, 5 skipped, 5 deselected)
  • Tests start real train and evaluation workers, call shutdown twice, and verify every worker exits
  • pre-commit run --all-files
  • git diff --check upstream/main

@bvolpato
bvolpato marked this pull request as ready for review August 6, 2026 15:18

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request ensures that trainer resources, specifically DataLoader worker processes, are properly shut down when training finishes or encounters an error. It introduces a shutdown_dataloader utility function to stop worker processes owned by a StatefulDataLoader, and integrates this into the shutdown methods of both SFTTrainer and Trainer. Additionally, finally blocks are added to the entry points to guarantee that shutdown is always executed. Feedback suggests improving the readability of the dataloader shutdown logic in SFTTrainer by separating the train and evaluation dataloader shutdown blocks instead of using an intermediate list.

Comment thread skyrl/train/sft_trainer.py Outdated
@bvolpato
bvolpato force-pushed the bvolpato/shutdown-dataloader-workers branch 2 times, most recently from d74713a to 1a025cd Compare August 7, 2026 00:21
@bvolpato
bvolpato force-pushed the bvolpato/shutdown-dataloader-workers branch from 1a025cd to 366d401 Compare August 13, 2026 01:21
@bvolpato
bvolpato force-pushed the bvolpato/shutdown-dataloader-workers branch from 366d401 to b1a262c Compare August 14, 2026 04:41
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