feat(sglang): gate chat-shaped Prometheus collectors on embedding worker (DIS-2107) - #9830
Merged
Conversation
Closes the first half of DIS-2107. The SGLang publisher previously registered chat-shaped collectors (SGLang's multiprocess ``sglang:*`` metrics, the Dynamo ``LLMBackendMetrics`` gauges ``total_blocks`` / ``gpu_cache_usage`` / ``model_load_time``, the ``DynamoSglangPublisher`` with its ZMQ scheduler pull + KV-event + FPM-relay wiring) on every worker, including embedding workers where they emit zeros forever. Operators monitoring an embedding fleet saw noise instead of signal. ``setup_sgl_metrics`` now short-circuits when ``config.dynamo_args.embedding_worker`` is set: no chat-shaped constructors run. It still returns a uniform ``(publisher_or_None, asyncio.Task, metrics_labels)`` shape so the embedding init path can keep the same ``await + cancel()`` cleanup sequence as the chat path. The corresponding ``dynamo_embedding_batch_size`` / ``dynamo_embedding_input_tokens`` collectors that DO belong on an embedding worker are registered separately by ``init_embedding.py`` via ``init_embedding_metrics`` (see PR #9753). Tests ----- ``test_sglang_publisher.py`` gains two new cases: - ``test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker`` — patches ``setup_prometheus_registry``, ``register_engine_metrics_callback``, ``LLMBackendMetrics``, and ``DynamoSglangPublisher`` to ``raise`` if invoked. Setting ``embedding_worker=True`` must not trigger any of them, the return tuple must have ``publisher is None``, the task must be cancellable, and the model label still propagates. - ``test_setup_sgl_metrics_returns_publisher_for_chat_worker`` — sibling check that the chat path is unchanged: all chat-shaped constructors fire and a publisher comes back. vLLM publisher gating is intentionally NOT in this PR — vLLM has no embedding-worker code path on ``main`` today. The mirror change lands once the vLLM-side embedding worker (DIS-2092) merges and the vLLM-side embedding observation PR (DIS-2094 part 3) is on deck. Refs DIS-2107. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
Contributor
WalkthroughThe PR adds embedding worker support to ChangesEmbedding worker metrics bypass
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
components/src/dynamo/sglang/publisher.py (1)
493-493: ⚡ Quick winReplace defensive
getattr(..., "embedding_worker", False)with directconfig.dynamo_args.embedding_worker
embedding_workeris declared onDynamoSGLangConfigand is populated from parsed CLI args (with--embedding-workerdefaulting toFalse), so thegetattr(..., default)silently masks missing/incorrect config. Use direct access to fail fast.📝 Proposed change
- if getattr(config.dynamo_args, "embedding_worker", False): + if config.dynamo_args.embedding_worker:🤖 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/sglang/publisher.py` at line 493, Replace the defensive getattr call with direct attribute access: where the code checks getattr(config.dynamo_args, "embedding_worker", False) in publisher.py, use config.dynamo_args.embedding_worker instead so missing/invalid configs raise immediately; this targets the attribute on the DynamoSGLangConfig-backed dynamo_args object and ensures the code fails fast rather than silently using a default.
🤖 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/sglang/publisher.py`:
- Around line 478-481: Remove the internal Linear ID from the docstring: replace
the `DIS-2094` reference with either a public GitHub issue reference (e.g.,
`#NNNN` or `GH-NNNN`) or reword to describe the cross-reference without a ticket
ID; update the sentence mentioning `init_embedding.py` and
`init_embedding_metrics` and the cleanup note about `metrics_task.cancel()` so
it reads the same but contains no `DIS-2094` token (e.g., "see init_embedding.py
/ init_embedding_metrics for embedding-shaped metrics registration" or "see
GH-NNNN" if you have a public issue).
In `@components/src/dynamo/sglang/tests/test_sglang_publisher.py`:
- Line 482: The inline comment contains an internal Linear ticket ID ("# ----
DIS-2107: per-worker metric gating ----"); remove the DIS-2107 reference and
replace it with an approved identifier or descriptive label (for example change
to a GitHub issue/PR number like "GH-1234" or simply "# ---- per-worker metric
gating ----") so the comment no longer contains internal Linear IDs; update the
comment text where it appears to use the new label.
- Around line 505-583: Re-run the code formatter so the test file meets the
project's Black settings: run the pre-commit black hook (e.g., pre-commit run
--all-files or black on the test file) and commit the reformatted file;
specifically wrap the long lines in the test function
test_setup_sgl_metrics_returns_publisher_for_chat_worker (e.g., the
_StubPublisher.metrics_publisher.create_endpoint lambda and any long
monkeypatch.setattr calls) so they comply with Black's line-length rules and
ensure the file passes pre-commit.
---
Nitpick comments:
In `@components/src/dynamo/sglang/publisher.py`:
- Line 493: Replace the defensive getattr call with direct attribute access:
where the code checks getattr(config.dynamo_args, "embedding_worker", False) in
publisher.py, use config.dynamo_args.embedding_worker instead so missing/invalid
configs raise immediately; this targets the attribute on the
DynamoSGLangConfig-backed dynamo_args object and ensures the code fails fast
rather than silently using a default.
🪄 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: Pro
Run ID: a5d91669-1fc7-4f95-82fb-ce825e142061
📒 Files selected for processing (2)
components/src/dynamo/sglang/publisher.pycomponents/src/dynamo/sglang/tests/test_sglang_publisher.py
Addresses three findings on PR #9830: 1. mypy: ``setup_sgl_metrics`` now declares ``Optional[DynamoSglangPublisher]`` in its return tuple, but the chat-worker call sites (init_llm.py x2, init_diffusion.py x1) pass the publisher to functions that expect the non-Optional type. Add ``assert publisher is not None`` at those three call sites — those code paths are reached only when ``embedding_worker=False``, so the assertion is correct by construction and narrows the type for mypy. 2. mypy: the embedding-worker no-op task used ``asyncio.create_task(asyncio.Event().wait())``. ``Event.wait()`` returns ``Literal[True]`` when set, which mypy then conflicts with the chat path's ``publisher.run() -> None`` return at the second ``create_task`` call. Wrap the wait in a local ``async _idle() -> None`` so both branches feed ``Coroutine[Any, Any, None]`` into ``create_task``. 3. Repository policy: scrub internal Linear ticket IDs from source. Removed ``DIS-2094`` from the publisher docstring and ``DIS-2107`` from the test-file section header. 4. Black: re-format the two files mentioned by pre-commit. No behavioral changes. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com>
dynamo-ops
approved these changes
May 21, 2026
tzulingk
enabled auto-merge (squash)
May 22, 2026 22:54
tzulingk
added a commit
that referenced
this pull request
May 27, 2026
Addresses three findings on PR #9830: 1. mypy: ``setup_sgl_metrics`` now declares ``Optional[DynamoSglangPublisher]`` in its return tuple, but the chat-worker call sites (init_llm.py x2, init_diffusion.py x1) pass the publisher to functions that expect the non-Optional type. Add ``assert publisher is not None`` at those three call sites — those code paths are reached only when ``embedding_worker=False``, so the assertion is correct by construction and narrows the type for mypy. 2. mypy: the embedding-worker no-op task used ``asyncio.create_task(asyncio.Event().wait())``. ``Event.wait()`` returns ``Literal[True]`` when set, which mypy then conflicts with the chat path's ``publisher.run() -> None`` return at the second ``create_task`` call. Wrap the wait in a local ``async _idle() -> None`` so both branches feed ``Coroutine[Any, Any, None]`` into ``create_task``. 3. Repository policy: scrub internal Linear ticket IDs from source. Removed ``DIS-2094`` from the publisher docstring and ``DIS-2107`` from the test-file section header. 4. Black: re-format the two files mentioned by pre-commit. No behavioral changes. Signed-off-by: Tzu-Ling <tzulingk@nvidia.com> (cherry picked from commit 2029e20)
dillon-cullinan
approved these changes
May 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview:
Closes the SGLang half of DIS-2107. Operators monitoring an embedding fleet today see chat-shaped collectors emit zeros forever: KV gauges, prefill/decode counters, SGLang's
sglang:*multiproc metrics, and the DynamoLLMBackendMetricsseries all live on every worker regardless of whether the engine even has a KV cache. This PR makes the SGLang publisher skip all of that wiring for embedding workers.Details:
setup_sgl_metrics(engine, config, generate_endpoint)now short-circuits whenconfig.dynamo_args.embedding_workeris True:setup_prometheus_registry()(SGLangsglang:*multiproc metrics),LLMBackendMetrics()(total_blocks, gpu_cache_usage, model_load_time),DynamoSglangPublisher()(ZMQ scheduler pull + KV-event publisher + FPM relay).(publisher_or_None, asyncio.Task, metrics_labels)shape soinit_embedding.pycan keep the sameawait metrics_task+metrics_task.cancel()cleanup as the chat-worker path. The task is a never-completing waiter for uniformity.init_embedding.py: the embedding-shaped collectors (dynamo_embedding_batch_size,dynamo_embedding_input_tokens— see feat(sglang): Prometheus metrics for embedding workload shape (DIS-2094 part 1) #9753).The result: on an embedding worker's
/metrics, you see only the metrics that actually move on a pooling engine, not the always-zero chat noise.Tests:
Two new cases in
components/src/dynamo/sglang/tests/test_sglang_publisher.py:test_setup_sgl_metrics_skips_chat_pipeline_for_embedding_worker— patches each chat-shaped constructor with a function that raises if invoked. Withconfig.dynamo_args.embedding_worker=True, none of them must fire. Verifies the return tuple too:publisher is None,taskis a cancellableasyncio.Task,metrics_labelsstill carries("model", served_name).test_setup_sgl_metrics_returns_publisher_for_chat_worker— sibling check that the chat-worker path is unchanged:register_engine_metrics_callback,LLMBackendMetrics, andDynamoSglangPublisherall fire exactly once. Confirms the gating doesn't accidentally short-circuit the default code path.Why not vLLM in this PR?
vLLM's embedding-worker shape (DIS-2092 / PR #9713) hasn't merged yet, and the vLLM-side embedding observation PR (DIS-2094 part 3) is still pending. Once both land, the same gating will mirror to
components/src/dynamo/vllm/publisher.pyin a follow-up PR.Where should the reviewer start?
components/src/dynamo/sglang/publisher.py— the early-exit insetup_sgl_metrics(). The shape of the return tuple is preserved by design.components/src/dynamo/sglang/tests/test_sglang_publisher.py— the two new tests at the bottom of the file.Linear: https://linear.app/nvidia/issue/DIS-2107
Summary by CodeRabbit
New Features
Tests