Repository navigation
fix(vllm): publish KV gauges from prefill workers - #12559
Conversation
GitHub Actions
Inactive
GitHub Actions
Inactive
|
👋 Hi gitover22! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
Bot
added
external-contribution
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
WalkthroughChangesPrefill metrics
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 winAssert the published gauge values.
The mock assertions verify call plumbing only. They do not prove that
total_blocksandgpu_cache_usage_percenthave 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
📒 Files selected for processing (3)
components/src/dynamo/vllm/tests/test_vllm_worker_factory.pycomponents/src/dynamo/vllm/worker_factory.pytests/serve/test_vllm.py
![dynamo-review-agent[bot]](https://avatars.githubusercontent.com/u/201626793?s=60&v=4){kind=small}
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
|
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. |
7285a7c to
28d26b8
Compare
GitHub Actions
Inactive
|
@jthomson04 the consolidation requested in #11919 is ready on current head Could you please vet this head and comment |
|
/ok to test 28d26b8 |
Signed-off-by: gitover22 <qidizou88@gmail.com>
28d26b8 to
7331f9d
Compare
GitHub Actions
Inactive
|
/ok to test 7331f9d |
|
@jthomson04 @galletas1712 hi, may I ask if there are any remaining concerns with this PR? Why hasn't it been merged? |
Overview
Fix missing
dynamo_component_total_blocksanddynamo_component_gpu_cache_usage_percentsamples on ordinary disaggregated vLLM prefill workers.This consolidates the remaining fix for #11919 on current
mainwhile preserving the retained per-rank snapshot logger lifecycle merged in #14258.Details
StatLoggerFactoryduring ordinary prefill engine setup.SYSTEM2metrics 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, thencomponents/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.pypython -m pytest -xvv -s 'tests/serve/test_vllm.py::test_serve_deployment[disaggregated-2]'dynamo_component_total_blocks = 49888anddynamo_component_gpu_cache_usage_percent = 0.0after successful chat and completion requests through prefill, NIXL transfer, and decode.python3 -m pre_commit run --files <changed files>git diff --checkRelated issues
Closes #11919