Skip to content

[Frontend][RL] Track HTTP weight operation outcomes and concurrency - #55781

Merged
aoshen02 merged 34 commits into
vllm-project:mainfrom
Ronald1995:codex/rl-weight-operation-metrics
Sep 30, 2026
Merged

aoshen02 merged 34 commits into
vllm-project:mainfrom
Ronald1995:codex/rl-weight-operation-metrics

Conversation

@Ronald1995

@Ronald1995 Ronald1995 commented Sep 8, 2026 •

Copy link
Copy Markdown
Collaborator

Purpose

Weight sync is on the critical path of every RL step, but the frontend exposes no signal for how long each weight-transfer operation takes or whether one is stuck. The generic HTTP metrics already count every weight route by status (http_requests_total{handler,status}), but their per-handler histogram tops out at 1 s, and nothing is recorded until a request completes, so a hung NCCL transfer is invisible.

This PR adds two metrics, with the same names, labels and buckets in the Python and Rust frontends:

Metric Type Labels
vllm:rl_weight_update_operation_duration_seconds Histogram, buckets 0.01–600 s operation
vllm:rl_weight_update_operations_in_flight Gauge (livesum in multiprocess mode) operation
Endpoint operation
/init_weight_transfer_engine init
/start_weight_update start
/start_draft_weight_update start_draft
/update_weights update (one per chunk when packed)
/finish_weight_update finish (includes the optional version update)
  • Each route wraps only its engine call, so requests rejected by validation are not recorded.
  • The in-flight gauge is released and the duration observed on success, error and cancellation. Python uses a context manager; Rust uses a guard dropped at the end of the handler.
  • Python collectors are created lazily on the default registry, after multiprocess Prometheus setup.
  • No request counter is added: http_requests_total already covers counts and status. This matches the sleep-mode metrics in [Frontend][RL] Align sleep-mode API responses and operation metrics #52864.

Test Plan

pytest tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py
pytest tests/entrypoints/serve/dev/rlhf/test_weight_operation_metrics.py  # 1 GPU
cd rust && cargo test -p vllm-metrics weight_operation
cd rust && cargo test -p vllm-server --lib weight_transfer
  • E2E (1 GPU): a real vllm serve --api-server-count 2 with the IPC weight-transfer backend. An HF trainer runs eight full syncs over HTTP. /metrics must report the exact summed duration counts per operation and zero in-flight, and the multiprocess files must show both API servers recorded, so the counts are cross-process sums.
  • Multiprocess (CPU): two real processes each hold an operation open; the registry that /metrics serves (get_prometheus_registry()) reports in-flight 2, then 0 and two duration samples after both finish. Fails if livesum is dropped.
  • Cancellation (CPU): cancelling a route task while the engine call is pending releases in-flight and still observes the duration. Fails if the recorder only handles Exception.
  • Also (CPU): recorder success/exception paths, each of the five routes, rejected input records nothing, lazy registration on the default registry.
  • Rust: the recorder guard tracks in-flight and duration; the HTTP lifecycle test checks update/finish duration counts and in-flight; invalid payloads record no duration.

Test Result

  • Python unit: 11 passed. E2E: 1 passed on GB200.
  • Rust: vllm-metrics 1 passed, vllm-server weight_transfer routes 3 passed.
  • pre-commit (ruff, mypy, rustfmt, typos, test tethering) clean.
  • No model output path changes, so no accuracy eval.

AI assistance (Claude Code) was used for this change; the submitter reviewed every changed line and ran the tests above.

🤖 Generated with Claude Code

@Ronald1995

Copy link
Copy Markdown
Collaborator Author

Scope correction pushed in fe7e73affd1bee4dd2c8047dfa038dddc067c556.

The earlier mainline synchronization accidentally included unrelated changes from a local tracking reference that did not match the PR's actual upstream base. Those changes have now been removed from the final PR diff:

  • tests/kernels/moe/test_b12x.py
  • tests/models/test_initialization.py
  • tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
  • vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
  • vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

I verified the five files against the immutable upstream merge-base 6fbb00b18874e27ba7d7adc0a3b8e93fee763ab1 and then checked GitHub's actual head SHA and Files changed list after pushing: this PR now contains exactly 3 intended files. The RL changes themselves are unchanged. This correction uses an additive cleanup commit; it does not rewrite contributor history.

This corrects the earlier scope-validation claim. No new GPU execution results are claimed.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995
Ronald1995 force-pushed the codex/rl-weight-operation-metrics branch 2 times, most recently from 777dcb0 to 46a0704 Compare September 22, 2026 03:55
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995
Ronald1995 force-pushed the codex/rl-weight-operation-metrics branch 3 times, most recently from dbf7e83 to 1f0262b Compare September 22, 2026 03:58
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995
Ronald1995 force-pushed the codex/rl-weight-operation-metrics branch from 1f0262b to 54ce10e Compare September 22, 2026 03:59
Add bounded-label telemetry around dispatched frontend weight operations: init,
start, start_draft, update, finish and set_version. Counters report
success/error, histograms measure one dispatched operation, and an in-flight
gauge tracks concurrency without one request clearing another. Input validation
stays outside the recorder.

Both frontends expose the same names and labels: the Python RL HTTP routes use
an independently owned `WeightOperationMetrics` created on first use on
`get_prometheus_registry()`, and the Rust HTTP weight-transfer routes use a
drop-safe RAII recorder.

`finish` and the weight-version handshake are recorded as separate operations so
that a version failure is not reported as a failed finish; the Python route now
calls `finish_weight_update()` and `update_weight_version()` separately to match
the Rust frontend, and `/update_weight_version` is instrumented on both.

Tests: recorder outcome/overlap tests, route-level in-flight and rejected-input
tests, a registry-binding regression test, and Rust recorder plus HTTP
lifecycle and rejection metric assertions.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

Scope split

The two changes that were bundled here are now separated, so this PR keeps only the RL weight-metric feature and the shared logger-ownership fix can be reviewed by its own owners:

  • RL weight-operation telemetry (this PR) — #11
  • PrometheusStatLogger ownership fix — #10

Both are currently based on codex/rl-weight-operation-metrics and touch disjoint files, so their diffs here are clean. They can be reviewed and landed in either order; once #10 lands, #11 rebases onto main and its diff shrinks to the 7 RL files.

What #11 changes relative to this branch

  1. finish and the weight-version handshake are counted as separate operations. Today the Rust recorder spans finish_weight_update() and set_weight_version(), so a version failure is reported as a failed finish with a duration that includes the version round-trip, while the Python route counts them as one finish. The same metric name therefore means different things per frontend. Add cache watermark to avoid frequent cache eviction #11 records finish when the engine finishes, and counts the handshake as set_version on both sides.
  2. /update_weight_version gets instrumented on both frontends (currently counted nowhere).
  3. The Python recorder is created on first use on get_prometheus_registry(). The module-level singleton is created at import time on the default REGISTRY while /metrics serves get_prometheus_registry(); that only works because prometheus_client picks the mmap-backed value class when PROMETHEUS_MULTIPROC_DIR happens to be set at creation time. A regression test now fails if the binding moves back to import time.
  4. Test gaps: rejected input asserts the in-flight gauge is untouched (Python and Rust), the Rust lifecycle test asserts finish / set_version counts, and a new Rust test covers "finish succeeds, version fails" (500, finish = success, set_version = error).
  5. The metrics documentation states that a cancelled operation is recorded as status="error".

Test evidence on this branch

H20 cluster (vllm-h20-head -> Slurm h20 -> vllm-h20-02), Python 3.12.14 venv with torch 2.14.0+cpu + requirements/common.txt, real vllm import chain:

13 passed, 1 warning in 0.83s

cargo test -p vllm-metrics on the same cluster: 3 passed; 0 failed (includes weight_operation_recorder_tracks_success_and_error).

cargo test -p vllm-server --lib weight_transfer is still unverified: the H20 compute-node image has no OpenSSL development headers (Failed to find OpenSSL development headers, SERVER_RC=101), the account has no sudo/apt, and the CUDA devel container lacks libssl-dev too. The route-level assertions added in #11 need a run on a host with pkg-config libssl-dev.

Review logistics

@Ronald1995 — since these branches sit on top of this PR's branch, the intended sequence is: review/land #10 (shared fix, separate owners) and #11 (RL metric, RL maintainers), then this branch can be retired or fast-forwarded. Nothing here needs to be re-litigated; the RL-specific design decisions are unchanged apart from the finish/set_version split described above.

The per-collector ownership tracking (`PrometheusStatLoggerRegistry`,
`unregister_prometheus_stat_logger_metrics`) and the connector /
speculative-decoding / perf metric-class plumbing belong with the shared
`PrometheusStatLogger` lifetime fix, not with the RL weight-operation metrics:
the same bug also silently drops the tool-call parser metrics, the WebSocket
connection metrics and the prometheus-fastapi-instrumentator HTTP metrics.

- Restore `vllm/v1/metrics/{loggers,prometheus,ray_wrappers,perf}.py`,
  `vllm/v1/spec_decode/metrics.py` and the KV/EC connector metric modules to
  their state on the PR base; this PR no longer touches them.
- Drop `tests/v1/metrics/test_prometheus_logger_ownership.py`, which tests that
  ownership behaviour and moves with it.
- Keep everything that is specific to this PR: the recorders in both frontends,
  the `finish` / `set_version` split, the lazy binding of the Python recorder to
  `get_prometheus_registry()`, and all the tests for them.

Known limitation after this revert: in single-process deployments
(`api_server_count == 1`, the default) a `PrometheusStatLogger` rebuild removes
every `vllm:*` collector from the default registry, so these frontend metrics
stop being exported from `/metrics` until the logger-ownership fix lands. The
recorder object stays alive and keeps counting; only the export is affected.
This is the pre-existing behaviour for all frontend-owned metrics.

Tested after the revert with the real modules: 11/11 weight-operation tests pass
(recorder outcomes, overlap, lazy registry binding, route in-flight/rejection,
finish-succeeds/version-fails, unversioned finish, update_weight_version).

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

Scope narrowed: the logger-ownership change is reverted out of this PR

Pushed 9368e3a4b2 (Revert the logger-ownership change out of this PR). This PR now contains 7 files / +741 / -11, all weight-operation metrics:

rust/src/metrics/src/api_server.rs
rust/src/server/src/routes/tests.rs
rust/src/server/src/routes/weight_transfer.rs
tests/entrypoints/unit_tests/test_rl_weight_operation_metrics.py
tests/entrypoints/unit_tests/test_rl_weight_operation_routes.py
vllm/entrypoints/serve/dev/rlhf/api_router.py
vllm/entrypoints/serve/dev/rlhf/metrics.py

Reverted to the PR base (this PR no longer touches them):

  • vllm/v1/metrics/loggers.py — back to unregister_vllm_metrics(); the _use_prometheus_registry flag and the functools.partial rebinding are gone.
  • vllm/v1/metrics/prometheus.py — PrometheusStatLoggerRegistry and unregister_prometheus_stat_logger_metrics() removed; the file is now byte-identical to base.
  • vllm/v1/metrics/ray_wrappers.py, vllm/v1/metrics/perf.py, vllm/v1/spec_decode/metrics.py.
  • vllm/distributed/kv_transfer/kv_connector/v1/metrics.py and vllm/distributed/ec_transfer/ec_connector/metrics.py.
  • tests/v1/metrics/test_prometheus_logger_ownership.py — deleted, it tests that behaviour and moves with it.

Kept, because it is specific to this PR: both recorders, the finish / set_version split, set_version coverage on both frontends, and the lazy binding of the Python recorder to get_prometheus_registry().

Known limitation (please read before reviewing)

In single-process deployments (api_server_count == 1, the default), a PrometheusStatLogger rebuild calls unregister_vllm_metrics(), which removes every vllm:* collector from the default registry. These frontend metrics then stop being exported from /metrics for the rest of the process. The recorder object stays alive and keeps counting; only the export is affected. Measured on a scratch build:

SINGLE PROCESS
  before unregister, exposed: True
  after  unregister, exposed: False     <-- export gone
  metric object still usable: 6.0       <-- counting continues
MULTIPROCESS (api_server_count > 1)
  before unregister, exposed: True
  after  unregister, exposed: True      <-- survives via the mmap file

This is the pre-existing behaviour for every frontend-owned metric in the tree (the tool-call parser metrics and the WebSocket connection metrics have the same exposure problem today), and it is fixed separately in a follow-up PR rather than here, so that this PR stays focused on the RL weight metrics.

Test status after the revert

Re-ran the weight-operation tests against the real modules: 11/11 pass (recorder success/error/cancellation, overlapping requests, lazy registry binding, recorder-created-once, route in-flight accounting, rejected input, finish succeeds while the version handshake fails, unversioned finish, /update_weight_version). ruff check and ruff format --check pass; check_test_tethering.py passes for both new test files.

The Rust route-level assertions (cargo test -p vllm-server --lib weight_transfer) are still unexecuted: the H20 compute-node image has no OpenSSL development headers (Failed to find OpenSSL development headers, SERVER_RC=101), the account has no sudo/apt, and the CUDA devel container lacks libssl-dev as well. cargo test -p vllm-metrics does pass on that cluster (3 passed; 0 failed). I will confirm the route tests on a host with pkg-config libssl-dev before this leaves draft.

…emantics

Rename and re-scope the frontend weight-operation telemetry so that the metrics
describe logical frontend weight operations, not HTTP requests:

- `vllm:rl_weight_update_operations_total` (was `..._requests_total`)
- `vllm:rl_weight_update_operation_duration_seconds` (was `..._request_...`)
- `vllm:rl_weight_update_operations_in_flight` (was `..._requests_in_flight`)

One HTTP request maps to zero or more operations: input validation happens
outside the recorder, and `/finish_weight_update` with a `weight_version`
records `finish` for `engine.finish_weight_update()` plus `set_version` for
`engine.update_weight_version()`. `finish{success}` together with
`set_version{error}` (HTTP 500 overall) is a valid outcome and is what makes the
failure attributable; `/update_weight_version` records `set_version` as well.
Help text, docstrings and comments no longer describe these as HTTP requests.

Metric creation no longer targets the scrape registry. The collector is created
on the process default registry with no `registry=` argument, which is the
`prometheus_client` contract for multiprocess mode: the values become
mmap-backed and the `/metrics` scrape registry aggregates them through
`MultiProcessCollector`. The lazy singleton stays - its purpose is that metric
objects exist only after multiprocess Prometheus setup has run - but it is no
longer described as "binding to the registry returned to /metrics". Tests may
still pass a private registry for isolation.

Tests:

- new `test_rl_weight_operation_multiprocess.py` uses real child processes and a
  real `PROMETHEUS_MULTIPROC_DIR`: two processes aggregating into
  `operations_total == 2`, `operation_duration_seconds_count == 2` and
  `operations_in_flight == 2`; per-operation series staying separate; and
  `mark_process_dead` clearing a dead process's stale `in_flight` while leaving a
  live process untouched (skipped on Windows, which locks the mmap files);
- `test_rl_weight_operation_routes.py` now parametrizes every operation
  (`init`, `start`, `start_draft`, `update`, `finish`, `set_version`) for both the
  success path (gauge 1 during the call, `operations_total{success}` and duration
  count +1, gauge back to 0) and the engine-exception path (`{error}` +1, no
  `{success}`, gauge released), plus rejected input touching no series at all;
- the recorder unit tests cover success/error/cancellation, overlapping
  operations, the `finish` + `set_version` split, and lazy creation.

Rust mirrors the same names, labels, buckets and boundaries; the counter is
registered without the `_total` suffix because the crate adds it when encoding.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

Semantics revision pushed: requests -> operations, and correct multiprocess registry usage

Latest head: cf3ee740b8. The PR now describes logical frontend weight operations everywhere (names, help text, docstrings, comments) instead of HTTP requests, and the metric creation registry follows the prometheus_client multiprocess contract.

Renames

Before After
vllm:rl_weight_update_requests_total vllm:rl_weight_update_operations_total
vllm:rl_weight_update_request_duration_seconds vllm:rl_weight_update_operation_duration_seconds
vllm:rl_weight_update_requests_in_flight vllm:rl_weight_update_operations_in_flight

The names now say operations because one HTTP request maps to zero or more of them (/finish_weight_update with a version records finish and set_version). The help text states this, and the PR body documents that summing across the operation label is not meaningful: update is per weight chunk while finish/set_version/init/start/start_draft are lifecycle events. So sum(..._operations_total) is not "number of weight updates" - query per operation instead.

Registry usage

The recorder is no longer created with get_prometheus_registry() as its creation registry. In multiprocess mode that returns a CollectorRegistry whose purpose is scrape/aggregation (MultiProcessCollector), not metric creation. The collector is now created on the process default registry with no registry= argument, which is what prometheus_client expects: the values become mmap-backed and the /metrics scrape registry aggregates them. The lazy singleton stays, but its stated purpose is now "metric objects are instantiated after multiprocess Prometheus setup has occurred" rather than "bind to the registry returned to /metrics".

New real multiprocess regression test

tests/entrypoints/unit_tests/test_rl_weight_operation_multiprocess.py uses real child processes and a real PROMETHEUS_MULTIPROC_DIR (skipped on Windows, where the mmap files are locked):

  • two processes each completing one operation: operations_total == 2, operation_duration_seconds_count == 2;
  • two processes each holding one operation open: operations_in_flight == 2, which is what actually validates multiprocess_mode="livesum";
  • different operation labels stay separate across processes;
  • mark_process_dead clears a dead process's stale in-flight gauge while a live process is unaffected;
  • a subprocess asserts the module import creates no collectors and that the counter/histogram/livesum mmap files appear only after first use.

The old test that mocked get_prometheus_registry() and asserted registration into the mock registry is gone - it did not test multiprocess behaviour, and test_metrics_are_created_lazily now covers the lazy contract directly.

Test evidence (H20)

29 passed, 1 warning in 29.18s

That is the recorder tests, the route wiring tests parametrized over all six operations (success and engine-exception paths, plus rejected input touching no series), and the four multiprocess tests. cargo test -p vllm-metrics on the same cluster: 3 passed; 0 failed.

cargo test -p vllm-server --lib weight_transfer is still blocked by the missing OpenSSL development headers on that image (no sudo/apt; the CUDA devel container has no libssl-dev either). The Rust route assertions use the new metric names and will be confirmed on a host with pkg-config libssl-dev before this leaves draft.

Known limitation, unchanged and still handled separately: in single-process deployments a PrometheusStatLogger rebuild calls unregister_vllm_metrics(), which removes these collectors from the default registry, so they stop being exported until the logger-ownership fix (pr-55781-ownership-fix) lands. That ownership change is deliberately not part of this PR.

Three follow-up commits were needed to get the multiprocess harness honest - a child stderr pipe that deadlocked on vLLM import warnings, a double increment of the held gauge, and children that never completed an operation so no counter was written. All three were test-harness bugs, not product bugs; the product code is unchanged apart from the rename and the registry change above.

@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

❌ This PR is 1 commit behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #92011 for commit ef0da41a5338.

@mergify

mergify Bot commented Sep 30, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Ronald1995.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Sep 30, 2026
Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@mergify mergify Bot removed the needs-rebase label Sep 30, 2026
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

❌ This PR is 3 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #92045 for commit 5a91e51b8eb9.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #92064 for commit 4f2f09d79296.

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

❌ This PR is 2 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

@Ronald1995

Copy link
Copy Markdown
Collaborator Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #92072 for commit 5103cca02ade.

@aoshen02
aoshen02 enabled auto-merge (squash) September 30, 2026 07:54
@aoshen02
aoshen02 merged commit 713ec07 into vllm-project:main Sep 30, 2026
109 checks passed
}


def test_weight_sync_metrics_aggregate_across_api_servers(tmp_path):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the help, will avoid it next time.

@MatthewBonanni

Copy link
Copy Markdown
Member

The test is also failing on main:

Build 92138: build · job
Build 92174: build · job
Build 92180: build · job

@AndreasKaratzas

Copy link
Copy Markdown
Member

@aoshen02 this is still flaky. Please review this one:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

frontend kv-connector ready ONLY add when PR is ready to merge/full CI is needed rust speculative-decoding

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants