Skip to content

[trainer] fix: gracefully finalize tracking and dataloader workers at fit() exits#6917

Open
seokhyunan wants to merge 1 commit into
verl-project:mainfrom
seokhyunan:fix/trainer-teardown
Open

[trainer] fix: gracefully finalize tracking and dataloader workers at fit() exits#6917
seokhyunan wants to merge 1 commit into
verl-project:mainfrom
seokhyunan:fix/trainer-teardown

Conversation

@seokhyunan

@seokhyunan seokhyunan commented Jul 2, 2026

Copy link
Copy Markdown

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 single try/finally around the training loop, so cleanup runs on normal completion, val_only, and exceptions:

  • W&B dashboards show only the "System" charts. Tracking finalized wandb only in __del__, which fires after the wandb-core service connection has already closed, so wandb.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: Killed printed at process exit when data.dataloader_num_workers > 0 — the live StatefulDataLoader worker iterator is never shut down before fit() 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 all fit() 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

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:

python3 examples/data_preprocess/gsm8k.py --local_dir ~/data/gsm8k

python3 -m verl.trainer.main_ppo \
    algorithm.adv_estimator=grpo \
    data.train_files=$HOME/data/gsm8k/train.parquet \
    data.val_files=$HOME/data/gsm8k/test.parquet \
    data.train_batch_size=8 \
    data.max_prompt_length=512 \
    data.max_response_length=256 \
    data.dataloader_num_workers=8 \
    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.ppo_mini_batch_size=8 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.rollout.n=4 \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \
    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \
    trainer.logger='["console","wandb"]' \
    trainer.project_name=verl_teardown_repro \
    trainer.experiment_name=grpo_2step \
    trainer.n_gpus_per_node=1 \
    trainer.nnodes=1 \
    trainer.val_before_train=False \
    trainer.test_freq=-1 \
    trainer.save_freq=-1 \
    trainer.total_training_steps=2
  • Before: W&B shows only System charts; the local run-*.wandb datastore has 0 history / 0 summary / 0 exit records; at shutdown teardown_atexit intermittently raises HandleAbandonedError / BrokenPipeError; and (with workers > 0) the DataLoader worker ... killed by signal: Killed traceback is printed.
  • After: training/* metrics appear on W&B; run-*.wandb has history ≥ 2, summary == 1, exit == 1; no teardown_atexit traceback; no DataLoader SIGKILL traceback; exit code 0. (Re-parse the datastore offline with wandb.sdk.internal.datastore to confirm the record counts.)

API and Usage Example

No CLI / config changes. One additive API: Tracking.finish() (idempotent, exception-safe) is now public, and Tracking.__del__ delegates to it. The V1 trainer calls it automatically at the end of fit(); other callers may invoke it explicitly:

tracking = Tracking(project_name, experiment_name, default_backend=["console", "wandb"])
try:
    ...  # training loop, tracking.log(...)
finally:
    tracking.finish()  # flush + finalize backends; safe to call once, a repeat call is a no-op

Design & Code Changes

  • verl/utils/tracking.py: add idempotent, exception-safe Tracking.finish() (per-backend try/except; _finished guard); __del__ now delegates to it (fallback only). After wandb.finish(exit_code=0), also call wandb.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's multiprocessing module" (Ray). Doing it in-band, while the core is alive, unregisters that atexit hook and closes the service cleanly, which removes the intermittent BrokenPipeError / HandleAbandonedError at shutdown. Guarded against SystemExit (wandb's teardown() calls sys.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() wraps self._fit() in try/finally and, in the finally, calls self._shutdown_dataloaders() then self.logger.finish(). The existing self._shutdown_dump_executor() stays inline at the return sites (it re-raises f.result(), so it must not run in the finally, where it would mask an in-flight exception). With torchdata's StatefulDataLoader, dataloader._iterator is the same object as self.train_dataloader_it, so this reaches the live train workers; the validation loader is already drained by its for-loop.

Checklist Before Submitting

  • Read the Contribute Guide.
  • Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
  • Add / Update documentation — N/A (internal end-of-run finalization; no user-facing API or config change).
  • Add unit or end-to-end test(s) to the CI workflow. If not feasible, explain why: the change only affects end-of-run finalization ordering, and its symptoms (the wandb-core atexit service race and the DataLoader worker SIGKILL) surface only in a live multi-process Ray + wandb-core shutdown, which the CI unit lanes do not reproduce deterministically. Validated end-to-end via the 2-step GRPO run in the Test section.
  • Once your PR is ready for CI, send a message in the ci-request channel.
  • Recipe submodule — N/A (not touched).

@CLAassistant

CLAassistant commented Jul 2, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@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 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.

Comment thread verl/trainer/ppo/v1/trainer_base.py
Comment thread verl/trainer/ppo/v1/trainer_base.py
Comment thread verl/utils/dataloader.py
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.

DataLoader worker (pid 10855) is killed by signal: Killed. [Bug] verl does not log the final validation stats

2 participants