[trainer] fix: gracefully finalize tracking and dataloader workers at fit() exits#6917
[trainer] fix: gracefully finalize tracking and dataloader workers at fit() exits#6917seokhyunan wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a graceful teardown mechanism to stop DataLoader worker processes and finalize tracking backends upon trainer exit, preventing spurious post-run SIGKILL errors. The review feedback suggests wrapping the dataloader shutdown in a try-except block to ensure tracking finalization always runs, and enhancing the dataloader shutdown utility to support standard PyTorch DataLoaders by accepting active iterators directly.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
What does this PR do?
The V1 PPO trainer finalizes two end-of-run resources only via
__del__/ garbage collection, which under Ray / interpreter teardown runs too late. This PR finalizes both explicitly, in a singletry/finallyaround the training loop, so cleanup runs on normal completion,val_only, and exceptions:Trackingfinalized wandb only in__del__, which fires after the wandb-core service connection has already closed, sowandb.finish()never delivers the exit record and history/summary/exit are never flushed. Fixes [Bug] verl does not log the final validation stats #1733.RuntimeError: DataLoader worker (pid ...) is killed by signal: Killedprinted at process exit whendata.dataloader_num_workers > 0— the liveStatefulDataLoaderworker iterator is never shut down beforefit()returns, so the workers are SIGKILLed during teardown. Fixes DataLoader worker (pid 10855) is killed by signal: Killed. #5922; supersedes Fix PPO dataloader worker teardown #6292 (closed for a CLA/authorship issue, not on technical merit — its helper is reused here) and, per that PR's review, covers allfit()exit paths.AI assistance was used to isolate the root causes and design the fix; the submitter has reviewed and tested the change.
Checklist Before Starting
DataLoader worker killed by signal→ DataLoader worker (pid 10855) is killed by signal: Killed. #5922dataloader worker teardown→ Fix PPO dataloader worker teardown #6292 (closed for CLA, not merit)wandb final validation stats→ [Bug] verl does not log the final validation stats #1733[{modules}] {type}: {description}→[trainer] fix: gracefully finalize tracking and dataloader workers at fit() exits(additive only, no[BREAKING]).Test
Not covered by CI (see Checklist Before Submitting). Validated with a minimal 2-step GRPO run (single GPU,
dataloader_num_workers=8, wandb logging) that reproduces both symptoms:run-*.wandbdatastore has 0 history / 0 summary / 0 exit records; at shutdownteardown_atexitintermittently raisesHandleAbandonedError/BrokenPipeError; and (with workers > 0) theDataLoader worker ... killed by signal: Killedtraceback is printed.training/*metrics appear on W&B;run-*.wandbhashistory ≥ 2,summary == 1,exit == 1; noteardown_atexittraceback; no DataLoader SIGKILL traceback; exit code 0. (Re-parse the datastore offline withwandb.sdk.internal.datastoreto confirm the record counts.)API and Usage Example
No CLI / config changes. One additive API:
Tracking.finish()(idempotent, exception-safe) is now public, andTracking.__del__delegates to it. The V1 trainer calls it automatically at the end offit(); other callers may invoke it explicitly:Design & Code Changes
verl/utils/tracking.py: add idempotent, exception-safeTracking.finish()(per-backendtry/except;_finishedguard);__del__now delegates to it (fallback only). Afterwandb.finish(exit_code=0), also callwandb.teardown(exit_code=0)in-band — wandb's own docs note its atexit service teardown "is not reliable in certain setups such as when using Python'smultiprocessingmodule" (Ray). Doing it in-band, while the core is alive, unregisters that atexit hook and closes the service cleanly, which removes the intermittentBrokenPipeError/HandleAbandonedErrorat shutdown. Guarded againstSystemExit(wandb'steardown()callssys.exit()if the core reports a non-zero exit code).verl/utils/dataloader.py(new):shutdown_dataloader_workers/shutdown_dataloaders— best-effort, exception-safe shutdown of a DataLoader's active worker iterator (calls the iterator's_shutdown_workers()while still in normal control flow, then clears_iterator). Reused from Fix PPO dataloader worker teardown #6292.verl/trainer/ppo/v1/trainer_base.py: extract the training body into_fit();fit()wrapsself._fit()intry/finallyand, in thefinally, callsself._shutdown_dataloaders()thenself.logger.finish(). The existingself._shutdown_dump_executor()stays inline at the return sites (it re-raisesf.result(), so it must not run in thefinally, where it would mask an in-flight exception). With torchdata'sStatefulDataLoader,dataloader._iteratoris the same object asself.train_dataloader_it, so this reaches the live train workers; the validation loader is already drained by itsfor-loop.Checklist Before Submitting
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=alwaysci-requestchannel.