diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index 09eabb3910..4c79926083 100644 --- a/skyrl/train/fully_async_trainer.py +++ b/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,10 +538,10 @@ 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 @@ -545,7 +549,7 @@ 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 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, @@ -553,14 +557,14 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don ) # 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). diff --git a/skyrl/train/utils/metrics.py b/skyrl/train/utils/metrics.py new file mode 100644 index 0000000000..a7727103d5 --- /dev/null +++ b/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/. 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/`` and mark ``name`` the active phase for its duration.""" + with Timer(name, timings), self.phase(name): + yield diff --git a/tests/train/utils/test_metrics.py b/tests/train/utils/test_metrics.py new file mode 100644 index 0000000000..bd69e979c2 --- /dev/null +++ b/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