diff --git a/skyrl/train/generators/base.py b/skyrl/train/generators/base.py index 26d95868b9..36eab8d581 100644 --- a/skyrl/train/generators/base.py +++ b/skyrl/train/generators/base.py @@ -46,6 +46,10 @@ 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]] + # Engine and env time splits of ``trajectory_generation_times``, one list entry per trajectory, + # 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 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..322f68ef7b 100644 --- a/skyrl/train/generators/skyrl_gym_generator.py +++ b/skyrl/train/generators/skyrl_gym_generator.py @@ -56,6 +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 seconds spent in each phase. "llm" is time in inference-engine calls, "env" is + # time in env.step. Field is None if any loop did not record a split. + time_splits: Optional[Dict[str, float]] = None @dataclass @@ -66,6 +69,9 @@ 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 spent in each phase. "llm" is time in inference-engine calls, "env" is + # time in env.step. Field is None if any loop did not record a split. + time_splits: Optional[Dict[str, float]] = None @dataclass @@ -140,6 +146,13 @@ def get_turn_rollout_logprobs(self) -> Optional[List[float]]: return self.output_logprobs + [0.0] * len(self.obs_ids) +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): def __init__( self, @@ -321,6 +334,7 @@ async def agent_loop( rollout_logprobs: Optional[List[float]] """ agent_loop_start_time = time.monotonic() + 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 @@ -409,7 +423,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) + 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] @@ -443,7 +459,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) + 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"] @@ -606,6 +624,7 @@ async def agent_loop( trajectory_id, ) agent_loop_output.e2e_time = time.monotonic() - agent_loop_start_time + agent_loop_output.time_splits = time_splits return agent_loop_output finally: @@ -873,6 +892,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False 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 = [] @@ -885,6 +905,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False out_trajectory_ids = [] out_env_classes = [] out_trajectory_generation_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) @@ -896,11 +917,13 @@ 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_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 + 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] @@ -913,6 +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_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 = ( @@ -953,6 +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_time_splits=trajectory_time_splits_per_prompt, ) if self.generator_cfg.zero_reward_on_non_stop: @@ -974,6 +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_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 4c5228f189..ec025c31d2 100644 --- a/skyrl/train/generators/utils.py +++ b/skyrl/train/generators/utils.py @@ -240,6 +240,31 @@ def _flatten_field(generator_outputs: List[GeneratorOutput], key: str) -> list: return flat +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 _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] + + 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. @@ -264,17 +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": ( - _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 - ), - "trajectory_generation_times": ( - _flatten_field(generator_outputs, "trajectory_generation_times") - if first.get("trajectory_generation_times") is not None - else None - ), + "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 @@ -284,19 +302,18 @@ 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_time_splits = _last_step_only(result.get("trajectory_time_splits"), is_last_step) - # Re-aggregate rollout metrics + # 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"], loss_masks=result.get("loss_masks"), trajectory_completion_times=trajectory_generation_times, + trajectory_time_splits=trajectory_time_splits, ) # Preserve generator-specific metrics from per-group rollout_metrics. get_rollout_metrics only @@ -377,6 +394,20 @@ 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]]) -> None: + """Add mean/p90/max stats for a per-trajectory time list under generate/trajectory_time__*.""" + if not times: + return + arr = np.array(times, dtype=np.float64) + rollout_metrics.update( + { + 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(), + } + ) + + def get_rollout_metrics( responses: List[List[int]], rewards: Union[List[float], List[List[float]]], @@ -384,6 +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_time_splits: Optional[Dict[str, List[float]]] = None, ): """ Computes rollout metrics including token statistics and optional environment-specific metrics. @@ -396,6 +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_time_splits: Optional per-component split of the completion times, e.g. + {"llm": [...], "env": [...]}, one entry per trajectory Returns: Dictionary of aggregated metrics @@ -450,15 +484,16 @@ 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(), - } - ) + _add_time_stats(rollout_metrics, "completion", trajectory_completion_times) + for name, times in (trajectory_time_splits or {}).items(): + _add_time_stats(rollout_metrics, name, times) + + if trajectory_completion_times and trajectory_time_splits: + # 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(): + 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 7546858bb8..877ea18091 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 @@ -32,6 +33,7 @@ def test_generator_output_concatenation(): # optional but present in the signature "trajectory_ids", "trajectory_generation_times", + "trajectory_time_splits", "is_last_step", "env_metrics", "pixel_values", @@ -93,6 +95,73 @@ 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_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_time_other_mean"] == pytest.approx(2.5) + assert metrics["generate/trajectory_time_other_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_time_llm_p90"] == pytest.approx( + np.percentile([4.0, 8.0, 12.0, 16.0], 90) + ) + assert concat_metrics["generate/trajectory_time_other_mean"] == pytest.approx(2.5) + + +def test_time_splits_concatenation_partial_is_none(): + """A batch that recorded no splits drops the whole time-split aggregate to None.""" + + 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) + # 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 + # 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 = { diff --git a/tests/train/generators/test_skyrl_gym_generator.py b/tests/train/generators/test_skyrl_gym_generator.py index c29afbe920..1b8f303d06 100644 --- a/tests/train/generators/test_skyrl_gym_generator.py +++ b/tests/train/generators/test_skyrl_gym_generator.py @@ -1740,14 +1740,123 @@ 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"] + 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] + + +@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): + """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 + + 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.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() + 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)], + } + + generator_output: GeneratorOutput = await generator.generate(input_batch) + + 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 + + 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 + assert llm_t + env_t <= e2e_t + 1e-6