Conversation
|
👋 Hi gitover22! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
| factory = StatLoggerFactory( | ||
| endpoint=generate_endpoint, | ||
| component_gauges=component_gauges, | ||
| ) |
There was a problem hiding this comment.
🔍 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.
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
🔍 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).
Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
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.
WalkthroughChangesPrefill metrics
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Comment |
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
| dp_size, | ||
| ) | ||
| factory.set_num_gpu_blocks_all(per_rank_num_gpu_blocks or 0) | ||
| factory.init_publish() |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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).
|
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
|
@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
|
/ok to test 7331f9d |
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