diff --git a/skyrl/train/fully_async_trainer.py b/skyrl/train/fully_async_trainer.py index 99fa0ff31b..c47e51a182 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 ScalarGauges, TrainingPhaseGauge from skyrl.train.utils.trainer_utils import ( ResumeMode, build_dataloader, @@ -456,9 +457,29 @@ async def train(self): with Timer("sync_weights_to_inference_engines"): await self.dispatch.save_weights_for_sampler() + # Per-step GPU utilization (-> tracker) and a training-phase gauge (-> Prometheus, joinable + # with node GPU metrics). The base trainer wires these for the synchronous loop; the async + # loop overrides train() and must start/flush/stop them itself. + if self._ray_gpu_monitor is not None: + self._ray_gpu_monitor.start() + self._phase_gauge = TrainingPhaseGauge() + # Buffer depth to Prometheus, so it joins the phase gauge and node GPU metrics in one store + # (the per-step tracker logging below only reaches the experiment tracker). keep_rate exposes + # the zero-variance drop dynamics under sample_full_batch. + self._loop_gauges = ScalarGauges( + descriptions={ + "skyrl_gen_buffer_qsize": "Completed generation groups buffered at step start.", + "skyrl_gen_buffer_maxsize": "Staleness-bounded generation-buffer capacity " + "(mini_batch_size * (max_staleness_steps + 1)).", + "skyrl_mini_batch_size": "Generation groups consumed per training step.", + "skyrl_gen_group_keep_rate": "Fraction of drained groups kept (not dropped as " + "zero-variance) while collecting the last mini-batch.", + } + ) + # Eval before training if self.cfg.trainer.eval_interval > 0 and self.cfg.trainer.eval_before_train: - with Timer("eval", self.all_timings): + with Timer("eval", self.all_timings), self._phase_gauge.phase("eval"): eval_metrics = await self.eval() self.tracker.log(eval_metrics, step=self.global_step, commit=True) @@ -503,15 +524,26 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don trained_steps_this_epoch = self.async_train_dataloader.num_trained() // self.mini_batch_size for _step_idx in range(self.global_step, (1 + epoch) * self.num_steps_per_epoch + 1): with Timer("step", self.all_timings): + # Buffer depth at step start: how many completed rollout batches are queued + # for the trainer. Near maxsize means generation is paused at the staleness + # cap; near zero means the trainer is starving. Previously only shown live in + # a tqdm postfix. + self.all_metrics["async/gen_buffer_qsize"] = generation_output_group_buffer.qsize() + self.all_metrics["async/gen_buffer_maxsize"] = generation_output_group_buffer.maxsize + self._loop_gauges.set("skyrl_gen_buffer_qsize", generation_output_group_buffer.qsize()) + self._loop_gauges.set("skyrl_gen_buffer_maxsize", generation_output_group_buffer.maxsize) + self._loop_gauges.set("skyrl_mini_batch_size", self.mini_batch_size) + # 1. Wait until we have a full mini-batch buffered (dropping zero-variance groups if # sample_full_batch). - ( - cur_generation_group_mini_batch, - cur_dropped_groups, - epoch_exhausted, - ) = await self._collect_generation_mini_batch( - generation_output_group_buffer, all_generators_done - ) + with self._phase_gauge.phase("waiting_for_buffer"): + ( + cur_generation_group_mini_batch, + cur_dropped_groups, + epoch_exhausted, + ) = await self._collect_generation_mini_batch( + generation_output_group_buffer, all_generators_done + ) if epoch_exhausted: # Exhausted mid mini-batch: discard the partial batch (marked consumed so it @@ -537,10 +569,18 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don break if self.sample_full_batch: - self.all_metrics["async/num_groups_dropped"] = len(cur_dropped_groups) + num_kept = len(cur_generation_group_mini_batch) + num_dropped = len(cur_dropped_groups) + self.all_metrics["async/num_groups_dropped"] = num_dropped + drained = num_kept + num_dropped + if drained > 0: + self._loop_gauges.set("skyrl_gen_group_keep_rate", num_kept / drained) # 2. Post-process the generated groups, aggregating to a single GeneratorOutput, and convert to training format. - with Timer("convert_to_training_input", self.all_timings): + with ( + Timer("convert_to_training_input", self.all_timings), + self._phase_gauge.phase("converting"), + ): training_input = await asyncio.to_thread( self.convert_generation_group_mini_batch_to_training_input, cur_generation_group_mini_batch, @@ -548,14 +588,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 +617,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 +625,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 +633,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 +699,8 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don # End of an epoch. finally: self._profiler_stop() + if self._ray_gpu_monitor is not None: + self._ray_gpu_monitor.stop() pbar.close() @@ -667,7 +711,7 @@ async def _watch_generators_done(tasks=generator_tasks, event=all_generators_don # safety net: always save final checkpoint at end of training. if self.cfg.trainer.ckpt_interval > 0: - with Timer("save_checkpoints", self.all_timings): + with Timer("save_checkpoints", self.all_timings), self._phase_gauge.phase("checkpoint"): await asyncio.to_thread(self.save_checkpoints) logger.info("Saved final checkpoint.") if self.cfg.trainer.hf_save_interval > 0: 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 cc9bc95e80..156fcf8599 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 = [] @@ -885,6 +917,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False out_trajectory_ids = [] out_env_classes = [] out_trajectory_generation_times = [] + out_trajectory_llm_times = [] + out_trajectory_env_times = [] for i, output in enumerate(all_outputs): for j, step_output in enumerate(output.step_outputs): responses.append(step_output.response_ids) @@ -896,11 +930,17 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False is_last_step.append(j == len(output.step_outputs) - 1) out_trajectory_ids.append(trajectory_ids[i]) out_env_classes.append(env_classes[i]) - # For trajectory completion per turn we just use the trajectory level e2e time + # For trajectory completion per turn we just use the trajectory level times out_trajectory_generation_times.append(getattr(output, "e2e_time", None)) + out_trajectory_llm_times.append(getattr(output, "llm_time", None)) + out_trajectory_env_times.append(getattr(output, "env_time", None)) # Keep aligned with the per-prompt None handling: if not trajectory_generation_times_per_prompt: out_trajectory_generation_times = None + if not trajectory_llm_times_per_prompt: + out_trajectory_llm_times = None + if not trajectory_env_times_per_prompt: + out_trajectory_env_times = None env_classes = out_env_classes else: responses = [output.response_ids for output in all_outputs] @@ -913,6 +953,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False out_trajectory_ids = None # One time per trajectory, already aligned 1:1 with responses (None if not all recorded). out_trajectory_generation_times = trajectory_generation_times_per_prompt + out_trajectory_llm_times = trajectory_llm_times_per_prompt + out_trajectory_env_times = trajectory_env_times_per_prompt has_vision_features = any(getattr(output, "pixel_values", None) is not None for output in all_outputs) pixel_values = ( @@ -953,6 +995,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False # NOTE: we only use trajectory completion times per prompt for # metrics, to avoid duplicate entries with step-wise training trajectory_completion_times=trajectory_generation_times_per_prompt, + trajectory_llm_times=trajectory_llm_times_per_prompt, + trajectory_env_times=trajectory_env_times_per_prompt, ) if self.generator_cfg.zero_reward_on_non_stop: @@ -974,6 +1018,8 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False "trajectory_ids": out_trajectory_ids, # NOTE: for completion metrics, we output the completion time "trajectory_generation_times": out_trajectory_generation_times, + "trajectory_llm_times": out_trajectory_llm_times, + "trajectory_env_times": out_trajectory_env_times, "rollout_expert_indices": rollout_expert_indices, "is_last_step": is_last_step, "env_metrics": env_metrics, diff --git a/skyrl/train/generators/utils.py b/skyrl/train/generators/utils.py index 4c5228f189..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 @@ -384,6 +398,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 +412,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 +480,35 @@ def get_rollout_metrics( } ) + 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(), + } + ) + + 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 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/skyrl/train/utils/phase_metrics.py b/skyrl/train/utils/phase_metrics.py new file mode 100644 index 0000000000..4aaf510f4d --- /dev/null +++ b/skyrl/train/utils/phase_metrics.py @@ -0,0 +1,111 @@ +"""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_``) 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/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index c29afbe920..e7953092d3 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -1751,3 +1751,144 @@ 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 + + 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 + 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 + + # 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()) diff --git a/tests/train/utils/test_phase_metrics.py b/tests/train/utils/test_phase_metrics.py new file mode 100644 index 0000000000..f3d9895814 --- /dev/null +++ b/tests/train/utils/test_phase_metrics.py @@ -0,0 +1,88 @@ +""" +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