Skip to content

fix(vllm): publish KV gauges from prefill workers - #12559

Open
gitover22 wants to merge 1 commit into
ai-dynamo:mainfrom
gitover22:fix/vllm-prefill-kv-gauges
Open

gitover22 wants to merge 1 commit into
ai-dynamo:mainfrom
gitover22:fix/vllm-prefill-kv-gauges

Conversation

@gitover22

@gitover22 gitover22 commented Aug 3, 2026

Copy link
Copy Markdown

Overview

Fix missing dynamo_component_total_blocks and dynamo_component_gpu_cache_usage_percent samples on ordinary disaggregated vLLM prefill workers.

This consolidates the remaining fix for #11919 on current main while preserving the retained per-rank snapshot logger lifecycle merged in #14258.

Details

  • Create and pass a StatLoggerFactory during ordinary prefill engine setup.
  • Share per-rank KV capacity initialization and initial publication between ordinary and restored prefill workers.
  • Publish configured total capacity at startup while keeping used blocks and cache utilization at zero until requests arrive.
  • Add deterministic nonzero-DP-rank assertions for post-bind startup, active, and idle publication, including fresh publisher state and the disabled-stat no-op path.
  • Extend the two-GPU disaggregated serve scenario to scrape the prefill SYSTEM2 metrics endpoint.

No snapshot state format or restored scheduler/FPM identity plumbing is changed.

Where should the reviewer start?

Start with components/src/dynamo/vllm/worker_factory.py, then components/src/dynamo/vllm/publisher.py. The focused regression coverage is in their adjacent test files.

Validation

  • python -m pytest -q components/src/dynamo/vllm/tests/test_vllm_worker_factory.py components/src/dynamo/vllm/tests/test_vllm_publisher.py components/src/dynamo/vllm/tests/test_vllm_snapshot.py
    • 65 passed with vLLM 0.28.0 and PyTorch 2.13.0+cu130.
  • python -m pytest -xvv -s 'tests/serve/test_vllm.py::test_serve_deployment[disaggregated-2]'
    • 1 passed in 240.40 seconds on two NVIDIA H20 GPUs.
    • The prefill scrape reported dynamo_component_total_blocks = 49888 and dynamo_component_gpu_cache_usage_percent = 0.0 after successful chat and completion requests through prefill, NIXL transfer, and decode.
  • python3 -m pre_commit run --files <changed files>
    • All hooks passed.
  • git diff --check
    • Passed.

Related issues

Closes #11919

@gitover22
gitover22 requested review from a team as code owners August 3, 2026 08:28
@copy-pr-bot

copy-pr-bot Bot commented Aug 3, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@gitover22
gitover22 temporarily deployed to external_collaborator August 3, 2026 08:28 — with GitHub Actions Inactive
@gitover22
gitover22 temporarily deployed to external_collaborator August 3, 2026 08:28 — with GitHub Actions Inactive
@github-actions github-actions Bot added the fix label Aug 3, 2026
@github-actions

github-actions Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

👋 Hi gitover22! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

@github-actions github-actions Bot added external-contribution Pull request is from an external contributor backend::vllm Relates to the vllm backend labels Aug 3, 2026

@devin-ai-integration devin-ai-integration Bot left a comment

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.

Devin Review found 2 potential issues.

Open in Devin Review

Comment on lines +1284 to +1287
factory = StatLoggerFactory(
endpoint=generate_endpoint,
component_gauges=component_gauges,
)

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.

🔍 Checkpoint/snapshot prefill path still ends up with no stat logger

In the snapshot branch the engine was already built by snapshot.py:41 (setup_vllm_engine(config) with no stat logger), so vLLM never called create_stat_logger() on this new factory and StatLoggerFactory.created_logger stays None (components/src/dynamo/vllm/publisher.py:144). Consequently set_num_gpu_blocks_all() and init_publish() at components/src/dynamo/vllm/worker_factory.py:1313-1314 are silent no-ops, and prefill workers restored from a checkpoint still emit no total_blocks / gpu_cache_usage_percent samples. This mirrors the pre-existing decode behavior (worker_factory.py:1001-1004), so it is not a regression, but the PR's stated goal is only achieved on the non-snapshot path.

Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Good catch. Fixed in b9bf5a1: snapshot setup now installs an unbound stat logger before capture, preserves it with the engine state, and binds the runtime endpoint only after restore. The restored prefill path reuses that factory before setting capacity and seeding gauges. Added regression coverage for snapshot setup, deferred binding, and the restored prefill path.

# a restart of the EngineCore process.
os.environ[ENV_FPM_WORKER_ID] = fpm_worker_id
else:
factory = StatLoggerFactory(endpoint=generate_endpoint)

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.

🔍 Prefill worker now advertises worker load metrics over NATS

Installing the stat logger also constructs a WorkerMetricsPublisher and creates its NATS endpoint on the prefill generate endpoint (components/src/dynamo/vllm/publisher.py:31-42), and init_publish() immediately publishes kv_used_blocks=0. Beyond Prometheus gauges (the PR's stated goal) this makes prefill workers visible to KV/load-aware routing consumers that previously saw nothing from them. Worth confirming the router treats these new prefill samples as intended, especially if vllm_config.cache_config.num_gpu_blocks is None (the or 0 fallback then pins reported used blocks to 0 forever).

Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Confirmed this is intended. KvWorkerMonitor creates a separate subscriber for the exact prefill endpoint and validates endpoint membership before accepting events, so these samples cannot be treated as decode-worker samples. They provide the prefill worker load signal used by overload handling. vLLM applies EngineCoreReadyResponse before setup returns and populates cache_config.num_gpu_blocks there; None is only a pre-initialization/shutdown state. The two-H20 disaggregated test passes with the new publisher active.

@coderabbitai

coderabbitai Bot commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

Changes

Prefill metrics

Layer / File(s) Summary
Prefill stat logger wiring
components/src/dynamo/vllm/worker_factory.py
The prefill worker creates StatLoggerFactory, configures per-rank GPU KV blocks, passes it to engine setup, and initializes metric publishing.
Metrics wiring validation
components/src/dynamo/vllm/tests/test_vllm_worker_factory.py, tests/serve/test_vllm.py
Tests verify logger creation, engine setup wiring, GPU block calculation, metric publication, and a disaggregated vLLM metrics request through SYSTEM2.

Estimated code review effort: 3 (Moderate) | ~20 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Linked Issues check ✅ Passed The changes satisfy issue #11919 by wiring stat logging and publishing seeded KV cache gauges for disaggregated prefill workers.
Out of Scope Changes check ✅ Passed All modified runtime code and tests directly support KV gauge publishing for disaggregated vLLM prefill workers.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check ✅ Passed The title clearly and concisely describes the main change: publishing KV cache gauges from vLLM prefill workers.
Description check ✅ Passed The description explains the change, identifies review starting points, documents validation, and includes the required linked issue with "Closes #11919". It is complete despite minor heading-format d…

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

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.

Actionable comments posted: 3

🧹 Nitpick comments (1)
components/src/dynamo/vllm/tests/test_vllm_worker_factory.py (1)

808-813: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Assert the published gauge values.

The mock assertions verify call plumbing only. They do not prove that total_blocks and gpu_cache_usage_percent have the expected values. Add a focused publisher/component-gauge assertion or a metrics scrape assertion for the pre-request state.

As per path instructions, tests should assert metric initialization and publication deterministically.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@components/src/dynamo/vllm/tests/test_vllm_worker_factory.py` around lines
808 - 813, Extend the test around setup_vllm_engine and stat_logger.init_publish
to assert the pre-request published gauge values for total_blocks and
gpu_cache_usage_percent, rather than only verifying call plumbing. Use a focused
publisher/component-gauge or deterministic metrics-scrape assertion and confirm
both expected values during metric initialization.

Source: Path instructions

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@components/src/dynamo/vllm/tests/test_vllm_worker_factory.py`:
- Around line 730-736: Update the test constructing engine_tuple to replace the
hardcoded "/tmp/prom" value with an isolated temporary path supplied by pytest's
tmp_path fixture or Python's tempfile module, while preserving the expected path
type and test behavior.

In `@components/src/dynamo/vllm/worker_factory.py`:
- Around line 1282-1287: Preserve the original StatLoggerFactory or publisher
when creating the snapshot EngineSetupResult, then reuse that instance in the
snapshot_engine branch instead of constructing a new factory around
engine_client. Update the prefill worker’s logger calls at the referenced lines
to operate on the preserved logger, and add a regression test covering
configured capacity and initial gauges for snapshot engines.

In `@tests/serve/test_vllm.py`:
- Around line 372-376: Add an explicit pytest timeout marker to the
disaggregated serve test containing metric_payload_default, using a value three
times the measured normal runtime. Keep the existing test behavior unchanged and
ensure the marker covers worker startup and network requests.

---

Nitpick comments:
In `@components/src/dynamo/vllm/tests/test_vllm_worker_factory.py`:
- Around line 808-813: Extend the test around setup_vllm_engine and
stat_logger.init_publish to assert the pre-request published gauge values for
total_blocks and gpu_cache_usage_percent, rather than only verifying call
plumbing. Use a focused publisher/component-gauge or deterministic
metrics-scrape assertion and confirm both expected values during metric
initialization.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 41c0ef53-d7ed-4652-83b8-7a419c8fec73

📥 Commits

Reviewing files that changed from the base of the PR and between c574e4d and 4a4b555.

📒 Files selected for processing (3)
  • components/src/dynamo/vllm/tests/test_vllm_worker_factory.py
  • components/src/dynamo/vllm/worker_factory.py
  • tests/serve/test_vllm.py

Comment thread components/src/dynamo/vllm/tests/test_vllm_worker_factory.py
Comment thread components/src/dynamo/vllm/worker_factory.py Outdated
Comment thread tests/serve/test_vllm.py
dp_size,
)
factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0)
factory.init_publish()

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Calling init_publish() after set_num_gpu_blocks_all() immediately overwrites the initialized total_blocks gauge with 0 because DynamoStatLoggerPublisher.init_publish() ignores self.num_gpu_block, so prefill worker /metrics reports zero capacity until the first scheduler stats record. Fix: initialize total_blocks from the configured per-rank block count while keeping used blocks and cache usage at zero.

🤖 AI Fix

In components/src/dynamo/vllm/publisher.py, update DynamoStatLoggerPublisher.init_publish to call self.component_gauges.set_total_blocks(dp_rank_str, self.num_gpu_block) instead of 0, leave self.inner.publish(... kv_used_blocks=0) and set_gpu_cache_usage(..., 0.0) unchanged, and add a unit test in components/src/dynamo/vllm/tests/test_vllm_publisher.py that sets num_gpu_block, calls init_publish, and asserts the total-block gauge receives the configured capacity.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Fixed in 8bf256e. init_publish now seeds total_blocks from the configured per-rank capacity while keeping kv_used_blocks and gpu_cache_usage at zero. The deterministic publisher test verifies an initial capacity of 48 and the subsequent scheduler update.

@gitover22
gitover22 requested review from a team as code owners August 3, 2026 08:49
@gitover22
gitover22 temporarily deployed to external_collaborator August 3, 2026 08:49 — with GitHub Actions Inactive
@pull-request-size pull-request-size Bot added size/L and removed size/M labels Aug 3, 2026
@gitover22
gitover22 temporarily deployed to external_collaborator August 3, 2026 08:56 — with GitHub Actions Inactive
Comment thread components/src/dynamo/vllm/publisher.py Outdated
raise RuntimeError("vLLM stat logger factory endpoint is already bound")
self.endpoint = endpoint
if self.created_logger is not None:
self.created_logger.bind_endpoint(endpoint)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

bind_endpoint() only binds self.created_logger, but vLLM calls this factory once per DP rank and created_logger is overwritten each time, so snapshot workers with multiple DP ranks leave every earlier rank's logger without a NATS endpoint. Fix: retain all created DynamoStatLoggerPublisher instances and apply endpoint binding, capacity seeding, and initial publish to each one.

🤖 AI Fix

In components/src/dynamo/vllm/publisher.py, change StatLoggerFactory to store self.created_loggers: list[DynamoStatLoggerPublisher], append each DynamoStatLoggerPublisher in create_stat_logger, update bind_endpoint, set_num_gpu_blocks_all, and init_publish to iterate over all stored loggers, and update tests to create two stat loggers and assert both are bound and initialized.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Fixed in 7285a7c. StatLoggerFactory now retains every per-rank publisher and applies delayed endpoint binding, capacity configuration, and initial publication to all of them. The regression test constructs ranks 0 and 1 and verifies both publishers receive endpoint binding, capacity, initial gauges, and subsequent usage updates. The worker/publisher suite passes (35 tests).

@gitover22
gitover22 temporarily deployed to external_collaborator August 3, 2026 12:19 — with GitHub Actions Inactive
@github-actions

github-actions Bot commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

This PR is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

@gitover22

Copy link
Copy Markdown
Author

@jthomson04 the consolidation requested in #11919 is ready on current head
28d26b8 (GitHub Verified and DCO-compliant). The lightweight checks are
running, but the full pipeline is awaiting external-contribution approval.

Could you please vet this head and comment /ok to test 28d26b8 when ready?
The current head has also passed the focused 65-test suite and the two-H20
disaggregated vLLM 0.28.0 serve test locally, including the prefill SYSTEM2
metrics scrape.

@jthomson04

Copy link
Copy Markdown
Contributor

/ok to test 28d26b8

Signed-off-by: gitover22 <qidizou88@gmail.com>
@jthomson04
jthomson04 force-pushed the fix/vllm-prefill-kv-gauges branch from 28d26b8 to 7331f9d Compare September 11, 2026 17:54
@jthomson04
jthomson04 deployed to external_collaborator September 11, 2026 17:54 — with GitHub Actions Active
@jthomson04

Copy link
Copy Markdown
Contributor

/ok to test 7331f9d

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

Labels

backend::vllm Relates to the vllm backend external-contribution Pull request is from an external contributor fix size/L Stale

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Disaggregated prefill worker does not populate KV gauges (total_blocks / gpu_cache_usage_percent)

3 participants