From c96b3ef648a6f624a94306277ba34a90c7aeeb51 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <eicherseiji@users.noreply.github.com>
Date: Tue, 14 Jul 2026 20:38:19 +0000
Subject: [PATCH 01/11] [logging] Split rollout trajectory time into
 inference-engine wait vs env.step

Rollout timing is currently reported only as a total per trajectory
(`generate/trajectory_completion_time_*`, added in #1804). That total cannot
distinguish an engine-bound rollout from an environment-bound one, which is the
first question to ask when GPU utilization sags during generation.

Accumulate, per trajectory in `agent_loop`, the wall-clock time spent awaiting
`inference_engine_client.generate()` and the time spent in `env.step()`, and
report:

  generate/trajectory_llm_time_{mean,p90,max}
  generate/trajectory_env_time_{mean,p90,max}
  generate/frac_time_in_env

`frac_time_in_env` is time-weighted (sum over sum) so long trajectories count
proportionally rather than each trajectory contributing equally.

The two components need not sum to `e2e_time`; the remainder is tokenization,
chat-template rendering and other in-loop bookkeeping. Both are plumbed like the
existing `e2e_time`: optional, and omitted entirely if any trajectory in the
batch did not record them.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 skyrl/train/generators/skyrl_gym_generator.py |  34 +++++
 skyrl/train/generators/utils.py               |  36 +++++
 .../generators/test_skyrl_gym_generator.py    | 125 ++++++++++++++++++
 3 files changed, 195 insertions(+)

diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index cc9bc95e80..e01128ddee 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/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 = []
@@ -953,6 +985,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:
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 4c5228f189..4ef3abcc87 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -384,6 +384,8 @@ def get_rollout_metrics(
     env_classes: Optional[List[str]] = None,
     loss_masks: Optional[List[List[int]]] = None,
     trajectory_completion_times: Optional[List[float]] = None,
+    trajectory_llm_times: Optional[List[float]] = None,
+    trajectory_env_times: Optional[List[float]] = None,
 ):
     """
     Computes rollout metrics including token statistics and optional environment-specific metrics.
@@ -396,6 +398,10 @@ def get_rollout_metrics(
         loss_masks: Optional list of per-token loss masks; used to compute assistant-only token counts
         trajectory_completion_times: Optional per-trajectory end-to-end generation times (seconds);
             used to compute aggregate trajectory completion-time stats (mean / p90 / max)
+        trajectory_llm_times: Optional per-trajectory time (seconds) spent awaiting the inference
+            engine, summed over turns
+        trajectory_env_times: Optional per-trajectory time (seconds) spent in ``env.step()``, summed
+            over turns
 
     Returns:
         Dictionary of aggregated metrics
@@ -460,6 +466,36 @@ def get_rollout_metrics(
             }
         )
 
+    if trajectory_llm_times:
+        llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64)
+        rollout_metrics.update(
+            {
+                "generate/trajectory_llm_time_mean": np.mean(llm_times_arr).item(),
+                "generate/trajectory_llm_time_p90": np.percentile(llm_times_arr, 90).item(),
+                "generate/trajectory_llm_time_max": np.max(llm_times_arr).item(),
+            }
+        )
+
+    if trajectory_env_times:
+        env_times_arr = np.array(trajectory_env_times, dtype=np.float64)
+        rollout_metrics.update(
+            {
+                "generate/trajectory_env_time_mean": np.mean(env_times_arr).item(),
+                "generate/trajectory_env_time_p90": np.percentile(env_times_arr, 90).item(),
+                "generate/trajectory_env_time_max": np.max(env_times_arr).item(),
+            }
+        )
+
+    if trajectory_llm_times and trajectory_env_times:
+        llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64)
+        env_times_arr = np.array(trajectory_env_times, dtype=np.float64)
+        # Batch-level share of rollout time spent in the environment rather than awaiting the
+        # inference engine. Time-weighted (sum over sum), so long trajectories count proportionally
+        # instead of every trajectory contributing equally.
+        total_busy_time = np.sum(llm_times_arr) + np.sum(env_times_arr)
+        if total_busy_time > 0:
+            rollout_metrics["generate/frac_time_in_env"] = (np.sum(env_times_arr) / total_busy_time).item()
+
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)
         for i, metrics in enumerate(env_metrics):
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index c29afbe920..67465e70cc 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1751,3 +1751,128 @@ def step(self, action):
         np.percentile(expected, 90).item()
     )
     assert rollout_metrics["generate/trajectory_completion_time_max"] == pytest.approx(np.max(expected).item())
+
+
+@pytest.mark.asyncio
+@patch("skyrl_gym.make")
+async def test_llm_vs_env_time_split_metrics(mock_make, mock_tokenizer, mock_llm, mock_env_cfg):
+    """The rollout time split attributes engine wait and ``env.step()`` time separately.
+
+    A trajectory's wall-clock time is split between awaiting the inference engine and executing
+    the environment. Only the total was previously recorded (``trajectory_completion_time_*``),
+    which cannot distinguish an engine-bound rollout from an environment-bound one.
+
+    Uses an environment that is deliberately slower than the engine, so a split that mixed the two
+    up (or attributed everything to one side) fails.
+    """
+    import asyncio
+    import time
+
+    from skyrl.train.generators import utils as generator_utils
+    from skyrl.train.generators.base import TrajectoryID
+
+    llm_sleep_s = 0.02
+    env_sleep_s = 0.06
+    num_turns = 2
+
+    mock_tokenizer.eos_token_id = 4
+
+    def apply_chat_template_side_effect(messages, **kwargs):
+        if kwargs.get("tokenize", True):
+            return [201, 202]
+        return "".join([m.get("content", "") for m in messages])
+
+    mock_tokenizer.apply_chat_template.side_effect = apply_chat_template_side_effect
+
+    async def llm_generate_side_effect(input_batch, model=None):
+        await asyncio.sleep(llm_sleep_s)
+        num = len(input_batch["prompt_token_ids"]) if "prompt_token_ids" in input_batch else len(input_batch["prompts"])
+        return {
+            "responses": ["step"] * num,
+            "stop_reasons": ["stop"] * num,
+            "response_logprobs": None,
+            "response_ids": [[10, 11, 12, mock_tokenizer.eos_token_id] for _ in range(num)],
+        }
+
+    mock_llm.generate = AsyncMock(side_effect=llm_generate_side_effect)
+
+    class SlowEnv(BaseTextEnv):
+        def __init__(self):
+            super().__init__()
+            self.turns = 0
+
+        def init(self, prompt):
+            return prompt, {}
+
+        def step(self, action):
+            time.sleep(env_sleep_s)
+            self.turns += 1
+            if self.turns < num_turns:
+                return BaseTextEnvStepOutput(
+                    observations=[{"role": "user", "content": "obs"}], reward=0.0, done=False, metadata={}
+                )
+            return BaseTextEnvStepOutput(observations=[], reward=1.0, done=True, metadata={})
+
+    mock_make.side_effect = lambda *args, **kwargs: SlowEnv()
+
+    # Run env steps on the executor (as in production, where ``max_env_workers`` defaults to 32).
+    # With inline execution a blocking ``env.step`` stalls the shared event loop, so a sibling
+    # trajectory's ``await generate()`` would absorb that stall and inflate its measured LLM time.
+    mock_env_cfg.max_env_workers = 4
+
+    cfg = GeneratorConfig()
+    cfg.sampling_params.max_generate_length = 50
+    cfg.sampling_params.logprobs = None
+    cfg.apply_overlong_filtering = False
+    cfg.max_input_length = 512
+    cfg.batched = False
+    cfg.max_turns = 10
+    cfg.zero_reward_on_non_stop = False
+    cfg.use_conversation_multi_turn = True
+    cfg.chat_template = ChatTemplateConfig(source="name", name_or_path=None)
+
+    generator = SkyRLGymGenerator(
+        generator_cfg=cfg,
+        skyrl_gym_cfg=mock_env_cfg,
+        inference_engine_client=mock_llm,
+        tokenizer=mock_tokenizer,
+    )
+    generator.base_conversation_token_ids = []
+
+    num_trajectories = 2
+    prompts = [[{"role": "user", "content": f"Q{i}?"}] for i in range(num_trajectories)]
+    input_batch: GeneratorInput = {
+        "prompts": prompts,
+        "env_extras": [{} for _ in prompts],
+        "env_classes": [mock_env_cfg.env_class for _ in prompts],
+        "trajectory_ids": [TrajectoryID(instance_id=f"uid{i}", repetition_id=0) for i in range(num_trajectories)],
+    }
+
+    spy = MagicMock(side_effect=generator_utils.get_rollout_metrics)
+    with patch("skyrl.train.generators.skyrl_gym_generator.get_rollout_metrics", spy):
+        generator_output: GeneratorOutput = await generator.generate(input_batch)
+
+    llm_times = spy.call_args.kwargs["trajectory_llm_times"]
+    env_times = spy.call_args.kwargs["trajectory_env_times"]
+    assert llm_times is not None and env_times is not None
+    assert len(llm_times) == num_trajectories
+    assert len(env_times) == num_trajectories
+
+    # Each side is at least its per-turn sleep summed over turns, and neither swallows the other.
+    for llm_t, env_t in zip(llm_times, env_times):
+        assert llm_t >= llm_sleep_s * num_turns
+        assert env_t >= env_sleep_s * num_turns
+        # The environment is ~3x slower than the engine here, so a swapped attribution would flip this.
+        assert env_t > llm_t
+
+    metrics = generator_output["rollout_metrics"]
+    assert metrics["generate/trajectory_llm_time_mean"] >= llm_sleep_s * num_turns
+    assert metrics["generate/trajectory_env_time_mean"] >= env_sleep_s * num_turns
+
+    # env_sleep / (env_sleep + llm_sleep) = 0.75; allow generous headroom for scheduling overhead
+    # attributed to the engine wait, but it must clearly indicate an environment-bound rollout.
+    assert 0.5 < metrics["generate/frac_time_in_env"] < 1.0
+
+    # The split accounts for real time inside the trajectory and never exceeds its end-to-end time.
+    for llm_t, env_t, e2e_t in zip(llm_times, env_times, generator_output["trajectory_generation_times"]):
+        assert llm_t + env_t <= e2e_t + 1e-6

From 896911cfeaf3aa624f3f21ff209269d81413c113 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <eicherseiji@users.noreply.github.com>
Date: Wed, 15 Jul 2026 00:30:30 +0000
Subject: [PATCH 02/11] [logging] Async trainer: GPU util per step,
 training-phase Prometheus gauge, buffer depth

Three observability additions to the fully-async RL trainer, all motivated by a
post-mortem where reconstructing GPU idle time required correlating Prometheus
(GPU/disk, wall-clock keyed) against W&B (phase timings, step keyed) by hand.

1. Wire RayGpuMonitor into the async loop. It is constructed in the base
   RayPPOTrainer and flushed in the base loop, but FullyAsyncRayPPOTrainer
   overrides train() and never started or flushed it -- so runs on the async
   trainer logged zero `ray/` GPU keys despite `enable_ray_gpu_monitor=True`.
   Start it at loop entry, flush per step into the committed payload, stop it in
   the finally.

2. Publish a `skyrl_training_phase` gauge via `ray.util.metrics`
   (skyrl/train/utils/phase_metrics.py). Ray exports these to the same
   Prometheus that scrapes node GPU metrics, so
   `avg(ray_node_gpus_utilization) by (Phase)` becomes a single-store query
   instead of a manual cross-tracker correlation -- and it works after a cluster
   restart, when the experiment tracker may be unreachable. Set at each existing
   Timer() phase boundary (waiting_for_buffer / converting / training /
   weight_sync / eval / checkpoint); best-effort, never raises.

3. Log `async/gen_buffer_qsize` and `async/gen_buffer_maxsize` at step start.
   Run-ahead depth was previously only in a tqdm postfix; as a metric it
   distinguishes a throughput-limited generator (buffer low) from a
   gate-limited one (buffer at the staleness-bounded maxsize).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 skyrl/train/fully_async_trainer.py      | 50 +++++++++++-----
 skyrl/train/utils/phase_metrics.py      | 77 +++++++++++++++++++++++++
 tests/train/utils/test_phase_metrics.py | 59 +++++++++++++++++++
 3 files changed, 172 insertions(+), 14 deletions(-)
 create mode 100644 skyrl/train/utils/phase_metrics.py
 create mode 100644 tests/train/utils/test_phase_metrics.py

diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py
index 99fa0ff31b..07bd555cab 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.phase_metrics import TrainingPhaseGauge
 from skyrl.train.utils.trainer_utils import (
     ResumeMode,
     build_dataloader,
@@ -456,9 +457,16 @@ 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()
+
         # 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 +511,22 @@ 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):
+                        # Run-ahead depth of the rollout buffer at step start: high => generation is
+                        # ahead (buffer near the staleness-bounded maxsize); low => the trainer is
+                        # starving for rollouts. Previously only visible 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
+
                         # 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
@@ -540,7 +555,10 @@ 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("converting"),
+                        ):
                             training_input = await asyncio.to_thread(
                                 self.convert_generation_group_mini_batch_to_training_input,
                                 cur_generation_group_mini_batch,
@@ -548,14 +566,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 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,7 +595,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 Timer("eval", self.all_timings), self._phase_gauge.phase("eval"):
                             eval_metrics = await self.eval()
                             self.all_metrics.update(eval_metrics)
 
@@ -585,7 +603,7 @@ 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 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:
@@ -593,6 +611,8 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                                 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 +677,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 +689,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:
diff --git a/skyrl/train/utils/phase_metrics.py b/skyrl/train/utils/phase_metrics.py
new file mode 100644
index 0000000000..af266747c1
--- /dev/null
+++ b/skyrl/train/utils/phase_metrics.py
@@ -0,0 +1,77 @@
+"""Publishes the training loop's current macro-phase to Prometheus via ``ray.util.metrics``.
+
+Ray exports metrics recorded through ``ray.util.metrics`` to the same Prometheus that scrapes
+cluster-wide node metrics (GPU/CPU/disk). Emitting the training-loop phase there lets a query join
+"what was the loop doing" against "how busy were the GPUs" in a single store and time base, e.g.
+
+    avg(ray_node_gpus_utilization) by (Phase)
+
+without correlating a separate experiment tracker by wall-clock. This is deliberately a Prometheus
+gauge rather than a tracker/W&B scalar: Prometheus is the store that has cluster-wide GPU data and
+survives a cluster restart, so it is the one place a post-hoc utilization breakdown can be computed.
+"""
+
+from contextlib import contextmanager
+
+from loguru import logger
+
+# Macro-phases of one async training step, plus the default "generating" state between blocks (the
+# trainer is not blocking, so generation and staleness control proceed in the background).
+PHASES = (
+    "generating",
+    "waiting_for_buffer",
+    "converting",
+    "training",
+    "weight_sync",
+    "eval",
+    "checkpoint",
+)
+
+
+class TrainingPhaseGauge:
+    """Sets a ``skyrl_training_phase`` gauge to 1.0 for the active phase and 0.0 for the rest.
+
+    Exactly one phase is 1.0 at any time, so ``ray_skyrl_training_phase{phase="eval"} == 1`` marks the
+    eval windows on the Prometheus timeline. Best-effort: if Ray metrics are unavailable the object
+    silently no-ops, so it is always safe to construct and call (including in unit tests without Ray).
+    """
+
+    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",),
+            )
+            self._emit("generating")
+        except Exception as e:
+            logger.warning(f"TrainingPhaseGauge disabled ({e}); the training-phase metric will not be published.")
+            self._gauge = None
+
+    def _emit(self, active: str) -> None:
+        if self._gauge is None:
+            return
+        try:
+            for phase in PHASES:
+                self._gauge.set(1.0 if phase == active else 0.0, tags={"phase": phase})
+        except Exception:
+            # Observability must never break training.
+            pass
+
+    def set_phase(self, name: str) -> None:
+        self._current = name
+        self._emit(name)
+
+    @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)
diff --git a/tests/train/utils/test_phase_metrics.py b/tests/train/utils/test_phase_metrics.py
new file mode 100644
index 0000000000..dee0bf0d28
--- /dev/null
+++ b/tests/train/utils/test_phase_metrics.py
@@ -0,0 +1,59 @@
+"""
+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 backed by a mock ray Gauge; return (obj, mock_gauge)."""
+    mock_gauge = MagicMock()
+    with patch("ray.util.metrics.Gauge", return_value=mock_gauge):
+        obj = TrainingPhaseGauge()
+    return obj, mock_gauge
+
+
+def _active_phase(mock_gauge):
+    """Return the single phase set to 1.0 in the most recent full emit (len(PHASES) set calls)."""
+    last = mock_gauge.set.call_args_list[-len(PHASES) :]
+    active = [c.kwargs["tags"]["phase"] for c in last if c.args[0] == 1.0]
+    assert len(active) == 1, f"expected exactly one active phase, got {active}"
+    # every phase must be written each emit (so stale phases are cleared to 0.0)
+    written = {c.kwargs["tags"]["phase"] for c in last}
+    assert written == set(PHASES)
+    return active[0]
+
+
+def test_construction_defaults_to_generating():
+    _, g = _make_gauge()
+    assert _active_phase(g) == "generating"
+
+
+def test_set_phase_marks_exactly_one_active():
+    obj, g = _make_gauge()
+    obj.set_phase("training")
+    assert _active_phase(g) == "training"
+    obj.set_phase("eval")
+    assert _active_phase(g) == "eval"
+
+
+def test_phase_context_manager_restores_prior_phase():
+    obj, g = _make_gauge()
+    obj.set_phase("training")
+    with obj.phase("checkpoint"):
+        assert _active_phase(g) == "checkpoint"
+    # restored to whatever was active before the block, not hard-coded to a default
+    assert _active_phase(g) == "training"
+
+
+def test_disabled_when_ray_metrics_unavailable():
+    # Gauge construction raising (e.g. Ray not initialized) must degrade to a silent no-op,
+    # never propagate, so training is never broken by observability.
+    with patch("ray.util.metrics.Gauge", side_effect=RuntimeError("ray not initialized")):
+        obj = TrainingPhaseGauge()
+    # all calls are safe no-ops
+    obj.set_phase("training")
+    with obj.phase("eval"):
+        pass

From d7fef66aa404b94c63dbcb2360e9efe9ed602ce0 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Thu, 16 Jul 2026 22:32:21 +0000
Subject: [PATCH 03/11] [logging] Aggregate the rollout time split across
 concatenation

Item 1's llm/env time split was computed only at the per-generate()
get_rollout_metrics call. Every logging path (async trainer, eval,
step-wise) re-aggregates through concatenate_generator_outputs, which had
no raw llm/env lists and fell back to the substring-based extra_keys
aggregator. There p90 and frac_time_in_env hit the sum() branch, so a
16-group mini-batch logged frac_time_in_env near 12 and p90 near 16x its
real value.

Store the raw per-trajectory llm/env lists on GeneratorOutput like
e2e_time, thread them through concatenate_generator_outputs with the same
step-wise last-step filter, and recompute all seven metrics there so the
extra_keys fallback leaves them alone.

Also in this commit:
- Extend test_llm_vs_env_time_split_metrics to assert the split on a
  concatenated output, the path that actually broke.
- Fix the phase_metrics module docstring query. The gauge has a lowercase
  `phase` tag and no shared label with the GPU metric, so
  `avg(ray_node_gpus_utilization) by (Phase)` does not run; replace it
  with a vector match against `ray_skyrl_training_phase{phase="eval"}`.
- Deduplicate the last-step filter and the frac-block array construction.
---
 skyrl/train/generators/base.py                |  5 ++
 skyrl/train/generators/skyrl_gym_generator.py | 16 +++++-
 skyrl/train/generators/utils.py               | 51 ++++++++++++-------
 skyrl/train/utils/phase_metrics.py            | 15 +++---
 .../generators/test_skyrl_gym_generator.py    | 16 ++++++
 5 files changed, 76 insertions(+), 27 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index 26d95868b9..96fb91d63d 100644
--- a/skyrl/train/generators/base.py
+++ b/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]]
diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index e01128ddee..156fcf8599 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -334,7 +334,7 @@ async def agent_loop(
         """
         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
+        # ``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
@@ -917,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)
@@ -928,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]
@@ -945,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 = (
@@ -1008,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,
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 4ef3abcc87..1dd0890d4b 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -240,6 +240,20 @@ def _flatten_field(generator_outputs: List[GeneratorOutput], key: str) -> list:
     return flat
 
 
+def _concat_field(generator_outputs: List[GeneratorOutput], key: str) -> Optional[list]:
+    """Flatten an optional per-trajectory field, keyed off the first output (None if absent)."""
+    if generator_outputs[0].get(key) is None:
+        return None
+    return _flatten_field(generator_outputs, key)
+
+
+def _last_step_only(values: Optional[list], is_last_step: Optional[List[bool]]) -> Optional[list]:
+    """Keep only last-step entries when step-wise, so one trajectory contributes one value."""
+    if values is None or not is_last_step:
+        return values
+    return [v for v, last in zip(values, is_last_step) if last]
+
+
 def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step_wise: bool = False) -> GeneratorOutput:
     """
     Concatenate the generator outputs of multiple batches. Then validate the concatenated result.
@@ -270,11 +284,9 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         "rollout_logprobs": (
             _flatten_field(generator_outputs, "rollout_logprobs") if first.get("rollout_logprobs") is not None else None
         ),
-        "trajectory_generation_times": (
-            _flatten_field(generator_outputs, "trajectory_generation_times")
-            if first.get("trajectory_generation_times") is not None
-            else None
-        ),
+        "trajectory_generation_times": _concat_field(generator_outputs, "trajectory_generation_times"),
+        "trajectory_llm_times": _concat_field(generator_outputs, "trajectory_llm_times"),
+        "trajectory_env_times": _concat_field(generator_outputs, "trajectory_env_times"),
     }
 
     # propagate additional keys with list values as-is
@@ -284,19 +296,21 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
     for key in additional_keys:
         result[key] = _flatten_field(generator_outputs, key)
 
-    # With step-wise training, only use the trajectory generation time from the last step
-    trajectory_generation_times = result.get("trajectory_generation_times")
-    if step_wise and trajectory_generation_times and result.get("is_last_step"):
-        trajectory_generation_times = [
-            t for t, is_last_step in zip(trajectory_generation_times, result.get("is_last_step")) if is_last_step
-        ]
+    # With step-wise training each trajectory spans multiple rows; keep only its last-step timing.
+    is_last_step = result.get("is_last_step") if step_wise else None
+    trajectory_generation_times = _last_step_only(result.get("trajectory_generation_times"), is_last_step)
+    trajectory_llm_times = _last_step_only(result.get("trajectory_llm_times"), is_last_step)
+    trajectory_env_times = _last_step_only(result.get("trajectory_env_times"), is_last_step)
 
-    # Re-aggregate rollout metrics
+    # Re-aggregate rollout metrics. The time splits must be recomputed here from the raw per-
+    # trajectory lists; the extra_keys fallback below cannot aggregate a p90 or a ratio correctly.
     rollout_metrics = get_rollout_metrics(
         result["response_ids"],
         result["rewards"],
         loss_masks=result.get("loss_masks"),
         trajectory_completion_times=trajectory_generation_times,
+        trajectory_llm_times=trajectory_llm_times,
+        trajectory_env_times=trajectory_env_times,
     )
 
     # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only
@@ -466,8 +480,10 @@ def get_rollout_metrics(
             }
         )
 
-    if trajectory_llm_times:
-        llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64)
+    llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64) if trajectory_llm_times else None
+    env_times_arr = np.array(trajectory_env_times, dtype=np.float64) if trajectory_env_times else None
+
+    if llm_times_arr is not None:
         rollout_metrics.update(
             {
                 "generate/trajectory_llm_time_mean": np.mean(llm_times_arr).item(),
@@ -476,8 +492,7 @@ def get_rollout_metrics(
             }
         )
 
-    if trajectory_env_times:
-        env_times_arr = np.array(trajectory_env_times, dtype=np.float64)
+    if env_times_arr is not None:
         rollout_metrics.update(
             {
                 "generate/trajectory_env_time_mean": np.mean(env_times_arr).item(),
@@ -486,9 +501,7 @@ def get_rollout_metrics(
             }
         )
 
-    if trajectory_llm_times and trajectory_env_times:
-        llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64)
-        env_times_arr = np.array(trajectory_env_times, dtype=np.float64)
+    if llm_times_arr is not None and env_times_arr is not None:
         # Batch-level share of rollout time spent in the environment rather than awaiting the
         # inference engine. Time-weighted (sum over sum), so long trajectories count proportionally
         # instead of every trajectory contributing equally.
diff --git a/skyrl/train/utils/phase_metrics.py b/skyrl/train/utils/phase_metrics.py
index af266747c1..3b42bf78d4 100644
--- a/skyrl/train/utils/phase_metrics.py
+++ b/skyrl/train/utils/phase_metrics.py
@@ -1,14 +1,17 @@
 """Publishes the training loop's current macro-phase to Prometheus via ``ray.util.metrics``.
 
 Ray exports metrics recorded through ``ray.util.metrics`` to the same Prometheus that scrapes
-cluster-wide node metrics (GPU/CPU/disk). Emitting the training-loop phase there lets a query join
-"what was the loop doing" against "how busy were the GPUs" in a single store and time base, e.g.
+cluster-wide node metrics (GPU/CPU/disk), prefixing them with ``ray_``. Emitting the training-loop
+phase there puts "what was the loop doing" on the same store and time base as "how busy were the
+GPUs", so the two can be overlaid without correlating a separate experiment tracker by wall-clock.
+``ray_skyrl_training_phase{phase="eval"} == 1`` selects the eval windows, and node GPU utilization
+during a phase follows from a vector match against that selector, for example:
 
-    avg(ray_node_gpus_utilization) by (Phase)
+    avg(ray_node_gpus_utilization) and on() (ray_skyrl_training_phase{phase="eval"} == 1)
 
-without correlating a separate experiment tracker by wall-clock. This is deliberately a Prometheus
-gauge rather than a tracker/W&B scalar: Prometheus is the store that has cluster-wide GPU data and
-survives a cluster restart, so it is the one place a post-hoc utilization breakdown can be computed.
+This is deliberately a Prometheus gauge rather than a tracker/W&B scalar: Prometheus is the store
+that has cluster-wide GPU data and survives a cluster restart, so it is the one place a post-hoc
+utilization breakdown can be computed.
 """
 
 from contextlib import contextmanager
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 67465e70cc..e7953092d3 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1768,6 +1768,8 @@ async def test_llm_vs_env_time_split_metrics(mock_make, mock_tokenizer, mock_llm
     import asyncio
     import time
 
+    import numpy as np
+
     from skyrl.train.generators import utils as generator_utils
     from skyrl.train.generators.base import TrajectoryID
 
@@ -1876,3 +1878,17 @@ def step(self, action):
     # The split accounts for real time inside the trajectory and never exceeds its end-to-end time.
     for llm_t, env_t, e2e_t in zip(llm_times, env_times, generator_output["trajectory_generation_times"]):
         assert llm_t + env_t <= e2e_t + 1e-6
+
+    # Regression: every logging path (async trainer, eval, step-wise) re-aggregates through
+    # concatenate_generator_outputs. The split must be recomputed there from the raw per-trajectory
+    # lists; when it was not, p90 and frac fell into a substring sum() fallback that inflated p90 by
+    # ~num_groups and summed the fractions past 1.0. Concatenating identical groups leaves the
+    # per-trajectory distribution unchanged, so the aggregates must match the single-group values.
+    concatenated = generator_utils.concatenate_generator_outputs([generator_output, generator_output])
+    concat_metrics = concatenated["rollout_metrics"]
+    # Concatenating two groups doubles the sample, so p90 is the percentile of the combined raw
+    # list, not the single-group p90. The sum() fallback instead added the two group p90s, exceeding
+    # the sample max, and summed the fractions past 1.0.
+    assert 0.5 < concat_metrics["generate/frac_time_in_env"] < 1.0
+    assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item())
+    assert concat_metrics["generate/trajectory_env_time_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item())

From e245bc0bba9ee03687ee525e2ce3c0e5cc4363c5 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Thu, 16 Jul 2026 22:56:02 +0000
Subject: [PATCH 04/11] [logging] Publish rollout-buffer levels as Prometheus
 gauges

Add ScalarGauges, a best-effort ray.util.metrics scalar-gauge helper that
no-ops without Ray, like TrainingPhaseGauge. Publish the async run-ahead
buffer levels through it so buffer pressure sits in the same Prometheus
store as the training-phase gauge and node GPU metrics, not only in the
experiment tracker:

  skyrl_gen_buffer_qsize     completed groups buffered at step start
  skyrl_gen_buffer_maxsize   staleness-bounded buffer capacity
  skyrl_mini_batch_size      groups consumed per training step
  skyrl_gen_group_keep_rate  kept / (kept + dropped) for the last
                             mini-batch, the zero-variance drop rate
                             under sample_full_batch
---
 skyrl/train/fully_async_trainer.py      | 25 +++++++++++++++++--
 skyrl/train/utils/phase_metrics.py      | 33 ++++++++++++++++++++++++-
 tests/train/utils/test_phase_metrics.py | 29 ++++++++++++++++++++++
 3 files changed, 84 insertions(+), 3 deletions(-)

diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py
index 07bd555cab..5b60a78f1a 100644
--- a/skyrl/train/fully_async_trainer.py
+++ b/skyrl/train/fully_async_trainer.py
@@ -39,7 +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.phase_metrics import ScalarGauges, TrainingPhaseGauge
 from skyrl.train.utils.trainer_utils import (
     ResumeMode,
     build_dataloader,
@@ -463,6 +463,19 @@ async def train(self):
         if self._ray_gpu_monitor is not None:
             self._ray_gpu_monitor.start()
         self._phase_gauge = TrainingPhaseGauge()
+        # Rollout-buffer levels to Prometheus, so run-ahead pressure 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:
@@ -516,6 +529,9 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                         # starving for rollouts. Previously only visible 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).
@@ -552,7 +568,12 @@ 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 (
diff --git a/skyrl/train/utils/phase_metrics.py b/skyrl/train/utils/phase_metrics.py
index 3b42bf78d4..4aaf510f4d 100644
--- a/skyrl/train/utils/phase_metrics.py
+++ b/skyrl/train/utils/phase_metrics.py
@@ -1,4 +1,5 @@
-"""Publishes the training loop's current macro-phase to Prometheus via ``ray.util.metrics``.
+"""Publishes training-loop state to Prometheus via ``ray.util.metrics``: the current macro-phase
+(``TrainingPhaseGauge``) and scalar loop levels such as buffer depth (``ScalarGauges``).
 
 Ray exports metrics recorded through ``ray.util.metrics`` to the same Prometheus that scrapes
 cluster-wide node metrics (GPU/CPU/disk), prefixing them with ``ray_``. Emitting the training-loop
@@ -15,6 +16,7 @@
 """
 
 from contextlib import contextmanager
+from typing import Dict, Optional
 
 from loguru import logger
 
@@ -78,3 +80,32 @@ def phase(self, name: str):
             yield
         finally:
             self.set_phase(prev)
+
+
+class ScalarGauges:
+    """Best-effort scalar gauges published to Prometheus via ``ray.util.metrics``.
+
+    Lazily creates a gauge the first time a name is ``set`` (Ray exports it as ``ray_<name>``) and
+    updates its value on each call. Like ``TrainingPhaseGauge`` this no-ops when Ray metrics are
+    unavailable, so it is safe to construct and call without Ray (including in unit tests).
+    """
+
+    def __init__(self, descriptions: Optional[Dict[str, str]] = None) -> None:
+        self._descriptions = descriptions or {}
+        self._gauges: Dict[str, object] = {}
+        self._enabled = True
+
+    def set(self, name: str, value: float) -> None:
+        if not self._enabled:
+            return
+        try:
+            gauge = self._gauges.get(name)
+            if gauge is None:
+                from ray.util.metrics import Gauge
+
+                gauge = Gauge(name, description=self._descriptions.get(name, name))
+                self._gauges[name] = gauge
+            gauge.set(float(value))
+        except Exception as e:
+            logger.warning(f"ScalarGauges disabled ({e}); scalar training metrics will not be published.")
+            self._enabled = False
diff --git a/tests/train/utils/test_phase_metrics.py b/tests/train/utils/test_phase_metrics.py
index dee0bf0d28..f3d9895814 100644
--- a/tests/train/utils/test_phase_metrics.py
+++ b/tests/train/utils/test_phase_metrics.py
@@ -57,3 +57,32 @@ def test_disabled_when_ray_metrics_unavailable():
     obj.set_phase("training")
     with obj.phase("eval"):
         pass
+
+
+def test_scalar_gauges_create_once_and_update():
+    from skyrl.train.utils.phase_metrics import ScalarGauges
+
+    created = {}
+
+    def fake_gauge(name, description=None):
+        created[name] = MagicMock()
+        return created[name]
+
+    with patch("ray.util.metrics.Gauge", side_effect=fake_gauge):
+        g = ScalarGauges()
+        g.set("skyrl_gen_buffer_qsize", 3)
+        g.set("skyrl_gen_buffer_qsize", 5)
+        g.set("skyrl_mini_batch_size", 8)
+
+    # A gauge is created once per name and reused, and values are coerced to float.
+    assert set(created) == {"skyrl_gen_buffer_qsize", "skyrl_mini_batch_size"}
+    created["skyrl_gen_buffer_qsize"].set.assert_called_with(5.0)
+    created["skyrl_mini_batch_size"].set.assert_called_with(8.0)
+
+
+def test_scalar_gauges_disabled_when_ray_metrics_unavailable():
+    from skyrl.train.utils.phase_metrics import ScalarGauges
+
+    with patch("ray.util.metrics.Gauge", side_effect=RuntimeError("ray not initialized")):
+        g = ScalarGauges()
+        g.set("skyrl_gen_buffer_qsize", 1)  # must not raise

From 77d283f16b8284799508b5c4523e3afc19644edf Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Thu, 16 Jul 2026 23:50:03 +0000
Subject: [PATCH 05/11] [logging] Call it buffer depth in the async-trainer
 comments

Rename the run-ahead phrasing to plain buffer depth in the buffer-metric and gauge comments; no behavior change.
---
 skyrl/train/fully_async_trainer.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py
index 5b60a78f1a..c47e51a182 100644
--- a/skyrl/train/fully_async_trainer.py
+++ b/skyrl/train/fully_async_trainer.py
@@ -463,9 +463,9 @@ async def train(self):
         if self._ray_gpu_monitor is not None:
             self._ray_gpu_monitor.start()
         self._phase_gauge = TrainingPhaseGauge()
-        # Rollout-buffer levels to Prometheus, so run-ahead pressure 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.
+        # 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.",
@@ -524,9 +524,10 @@ 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):
-                        # Run-ahead depth of the rollout buffer at step start: high => generation is
-                        # ahead (buffer near the staleness-bounded maxsize); low => the trainer is
-                        # starving for rollouts. Previously only visible in a tqdm postfix.
+                        # 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())

From 6a9c5d8a6386e7377c66cb140514f1a87e47c3d3 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Fri, 17 Jul 2026 22:26:03 +0000
Subject: [PATCH 06/11] [logging] Narrow this PR to the rollout time split

The async-trainer observability items bundled here (RayGpuMonitor
wiring, training-phase gauge, buffer depth) move to their own PRs so
each is a reviewable unit. This restores fully_async_trainer.py to base
and removes phase_metrics.py and its tests. The rollout LLM vs env time
split and its aggregation across concatenate_generator_outputs stay.
---
 skyrl/train/fully_async_trainer.py      |  74 ++++------------
 skyrl/train/utils/phase_metrics.py      | 111 ------------------------
 tests/train/utils/test_phase_metrics.py |  88 -------------------
 3 files changed, 15 insertions(+), 258 deletions(-)
 delete mode 100644 skyrl/train/utils/phase_metrics.py
 delete mode 100644 tests/train/utils/test_phase_metrics.py

diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py
index c47e51a182..99fa0ff31b 100644
--- a/skyrl/train/fully_async_trainer.py
+++ b/skyrl/train/fully_async_trainer.py
@@ -39,7 +39,6 @@
 )
 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,
@@ -457,29 +456,9 @@ 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), self._phase_gauge.phase("eval"):
+            with Timer("eval", self.all_timings):
                 eval_metrics = await self.eval()
                 self.tracker.log(eval_metrics, step=self.global_step, commit=True)
 
@@ -524,26 +503,15 @@ 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).
-                        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
-                            )
+                        (
+                            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
@@ -569,18 +537,10 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                             break
 
                         if self.sample_full_batch:
-                            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)
+                            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),
-                            self._phase_gauge.phase("converting"),
-                        ):
+                        with Timer("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,
@@ -588,14 +548,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), self._phase_gauge.phase("training"):
+                        with Timer("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), self._phase_gauge.phase("weight_sync"):
+                        with Timer("sync_weights", self.all_timings):
                             await self.dispatch.save_weights_for_sampler()
 
                     # A training step completed: count it for this epoch's bookkeeping.
@@ -617,7 +577,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), self._phase_gauge.phase("eval"):
+                        with Timer("eval", self.all_timings):
                             eval_metrics = await self.eval()
                             self.all_metrics.update(eval_metrics)
 
@@ -625,7 +585,7 @@ 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), self._phase_gauge.phase("checkpoint"):
+                            with Timer("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:
@@ -633,8 +593,6 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don
                                 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)
@@ -699,8 +657,6 @@ 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()
 
@@ -711,7 +667,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), self._phase_gauge.phase("checkpoint"):
+            with Timer("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:
diff --git a/skyrl/train/utils/phase_metrics.py b/skyrl/train/utils/phase_metrics.py
deleted file mode 100644
index 4aaf510f4d..0000000000
--- a/skyrl/train/utils/phase_metrics.py
+++ /dev/null
@@ -1,111 +0,0 @@
-"""Publishes training-loop state to Prometheus via ``ray.util.metrics``: the current macro-phase
-(``TrainingPhaseGauge``) and scalar loop levels such as buffer depth (``ScalarGauges``).
-
-Ray exports metrics recorded through ``ray.util.metrics`` to the same Prometheus that scrapes
-cluster-wide node metrics (GPU/CPU/disk), prefixing them with ``ray_``. Emitting the training-loop
-phase there puts "what was the loop doing" on the same store and time base as "how busy were the
-GPUs", so the two can be overlaid without correlating a separate experiment tracker by wall-clock.
-``ray_skyrl_training_phase{phase="eval"} == 1`` selects the eval windows, and node GPU utilization
-during a phase follows from a vector match against that selector, for example:
-
-    avg(ray_node_gpus_utilization) and on() (ray_skyrl_training_phase{phase="eval"} == 1)
-
-This is deliberately a Prometheus gauge rather than a tracker/W&B scalar: Prometheus is the store
-that has cluster-wide GPU data and survives a cluster restart, so it is the one place a post-hoc
-utilization breakdown can be computed.
-"""
-
-from contextlib import contextmanager
-from typing import Dict, Optional
-
-from loguru import logger
-
-# Macro-phases of one async training step, plus the default "generating" state between blocks (the
-# trainer is not blocking, so generation and staleness control proceed in the background).
-PHASES = (
-    "generating",
-    "waiting_for_buffer",
-    "converting",
-    "training",
-    "weight_sync",
-    "eval",
-    "checkpoint",
-)
-
-
-class TrainingPhaseGauge:
-    """Sets a ``skyrl_training_phase`` gauge to 1.0 for the active phase and 0.0 for the rest.
-
-    Exactly one phase is 1.0 at any time, so ``ray_skyrl_training_phase{phase="eval"} == 1`` marks the
-    eval windows on the Prometheus timeline. Best-effort: if Ray metrics are unavailable the object
-    silently no-ops, so it is always safe to construct and call (including in unit tests without Ray).
-    """
-
-    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",),
-            )
-            self._emit("generating")
-        except Exception as e:
-            logger.warning(f"TrainingPhaseGauge disabled ({e}); the training-phase metric will not be published.")
-            self._gauge = None
-
-    def _emit(self, active: str) -> None:
-        if self._gauge is None:
-            return
-        try:
-            for phase in PHASES:
-                self._gauge.set(1.0 if phase == active else 0.0, tags={"phase": phase})
-        except Exception:
-            # Observability must never break training.
-            pass
-
-    def set_phase(self, name: str) -> None:
-        self._current = name
-        self._emit(name)
-
-    @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)
-
-
-class ScalarGauges:
-    """Best-effort scalar gauges published to Prometheus via ``ray.util.metrics``.
-
-    Lazily creates a gauge the first time a name is ``set`` (Ray exports it as ``ray_<name>``) and
-    updates its value on each call. Like ``TrainingPhaseGauge`` this no-ops when Ray metrics are
-    unavailable, so it is safe to construct and call without Ray (including in unit tests).
-    """
-
-    def __init__(self, descriptions: Optional[Dict[str, str]] = None) -> None:
-        self._descriptions = descriptions or {}
-        self._gauges: Dict[str, object] = {}
-        self._enabled = True
-
-    def set(self, name: str, value: float) -> None:
-        if not self._enabled:
-            return
-        try:
-            gauge = self._gauges.get(name)
-            if gauge is None:
-                from ray.util.metrics import Gauge
-
-                gauge = Gauge(name, description=self._descriptions.get(name, name))
-                self._gauges[name] = gauge
-            gauge.set(float(value))
-        except Exception as e:
-            logger.warning(f"ScalarGauges disabled ({e}); scalar training metrics will not be published.")
-            self._enabled = False
diff --git a/tests/train/utils/test_phase_metrics.py b/tests/train/utils/test_phase_metrics.py
deleted file mode 100644
index f3d9895814..0000000000
--- a/tests/train/utils/test_phase_metrics.py
+++ /dev/null
@@ -1,88 +0,0 @@
-"""
-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 backed by a mock ray Gauge; return (obj, mock_gauge)."""
-    mock_gauge = MagicMock()
-    with patch("ray.util.metrics.Gauge", return_value=mock_gauge):
-        obj = TrainingPhaseGauge()
-    return obj, mock_gauge
-
-
-def _active_phase(mock_gauge):
-    """Return the single phase set to 1.0 in the most recent full emit (len(PHASES) set calls)."""
-    last = mock_gauge.set.call_args_list[-len(PHASES) :]
-    active = [c.kwargs["tags"]["phase"] for c in last if c.args[0] == 1.0]
-    assert len(active) == 1, f"expected exactly one active phase, got {active}"
-    # every phase must be written each emit (so stale phases are cleared to 0.0)
-    written = {c.kwargs["tags"]["phase"] for c in last}
-    assert written == set(PHASES)
-    return active[0]
-
-
-def test_construction_defaults_to_generating():
-    _, g = _make_gauge()
-    assert _active_phase(g) == "generating"
-
-
-def test_set_phase_marks_exactly_one_active():
-    obj, g = _make_gauge()
-    obj.set_phase("training")
-    assert _active_phase(g) == "training"
-    obj.set_phase("eval")
-    assert _active_phase(g) == "eval"
-
-
-def test_phase_context_manager_restores_prior_phase():
-    obj, g = _make_gauge()
-    obj.set_phase("training")
-    with obj.phase("checkpoint"):
-        assert _active_phase(g) == "checkpoint"
-    # restored to whatever was active before the block, not hard-coded to a default
-    assert _active_phase(g) == "training"
-
-
-def test_disabled_when_ray_metrics_unavailable():
-    # Gauge construction raising (e.g. Ray not initialized) must degrade to a silent no-op,
-    # never propagate, so training is never broken by observability.
-    with patch("ray.util.metrics.Gauge", side_effect=RuntimeError("ray not initialized")):
-        obj = TrainingPhaseGauge()
-    # all calls are safe no-ops
-    obj.set_phase("training")
-    with obj.phase("eval"):
-        pass
-
-
-def test_scalar_gauges_create_once_and_update():
-    from skyrl.train.utils.phase_metrics import ScalarGauges
-
-    created = {}
-
-    def fake_gauge(name, description=None):
-        created[name] = MagicMock()
-        return created[name]
-
-    with patch("ray.util.metrics.Gauge", side_effect=fake_gauge):
-        g = ScalarGauges()
-        g.set("skyrl_gen_buffer_qsize", 3)
-        g.set("skyrl_gen_buffer_qsize", 5)
-        g.set("skyrl_mini_batch_size", 8)
-
-    # A gauge is created once per name and reused, and values are coerced to float.
-    assert set(created) == {"skyrl_gen_buffer_qsize", "skyrl_mini_batch_size"}
-    created["skyrl_gen_buffer_qsize"].set.assert_called_with(5.0)
-    created["skyrl_mini_batch_size"].set.assert_called_with(8.0)
-
-
-def test_scalar_gauges_disabled_when_ray_metrics_unavailable():
-    from skyrl.train.utils.phase_metrics import ScalarGauges
-
-    with patch("ray.util.metrics.Gauge", side_effect=RuntimeError("ray not initialized")):
-        g = ScalarGauges()
-        g.set("skyrl_gen_buffer_qsize", 1)  # must not raise

From 78ac038c95d1265060ba3528c452483d33edfb17 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Fri, 17 Jul 2026 23:12:19 +0000
Subject: [PATCH 07/11] [logging] Simplify the time-split diff

Review feedback pass:
- Add trajectory_llm_times/env_times to the GeneratorOutput field
  guardrail test, which this branch had broken.
- Extract _add_time_stats for the mean/p90/max blocks (completion, llm,
  env) and reuse its sums for frac_time_in_env.
- Extract _optional_times for the omit-if-any-None per-prompt lists.
- Use _concat_field for stop_reasons and rollout_logprobs too.
- Cut duplicated comments; the llm/env split was explained in six
  places, now documented once per surface.
- Assert the step-wise per-step replication of the split in the
  existing step-wise timing test.
---
 skyrl/train/generators/base.py                |  5 +-
 skyrl/train/generators/skyrl_gym_generator.py | 39 ++++-------
 skyrl/train/generators/utils.py               | 70 +++++++------------
 .../generators/test_generator_output_utils.py |  2 +
 .../generators/test_skyrl_gym_generator.py    | 34 ++++-----
 5 files changed, 57 insertions(+), 93 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index 96fb91d63d..b2ccc102d6 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -46,9 +46,8 @@ 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.
+    # Split of ``trajectory_generation_times`` into engine-wait vs ``env.step()`` time (seconds).
+    # None if any trajectory did not record it.
     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]
diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index 156fcf8599..a317a1d6c0 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -56,11 +56,9 @@ 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.
+    # Wall-clock time (seconds) awaiting the inference engine, summed over turns. None if untracked.
     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.
+    # Wall-clock time (seconds) in ``env.step()``, summed over turns. None if untracked.
     env_time: Optional[float] = None
 
 
@@ -72,11 +70,8 @@ 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.
+    # Same split as TrajectoryOutput.llm_time / env_time.
     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
 
 
@@ -152,6 +147,12 @@ def get_turn_rollout_logprobs(self) -> Optional[List[float]]:
         return self.output_logprobs + [0.0] * len(self.obs_ids)
 
 
+def _optional_times(outputs, attr: str) -> Optional[List[float]]:
+    """One value per trajectory, or None if any trajectory did not record ``attr``."""
+    values = [getattr(o, attr, None) for o in outputs]
+    return None if any(v is None for v in values) else values
+
+
 class SkyRLGymGenerator(GeneratorInterface):
     def __init__(
         self,
@@ -333,9 +334,8 @@ 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.
+        # Engine wait vs. env.step() time, accumulated across turns. The gap to e2e_time is
+        # tokenization, chat templating, and event-loop scheduling.
         llm_time_s = 0.0
         env_time_s = 0.0
 
@@ -891,20 +891,9 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             disable=disable_tqdm,
         )
         # Per-trajectory end-to-end generation times (one entry per prompt, preserving input order).
-        # ``e2e_time`` is optional for agent loops; if any trajectory did not record it, we omit the
-        # field entirely rather than emit a partially-populated list.
-        trajectory_generation_times_per_prompt = [getattr(output, "e2e_time", None) for output in all_outputs]
-        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
+        trajectory_generation_times_per_prompt = _optional_times(all_outputs, "e2e_time")
+        trajectory_llm_times_per_prompt = _optional_times(all_outputs, "llm_time")
+        trajectory_env_times_per_prompt = _optional_times(all_outputs, "env_time")
 
         if self.generator_cfg.step_wise_trajectories:
             responses = []
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 1dd0890d4b..3137f5f319 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -278,12 +278,8 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         "response_ids": _flatten_field(generator_outputs, "response_ids"),
         "rewards": _flatten_field(generator_outputs, "rewards"),
         "loss_masks": _flatten_field(generator_outputs, "loss_masks"),
-        "stop_reasons": (
-            _flatten_field(generator_outputs, "stop_reasons") if first.get("stop_reasons") is not None else None
-        ),
-        "rollout_logprobs": (
-            _flatten_field(generator_outputs, "rollout_logprobs") if first.get("rollout_logprobs") is not None else None
-        ),
+        "stop_reasons": _concat_field(generator_outputs, "stop_reasons"),
+        "rollout_logprobs": _concat_field(generator_outputs, "rollout_logprobs"),
         "trajectory_generation_times": _concat_field(generator_outputs, "trajectory_generation_times"),
         "trajectory_llm_times": _concat_field(generator_outputs, "trajectory_llm_times"),
         "trajectory_env_times": _concat_field(generator_outputs, "trajectory_env_times"),
@@ -302,8 +298,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
     trajectory_llm_times = _last_step_only(result.get("trajectory_llm_times"), is_last_step)
     trajectory_env_times = _last_step_only(result.get("trajectory_env_times"), is_last_step)
 
-    # Re-aggregate rollout metrics. The time splits must be recomputed here from the raw per-
-    # trajectory lists; the extra_keys fallback below cannot aggregate a p90 or a ratio correctly.
+    # Re-aggregate rollout metrics; the extra_keys fallback below cannot aggregate a p90 or a ratio.
     rollout_metrics = get_rollout_metrics(
         result["response_ids"],
         result["rewards"],
@@ -391,6 +386,21 @@ def compute_turn_token_counts(loss_masks: List[List[int]]) -> List[int]:
     return turn_token_counts
 
 
+def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[List[float]]) -> Optional[float]:
+    """Add mean/p90/max for a per-trajectory time list; returns its sum, or None if absent."""
+    if not times:
+        return None
+    arr = np.array(times, dtype=np.float64)
+    rollout_metrics.update(
+        {
+            f"generate/trajectory_{name}_time_mean": np.mean(arr).item(),
+            f"generate/trajectory_{name}_time_p90": np.percentile(arr, 90).item(),
+            f"generate/trajectory_{name}_time_max": np.max(arr).item(),
+        }
+    )
+    return np.sum(arr).item()
+
+
 def get_rollout_metrics(
     responses: List[List[int]],
     rewards: Union[List[float], List[List[float]]],
@@ -470,44 +480,14 @@ def get_rollout_metrics(
                 }
             )
 
-    if trajectory_completion_times:
-        completion_times_arr = np.array(trajectory_completion_times, dtype=np.float64)
-        rollout_metrics.update(
-            {
-                "generate/trajectory_completion_time_mean": np.mean(completion_times_arr).item(),
-                "generate/trajectory_completion_time_p90": np.percentile(completion_times_arr, 90).item(),
-                "generate/trajectory_completion_time_max": np.max(completion_times_arr).item(),
-            }
-        )
-
-    llm_times_arr = np.array(trajectory_llm_times, dtype=np.float64) if trajectory_llm_times else None
-    env_times_arr = np.array(trajectory_env_times, dtype=np.float64) if trajectory_env_times else None
-
-    if llm_times_arr is not None:
-        rollout_metrics.update(
-            {
-                "generate/trajectory_llm_time_mean": np.mean(llm_times_arr).item(),
-                "generate/trajectory_llm_time_p90": np.percentile(llm_times_arr, 90).item(),
-                "generate/trajectory_llm_time_max": np.max(llm_times_arr).item(),
-            }
-        )
-
-    if env_times_arr is not None:
-        rollout_metrics.update(
-            {
-                "generate/trajectory_env_time_mean": np.mean(env_times_arr).item(),
-                "generate/trajectory_env_time_p90": np.percentile(env_times_arr, 90).item(),
-                "generate/trajectory_env_time_max": np.max(env_times_arr).item(),
-            }
-        )
+    _add_time_stats(rollout_metrics, "completion", trajectory_completion_times)
+    llm_sum = _add_time_stats(rollout_metrics, "llm", trajectory_llm_times)
+    env_sum = _add_time_stats(rollout_metrics, "env", trajectory_env_times)
 
-    if llm_times_arr is not None and env_times_arr is not None:
-        # Batch-level share of rollout time spent in the environment rather than awaiting the
-        # inference engine. Time-weighted (sum over sum), so long trajectories count proportionally
-        # instead of every trajectory contributing equally.
-        total_busy_time = np.sum(llm_times_arr) + np.sum(env_times_arr)
-        if total_busy_time > 0:
-            rollout_metrics["generate/frac_time_in_env"] = (np.sum(env_times_arr) / total_busy_time).item()
+    if llm_sum is not None and env_sum is not None and llm_sum + env_sum > 0:
+        # Time-weighted (sum over sum), not a mean of per-trajectory ratios, so long trajectories
+        # count proportionally.
+        rollout_metrics["generate/frac_time_in_env"] = env_sum / (llm_sum + env_sum)
 
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 7546858bb8..84e3a87aec 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -32,6 +32,8 @@ def test_generator_output_concatenation():
         # optional but present in the signature
         "trajectory_ids",
         "trajectory_generation_times",
+        "trajectory_llm_times",
+        "trajectory_env_times",
         "is_last_step",
         "env_metrics",
         "pixel_values",
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index e7953092d3..d9c0c1bc58 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1752,19 +1752,20 @@ def step(self, action):
     )
     assert rollout_metrics["generate/trajectory_completion_time_max"] == pytest.approx(np.max(expected).item())
 
+    # The llm/env split is stored per step and replicated like the completion times above.
+    for key in ("trajectory_llm_times", "trajectory_env_times"):
+        split_times = generator_output[key]
+        assert split_times is not None
+        assert len(split_times) == num_steps
+        assert split_times[0] == split_times[1] and split_times[2] == split_times[3]
+
 
 @pytest.mark.asyncio
 @patch("skyrl_gym.make")
 async def test_llm_vs_env_time_split_metrics(mock_make, mock_tokenizer, mock_llm, mock_env_cfg):
-    """The rollout time split attributes engine wait and ``env.step()`` time separately.
-
-    A trajectory's wall-clock time is split between awaiting the inference engine and executing
-    the environment. Only the total was previously recorded (``trajectory_completion_time_*``),
-    which cannot distinguish an engine-bound rollout from an environment-bound one.
-
-    Uses an environment that is deliberately slower than the engine, so a split that mixed the two
-    up (or attributed everything to one side) fails.
-    """
+    """Only the total rollout time was previously recorded, which cannot distinguish an engine-bound
+    rollout from an environment-bound one. Uses an env deliberately slower than the engine, so a
+    swapped or one-sided attribution fails."""
     import asyncio
     import time
 
@@ -1817,9 +1818,8 @@ def step(self, action):
 
     mock_make.side_effect = lambda *args, **kwargs: SlowEnv()
 
-    # Run env steps on the executor (as in production, where ``max_env_workers`` defaults to 32).
-    # With inline execution a blocking ``env.step`` stalls the shared event loop, so a sibling
-    # trajectory's ``await generate()`` would absorb that stall and inflate its measured LLM time.
+    # Run env.step on the executor as in production; a blocking step inline would stall the event
+    # loop and inflate a sibling trajectory's measured LLM time.
     mock_env_cfg.max_env_workers = 4
 
     cfg = GeneratorConfig()
@@ -1879,16 +1879,10 @@ def step(self, action):
     for llm_t, env_t, e2e_t in zip(llm_times, env_times, generator_output["trajectory_generation_times"]):
         assert llm_t + env_t <= e2e_t + 1e-6
 
-    # Regression: every logging path (async trainer, eval, step-wise) re-aggregates through
-    # concatenate_generator_outputs. The split must be recomputed there from the raw per-trajectory
-    # lists; when it was not, p90 and frac fell into a substring sum() fallback that inflated p90 by
-    # ~num_groups and summed the fractions past 1.0. Concatenating identical groups leaves the
-    # per-trajectory distribution unchanged, so the aggregates must match the single-group values.
+    # Regression: concatenate_generator_outputs (every logging path) must recompute the split from
+    # the raw per-trajectory lists, not sum per-group p90/frac past the sample max and 1.0.
     concatenated = generator_utils.concatenate_generator_outputs([generator_output, generator_output])
     concat_metrics = concatenated["rollout_metrics"]
-    # Concatenating two groups doubles the sample, so p90 is the percentile of the combined raw
-    # list, not the single-group p90. The sum() fallback instead added the two group p90s, exceeding
-    # the sample max, and summed the fractions past 1.0.
     assert 0.5 < concat_metrics["generate/frac_time_in_env"] < 1.0
     assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item())
     assert concat_metrics["generate/trajectory_env_time_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item())

From 48208fe8510dc3674107918f6b59561ca7ed2ef4 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Fri, 17 Jul 2026 23:28:29 +0000
Subject: [PATCH 08/11] [logging] Account for env setup and overhead in
 trajectory time

Closes the trajectory-time stack so everything inside
trajectory_completion_time is explicit:

  e2e = llm_time + env_time + env_setup_time + overhead_time

env_setup_time brackets skyrl_gym.make plus env.init, which run inside
the e2e clock but outside env_time; for sandboxed tool envs setup can be
seconds and was indistinguishable from tokenization overhead.
overhead_time is the derived remainder: tokenization, chat templating,
output assembly, and event-loop scheduling. Emitted as
generate/trajectory_{env_setup,overhead}_time_{mean,p90,max} on both the
per-generate and concatenated paths.
---
 skyrl/train/generators/base.py                |  5 +++--
 skyrl/train/generators/skyrl_gym_generator.py | 16 ++++++++++++++-
 skyrl/train/generators/utils.py               | 19 ++++++++++++++++++
 .../generators/test_generator_output_utils.py |  1 +
 .../generators/test_skyrl_gym_generator.py    | 20 +++++++++++++------
 5 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index b2ccc102d6..d017b833df 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -46,10 +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]]
-    # Split of ``trajectory_generation_times`` into engine-wait vs ``env.step()`` time (seconds).
-    # None if any trajectory did not record it.
+    # Split of ``trajectory_generation_times`` into engine-wait vs ``env.step()`` time (seconds),
+    # plus env construction/init time. None if any trajectory did not record it.
     trajectory_llm_times: Optional[List[float]]
     trajectory_env_times: Optional[List[float]]
+    trajectory_env_setup_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]]
diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index a317a1d6c0..56fd25915c 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -60,6 +60,8 @@ class TrajectoryOutput:
     llm_time: Optional[float] = None
     # Wall-clock time (seconds) in ``env.step()``, summed over turns. None if untracked.
     env_time: Optional[float] = None
+    # Wall-clock time (seconds) constructing and initializing the env. None if untracked.
+    env_setup_time: Optional[float] = None
 
 
 @dataclass
@@ -70,9 +72,10 @@ 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
-    # Same split as TrajectoryOutput.llm_time / env_time.
+    # Same split as TrajectoryOutput.llm_time / env_time / env_setup_time.
     llm_time: Optional[float] = None
     env_time: Optional[float] = None
+    env_setup_time: Optional[float] = None
 
 
 @dataclass
@@ -353,6 +356,7 @@ async def agent_loop(
             env_extras = self._setup_env_extras(env_class, env_extras, sampling_params, trajectory_id)
 
             env_config = getattr(self.skyrl_gym_cfg, env_class, dict())
+            env_setup_start_time = time.monotonic()
             env = skyrl_gym.make(env_class, env_config=env_config, extras=env_extras)
 
             # Instantiate chat_history and chat_end_index, which are only used if `retokenize_chat_history==True`.
@@ -361,6 +365,7 @@ async def agent_loop(
 
             # init() returns the first prompt to be given to the model, and optional metadata dict
             chat_history, _ = await self._run_in_executor_if_available(env.init, chat_history)
+            env_setup_time_s = time.monotonic() - env_setup_start_time
             initial_chat_history_length = len(chat_history)
             initial_input_ids = self.tokenizer.apply_chat_template(
                 chat_history,
@@ -629,6 +634,7 @@ async def agent_loop(
             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
+            agent_loop_output.env_setup_time = env_setup_time_s
             return agent_loop_output
 
         finally:
@@ -894,6 +900,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
         trajectory_generation_times_per_prompt = _optional_times(all_outputs, "e2e_time")
         trajectory_llm_times_per_prompt = _optional_times(all_outputs, "llm_time")
         trajectory_env_times_per_prompt = _optional_times(all_outputs, "env_time")
+        trajectory_env_setup_times_per_prompt = _optional_times(all_outputs, "env_setup_time")
 
         if self.generator_cfg.step_wise_trajectories:
             responses = []
@@ -908,6 +915,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             out_trajectory_generation_times = []
             out_trajectory_llm_times = []
             out_trajectory_env_times = []
+            out_trajectory_env_setup_times = []
             for i, output in enumerate(all_outputs):
                 for j, step_output in enumerate(output.step_outputs):
                     responses.append(step_output.response_ids)
@@ -923,6 +931,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                     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))
+                    out_trajectory_env_setup_times.append(getattr(output, "env_setup_time", None))
             # Keep aligned with the per-prompt None handling:
             if not trajectory_generation_times_per_prompt:
                 out_trajectory_generation_times = None
@@ -930,6 +939,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                 out_trajectory_llm_times = None
             if not trajectory_env_times_per_prompt:
                 out_trajectory_env_times = None
+            if not trajectory_env_setup_times_per_prompt:
+                out_trajectory_env_setup_times = None
             env_classes = out_env_classes
         else:
             responses = [output.response_ids for output in all_outputs]
@@ -944,6 +955,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             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
+            out_trajectory_env_setup_times = trajectory_env_setup_times_per_prompt
 
         has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs)
         pixel_values = (
@@ -986,6 +998,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             trajectory_completion_times=trajectory_generation_times_per_prompt,
             trajectory_llm_times=trajectory_llm_times_per_prompt,
             trajectory_env_times=trajectory_env_times_per_prompt,
+            trajectory_env_setup_times=trajectory_env_setup_times_per_prompt,
         )
 
         if self.generator_cfg.zero_reward_on_non_stop:
@@ -1009,6 +1022,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
             "trajectory_generation_times": out_trajectory_generation_times,
             "trajectory_llm_times": out_trajectory_llm_times,
             "trajectory_env_times": out_trajectory_env_times,
+            "trajectory_env_setup_times": out_trajectory_env_setup_times,
             "rollout_expert_indices": rollout_expert_indices,
             "is_last_step": is_last_step,
             "env_metrics": env_metrics,
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 3137f5f319..fa30192003 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -283,6 +283,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         "trajectory_generation_times": _concat_field(generator_outputs, "trajectory_generation_times"),
         "trajectory_llm_times": _concat_field(generator_outputs, "trajectory_llm_times"),
         "trajectory_env_times": _concat_field(generator_outputs, "trajectory_env_times"),
+        "trajectory_env_setup_times": _concat_field(generator_outputs, "trajectory_env_setup_times"),
     }
 
     # propagate additional keys with list values as-is
@@ -297,6 +298,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
     trajectory_generation_times = _last_step_only(result.get("trajectory_generation_times"), is_last_step)
     trajectory_llm_times = _last_step_only(result.get("trajectory_llm_times"), is_last_step)
     trajectory_env_times = _last_step_only(result.get("trajectory_env_times"), is_last_step)
+    trajectory_env_setup_times = _last_step_only(result.get("trajectory_env_setup_times"), is_last_step)
 
     # Re-aggregate rollout metrics; the extra_keys fallback below cannot aggregate a p90 or a ratio.
     rollout_metrics = get_rollout_metrics(
@@ -306,6 +308,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         trajectory_completion_times=trajectory_generation_times,
         trajectory_llm_times=trajectory_llm_times,
         trajectory_env_times=trajectory_env_times,
+        trajectory_env_setup_times=trajectory_env_setup_times,
     )
 
     # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only
@@ -410,6 +413,7 @@ def get_rollout_metrics(
     trajectory_completion_times: Optional[List[float]] = None,
     trajectory_llm_times: Optional[List[float]] = None,
     trajectory_env_times: Optional[List[float]] = None,
+    trajectory_env_setup_times: Optional[List[float]] = None,
 ):
     """
     Computes rollout metrics including token statistics and optional environment-specific metrics.
@@ -426,6 +430,8 @@ def get_rollout_metrics(
             engine, summed over turns
         trajectory_env_times: Optional per-trajectory time (seconds) spent in ``env.step()``, summed
             over turns
+        trajectory_env_setup_times: Optional per-trajectory time (seconds) spent constructing and
+            initializing the env
 
     Returns:
         Dictionary of aggregated metrics
@@ -483,12 +489,25 @@ def get_rollout_metrics(
     _add_time_stats(rollout_metrics, "completion", trajectory_completion_times)
     llm_sum = _add_time_stats(rollout_metrics, "llm", trajectory_llm_times)
     env_sum = _add_time_stats(rollout_metrics, "env", trajectory_env_times)
+    _add_time_stats(rollout_metrics, "env_setup", trajectory_env_setup_times)
 
     if llm_sum is not None and env_sum is not None and llm_sum + env_sum > 0:
         # Time-weighted (sum over sum), not a mean of per-trajectory ratios, so long trajectories
         # count proportionally.
         rollout_metrics["generate/frac_time_in_env"] = env_sum / (llm_sum + env_sum)
 
+    # Everything in e2e not attributed to the engine, the env, or env setup: tokenization, chat
+    # templating, output assembly, and event-loop scheduling. Closes the trajectory-time stack.
+    if trajectory_completion_times and trajectory_llm_times and trajectory_env_times:
+        setup_times = trajectory_env_setup_times or [0.0] * len(trajectory_completion_times)
+        overhead = [
+            e2e - llm - env - setup
+            for e2e, llm, env, setup in zip(
+                trajectory_completion_times, trajectory_llm_times, trajectory_env_times, setup_times
+            )
+        ]
+        _add_time_stats(rollout_metrics, "overhead", overhead)
+
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)
         for i, metrics in enumerate(env_metrics):
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 84e3a87aec..827dd4510d 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -34,6 +34,7 @@ def test_generator_output_concatenation():
         "trajectory_generation_times",
         "trajectory_llm_times",
         "trajectory_env_times",
+        "trajectory_env_setup_times",
         "is_last_step",
         "env_metrics",
         "pixel_values",
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index d9c0c1bc58..5421f5b5ba 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1752,8 +1752,8 @@ def step(self, action):
     )
     assert rollout_metrics["generate/trajectory_completion_time_max"] == pytest.approx(np.max(expected).item())
 
-    # The llm/env split is stored per step and replicated like the completion times above.
-    for key in ("trajectory_llm_times", "trajectory_env_times"):
+    # The llm/env/setup split is stored per step and replicated like the completion times above.
+    for key in ("trajectory_llm_times", "trajectory_env_times", "trajectory_env_setup_times"):
         split_times = generator_output[key]
         assert split_times is not None
         assert len(split_times) == num_steps
@@ -1776,6 +1776,7 @@ async def test_llm_vs_env_time_split_metrics(mock_make, mock_tokenizer, mock_llm
 
     llm_sleep_s = 0.02
     env_sleep_s = 0.06
+    setup_sleep_s = 0.03
     num_turns = 2
 
     mock_tokenizer.eos_token_id = 4
@@ -1805,6 +1806,7 @@ def __init__(self):
             self.turns = 0
 
         def init(self, prompt):
+            time.sleep(setup_sleep_s)
             return prompt, {}
 
         def step(self, action):
@@ -1856,9 +1858,11 @@ def step(self, action):
 
     llm_times = spy.call_args.kwargs["trajectory_llm_times"]
     env_times = spy.call_args.kwargs["trajectory_env_times"]
-    assert llm_times is not None and env_times is not None
+    setup_times = spy.call_args.kwargs["trajectory_env_setup_times"]
+    assert llm_times is not None and env_times is not None and setup_times is not None
     assert len(llm_times) == num_trajectories
     assert len(env_times) == num_trajectories
+    assert all(t >= setup_sleep_s for t in setup_times)
 
     # Each side is at least its per-turn sleep summed over turns, and neither swallows the other.
     for llm_t, env_t in zip(llm_times, env_times):
@@ -1870,14 +1874,18 @@ def step(self, action):
     metrics = generator_output["rollout_metrics"]
     assert metrics["generate/trajectory_llm_time_mean"] >= llm_sleep_s * num_turns
     assert metrics["generate/trajectory_env_time_mean"] >= env_sleep_s * num_turns
+    assert metrics["generate/trajectory_env_setup_time_mean"] >= setup_sleep_s
+    assert metrics["generate/trajectory_overhead_time_mean"] >= 0.0
 
     # env_sleep / (env_sleep + llm_sleep) = 0.75; allow generous headroom for scheduling overhead
     # attributed to the engine wait, but it must clearly indicate an environment-bound rollout.
     assert 0.5 < metrics["generate/frac_time_in_env"] < 1.0
 
-    # The split accounts for real time inside the trajectory and never exceeds its end-to-end time.
-    for llm_t, env_t, e2e_t in zip(llm_times, env_times, generator_output["trajectory_generation_times"]):
-        assert llm_t + env_t <= e2e_t + 1e-6
+    # The bands never exceed the trajectory's end-to-end time; overhead is the exact remainder.
+    for llm_t, env_t, setup_t, e2e_t in zip(
+        llm_times, env_times, setup_times, generator_output["trajectory_generation_times"]
+    ):
+        assert llm_t + env_t + setup_t <= e2e_t + 1e-6
 
     # Regression: concatenate_generator_outputs (every logging path) must recompute the split from
     # the raw per-trajectory lists, not sum per-group p90/frac past the sample max and 1.0.

From c9ba6fcbba63ea265990cd20e2296a2f32fcc876 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Fri, 17 Jul 2026 23:50:57 +0000
Subject: [PATCH 09/11] [logging] Trajectory-time accounting review pass

- Require all four raw lists before computing overhead instead of
  zero-filling a missing env_setup list, which silently folded real
  setup time into overhead for any generator reporting only llm and env.
- Name env teardown in the overhead comment; env.close() runs inside the
  e2e window but outside the env_time and env_setup_time brackets.
- Vectorize the overhead computation to match the rest of the function.
- Assert setup_times length and the new bands on the concatenated path.
---
 skyrl/train/generators/utils.py               | 21 +++++++++----------
 .../generators/test_skyrl_gym_generator.py    |  5 +++++
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index fa30192003..7a29f8c9c1 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -496,17 +496,16 @@ def get_rollout_metrics(
         # count proportionally.
         rollout_metrics["generate/frac_time_in_env"] = env_sum / (llm_sum + env_sum)
 
-    # Everything in e2e not attributed to the engine, the env, or env setup: tokenization, chat
-    # templating, output assembly, and event-loop scheduling. Closes the trajectory-time stack.
-    if trajectory_completion_times and trajectory_llm_times and trajectory_env_times:
-        setup_times = trajectory_env_setup_times or [0.0] * len(trajectory_completion_times)
-        overhead = [
-            e2e - llm - env - setup
-            for e2e, llm, env, setup in zip(
-                trajectory_completion_times, trajectory_llm_times, trajectory_env_times, setup_times
-            )
-        ]
-        _add_time_stats(rollout_metrics, "overhead", overhead)
+    # Remainder of e2e not attributed to the engine, env steps, or env setup. Tokenization, chat
+    # templating, output assembly, env teardown, and event-loop scheduling.
+    if trajectory_completion_times and trajectory_llm_times and trajectory_env_times and trajectory_env_setup_times:
+        overhead = (
+            np.array(trajectory_completion_times, dtype=np.float64)
+            - np.array(trajectory_llm_times, dtype=np.float64)
+            - np.array(trajectory_env_times, dtype=np.float64)
+            - np.array(trajectory_env_setup_times, dtype=np.float64)
+        )
+        _add_time_stats(rollout_metrics, "overhead", overhead.tolist())
 
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 5421f5b5ba..1d5b445f47 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1862,6 +1862,7 @@ def step(self, action):
     assert llm_times is not None and env_times is not None and setup_times is not None
     assert len(llm_times) == num_trajectories
     assert len(env_times) == num_trajectories
+    assert len(setup_times) == num_trajectories
     assert all(t >= setup_sleep_s for t in setup_times)
 
     # Each side is at least its per-turn sleep summed over turns, and neither swallows the other.
@@ -1894,3 +1895,7 @@ def step(self, action):
     assert 0.5 < concat_metrics["generate/frac_time_in_env"] < 1.0
     assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item())
     assert concat_metrics["generate/trajectory_env_time_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item())
+    assert concat_metrics["generate/trajectory_env_setup_time_p90"] == pytest.approx(
+        np.percentile(setup_times * 2, 90).item()
+    )
+    assert concat_metrics["generate/trajectory_overhead_time_mean"] >= 0.0

From e9e7660e670418b244ed6d5aa4abff0dc9156d12 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Tue, 21 Jul 2026 21:00:59 +0000
Subject: [PATCH 10/11] [logging] Rename the trajectory time remainder band to
 "other"

The remainder of e2e after llm/env/env_setup is not pure overhead: it
holds tokenization and chat templating. Rename the band from overhead to
other so a large value reads as a problem to investigate rather than
expected slack. Metric generate/trajectory_overhead_time_* becomes
generate/trajectory_other_time_*.
---
 skyrl/train/generators/utils.py                    | 4 ++--
 tests/train/generators/test_skyrl_gym_generator.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 7a29f8c9c1..89fa00eed8 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -499,13 +499,13 @@ def get_rollout_metrics(
     # Remainder of e2e not attributed to the engine, env steps, or env setup. Tokenization, chat
     # templating, output assembly, env teardown, and event-loop scheduling.
     if trajectory_completion_times and trajectory_llm_times and trajectory_env_times and trajectory_env_setup_times:
-        overhead = (
+        other = (
             np.array(trajectory_completion_times, dtype=np.float64)
             - np.array(trajectory_llm_times, dtype=np.float64)
             - np.array(trajectory_env_times, dtype=np.float64)
             - np.array(trajectory_env_setup_times, dtype=np.float64)
         )
-        _add_time_stats(rollout_metrics, "overhead", overhead.tolist())
+        _add_time_stats(rollout_metrics, "other", other.tolist())
 
     if env_metrics is not None and env_classes is not None:
         env_to_metrics = defaultdict(list)
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 1d5b445f47..7eced24b20 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1876,13 +1876,13 @@ def step(self, action):
     assert metrics["generate/trajectory_llm_time_mean"] >= llm_sleep_s * num_turns
     assert metrics["generate/trajectory_env_time_mean"] >= env_sleep_s * num_turns
     assert metrics["generate/trajectory_env_setup_time_mean"] >= setup_sleep_s
-    assert metrics["generate/trajectory_overhead_time_mean"] >= 0.0
+    assert metrics["generate/trajectory_other_time_mean"] >= 0.0
 
     # env_sleep / (env_sleep + llm_sleep) = 0.75; allow generous headroom for scheduling overhead
     # attributed to the engine wait, but it must clearly indicate an environment-bound rollout.
     assert 0.5 < metrics["generate/frac_time_in_env"] < 1.0
 
-    # The bands never exceed the trajectory's end-to-end time; overhead is the exact remainder.
+    # The bands never exceed the trajectory's end-to-end time; "other" is the exact remainder.
     for llm_t, env_t, setup_t, e2e_t in zip(
         llm_times, env_times, setup_times, generator_output["trajectory_generation_times"]
     ):
@@ -1898,4 +1898,4 @@ def step(self, action):
     assert concat_metrics["generate/trajectory_env_setup_time_p90"] == pytest.approx(
         np.percentile(setup_times * 2, 90).item()
     )
-    assert concat_metrics["generate/trajectory_overhead_time_mean"] >= 0.0
+    assert concat_metrics["generate/trajectory_other_time_mean"] >= 0.0

From e15eb9202fd04e2a7245e08f56fea069ffe5d70d Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Tue, 21 Jul 2026 21:14:08 +0000
Subject: [PATCH 11/11] [logging] Group time bands under the trajectory_time_
 metric namespace

Emit generate/trajectory_time_<band>_<stat> instead of
generate/trajectory_<band>_time_<stat> so the bands (completion, llm,
env, env_setup, other) read as components of one trajectory_time total
and one glob pulls the whole breakdown.
---
 skyrl/train/generators/utils.py               |  6 ++--
 .../generators/test_skyrl_gym_generator.py    | 28 +++++++++----------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 89fa00eed8..98ae5c2138 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -396,9 +396,9 @@ def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[
     arr = np.array(times, dtype=np.float64)
     rollout_metrics.update(
         {
-            f"generate/trajectory_{name}_time_mean": np.mean(arr).item(),
-            f"generate/trajectory_{name}_time_p90": np.percentile(arr, 90).item(),
-            f"generate/trajectory_{name}_time_max": np.max(arr).item(),
+            f"generate/trajectory_time_{name}_mean": np.mean(arr).item(),
+            f"generate/trajectory_time_{name}_p90": np.percentile(arr, 90).item(),
+            f"generate/trajectory_time_{name}_max": np.max(arr).item(),
         }
     )
     return np.sum(arr).item()
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 7eced24b20..b20d986797 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1740,17 +1740,17 @@ def step(self, action):
     # Aggregate stats are present and computed over the per-prompt times.
     rollout_metrics = generator_output["rollout_metrics"]
     for key in (
-        "generate/trajectory_completion_time_mean",
-        "generate/trajectory_completion_time_p90",
-        "generate/trajectory_completion_time_max",
+        "generate/trajectory_time_completion_mean",
+        "generate/trajectory_time_completion_p90",
+        "generate/trajectory_time_completion_max",
     ):
         assert key in rollout_metrics, f"missing metric {key}"
     expected = np.array(metrics_times, dtype=np.float64)
-    assert rollout_metrics["generate/trajectory_completion_time_mean"] == pytest.approx(np.mean(expected).item())
-    assert rollout_metrics["generate/trajectory_completion_time_p90"] == pytest.approx(
+    assert rollout_metrics["generate/trajectory_time_completion_mean"] == pytest.approx(np.mean(expected).item())
+    assert rollout_metrics["generate/trajectory_time_completion_p90"] == pytest.approx(
         np.percentile(expected, 90).item()
     )
-    assert rollout_metrics["generate/trajectory_completion_time_max"] == pytest.approx(np.max(expected).item())
+    assert rollout_metrics["generate/trajectory_time_completion_max"] == pytest.approx(np.max(expected).item())
 
     # The llm/env/setup split is stored per step and replicated like the completion times above.
     for key in ("trajectory_llm_times", "trajectory_env_times", "trajectory_env_setup_times"):
@@ -1873,10 +1873,10 @@ def step(self, action):
         assert env_t > llm_t
 
     metrics = generator_output["rollout_metrics"]
-    assert metrics["generate/trajectory_llm_time_mean"] >= llm_sleep_s * num_turns
-    assert metrics["generate/trajectory_env_time_mean"] >= env_sleep_s * num_turns
-    assert metrics["generate/trajectory_env_setup_time_mean"] >= setup_sleep_s
-    assert metrics["generate/trajectory_other_time_mean"] >= 0.0
+    assert metrics["generate/trajectory_time_llm_mean"] >= llm_sleep_s * num_turns
+    assert metrics["generate/trajectory_time_env_mean"] >= env_sleep_s * num_turns
+    assert metrics["generate/trajectory_time_env_setup_mean"] >= setup_sleep_s
+    assert metrics["generate/trajectory_time_other_mean"] >= 0.0
 
     # env_sleep / (env_sleep + llm_sleep) = 0.75; allow generous headroom for scheduling overhead
     # attributed to the engine wait, but it must clearly indicate an environment-bound rollout.
@@ -1893,9 +1893,9 @@ def step(self, action):
     concatenated = generator_utils.concatenate_generator_outputs([generator_output, generator_output])
     concat_metrics = concatenated["rollout_metrics"]
     assert 0.5 < concat_metrics["generate/frac_time_in_env"] < 1.0
-    assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item())
-    assert concat_metrics["generate/trajectory_env_time_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item())
-    assert concat_metrics["generate/trajectory_env_setup_time_p90"] == pytest.approx(
+    assert concat_metrics["generate/trajectory_time_llm_p90"] == pytest.approx(np.percentile(llm_times * 2, 90).item())
+    assert concat_metrics["generate/trajectory_time_env_p90"] == pytest.approx(np.percentile(env_times * 2, 90).item())
+    assert concat_metrics["generate/trajectory_time_env_setup_p90"] == pytest.approx(
         np.percentile(setup_times * 2, 90).item()
     )
-    assert concat_metrics["generate/trajectory_other_time_mean"] >= 0.0
+    assert concat_metrics["generate/trajectory_time_other_mean"] >= 0.0