[logging] Fix: RayGpuMonitor never ran on the fully async trainer#1922
Conversation
RayGpuMonitor is constructed in RayPPOTrainer.__init__ and default-on, and the base loop starts, flushes, and stops it. FullyAsyncRayPPOTrainer overrides train() and never called any of the three, so async runs logged zero ray/ GPU keys despite enable_ray_gpu_monitor=True. Start the monitor at loop entry, flush it into the per-step timing payload, and stop it in the finally.
There was a problem hiding this comment.
Code Review
This pull request integrates _ray_gpu_monitor into the fully asynchronous trainer to track per-step GPU utilization. The monitor is started before training, flushed during step logging, and stopped in the finally block. Feedback suggests wrapping the monitor's lifecycle in an outer try...finally block to prevent potential thread leaks and metric corruption if an exception occurs before entering the main training loop.
| # Per-step GPU utilization to the tracker. The base loop starts, flushes, and stops the | ||
| # monitor itself. The async loop overrides train() and must wire it here. | ||
| if self._ray_gpu_monitor is not None: | ||
| self._ray_gpu_monitor.start() |
There was a problem hiding this comment.
Potential Resource Leak and Metric Corruption on Failure
If an exception occurs during the initial evaluation (e.g., in self.eval() at line 467) or during progress bar/profiler initialization before entering the main try block at line 476, the finally block at line 665 will not be executed.
This leads to two issues:
- Thread Leak: The background daemon thread of
RayGpuMonitorwill continue running indefinitely. - Metric Corruption: Since the monitor is never stopped, it will keep collecting idle GPU metrics in its internal buffer. If
train()is called again in the same process (e.g., in a notebook, multi-run script, or unit test), the monitor will not start a new thread (asself._runningis stillTrue), and the firstflush()call will average the idle metrics collected between the runs, corrupting the step metrics.
Suggested Fix
To guarantee that the monitor is stopped under any failure scenario, you can wrap the entire training process (from monitor start onwards) in an outer try...finally block. For example:
if self._ray_gpu_monitor is not None:
self._ray_gpu_monitor.start()
try:
# Eval before training
if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train:
...
# main training loop
...
try:
for epoch in range(start_epoch, self.cfg.trainer.epochs):
...
finally:
self._profiler_stop()
finally:
if self._ray_gpu_monitor is not None:
self._ray_gpu_monitor.stop()
What does this PR do?
Fixes a wiring bug:
RayGpuMonitornever ran on the fully async trainer.RayGpuMonitor(#1712) is constructed inRayPPOTrainer.__init__and enabled by default withtrainer.enable_ray_gpu_monitor=True. The base loop starts it, flushes it into the per-step log payload, and stops it.FullyAsyncRayPPOTraineroverridestrain()and never called any of the three, so fully async runs logged noray/GPU keys at all. Confirmed on real runs.The async loop now starts the monitor at loop entry, flushes it into the per-step timing payload, and stops it in the
finally, mirroring the base loop wiring.Tests
Wiring only, covered by lint and existing tests. Exercising it end to end needs a GPU cluster. A structural lifecycle test that catches any base-loop hook dropped by an overriding trainer is a planned follow-up.
Split out of #1900.