[logging] Account for env setup and overhead in trajectory time#1925
Draft
eicherseiji wants to merge 9 commits into
Draft
[logging] Account for env setup and overhead in trajectory time#1925eicherseiji wants to merge 9 commits into
eicherseiji wants to merge 9 commits into
Conversation
… env.step Rollout timing is currently reported only as a total per trajectory (`generate/trajectory_completion_time_*`, added in NovaSky-AI#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>
… 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>
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.
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
Rename the run-ahead phrasing to plain buffer depth in the buffer-metric and gauge comments; no behavior change.
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.
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.
Closes the trajectory-time stack so everything inside
trajectory_completion_time is explicit:
e2e = llm_time + env_time + env_setup_time + overhead_time
env_setup_time brackets skyrl_gym.make plus env.init, which run inside
the e2e clock but outside env_time; for sandboxed tool envs setup can be
seconds and was indistinguishable from tokenization overhead.
overhead_time is the derived remainder: tokenization, chat templating,
output assembly, and event-loop scheduling. Emitted as
generate/trajectory_{env_setup,overhead}_time_{mean,p90,max} on both the
per-generate and concatenated paths.
- Require all four raw lists before computing overhead instead of zero-filling a missing env_setup list, which silently folded real setup time into overhead for any generator reporting only llm and env. - Name env teardown in the overhead comment; env.close() runs inside the e2e window but outside the env_time and env_setup_time brackets. - Vectorize the overhead computation to match the rest of the function. - Assert setup_times length and the new bands on the concatenated path.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Accounts for everything inside
generate/trajectory_completion_timeexplicitly:Stacked on #1900, which adds the llm/env split this PR extends. Until #1900 merges this diff includes its base; the reviewable delta is the last commit.
Two new bands:
generate/trajectory_env_setup_time_{mean,p90,max}: a direct bracket aroundskyrl_gym.makeplusenv.init. Both run inside the e2e clock but outsideenv_time, which only sumsenv.step(). For sandboxed tool envs setup can take seconds and was previously indistinguishable from tokenization overhead.generate/trajectory_overhead_time_{mean,p90,max}: the derived remainder. Tokenization, chat templating, output assembly, env teardown (env.close()runs inside the e2e window but outside both env brackets), and event-loop scheduling. Computed only when all four raw lists are present, so a generator reporting a partial split never silently inflates it. With this band the stack sums to the trajectory's end-to-end time, so a growing unaccounted cost is visible instead of silent.Both are computed on the per-
generate()path and recomputed from raw per-trajectory lists inconcatenate_generator_outputs, the path every logger uses.The engine-side subdivision of
llm_time(queue, prefill, decode per turn) intentionally stays out: vLLM already exports those per request to Prometheus, so they are a dashboard join, not new trainer fields.Tests
test_llm_vs_env_time_split_metrics: the env now sleeps ininit, and the test assertsenv_setup_timecatches it,overhead_timeis non-negative, andllm + env + setup <= e2eper trajectory.GeneratorOutputfield guardrail cover the new field.uv run --isolated --extra dev --extra skyrl-train pytest tests/train/generators/passes, 49 tests. ruff and black clean.