[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.phase_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 phase_metrics.py.
+        self._phase_gauge = TrainingPhaseGauge()
+
         assert (
             # otherwise wasted throughput
             self.mini_batch_size <= self.num_parallel_generation_workers
@@ -458,7 +462,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 Timer("eval", self.all_timings), self._phase_gauge.phase("eval"):
                 eval_metrics = await self.eval()
                 self.tracker.log(eval_metrics, step=self.global_step, commit=True)
 
@@ -505,13 +509,14 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                     with Timer("step", self.all_timings):
                         # 1. Wait until we have a full mini-batch buffered (dropping zero-variance groups if
                         # sample_full_batch).
-                        (
-                            cur_generation_group_mini_batch,
-                            cur_dropped_groups,
-                            epoch_exhausted,
-                        ) = await self._collect_generation_mini_batch(
-                            generation_output_group_buffer, all_generators_done
-                        )
+                        with self._phase_gauge.phase("wait_for_generation_buffer"):
+                            (
+                                cur_generation_group_mini_batch,
+                                cur_dropped_groups,
+                                epoch_exhausted,
+                            ) = await self._collect_generation_mini_batch(
+                                generation_output_group_buffer, all_generators_done
+                            )
 
                         if epoch_exhausted:
                             # Exhausted mid mini-batch: discard the partial batch (marked consumed so it
@@ -540,22 +545,25 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                             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 (
+                            Timer("convert_to_training_input", self.all_timings),
+                            self._phase_gauge.phase("convert_to_training_input"),
+                        ):
                             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 Timer("run_training", self.all_timings), self._phase_gauge.phase("run_training"):
                             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 Timer("sync_weights", self.all_timings), self._phase_gauge.phase("sync_weights"):
                             await self.dispatch.save_weights_for_sampler()
 
                     # A training step completed: count it for this epoch's bookkeeping.
@@ -577,15 +585,18 @@ 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 Timer("eval", self.all_timings), self._phase_gauge.phase("eval"):
                             eval_metrics = await self.eval()
                             self.all_metrics.update(eval_metrics)
 
                     # 7. Checkpointing. At interval and at the last step of each epoch.
                     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 (
+                                Timer("save_checkpoints", self.all_timings),
+                                self._phase_gauge.phase("save_checkpoints"),
+                            ):
                                 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:
@@ -667,7 +678,7 @@ 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 Timer("save_checkpoints", self.all_timings), self._phase_gauge.phase("save_checkpoints"):
                 await asyncio.to_thread(self.save_checkpoints)
                 logger.info("Saved final checkpoint.")
         if self.cfg.trainer.hf_save_interval > 0:

skyrl/train/utils/phase_metrics.py

@@ -0,0 +1,73 @@
+"""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 loguru import logger
+
+# Macro-phases of one async training step. Names match the paired Timer keys so the wandb timing/*
+# 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",
+    "wait_for_generation_buffer",
+    "convert_to_training_input",
+    "run_training",
+    "sync_weights",
+    "eval",
+    "save_checkpoints",
+)
+
+
+class TrainingPhaseGauge:
+    """Sets ``skyrl_training_phase{phase=...}`` to 1.0 for the active phase and 0.0 for the rest.
+
+    Best-effort: silently no-ops when Ray metrics are unavailable, so it is always safe to
+    construct and call.
+    """
+
+    def __init__(self) -> None:
+        self._gauge = None
+        self._current = "generating"
+        try:
+            from ray.util.metrics import Gauge
+
+            self._gauge = Gauge(
+                "skyrl_training_phase",
+                description="1.0 for the active training-loop macro-phase, 0.0 otherwise.",
+                tag_keys=("phase",),
+            )
+            # Seed every series once so PromQL selectors never hit a missing series.
+            for phase in PHASES:
+                self._gauge.set(1.0 if phase == self._current else 0.0, tags={"phase": phase})
+        except Exception as e:
+            logger.warning(f"TrainingPhaseGauge disabled ({e}); the training-phase metric will not be published.")
+            self._gauge = None
+
+    def set_phase(self, name: str) -> None:
+        if name not in PHASES:
+            logger.warning(f"TrainingPhaseGauge: unknown phase {name!r}; keeping {self._current!r}.")
+            return
+        prev, self._current = self._current, name
+        if self._gauge is None or prev == name:
+            return
+        try:
+            self._gauge.set(0.0, tags={"phase": prev})
+            self._gauge.set(1.0, tags={"phase": name})
+        except Exception:
+            pass
+
+    @contextmanager
+    def phase(self, name: str):
+        """Mark ``name`` active for the duration of the block, restoring the prior phase on exit."""
+        prev = self._current
+        self.set_phase(name)
+        try:
+            yield
+        finally:
+            self.set_phase(prev)

tests/train/utils/test_phase_metrics.py

@@ -0,0 +1,63 @@
+"""
+uv run --isolated --extra dev --extra skyrl-train pytest tests/train/utils/test_phase_metrics.py
+"""
+
+from unittest.mock import MagicMock, patch
+
+from skyrl.train.utils.phase_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_phase(values):
+    # every phase series exists (seeded at construction) and exactly one is 1.0
+    assert set(values) == set(PHASES)
+    active = [p for p, v in values.items() if v == 1.0]
+    assert len(active) == 1, f"expected exactly one active phase, got {active}"
+    return active[0]
+
+
+def test_construction_defaults_to_generating():
+    _, values = _make_gauge()
+    assert _active_phase(values) == "generating"
+
+
+def test_set_phase_marks_exactly_one_active():
+    obj, values = _make_gauge()
+    obj.set_phase("run_training")
+    assert _active_phase(values) == "run_training"
+    obj.set_phase("eval")
+    assert _active_phase(values) == "eval"
+
+
+def test_phase_context_manager_restores_prior_phase():
+    obj, values = _make_gauge()
+    obj.set_phase("run_training")
+    with obj.phase("save_checkpoints"):
+        assert _active_phase(values) == "save_checkpoints"
+    # restored to whatever was active before the block, not hard-coded to a default
+    assert _active_phase(values) == "run_training"
+
+
+def test_unknown_phase_is_rejected():
+    obj, values = _make_gauge()
+    obj.set_phase("run_training")
+    obj.set_phase("bogus_phase")
+    assert _active_phase(values) == "run_training"
+
+
+def test_disabled_when_ray_metrics_unavailable():
+    # Gauge() raising must not propagate
+    with patch("ray.util.metrics.Gauge", side_effect=RuntimeError("ray not initialized")):
+        obj = TrainingPhaseGauge()
+    obj.set_phase("run_training")
+    with obj.phase("eval"):
+        pass