[logging] Publish the training loop phase as a Prometheus gauge#1923
[logging] Publish the training loop phase as a Prometheus gauge#1923eicherseiji wants to merge 5 commits into
Conversation
Adds skyrl/train/utils/phase_metrics.py with TrainingPhaseGauge, which sets a skyrl_training_phase gauge to 1.0 for the active macro-phase and 0.0 for the rest via ray.util.metrics. Ray exports it to the same Prometheus that scrapes node GPU metrics, so phase windows join GPU utilization on one wall-clock axis without correlating the experiment tracker by hand. The async loop sets the phase at each existing Timer boundary. Best-effort: the gauge silently no-ops when Ray metrics are unavailable, so it never breaks training or tests.
There was a problem hiding this comment.
Code Review
This pull request introduces a TrainingPhaseGauge to publish the training loop's current macro-phase (such as eval, training, weight synchronization, and checkpointing) to Prometheus via Ray metrics. This allows training loop phases to be overlaid with cluster-wide GPU utilization metrics. The feedback suggests validating phase names against the defined list of allowed phases to prevent developer typos and silent failures, along with adding a corresponding unit test for this validation.
- Emit only the two changed series per transition instead of sweeping all phases; the full sweep runs once at construction to seed every series for PromQL selectors. - Warn and keep the current phase on an unknown phase name instead of silently zeroing everything. - Rename phases to match their paired Timer keys so wandb timing/* keys and the Prometheus phase label share one vocabulary. - Construct the gauge in __init__ behind trainer.enable_training_phase_gauge, matching the RayGpuMonitor lifecycle. - Deduplicate docstrings and comments; add unknown-phase and disabled-by-flag tests.
The gauge writes two in-process values per phase transition with no thread or I/O, so there is nothing an operator needs to switch off. Construct it unconditionally; it already no-ops without Ray.
ScalarGauges had no functional dependency on TrainingPhaseGauge; the stack existed only because both classes shared phase_metrics.py. Move ScalarGauges to its own module (scalar_gauges.py, one module per concern) and revert the phase-gauge content from this branch's tree, so this PR is independently mergeable and its diff is buffer-only. Forward revert, no history rewrite; NovaSky-AI#1923 still owns the phase gauge.
| # keys and the Prometheus phase label share one vocabulary. "generating" is the default between | ||
| # blocks, when the trainer is not blocking and generation proceeds in the background. | ||
| PHASES = ( | ||
| "generating", |
There was a problem hiding this comment.
There is currently no equivalent to generating in the wandb timing/* metrics, right? Is there any way to align the metrics better?
There was a problem hiding this comment.
Ah true; fixed. Dropped generating. In fully async RL it doesn't really align well.
pcmoritz
left a comment
There was a problem hiding this comment.
Two more suggestions:
- Put skyrl/train/utils/phase_metrics.py in skyrl/train/utils/metrics.py so the other metrics can also be put into this file going forward
- You could unify the
Timerandphasecontext manager into a combined util context manager that internally calls both of them
eicherseiji
left a comment
There was a problem hiding this comment.
Thanks @pcmoritz! Responded to comments
| # keys and the Prometheus phase label share one vocabulary. "generating" is the default between | ||
| # blocks, when the trainer is not blocking and generation proceeds in the background. | ||
| PHASES = ( | ||
| "generating", |
There was a problem hiding this comment.
Ah true; fixed. Dropped generating. In fully async RL it doesn't really align well.
… metrics.py Remove the generating phase. Generation runs in the background across every phase, so it was never a real phase, just a default label with no timing key. The gauge now lights a phase on enter and clears it on exit, reading nothing between phases. Add save_hf_model as a phase so a real timed block is no longer unlabeled. Unify Timer and phase into TrainingPhaseGauge.timed_phase, so every phase is a timed block whose name is its timing/<name> wandb key. Rename phase_metrics.py to metrics.py. Drop the try/except around Gauge construction; it does not raise on the pinned ray 2.56.0.
…gauge # Conflicts: # skyrl/train/fully_async_trainer.py
What does this PR do?
Publishes which macro-phase the fully async training loop is in as a Prometheus gauge. Today, joining "what was the loop doing" to "how busy were the GPUs" means hand-correlating the experiment tracker against Prometheus by wall-clock.
New metric:
skyrl_training_phase{phase=...}: 1.0 for the active phase, 0.0 for the rest. Phases match the loop'sTimer()keys (wait_for_generation_buffer,convert_to_training_input,run_training,sync_weights,eval,save_checkpoints) plus the defaultgenerating, so the wandbtiming/*keys and the Prometheus label share one vocabulary.Ray exports it to the same Prometheus that scrapes node GPU metrics, so per-phase GPU utilization is a single-store query:
Prometheus also survives a cluster restart, when the tracker may be unreachable.
How
New module
skyrl/train/utils/phase_metrics.pywithTrainingPhaseGauge, following the one-module-per-concern layout ofray_gpu_monitor.py. The trainer constructs it in__init__and sets the phase at each existingTimer()boundary. Each transition writes only the two changed series; construction seeds all series so PromQL selectors never hit a missing one. Best-effort: it no-ops when Ray metrics are unavailable and warns on unknown phase names, so it never breaks training or tests.Tests
tests/train/utils/test_phase_metrics.py(new): exactly one phase active, context-manager restore, unknown-phase rejection, no-op when Ray is unavailable. 5 passed.