[logging] Rollout time split + async-trainer GPU/phase/buffer observability#1900
Draft
eicherseiji wants to merge 2 commits into
Draft
[logging] Rollout time split + async-trainer GPU/phase/buffer observability#1900eicherseiji wants to merge 2 commits into
eicherseiji wants to merge 2 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>
CharlieFRuan
changed the title
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?
Adds observability for the async RL trainer, motivated by a post-mortem where reconstructing where GPU time went required correlating Prometheus (GPU/disk, wall-clock keyed) against W&B (phase timings, step keyed) by hand — the two never joined. Four additions, each closing part of that gap.
1. Split rollout trajectory time into inference-engine wait vs
env.step()generate/trajectory_{llm,env}_time_{mean,p90,max}andgenerate/frac_time_in_env. Only the total per-trajectory time existed (trajectory_completion_time_*, #1804); it cannot tell an engine-bound rollout from an environment-bound one — the first question for multi-turn/agentic envs. Plumbed like the existinge2e_time(optional, omitted if any trajectory didn't record it). (Original scope of this PR.)2. Fix:
RayGpuMonitorwas never running on the async trainerRayGpuMonitor(#1712) is constructed inRayPPOTrainer.__init__and flushed in the base loop, butFullyAsyncRayPPOTraineroverridestrain()and never calls.start()/.flush()/.stop(). Net effect: async RL runs logged zeroray/GPU keys despiteenable_ray_gpu_monitor=True(default) — verified against real runs. Now started at loop entry, flushed per step into the committed payload, stopped in thefinally.3.
skyrl_training_phasegauge → Prometheus (skyrl/train/utils/phase_metrics.py)Emits the loop's current macro-phase (
waiting_for_buffer/converting/training/weight_sync/eval/checkpoint, defaultgenerating) throughray.util.metrics, which Ray exports to the same Prometheus that scrapes node GPU metrics. That makesavg(ray_node_gpus_utilization) by (Phase)a single-store query instead of a manual cross-tracker correlation — and it still works after a cluster restart, when the tracker may be unreachable. Set at each existingTimer()boundary; best-effort (silently no-ops if Ray metrics are unavailable, so it never breaks training or tests).4. Buffer run-ahead depth as a metric
async/gen_buffer_qsize+async/gen_buffer_maxsize, sampled at step start (previously only in a tqdm postfix). Distinguishes a throughput-limited generator (buffer low) from a gate-limited one (buffer at the staleness-bounded maxsize) — different fixes.Scope note
Began as (1); (2)–(4) broaden it into a small async-observability PR — happy to split (2)–(4) into a follow-up if preferred. Deliberately not included: a
skyrl_node_role{train|inference}gauge (inference node ids aren't cleanly reachable from the trainer) — noted as a follow-up.Tests
tests/train/utils/test_phase_metrics.py(new): exactly-one-active-phase invariant, context-manager restore, no-op-when-Ray-unavailable. 4 passed.tests/train/generators/test_skyrl_gym_generator.py: 30 passed (incl.test_llm_vs_env_time_split_metrics).pre-commit(ruff / black / gitleaks) clean on changed files.Trainer wiring (items 2–4) mirrors the base-trainer pattern and is verified by unit tests + lint; not exercised end-to-end here (needs a GPU cluster).
🤖 Generated with Claude Code