[logging] Rollout time split + async-trainer GPU/phase/buffer observability

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 ScalarGauges, TrainingPhaseGauge
 from skyrl.train.utils.trainer_utils import (
     ResumeMode,
     build_dataloader,
@@ -456,9 +457,29 @@ async def train(self):
         with Timer("sync_weights_to_inference_engines"):
             await self.dispatch.save_weights_for_sampler()
 
+        # Per-step GPU utilization (-> tracker) and a training-phase gauge (-> Prometheus, joinable
+        # with node GPU metrics). The base trainer wires these for the synchronous loop; the async
+        # loop overrides train() and must start/flush/stop them itself.
+        if self._ray_gpu_monitor is not None:
+            self._ray_gpu_monitor.start()
+        self._phase_gauge = TrainingPhaseGauge()
+        # Buffer depth to Prometheus, so it joins the phase gauge and node GPU metrics in one store
+        # (the per-step tracker logging below only reaches the experiment tracker). keep_rate exposes
+        # the zero-variance drop dynamics under sample_full_batch.
+        self._loop_gauges = ScalarGauges(
+            descriptions={
+                "skyrl_gen_buffer_qsize": "Completed generation groups buffered at step start.",
+                "skyrl_gen_buffer_maxsize": "Staleness-bounded generation-buffer capacity "
+                "(mini_batch_size * (max_staleness_steps + 1)).",
+                "skyrl_mini_batch_size": "Generation groups consumed per training step.",
+                "skyrl_gen_group_keep_rate": "Fraction of drained groups kept (not dropped as "
+                "zero-variance) while collecting the last mini-batch.",
+            }
+        )
+
         # 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)
 
@@ -503,15 +524,26 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                 trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size
                 for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1):
                     with Timer("step", self.all_timings):
+                        # Buffer depth at step start: how many completed rollout batches are queued
+                        # for the trainer. Near maxsize means generation is paused at the staleness
+                        # cap; near zero means the trainer is starving. Previously only shown live in
+                        # a tqdm postfix.
+                        self.all_metrics["async/gen_buffer_qsize"] = generation_output_group_buffer.qsize()
+                        self.all_metrics["async/gen_buffer_maxsize"] = generation_output_group_buffer.maxsize
+                        self._loop_gauges.set("skyrl_gen_buffer_qsize", generation_output_group_buffer.qsize())
+                        self._loop_gauges.set("skyrl_gen_buffer_maxsize", generation_output_group_buffer.maxsize)
+                        self._loop_gauges.set("skyrl_mini_batch_size", self.mini_batch_size)
+
                         # 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("waiting_for_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
@@ -537,25 +569,33 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                             break
 
                         if self.sample_full_batch:
-                            self.all_metrics["async/num_groups_dropped"] = len(cur_dropped_groups)
+                            num_kept = len(cur_generation_group_mini_batch)
+                            num_dropped = len(cur_dropped_groups)
+                            self.all_metrics["async/num_groups_dropped"] = num_dropped
+                            drained = num_kept + num_dropped
+                            if drained > 0:
+                                self._loop_gauges.set("skyrl_gen_group_keep_rate", num_kept / drained)
 
                         # 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("converting"),
+                        ):
                             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("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("weight_sync"):
                             await self.dispatch.save_weights_for_sampler()
 
                     # A training step completed: count it for this epoch's bookkeeping.
@@ -577,22 +617,24 @@ 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("checkpoint"):
                                 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):
                                 await asyncio.to_thread(self.save_models)
 
                     timing_payload = {"timing/" + k: v for k, v in self.all_timings.items()}
+                    if self._ray_gpu_monitor is not None:
+                        timing_payload.update(self._ray_gpu_monitor.flush())
                     if self._vllm_metrics_scraper is not None:
                         timing_payload.update(await self._vllm_metrics_scraper.sample())
                     self.tracker.log(timing_payload, step=self.global_step, commit=True)
@@ -657,6 +699,8 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                 # End of an epoch.
         finally:
             self._profiler_stop()
+            if self._ray_gpu_monitor is not None:
+                self._ray_gpu_monitor.stop()
 
         pbar.close()
 
@@ -667,7 +711,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("checkpoint"):
                 await asyncio.to_thread(self.save_checkpoints)
                 logger.info("Saved final checkpoint.")
         if self.cfg.trainer.hf_save_interval > 0:

skyrl/train/generators/base.py

@@ -46,6 +46,11 @@ class GeneratorOutput(TypedDict):
     # trajectory in the input batch (i.e. per ``agent_loop`` call). Used by the fully
     # async trainer to compute per-group / intra-group completion-time metrics.
     trajectory_generation_times: Optional[List[float]]
+    # Per-trajectory split of ``trajectory_generation_times`` into inference-engine wait and
+    # ``env.step()`` time (seconds), one entry per trajectory. Re-aggregated into rollout metrics
+    # by ``concatenate_generator_outputs``. None if any trajectory did not record the split.
+    trajectory_llm_times: Optional[List[float]]
+    trajectory_env_times: Optional[List[float]]
     rollout_expert_indices: Optional[List[List[List[List[int]]]]]  # [batch_size, seq_len, layer_num, topk]
     # Applicable only for step-wise training
     is_last_step: Optional[List[bool]]

skyrl/train/generators/skyrl_gym_generator.py

@@ -56,6 +56,12 @@ class TrajectoryOutput:
     # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
+    # Wall-clock time (seconds) this trajectory spent awaiting the inference engine, summed over
+    # turns. Optional: agent loops may leave this as None if they do not track timing.
+    llm_time: Optional[float] = None
+    # Wall-clock time (seconds) this trajectory spent in ``env.step()``, summed over turns.
+    # Optional: agent loops may leave this as None if they do not track timing.
+    env_time: Optional[float] = None
 
 
 @dataclass
@@ -66,6 +72,12 @@ class StepWiseOutput:
     # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: agent loops may
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
+    # Wall-clock time (seconds) this trajectory spent awaiting the inference engine, summed over
+    # turns. Optional: agent loops may leave this as None if they do not track timing.
+    llm_time: Optional[float] = None
+    # Wall-clock time (seconds) this trajectory spent in ``env.step()``, summed over turns.
+    # Optional: agent loops may leave this as None if they do not track timing.
+    env_time: Optional[float] = None
 
 
 @dataclass
@@ -321,6 +333,11 @@ async def agent_loop(
             rollout_logprobs: Optional[List[float]]
         """
         agent_loop_start_time = time.monotonic()
+        # Wall-clock split of the trajectory: time awaiting the inference engine vs. time in
+        # ``env.step()``, accumulated across turns. The two need not sum to ``e2e_time``. The
+        # remainder is tokenization, chat-template rendering and other in-loop bookkeeping.
+        llm_time_s = 0.0
+        env_time_s = 0.0
 
         session_id = (
             f"{trajectory_id.instance_id}_{trajectory_id.repetition_id}" if trajectory_id is not None else uuid4().hex
@@ -409,7 +426,9 @@ async def agent_loop(
                     sampling_params=sampling_params,
                     cache_salt=cache_salt,
                 )
+                llm_call_start_time = time.monotonic()
                 engine_output = await self.inference_engine_client.generate(engine_input, model=self.policy_model_name)
+                llm_time_s += time.monotonic() - llm_call_start_time
                 output = engine_output["responses"][0]
                 output_ids = engine_output["response_ids"][0]
                 stop_reason = engine_output["stop_reasons"][0]
@@ -443,7 +462,9 @@ async def agent_loop(
                         added_eos = True
 
                 # 2. Environment step
+                env_step_start_time = time.monotonic()
                 env_step_output: BaseTextEnvStepOutput = await self._run_in_executor_if_available(env.step, output)
+                env_time_s += time.monotonic() - env_step_start_time
                 new_obs = env_step_output["observations"]
                 step_reward: float = env_step_output["reward"]
                 agent_loop_state.done = env_step_output["done"]
@@ -606,6 +627,8 @@ async def agent_loop(
                 trajectory_id,
             )
             agent_loop_output.e2e_time = time.monotonic() - agent_loop_start_time
+            agent_loop_output.llm_time = llm_time_s
+            agent_loop_output.env_time = env_time_s
             return agent_loop_output
 
         finally:
@@ -874,6 +897,15 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         if any(t is None for t in trajectory_generation_times_per_prompt):
             trajectory_generation_times_per_prompt = None
 
+        # Per-trajectory split of that end-to-end time into inference-engine wait vs. environment
+        # execution. Handled like ``e2e_time``: omitted entirely if any trajectory did not record it.
+        trajectory_llm_times_per_prompt = [getattr(output, "llm_time", None) for output in all_outputs]
+        if any(t is None for t in trajectory_llm_times_per_prompt):
+            trajectory_llm_times_per_prompt = None
+        trajectory_env_times_per_prompt = [getattr(output, "env_time", None) for output in all_outputs]
+        if any(t is None for t in trajectory_env_times_per_prompt):
+            trajectory_env_times_per_prompt = None
+
         if self.generator_cfg.step_wise_trajectories:
             responses = []
             rewards = []
@@ -885,6 +917,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_ids = []
             out_env_classes = []
             out_trajectory_generation_times = []
+            out_trajectory_llm_times = []
+            out_trajectory_env_times = []
             for i, output in enumerate(all_outputs):
                 for j, step_output in enumerate(output.step_outputs):
                     responses.append(step_output.response_ids)
@@ -896,11 +930,17 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                     is_last_step.append(j == len(output.step_outputs) - 1)
                     out_trajectory_ids.append(trajectory_ids[i])
                     out_env_classes.append(env_classes[i])
-                    # For trajectory completion per turn we just use the trajectory level e2e time
+                    # For trajectory completion per turn we just use the trajectory level times
                     out_trajectory_generation_times.append(getattr(output, "e2e_time", None))
+                    out_trajectory_llm_times.append(getattr(output, "llm_time", None))
+                    out_trajectory_env_times.append(getattr(output, "env_time", None))
             # Keep aligned with the per-prompt None handling:
             if not trajectory_generation_times_per_prompt:
                 out_trajectory_generation_times = None
+            if not trajectory_llm_times_per_prompt:
+                out_trajectory_llm_times = None
+            if not trajectory_env_times_per_prompt:
+                out_trajectory_env_times = None
             env_classes = out_env_classes
         else:
             responses = [output.response_ids for output in all_outputs]
@@ -913,6 +953,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_ids = None
             # One time per trajectory, already aligned 1:1 with responses (None if not all recorded).
             out_trajectory_generation_times = trajectory_generation_times_per_prompt
+            out_trajectory_llm_times = trajectory_llm_times_per_prompt
+            out_trajectory_env_times = trajectory_env_times_per_prompt
 
         has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs)
         pixel_values = (
@@ -953,6 +995,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             # NOTE: we only use trajectory completion times per prompt for
             # metrics, to avoid duplicate entries with step-wise training
             trajectory_completion_times=trajectory_generation_times_per_prompt,
+            trajectory_llm_times=trajectory_llm_times_per_prompt,
+            trajectory_env_times=trajectory_env_times_per_prompt,
         )
 
         if self.generator_cfg.zero_reward_on_non_stop:
@@ -974,6 +1018,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             "trajectory_ids": out_trajectory_ids,
             # NOTE: for completion metrics, we output the completion time
             "trajectory_generation_times": out_trajectory_generation_times,
+            "trajectory_llm_times": out_trajectory_llm_times,
+            "trajectory_env_times": out_trajectory_env_times,
             "rollout_expert_indices": rollout_expert_indices,
             "is_last_step": is_last_step,
             "env_metrics": env_metrics,