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/17] [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/17] [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/17] [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/17] [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/17] [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/17] [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/17] [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 15d0c579a674a464231b16f90f5e01a7b28c0c81 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:03:00 +0000
Subject: [PATCH 08/17] [logging] Carry the time split as one dict instead of
 parallel fields

Review feedback: the llm and env times traveled as parallel named fields
(llm_time/env_time on two dataclasses, trajectory_llm_times/
trajectory_env_times on GeneratorOutput), so every hop re-documented and
re-plumbed each field, and the omit-if-any-None handling needed its own
helper per field.

One time_splits dict per trajectory (keys "llm", "env") and one
trajectory_time_splits dict-of-lists on GeneratorOutput replace them.
The split semantics are documented once, on TrajectoryOutput.time_splits.
_optional_times is gone; get_rollout_metrics loops the dict through
_add_time_stats generically, so a future component is a new key rather
than a new field at every layer. Metric names are unchanged.
---
 skyrl/train/generators/base.py                |  7 +--
 skyrl/train/generators/skyrl_gym_generator.py | 63 ++++++++-----------
 skyrl/train/generators/utils.py               | 34 +++++-----
 .../generators/test_generator_output_utils.py |  3 +-
 .../generators/test_skyrl_gym_generator.py    | 14 ++---
 5 files changed, 58 insertions(+), 63 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index b2ccc102d6..b7b4973be8 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -46,10 +46,9 @@ 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.
-    trajectory_llm_times: Optional[List[float]]
-    trajectory_env_times: Optional[List[float]]
+    # Engine and env time splits of ``trajectory_generation_times``, one list entry per trajectory,
+    # e.g. {"llm": [...], "env": [...]}. None if any trajectory did not record them.
+    trajectory_time_splits: Optional[Dict[str, 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..e7a6eb4bf9 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -56,10 +56,10 @@ 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) awaiting the inference engine, summed over turns. None if untracked.
-    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 split of ``e2e_time`` by component (seconds), summed over turns: "llm" is time in
+    # the inference-engine call, "env" is time in ``env.step()``. The remainder of ``e2e_time`` is
+    # tokenization, chat templating, and event-loop scheduling. None if untracked.
+    time_splits: Optional[Dict[str, float]] = None
 
 
 @dataclass
@@ -70,9 +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
-    # Same split as TrajectoryOutput.llm_time / env_time.
-    llm_time: Optional[float] = None
-    env_time: Optional[float] = None
+    # Same as TrajectoryOutput.time_splits.
+    time_splits: Optional[Dict[str, float]] = None
 
 
 @dataclass
@@ -147,10 +146,11 @@ 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
+def _split_lists(time_splits: List[Optional[Dict[str, float]]]) -> Optional[Dict[str, List[float]]]:
+    """Per-component lists from per-trajectory time splits, or None if any trajectory lacks them."""
+    if not time_splits or any(s is None for s in time_splits):
+        return None
+    return {name: [s[name] for s in time_splits] for name in time_splits[0]}
 
 
 class SkyRLGymGenerator(GeneratorInterface):
@@ -334,10 +334,7 @@ async def agent_loop(
             rollout_logprobs: Optional[List[float]]
         """
         agent_loop_start_time = time.monotonic()
-        # 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
+        time_splits = {"llm": 0.0, "env": 0.0}
 
         session_id = (
             f"{trajectory_id.instance_id}_{trajectory_id.repetition_id}" if trajectory_id is not None else uuid4().hex
@@ -428,7 +425,7 @@ async def agent_loop(
                 )
                 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
+                time_splits["llm"] += 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]
@@ -464,7 +461,7 @@ async def agent_loop(
                 # 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
+                time_splits["env"] += 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"]
@@ -627,8 +624,7 @@ 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
+            agent_loop_output.time_splits = time_splits
             return agent_loop_output
 
         finally:
@@ -891,9 +887,12 @@ 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).
-        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")
+        # ``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
+        trajectory_time_splits_per_prompt = _split_lists([getattr(o, "time_splits", None) for o in all_outputs])
 
         if self.generator_cfg.step_wise_trajectories:
             responses = []
@@ -906,8 +905,7 @@ 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 = []
+            out_step_time_splits = []
             for i, output in enumerate(all_outputs):
                 for j, step_output in enumerate(output.step_outputs):
                     responses.append(step_output.response_ids)
@@ -921,15 +919,11 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False
                     out_env_classes.append(env_classes[i])
                     # 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))
+                    out_step_time_splits.append(getattr(output, "time_splits", 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
+            out_trajectory_time_splits = _split_lists(out_step_time_splits)
             env_classes = out_env_classes
         else:
             responses = [output.response_ids for output in all_outputs]
@@ -942,8 +936,7 @@ 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
+            out_trajectory_time_splits = trajectory_time_splits_per_prompt
 
         has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs)
         pixel_values = (
@@ -984,8 +977,7 @@ 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,
+            trajectory_time_splits=trajectory_time_splits_per_prompt,
         )
 
         if self.generator_cfg.zero_reward_on_non_stop:
@@ -1007,8 +999,7 @@ 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,
+            "trajectory_time_splits": out_trajectory_time_splits,
             "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..fea2a9e410 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -247,6 +247,14 @@ def _concat_field(generator_outputs: List[GeneratorOutput], key: str) -> Optiona
     return _flatten_field(generator_outputs, key)
 
 
+def _concat_time_splits(generator_outputs: List[GeneratorOutput]) -> Optional[Dict[str, List[float]]]:
+    """Concatenate the per-trajectory time-split lists across outputs, per component."""
+    first = generator_outputs[0].get("trajectory_time_splits")
+    if first is None:
+        return None
+    return {name: [t for go in generator_outputs for t in go["trajectory_time_splits"][name]] for name in first}
+
+
 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:
@@ -281,8 +289,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         "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"),
+        "trajectory_time_splits": _concat_time_splits(generator_outputs),
     }
 
     # propagate additional keys with list values as-is
@@ -295,8 +302,9 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], 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)
+    trajectory_time_splits = result.get("trajectory_time_splits")
+    if trajectory_time_splits is not None and is_last_step:
+        trajectory_time_splits = {k: _last_step_only(v, is_last_step) for k, v in trajectory_time_splits.items()}
 
     # Re-aggregate rollout metrics; the extra_keys fallback below cannot aggregate a p90 or a ratio.
     rollout_metrics = get_rollout_metrics(
@@ -304,8 +312,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], step
         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,
+        trajectory_time_splits=trajectory_time_splits,
     )
 
     # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only
@@ -408,8 +415,7 @@ 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,
+    trajectory_time_splits: Optional[Dict[str, List[float]]] = None,
 ):
     """
     Computes rollout metrics including token statistics and optional environment-specific metrics.
@@ -422,10 +428,8 @@ 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
+        trajectory_time_splits: Optional per-component split of the completion times, e.g.
+            {"llm": [...], "env": [...]}, one entry per trajectory
 
     Returns:
         Dictionary of aggregated metrics
@@ -481,8 +485,10 @@ 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)
+    split_sums = {
+        name: _add_time_stats(rollout_metrics, name, times) for name, times in (trajectory_time_splits or {}).items()
+    }
+    llm_sum, env_sum = split_sums.get("llm"), split_sums.get("env")
 
     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
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 84e3a87aec..a45eb84bc2 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -32,8 +32,7 @@ def test_generator_output_concatenation():
         # optional but present in the signature
         "trajectory_ids",
         "trajectory_generation_times",
-        "trajectory_llm_times",
-        "trajectory_env_times",
+        "trajectory_time_splits",
         "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..47e4291151 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1752,10 +1752,10 @@ 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
+    # The time splits are stored per step and replicated like the completion times above.
+    time_splits = generator_output["trajectory_time_splits"]
+    assert time_splits is not None and set(time_splits) == {"llm", "env"}
+    for split_times in time_splits.values():
         assert len(split_times) == num_steps
         assert split_times[0] == split_times[1] and split_times[2] == split_times[3]
 
@@ -1854,9 +1854,9 @@ def step(self, action):
     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
+    time_splits = spy.call_args.kwargs["trajectory_time_splits"]
+    assert time_splits is not None and set(time_splits) == {"llm", "env"}
+    llm_times, env_times = time_splits["llm"], time_splits["env"]
     assert len(llm_times) == num_trajectories
     assert len(env_times) == num_trajectories
 

From d7a6f04a4a0a7610a223dc66bac57634c85ded4d Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:12:45 +0000
Subject: [PATCH 09/17] [logging] Drop frac_time_in_env, add the overhead band

frac_time_in_env is derivable exactly as env_mean / (llm_mean + env_mean)
from metrics already logged, so it added convenience, not information.

Replace the "need not sum to e2e" caveat with the metric that closes the
stack: generate/trajectory_overhead_time_{mean,p90,max}, computed as e2e
minus every recorded split component. Generic over the split keys, so a
future component reallocates time out of overhead without touching this
code. The test asserts overhead is the exact remainder and survives
concatenation.
---
 skyrl/train/generators/utils.py               | 27 ++++++++++---------
 .../generators/test_skyrl_gym_generator.py    | 16 ++++++-----
 2 files changed, 23 insertions(+), 20 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index fea2a9e410..3d9614e708 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -393,10 +393,10 @@ 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."""
+def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[List[float]]) -> None:
+    """Add mean/p90/max stats for a per-trajectory time list under generate/trajectory_<name>_time_*."""
     if not times:
-        return None
+        return
     arr = np.array(times, dtype=np.float64)
     rollout_metrics.update(
         {
@@ -405,7 +405,6 @@ def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[
             f"generate/trajectory_{name}_time_max": np.max(arr).item(),
         }
     )
-    return np.sum(arr).item()
 
 
 def get_rollout_metrics(
@@ -485,15 +484,17 @@ def get_rollout_metrics(
             )
 
     _add_time_stats(rollout_metrics, "completion", trajectory_completion_times)
-    split_sums = {
-        name: _add_time_stats(rollout_metrics, name, times) for name, times in (trajectory_time_splits or {}).items()
-    }
-    llm_sum, env_sum = split_sums.get("llm"), split_sums.get("env")
-
-    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)
+    for name, times in (trajectory_time_splits or {}).items():
+        _add_time_stats(rollout_metrics, name, times)
+
+    # Overhead is e2e minus every recorded split component: env construction and teardown,
+    # tokenization, chat templating, and event-loop scheduling. It closes the stack, so
+    # unattributed time is visible instead of silent.
+    if trajectory_completion_times and trajectory_time_splits:
+        overhead = np.array(trajectory_completion_times, dtype=np.float64)
+        for times in trajectory_time_splits.values():
+            overhead = overhead - np.array(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 47e4291151..12d32868ae 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1870,19 +1870,21 @@ 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_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.
+    # The splits never exceed the trajectory's end-to-end time; overhead is the exact remainder.
+    overhead_times = []
     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
+        overhead_times.append(e2e_t - llm_t - env_t)
+    assert metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(np.mean(overhead_times).item())
 
     # 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.
+    # the raw per-trajectory lists, not sum per-group aggregates past the sample max.
     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_overhead_time_p90"] == pytest.approx(
+        np.percentile(overhead_times * 2, 90).item()
+    )

From bea12c69fc12671b6b505b200c22b9c1fccc7162 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:18:31 +0000
Subject: [PATCH 10/17] [logging] Split the timing test by layer

The generate() harness test conflated timing attribution with metric
math. Attribution keeps the integration test, now reading the splits
straight off GeneratorOutput instead of spying on get_rollout_metrics.
The metric math and concatenation move to exact-value tests over
hand-built outputs in test_generator_output_utils.py. Also drop the
overhead code comment.
---
 skyrl/train/generators/utils.py               |  3 --
 .../generators/test_generator_output_utils.py | 43 +++++++++++++++++++
 .../generators/test_skyrl_gym_generator.py    | 38 +++-------------
 3 files changed, 48 insertions(+), 36 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 3d9614e708..1d81a3ca87 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -487,9 +487,6 @@ def get_rollout_metrics(
     for name, times in (trajectory_time_splits or {}).items():
         _add_time_stats(rollout_metrics, name, times)
 
-    # Overhead is e2e minus every recorded split component: env construction and teardown,
-    # tokenization, chat templating, and event-loop scheduling. It closes the stack, so
-    # unattributed time is visible instead of silent.
     if trajectory_completion_times and trajectory_time_splits:
         overhead = np.array(trajectory_completion_times, dtype=np.float64)
         for times in trajectory_time_splits.values():
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index a45eb84bc2..6ea508189d 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -12,6 +12,7 @@
     compute_turn_token_counts,
     concatenate_generator_outputs,
     get_metrics_from_generator_output,
+    get_rollout_metrics,
     merge_stepwise_output,
 )
 from skyrl.train.utils.utils import validate_cfg
@@ -94,6 +95,48 @@ def test_generator_output_concatenation():
         np.testing.assert_allclose(concatenated_output["rollout_metrics"][key], value)
 
 
+def test_time_split_rollout_metrics():
+    metrics = get_rollout_metrics(
+        responses=[[1, 2]] * 4,
+        rewards=[1.0] * 4,
+        trajectory_completion_times=[10.0, 20.0, 30.0, 40.0],
+        trajectory_time_splits={"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]},
+    )
+    assert metrics["generate/trajectory_llm_time_mean"] == 10.0
+    assert metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile([4.0, 8.0, 12.0, 16.0], 90))
+    assert metrics["generate/trajectory_llm_time_max"] == 16.0
+    assert metrics["generate/trajectory_env_time_mean"] == 12.5
+    # Overhead is the exact per-trajectory remainder, here [1.0, 2.0, 3.0, 4.0].
+    assert metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(2.5)
+    assert metrics["generate/trajectory_overhead_time_max"] == pytest.approx(4.0)
+
+
+def test_time_splits_concatenation():
+    def make_output(times, splits) -> GeneratorOutput:
+        return {
+            "prompt_token_ids": [[1]] * len(times),
+            "response_ids": [[1, 2]] * len(times),
+            "rewards": [1.0] * len(times),
+            "loss_masks": [[1, 1]] * len(times),
+            "stop_reasons": ["stop"] * len(times),
+            "rollout_logprobs": None,
+            "trajectory_generation_times": times,
+            "trajectory_time_splits": splits,
+        }
+
+    out1 = make_output([10.0, 20.0], {"llm": [4.0, 8.0], "env": [5.0, 10.0]})
+    out2 = make_output([30.0, 40.0], {"llm": [12.0, 16.0], "env": [15.0, 20.0]})
+    concatenated = concatenate_generator_outputs([out1, out2])
+
+    assert concatenated["trajectory_time_splits"] == {"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]}
+    # Aggregates are recomputed over the combined sample, not combined from per-group aggregates.
+    concat_metrics = concatenated["rollout_metrics"]
+    assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(
+        np.percentile([4.0, 8.0, 12.0, 16.0], 90)
+    )
+    assert concat_metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(2.5)
+
+
 def test_get_metrics_from_generator_output():
     # Per trajectory rewards, where rewards are List[float]
     generator_output: GeneratorOutput = {
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 12d32868ae..193b7e5e68 100644
--- a/tests/train/generators/test_skyrl_gym_generator.py
+++ b/tests/train/generators/test_skyrl_gym_generator.py
@@ -1763,15 +1763,11 @@ def step(self, action):
 @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):
-    """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."""
+    """agent_loop attributes engine sleep to "llm" and env sleep to "env", not swapped or merged.
+    Metric math over the splits is covered by exact-value tests in test_generator_output_utils.py."""
     import asyncio
     import time
 
-    import numpy as np
-
-    from skyrl.train.generators import utils as generator_utils
     from skyrl.train.generators.base import TrajectoryID
 
     llm_sleep_s = 0.02
@@ -1850,41 +1846,17 @@ def step(self, action):
         "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)
+    generator_output: GeneratorOutput = await generator.generate(input_batch)
 
-    time_splits = spy.call_args.kwargs["trajectory_time_splits"]
+    time_splits = generator_output["trajectory_time_splits"]
     assert time_splits is not None and set(time_splits) == {"llm", "env"}
     llm_times, env_times = time_splits["llm"], time_splits["env"]
     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):
+    for llm_t, env_t, e2e_t in zip(llm_times, env_times, generator_output["trajectory_generation_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
-    assert metrics["generate/trajectory_overhead_time_mean"] >= 0.0
-
-    # The splits never exceed the trajectory's end-to-end time; overhead is the exact remainder.
-    overhead_times = []
-    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
-        overhead_times.append(e2e_t - llm_t - env_t)
-    assert metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(np.mean(overhead_times).item())
-
-    # Regression: concatenate_generator_outputs (every logging path) must recompute the split from
-    # the raw per-trajectory lists, not sum per-group aggregates past the sample max.
-    concatenated = generator_utils.concatenate_generator_outputs([generator_output, generator_output])
-    concat_metrics = concatenated["rollout_metrics"]
-    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_overhead_time_p90"] == pytest.approx(
-        np.percentile(overhead_times * 2, 90).item()
-    )

From 3d3171fc4034564c9345e6d553ecfcb42c76d498 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:24:34 +0000
Subject: [PATCH 11/17] [logging] Document the time split once, on the
 GeneratorOutput contract

The dataclass fields carry no comment; base.py states the keys and that the whole field is None (never partially populated) when any trajectory did not record its split.
---
 skyrl/train/generators/base.py                | 3 ++-
 skyrl/train/generators/skyrl_gym_generator.py | 4 ----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index b7b4973be8..a8a0485d89 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -47,7 +47,8 @@ class GeneratorOutput(TypedDict):
     # async trainer to compute per-group / intra-group completion-time metrics.
     trajectory_generation_times: Optional[List[float]]
     # Engine and env time splits of ``trajectory_generation_times``, one list entry per trajectory,
-    # e.g. {"llm": [...], "env": [...]}. None if any trajectory did not record them.
+    # e.g. {"llm": [...], "env": [...]}. This whole field is None if any trajectory did not record
+    # its split; the per-key lists are never partially populated.
     trajectory_time_splits: Optional[Dict[str, List[float]]]
     rollout_expert_indices: Optional[List[List[List[List[int]]]]]  # [batch_size, seq_len, layer_num, topk]
     # Applicable only for step-wise training
diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index e7a6eb4bf9..f3ea06e0e8 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -56,9 +56,6 @@ 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 split of ``e2e_time`` by component (seconds), summed over turns: "llm" is time in
-    # the inference-engine call, "env" is time in ``env.step()``. The remainder of ``e2e_time`` is
-    # tokenization, chat templating, and event-loop scheduling. None if untracked.
     time_splits: Optional[Dict[str, float]] = None
 
 
@@ -70,7 +67,6 @@ 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 as TrajectoryOutput.time_splits.
     time_splits: Optional[Dict[str, float]] = None
 
 

From 31c83885a3791b7f0fb73ef827daf7168925aae1 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:24:58 +0000
Subject: [PATCH 12/17] [logging] Simplify the trajectory_time_splits None
 wording

---
 skyrl/train/generators/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py
index a8a0485d89..36eab8d581 100644
--- a/skyrl/train/generators/base.py
+++ b/skyrl/train/generators/base.py
@@ -47,8 +47,8 @@ class GeneratorOutput(TypedDict):
     # async trainer to compute per-group / intra-group completion-time metrics.
     trajectory_generation_times: Optional[List[float]]
     # Engine and env time splits of ``trajectory_generation_times``, one list entry per trajectory,
-    # e.g. {"llm": [...], "env": [...]}. This whole field is None if any trajectory did not record
-    # its split; the per-key lists are never partially populated.
+    # e.g. {"llm": [...], "env": [...]}. trajectory_time_splits is None if any trajectory did not
+    # record its split.
     trajectory_time_splits: Optional[Dict[str, List[float]]]
     rollout_expert_indices: Optional[List[List[List[List[int]]]]]  # [batch_size, seq_len, layer_num, topk]
     # Applicable only for step-wise training

From 77bdf77caeedce53f046f08d37fb8a547c66221f Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:47:41 +0000
Subject: [PATCH 13/17] [logging] Concatenate optional fields with
 any-missing-is-None semantics

_concat_time_splits keyed off the first output, so concatenating a batch
that recorded splits with one that did not crashed when the missing
batch came second. Fold list and dict-of-lists handling into one
_concat_optional_field with any-missing-is-None semantics for every
optional field, and let _last_step_only recurse into dicts. Adds an
order-independence regression test.
---
 skyrl/train/generators/utils.py               | 41 ++++++++++---------
 .../generators/test_generator_output_utils.py | 25 +++++++++++
 2 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 1d81a3ca87..7d26fa79f5 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -240,25 +240,28 @@ 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:
+def _concat_optional_field(
+    generator_outputs: List[GeneratorOutput], key: str
+) -> Optional[Union[list, Dict[str, list]]]:
+    """Concatenate an optional per-trajectory field across outputs. None if any output lacks it.
+    Handles a flat list or a dict-of-lists, concatenating each component."""
+    values = [go.get(key) for go in generator_outputs]
+    if any(v is None for v in values):
         return None
+    if isinstance(values[0], dict):
+        return {name: [t for v in values for t in v[name]] for name in values[0]}
     return _flatten_field(generator_outputs, key)
 
 
-def _concat_time_splits(generator_outputs: List[GeneratorOutput]) -> Optional[Dict[str, List[float]]]:
-    """Concatenate the per-trajectory time-split lists across outputs, per component."""
-    first = generator_outputs[0].get("trajectory_time_splits")
-    if first is None:
-        return None
-    return {name: [t for go in generator_outputs for t in go["trajectory_time_splits"][name]] for name in first}
-
-
-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."""
+def _last_step_only(
+    values: Optional[Union[list, Dict[str, list]]], is_last_step: Optional[List[bool]]
+) -> Optional[Union[list, Dict[str, list]]]:
+    """Keep only last-step entries when step-wise, so one trajectory contributes one value.
+    Handles a flat list or a dict-of-lists. No-op when not step-wise."""
     if values is None or not is_last_step:
         return values
+    if isinstance(values, dict):
+        return {name: _last_step_only(v, is_last_step) for name, v in values.items()}
     return [v for v, last in zip(values, is_last_step) if last]
 
 
@@ -286,10 +289,10 @@ 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": _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_time_splits": _concat_time_splits(generator_outputs),
+        "stop_reasons": _concat_optional_field(generator_outputs, "stop_reasons"),
+        "rollout_logprobs": _concat_optional_field(generator_outputs, "rollout_logprobs"),
+        "trajectory_generation_times": _concat_optional_field(generator_outputs, "trajectory_generation_times"),
+        "trajectory_time_splits": _concat_optional_field(generator_outputs, "trajectory_time_splits"),
     }
 
     # propagate additional keys with list values as-is
@@ -302,9 +305,7 @@ def concatenate_generator_outputs(generator_outputs: List[GeneratorOutput], 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_time_splits = result.get("trajectory_time_splits")
-    if trajectory_time_splits is not None and is_last_step:
-        trajectory_time_splits = {k: _last_step_only(v, is_last_step) for k, v in trajectory_time_splits.items()}
+    trajectory_time_splits = _last_step_only(result.get("trajectory_time_splits"), is_last_step)
 
     # Re-aggregate rollout metrics; the extra_keys fallback below cannot aggregate a p90 or a ratio.
     rollout_metrics = get_rollout_metrics(
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 6ea508189d..af0b1c2141 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -137,6 +137,31 @@ def make_output(times, splits) -> GeneratorOutput:
     assert concat_metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(2.5)
 
 
+def test_time_splits_concatenation_partial_is_none():
+    """A batch that recorded no splits drops the whole aggregate to None rather than crashing."""
+
+    def make_output(times, splits) -> GeneratorOutput:
+        return {
+            "prompt_token_ids": [[1]] * len(times),
+            "response_ids": [[1, 2]] * len(times),
+            "rewards": [1.0] * len(times),
+            "loss_masks": [[1, 1]] * len(times),
+            "stop_reasons": ["stop"] * len(times),
+            "rollout_logprobs": None,
+            "trajectory_generation_times": times,
+            "trajectory_time_splits": splits,
+        }
+
+    recorded = make_output([10.0, 20.0], {"llm": [4.0, 8.0], "env": [5.0, 10.0]})
+    missing = make_output([30.0, 40.0], None)
+    # Order must not matter: keying off the first output would crash when the missing batch is second.
+    for outputs in ([recorded, missing], [missing, recorded]):
+        concatenated = concatenate_generator_outputs(outputs)
+        assert concatenated["trajectory_time_splits"] is None
+        # generation_times has its own None-ness, so it still concatenates fully.
+        assert sorted(concatenated["trajectory_generation_times"]) == [10.0, 20.0, 30.0, 40.0]
+
+
 def test_get_metrics_from_generator_output():
     # Per trajectory rewards, where rewards are List[float]
     generator_output: GeneratorOutput = {

From 60edb46e963f148bf6d60061e8448f234dfe6bad Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Sat, 18 Jul 2026 00:50:22 +0000
Subject: [PATCH 14/17] [logging] Drop archaeology from the time-split concat
 test

---
 tests/train/generators/test_generator_output_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index af0b1c2141..8ec78b797d 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -138,7 +138,7 @@ def make_output(times, splits) -> GeneratorOutput:
 
 
 def test_time_splits_concatenation_partial_is_none():
-    """A batch that recorded no splits drops the whole aggregate to None rather than crashing."""
+    """A batch that recorded no splits drops the whole time-split aggregate to None."""
 
     def make_output(times, splits) -> GeneratorOutput:
         return {
@@ -154,7 +154,7 @@ def make_output(times, splits) -> GeneratorOutput:
 
     recorded = make_output([10.0, 20.0], {"llm": [4.0, 8.0], "env": [5.0, 10.0]})
     missing = make_output([30.0, 40.0], None)
-    # Order must not matter: keying off the first output would crash when the missing batch is second.
+    # Both orders: the None batch must not depend on being first or last.
     for outputs in ([recorded, missing], [missing, recorded]):
         concatenated = concatenate_generator_outputs(outputs)
         assert concatenated["trajectory_time_splits"] is None

From 4c02371d0509621f22381b0729bf427762c77d53 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Tue, 21 Jul 2026 21:01:20 +0000
Subject: [PATCH 15/17] [logging] Address review: rename remainder band to
 "other", document time_splits keys

- Rename the e2e remainder band from overhead to other. The remainder
  holds tokenization and chat templating, so overhead overstated it; a
  large "other" should read as a problem to investigate.
- Document the time_splits dict keys on TrajectoryOutput and
  StepWiseOutput, matching the trajectory_time_splits contract field.
---
 skyrl/train/generators/skyrl_gym_generator.py         | 4 ++++
 skyrl/train/generators/utils.py                       | 7 ++++---
 tests/train/generators/test_generator_output_utils.py | 8 ++++----
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index f3ea06e0e8..210c0115c1 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -56,6 +56,8 @@ 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 seconds keyed by phase, e.g. {"llm": <inference-engine calls>, "env": <env.step>}.
+    # None if the agent loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None
 
 
@@ -67,6 +69,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 seconds keyed by phase, e.g. {"llm": <inference-engine calls>, "env": <env.step>}.
+    # None if the agent loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None
 
 
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 7d26fa79f5..17d4542387 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -489,10 +489,11 @@ def get_rollout_metrics(
         _add_time_stats(rollout_metrics, name, times)
 
     if trajectory_completion_times and trajectory_time_splits:
-        overhead = np.array(trajectory_completion_times, dtype=np.float64)
+        # Completion time not attributed to a named split, e.g. tokenization and chat templating.
+        other = np.array(trajectory_completion_times, dtype=np.float64)
         for times in trajectory_time_splits.values():
-            overhead = overhead - np.array(times, dtype=np.float64)
-        _add_time_stats(rollout_metrics, "overhead", overhead.tolist())
+            other = other - np.array(times, dtype=np.float64)
+        _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_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 8ec78b797d..9671460771 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -106,9 +106,9 @@ def test_time_split_rollout_metrics():
     assert metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile([4.0, 8.0, 12.0, 16.0], 90))
     assert metrics["generate/trajectory_llm_time_max"] == 16.0
     assert metrics["generate/trajectory_env_time_mean"] == 12.5
-    # Overhead is the exact per-trajectory remainder, here [1.0, 2.0, 3.0, 4.0].
-    assert metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(2.5)
-    assert metrics["generate/trajectory_overhead_time_max"] == pytest.approx(4.0)
+    # "other" is the exact per-trajectory remainder, here [1.0, 2.0, 3.0, 4.0].
+    assert metrics["generate/trajectory_other_time_mean"] == pytest.approx(2.5)
+    assert metrics["generate/trajectory_other_time_max"] == pytest.approx(4.0)
 
 
 def test_time_splits_concatenation():
@@ -134,7 +134,7 @@ def make_output(times, splits) -> GeneratorOutput:
     assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(
         np.percentile([4.0, 8.0, 12.0, 16.0], 90)
     )
-    assert concat_metrics["generate/trajectory_overhead_time_mean"] == pytest.approx(2.5)
+    assert concat_metrics["generate/trajectory_other_time_mean"] == pytest.approx(2.5)
 
 
 def test_time_splits_concatenation_partial_is_none():

From b7772e0f4cb3c978af609a32ffc6652fe76f17d1 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Tue, 21 Jul 2026 21:13:09 +0000
Subject: [PATCH 16/17] [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, other) read as components of one trajectory_time total and one glob
pulls the whole breakdown. Also reword the time_splits docstrings to name
the keys and say the whole field is None when unrecorded.
---
 skyrl/train/generators/skyrl_gym_generator.py    |  8 ++++----
 skyrl/train/generators/utils.py                  |  8 ++++----
 .../generators/test_generator_output_utils.py    | 16 ++++++++--------
 .../train/generators/test_skyrl_gym_generator.py | 12 ++++++------
 4 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index 210c0115c1..0f4076165f 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -56,8 +56,8 @@ 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 seconds keyed by phase, e.g. {"llm": <inference-engine calls>, "env": <env.step>}.
-    # None if the agent loop did not record a split.
+    # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
+    # time in env.step. The whole field is None if the agent loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None
 
 
@@ -69,8 +69,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 seconds keyed by phase, e.g. {"llm": <inference-engine calls>, "env": <env.step>}.
-    # None if the agent loop did not record a split.
+    # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
+    # time in env.step. The whole field is None if the agent loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None
 
 
diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py
index 17d4542387..ec025c31d2 100644
--- a/skyrl/train/generators/utils.py
+++ b/skyrl/train/generators/utils.py
@@ -395,15 +395,15 @@ def compute_turn_token_counts(loss_masks: List[List[int]]) -> List[int]:
 
 
 def _add_time_stats(rollout_metrics: Dict[str, Any], name: str, times: Optional[List[float]]) -> None:
-    """Add mean/p90/max stats for a per-trajectory time list under generate/trajectory_<name>_time_*."""
+    """Add mean/p90/max stats for a per-trajectory time list under generate/trajectory_time_<name>_*."""
     if not times:
         return
     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(),
         }
     )
 
diff --git a/tests/train/generators/test_generator_output_utils.py b/tests/train/generators/test_generator_output_utils.py
index 9671460771..877ea18091 100644
--- a/tests/train/generators/test_generator_output_utils.py
+++ b/tests/train/generators/test_generator_output_utils.py
@@ -102,13 +102,13 @@ def test_time_split_rollout_metrics():
         trajectory_completion_times=[10.0, 20.0, 30.0, 40.0],
         trajectory_time_splits={"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]},
     )
-    assert metrics["generate/trajectory_llm_time_mean"] == 10.0
-    assert metrics["generate/trajectory_llm_time_p90"] == pytest.approx(np.percentile([4.0, 8.0, 12.0, 16.0], 90))
-    assert metrics["generate/trajectory_llm_time_max"] == 16.0
-    assert metrics["generate/trajectory_env_time_mean"] == 12.5
+    assert metrics["generate/trajectory_time_llm_mean"] == 10.0
+    assert metrics["generate/trajectory_time_llm_p90"] == pytest.approx(np.percentile([4.0, 8.0, 12.0, 16.0], 90))
+    assert metrics["generate/trajectory_time_llm_max"] == 16.0
+    assert metrics["generate/trajectory_time_env_mean"] == 12.5
     # "other" is the exact per-trajectory remainder, here [1.0, 2.0, 3.0, 4.0].
-    assert metrics["generate/trajectory_other_time_mean"] == pytest.approx(2.5)
-    assert metrics["generate/trajectory_other_time_max"] == pytest.approx(4.0)
+    assert metrics["generate/trajectory_time_other_mean"] == pytest.approx(2.5)
+    assert metrics["generate/trajectory_time_other_max"] == pytest.approx(4.0)
 
 
 def test_time_splits_concatenation():
@@ -131,10 +131,10 @@ def make_output(times, splits) -> GeneratorOutput:
     assert concatenated["trajectory_time_splits"] == {"llm": [4.0, 8.0, 12.0, 16.0], "env": [5.0, 10.0, 15.0, 20.0]}
     # Aggregates are recomputed over the combined sample, not combined from per-group aggregates.
     concat_metrics = concatenated["rollout_metrics"]
-    assert concat_metrics["generate/trajectory_llm_time_p90"] == pytest.approx(
+    assert concat_metrics["generate/trajectory_time_llm_p90"] == pytest.approx(
         np.percentile([4.0, 8.0, 12.0, 16.0], 90)
     )
-    assert concat_metrics["generate/trajectory_other_time_mean"] == pytest.approx(2.5)
+    assert concat_metrics["generate/trajectory_time_other_mean"] == pytest.approx(2.5)
 
 
 def test_time_splits_concatenation_partial_is_none():
diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py
index 193b7e5e68..1b8f303d06 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 time splits are stored per step and replicated like the completion times above.
     time_splits = generator_output["trajectory_time_splits"]

From 9d051911a3a34f004ca2e27a24d63d37436cf2b8 Mon Sep 17 00:00:00 2001
From: Seiji Eicher <seiji@anyscale.com>
Date: Tue, 21 Jul 2026 21:16:55 +0000
Subject: [PATCH 17/17] [logging] Reword time_splits docstring

---
 skyrl/train/generators/skyrl_gym_generator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/skyrl/train/generators/skyrl_gym_generator.py b/skyrl/train/generators/skyrl_gym_generator.py
index 0f4076165f..322f68ef7b 100644
--- a/skyrl/train/generators/skyrl_gym_generator.py
+++ b/skyrl/train/generators/skyrl_gym_generator.py
@@ -57,7 +57,7 @@ class TrajectoryOutput:
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
     # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
-    # time in env.step. The whole field is None if the agent loop did not record a split.
+    # time in env.step. Field is None if any loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None
 
 
@@ -70,7 +70,7 @@ class StepWiseOutput:
     # leave this as None if they do not track timing.
     e2e_time: Optional[float] = None
     # Wall-clock seconds spent in each phase. "llm" is time in inference-engine calls, "env" is
-    # time in env.step. The whole field is None if the agent loop did not record a split.
+    # time in env.step. Field is None if any loop did not record a split.
     time_splits: Optional[Dict[str, float]] = None