[logging] Publish the training loop phase as a Prometheus gauge

skyrl/train/fully_async_trainer.py

@@ -39,6 +39,7 @@
 )
 from skyrl.train.trainer import RayPPOTrainer
 from skyrl.train.utils import Timer
+from skyrl.train.utils.metrics import TrainingPhaseGauge
 from skyrl.train.utils.trainer_utils import (
     ResumeMode,
     build_dataloader,
@@ -335,6 +336,9 @@ def __init__(self, *args, **kwargs):
         self.max_staleness_steps = cfg.trainer.fully_async.max_staleness_steps
         self.sample_full_batch = cfg.trainer.fully_async.sample_full_batch
 
+        # Publishes the loop macro-phase to Prometheus; see metrics.py.
+        self._phase_gauge = TrainingPhaseGauge()
+
         assert (
             # otherwise wasted throughput
             self.mini_batch_size <= self.num_parallel_generation_workers
@@ -463,7 +467,7 @@ async def train(self):
 
         # Eval before training
         if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train:
-            with Timer("eval", self.all_timings):
+            with self._phase_gauge.timed_phase("eval", self.all_timings):
                 eval_metrics = await self.eval()
                 self.tracker.log(eval_metrics, step=self.global_step, commit=True)
 
@@ -534,33 +538,33 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                             # Save the end-of-epoch checkpoint the normal is_epoch_end path would have, since
                             # we break before reaching it.
                             if self.cfg.trainer.ckpt_interval > 0:
-                                with Timer("save_checkpoints", self.all_timings):
+                                with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
                                     await asyncio.to_thread(self.save_checkpoints)
                             if self.cfg.trainer.hf_save_interval > 0:
-                                with Timer("save_hf_model", self.all_timings):
+                                with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
                                     await asyncio.to_thread(self.save_models)
                             break
 
                         if self.sample_full_batch:
                             self.all_metrics["async/num_groups_dropped"] = len(cur_dropped_groups)
 
                         # 2. Post-process the generated groups, aggregating to a single GeneratorOutput, and convert to training format.
-                        with Timer("convert_to_training_input", self.all_timings):
+                        with self._phase_gauge.timed_phase("convert_to_training_input", self.all_timings):
                             training_input = await asyncio.to_thread(
                                 self.convert_generation_group_mini_batch_to_training_input,
                                 cur_generation_group_mini_batch,
                                 cur_dropped_groups,
                             )
 
                         # 3. Run training and update consumed UIDs.
-                        with Timer("run_training", self.all_timings):
+                        with self._phase_gauge.timed_phase("run_training", self.all_timings):
                             status = await self._run_training(training_input)
                             await self.async_train_dataloader.mark_consumed_uids(
                                 [g.uid for g in cur_generation_group_mini_batch]
                             )
 
                         # 4. After training: pause generation, sync weights, resume.
-                        with Timer("sync_weights", self.all_timings):
+                        with self._phase_gauge.timed_phase("sync_weights", self.all_timings):
                             await self.dispatch.save_weights_for_sampler()
 
                     # A training step completed: count it for this epoch's bookkeeping.
@@ -580,7 +584,7 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                         self.global_step % self.cfg.trainer.eval_interval == 0
                         or self.global_step == self.total_training_steps
                     ):
-                        with Timer("eval", self.all_timings):
+                        with self._phase_gauge.timed_phase("eval", self.all_timings):
                             eval_metrics = await self.eval()
                             self.all_metrics.update(eval_metrics)
 
@@ -592,11 +596,11 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                     is_epoch_end = trained_steps_this_epoch == self.num_steps_per_epoch
                     if self.cfg.trainer.ckpt_interval > 0:
                         if is_epoch_end or self.global_step % self.cfg.trainer.ckpt_interval == 0:
-                            with Timer("save_checkpoints", self.all_timings):
+                            with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
                                 await asyncio.to_thread(self.save_checkpoints)
                     if self.cfg.trainer.hf_save_interval > 0:
                         if is_epoch_end or self.global_step % self.cfg.trainer.hf_save_interval == 0:
-                            with Timer("save_hf_model", self.all_timings):
+                            with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
                                 await asyncio.to_thread(self.save_models)
 
                     timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()}
@@ -678,11 +682,11 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
 
         # safety net: always save final checkpoint at end of training.
         if self.cfg.trainer.ckpt_interval > 0:
-            with Timer("save_checkpoints", self.all_timings):
+            with self._phase_gauge.timed_phase("save_checkpoints", self.all_timings):
                 await asyncio.to_thread(self.save_checkpoints)
                 logger.info("Saved final checkpoint.")
         if self.cfg.trainer.hf_save_interval > 0:
-            with Timer("save_hf_model", self.all_timings):
+            with self._phase_gauge.timed_phase("save_hf_model", self.all_timings):
                 await asyncio.to_thread(self.save_models)
                 logger.info("Saved final model.")
 
@@ -761,7 +765,7 @@ async def _collect_generation_mini_batch(
         epoch_exhausted = False
         # Buffer occupancy when this step starts waiting.
         self.all_metrics["async/gen_buffer_qsize_at_wait_start"] = generation_output_group_buffer.qsize()
-        with Timer("wait_for_generation_buffer", self.all_timings):
+        with self._phase_gauge.timed_phase("wait_for_generation_buffer", self.all_timings):
             buffer_pbar = tqdm(total=self.mini_batch_size, initial=0, desc="Generation Buffer Progress", position=1)
             while len(kept_groups) < self.mini_batch_size:
                 # We do finish-time FIFO here (not schedule-time FIFO).

skyrl/train/utils/metrics.py

@@ -0,0 +1,65 @@
+"""Publishes the training loop's current macro-phase to Prometheus via ``ray.util.metrics``.
+
+Ray exports these to the same Prometheus that scrapes node GPU metrics, under the ``ray_`` prefix,
+so loop phase joins GPU utilization on one wall-clock axis and survives a cluster restart, e.g.
+
+    avg(ray_node_gpus_utilization) and on() (ray_skyrl_training_phase{phase="eval"} == 1)
+"""
+
+from contextlib import contextmanager
+from typing import Dict
+
+from loguru import logger
+
+from skyrl.train.utils.utils import Timer
+
+# Macro-phases of one async training step. Each name is the Timer key for that block, so the same
+# phase shows up in wandb as timing/<name>. Generation is not a phase: it runs in the background
+# across all of them, and its activity is tracked separately by the buffer-depth gauges.
+PHASES = (
+    "wait_for_generation_buffer",
+    "convert_to_training_input",
+    "run_training",
+    "sync_weights",
+    "eval",
+    "save_checkpoints",
+    "save_hf_model",
+)
+
+
+class TrainingPhaseGauge:
+    """Sets ``skyrl_training_phase{phase=...}`` to 1.0 while the trainer is in that macro-phase.
+
+    Exactly one phase is 1.0 at a time, or none between phases.
+    """
+
+    def __init__(self) -> None:
+        from ray.util.metrics import Gauge
+
+        self._gauge = Gauge(
+            "skyrl_training_phase",
+            description="1.0 while the trainer is in this macro-phase, 0.0 otherwise.",
+            tag_keys=("phase",),
+        )
+        # Seed every series to 0.0 so PromQL selectors never hit a missing series.
+        for phase in PHASES:
+            self._gauge.set(0.0, tags={"phase": phase})
+
+    @contextmanager
+    def phase(self, name: str):
+        """Mark ``name`` active for the duration of the block."""
+        if name not in PHASES:
+            logger.warning(f"TrainingPhaseGauge: unknown phase {name!r}; not publishing it.")
+            yield
+            return
+        self._gauge.set(1.0, tags={"phase": name})
+        try:
+            yield
+        finally:
+            self._gauge.set(0.0, tags={"phase": name})
+
+    @contextmanager
+    def timed_phase(self, name: str, timings: Dict[str, float]):
+        """Time the block as wandb ``timing/<name>`` and mark ``name`` the active phase for its duration."""
+        with Timer(name, timings), self.phase(name):
+            yield

tests/train/utils/test_metrics.py

@@ -0,0 +1,49 @@
+"""
+uv run --isolated --extra dev --extra skyrl-train pytest tests/train/utils/test_metrics.py
+"""
+
+from unittest.mock import MagicMock, patch
+
+from skyrl.train.utils.metrics import PHASES, TrainingPhaseGauge
+
+
+def _make_gauge():
+    """Build a TrainingPhaseGauge over a mock ray Gauge; return (obj, {phase: latest value})."""
+    values = {}
+    mock_gauge = MagicMock()
+    mock_gauge.set.side_effect = lambda v, tags: values.__setitem__(tags["phase"], v)
+    with patch("ray.util.metrics.Gauge", return_value=mock_gauge):
+        obj = TrainingPhaseGauge()
+    return obj, values
+
+
+def _active(values):
+    return {p for p, v in values.items() if v == 1.0}
+
+
+def test_construction_seeds_every_phase_to_zero():
+    _, values = _make_gauge()
+    assert set(values) == set(PHASES)
+    assert _active(values) == set()
+
+
+def test_phase_marks_one_active_then_clears():
+    obj, values = _make_gauge()
+    with obj.phase("run_training"):
+        assert _active(values) == {"run_training"}
+    assert _active(values) == set()
+
+
+def test_unknown_phase_publishes_nothing():
+    obj, values = _make_gauge()
+    with obj.phase("bogus_phase"):
+        assert _active(values) == set()
+
+
+def test_timed_phase_records_timing_and_marks_phase():
+    obj, values = _make_gauge()
+    timings = {}
+    with obj.timed_phase("run_training", timings):
+        assert _active(values) == {"run_training"}
+    assert _active(values) == set()
+    assert "run_training" in timings