Skip to content

[logging] Fix: RayGpuMonitor never ran on the fully async trainer#1922

Merged
SumanthRH merged 1 commit into
NovaSky-AI:mainfrom
eicherseiji:seiji/async-gpu-monitor-fix
Jul 17, 2026
Merged

[logging] Fix: RayGpuMonitor never ran on the fully async trainer#1922
SumanthRH merged 1 commit into
NovaSky-AI:mainfrom
eicherseiji:seiji/async-gpu-monitor-fix

Conversation

@eicherseiji

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes a wiring bug: RayGpuMonitor never ran on the fully async trainer.

RayGpuMonitor (#1712) is constructed in RayPPOTrainer.__init__ and enabled by default with trainer.enable_ray_gpu_monitor=True. The base loop starts it, flushes it into the per-step log payload, and stops it. FullyAsyncRayPPOTrainer overrides train() and never called any of the three, so fully async runs logged no ray/ 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.

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.
@SumanthRH
SumanthRH merged commit a07f44e into NovaSky-AI:main Jul 17, 2026
4 of 6 checks passed

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

Comment on lines +459 to +462
# 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()

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.

medium

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:

  1. Thread Leak: The background daemon thread of RayGpuMonitor will continue running indefinitely.
  2. 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 (as self._running is still True), and the first flush() 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()

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.

2 participants