[metrics 5/5] Scope prefill/decode and external vLLM metrics

[SumanthRH]

What does this PR do?

Scope vLLM step metrics and fresh-run summaries separately for prefill and decode, including externally deployed SkyRL servers in the same Ray cluster.

TLDR: Resolve the servers' frontend WorkerIds and roles, then reuse the existing worker-filtered collectors and reductions for each role.

Why this is needed

Prefill and decode can count the same prompt, prefill can report a placeholder output token, and each role measures different latency boundaries. Combining their counters or histograms obscures queue pressure and produces ambiguous run comparisons. External clients also lack the server actor handles used for managed worker discovery.

Usage

Enable generator.inference_engine.enable_ray_prometheus_stats=true on both the serving and client configurations. External clients specify individual external_server_urls with run_engines_locally=false, optionally together with external_proxy_url. The same Ray cluster must expose the servers' metrics agents. PD roles come from the servers; the client does not need to repeat their engine settings.

Output	Namespace
Sync charts	vllm/{train,eval}/{prefill,decode}/*
Fully async charts	vllm/{prefill,decode}/*
W&B Summary	vllm_correct_aggregate/{train,eval,combined}/{prefill,decode}/*

How it works

Managed deployments resolve WorkerIds from their existing server groups. External clients query /get_metrics_worker_info on each backend for the exporting frontend's WorkerId and optional prefill/decode role. This endpoint is exposed only with Ray stats enabled and avoids vLLM's /metrics mount. Each role gets its own collector for samples, generation windows and finalization; sync role clocks begin after both baseline reads complete.

Keep each available metric in its role scope, including engine imbalance within that role. Gauges remain step snapshots. Fresh-run-only publication, missing/reset omission and the bounded terminal attempt remain; an unavailable role omits only that role's summary. Each new client takes its own baseline against a reused deployment, excluding earlier served work. Concurrent clients on the same measured engines are not distinguishable by engine counters.

Legacy external endpoints retain cluster-wide collection. Unknown-role external PD suppresses summaries when the client sets enable_pd=true. The metrics guide documents setup and engine-level token/latency interpretation.

The accompanying small simulation fix supplies the timing and checkpoint-finalization hooks now called unconditionally by the async trainer, allowing the simulation entrypoint to complete its run.

Test plan

PD cases cover sync/async role separation, weighted counters/histograms, reused baselines, gauges, missing role workers and managed setup. External setup cases cover regular serving, automatic PD role discovery, explicit PD and legacy fallback. A native vLLM HTTP-app test verifies worker metadata is not captured by the Prometheus route. Existing lifecycle cases cover unknown roles and checkpoint resumption.

Validation

180 focused tests passed, including metrics, HTTP inference, async training and collation. Ruff/Black, secret checks and the production docs build passed. --noconftest avoids the repository's automatic Ray lifecycle fixture touching the live cluster.

uv run --isolated --offline --frozen --extra fsdp --extra dev --with ray==2.58.0 pytest --noconftest \
  tests/train/test_grafana_annotations.py tests/train/test_metrics_lifecycle.py \
  tests/train/test_vllm_run_statistics.py tests/train/test_vllm_metrics_scraper.py \
  tests/train/test_vllm_window_statistics.py tests/train/test_vllm_engine_imbalance.py \
  tests/train/test_vllm_pd_metrics.py tests/train/test_tracking.py \
  tests/train/test_fully_async_trainer.py tests/train/test_async_batch_collation.py \
  tests/backends/skyrl_train/inference_servers/test_external_metrics.py \
  tests/backends/skyrl_train/inference_servers/test_vllm_sample_support.py \
  tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py -q

Real-inference validation used Qwen2.5-1.5B-Instruct/GSM8K, B200 GPUs, vLLM 0.30 and Ray 2.58, with graph execution enabled and NixlConnector for 1P1D. Two external simulation clients each completed two steps against a reused regular server; another two clients did the same against reused 1P1D servers. No gradient computation or model transfer is performed by the simulation.

bash /home/ray/hackskyrl/logs/1004/metrics_external_mode/run_regular.sh
bash /home/ray/hackskyrl/logs/1004/metrics_external_mode/run_pd.sh

Mode	W&B runs	Independent Prometheus aggregate matches
Regular	client 1, client 2	24 / 24
1P1D	client 1, client 2	52 / 52

Query the recorded baseline/terminal timestamps with each role's exact WorkerIds and ClusterId. Subtract observed counters, check resets, divide summed hits by summed queries, and reduce histogram sum/count deltas and merged bucket deltas. Throughput uses independently checked async observation seconds. Each client excludes earlier server work: regular/decode totals are 1,536 output tokens and 2,526 prompt tokens; prefill's 24 placeholder outputs stay separate.

W&B readback verifies all 120 PD step values and 42 regular step values. Regular client 1's second history row remains unavailable through the API, although its explicit run aggregates all match Prometheus. Earlier managed-PD evaluation validation also matched all 26 aggregates and 60 role-scoped chart values in this run.

Validation-only export waits establish visibility; PD async durations include a 25-second terminal idle wait. Production finalization has no export wait. These are compatibility checks, not a throughput benchmark.

Configs, raw observations, timestamps, queries, comparison tables and logs are retained under /home/ray/hackskyrl/logs/1004/metrics_external_mode/ and the earlier /home/ray/hackskyrl/logs/1004/metrics_pd/.

Existing cleanup issue observed during validation

Both PD clients completed successfully and finalized metrics. The serving driver then raised ActorDiedError during teardown: InferenceOnlyEntrypoint._teardown() lists its PD groups twice because _server_groups already contains both roles. This code is unchanged in this PR and present in main. All test servers were stopped and GPU memory release was verified. Existing aiohttp session warnings also appear after the synchronous harness closes training event loops.