feat(grpc-servicer): add Prometheus /metrics sidecar#1782
Conversation
📝 WalkthroughWalkthroughAdds a best-effort Prometheus ChangesPrometheus /metrics sidecar for gRPC servicers
Sequence Diagram(s)sequenceDiagram
participant CLI as TokenSpeed CLI
participant Serve as serve_grpc
participant Sidecar as MetricsSidecar
participant Servicer as SchedulerServicer
participant Scraper as Prometheus scraper
CLI->>Serve: metrics_port
Serve->>Sidecar: start()
Sidecar-->>Serve: bound port
Serve->>Servicer: store metrics_port
Scraper->>Sidecar: GET /metrics
Sidecar-->>Scraper: Prometheus exposition
Servicer-->>CLI: GetServerInfo.server_args
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Code Review
This pull request introduces a lightweight, best-effort Prometheus /metrics HTTP sidecar for the gRPC servicers (sglang and tokenspeed), allowing per-worker metrics to be scraped over plain HTTP. It adds a custom MetricsSidecar server using asyncio and a SchedulerLoadCollector to expose scheduler load snapshots. The feedback highlights critical improvements for the custom HTTP server implementation: adding read timeouts and EOF checks to prevent Slowloris-style resource exhaustion, awaiting writer.wait_closed() to avoid socket resource leaks, and wrapping the snapshot retrieval in a try-except block to ensure a single failure does not crash the entire metrics scrape.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3e9c6d3634
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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 `@grpc_servicer/smg_grpc_servicer/metrics.py`:
- Around line 45-58: The explicit metrics_port parameter is returned without
validation when it's not None, bypassing the range checks that are applied to
the environment variable input. Add the same port range validation (0 < port <
65536) to the explicit parameter before returning it, ensuring that invalid port
values cannot be advertised or passed into startup paths regardless of whether
they come from the explicit parameter or the METRICS_PORT_ENV environment
variable.
- Around line 69-71: The metrics_url function constructs URLs without accounting
for IPv6 addresses, which require bracket notation in URLs. Before building the
URL string with the host and port, add a check to determine if the host is an
IPv6 address. If it is an IPv6 literal, wrap it in brackets when constructing
the f-string. This ensures the resulting URL is valid and unambiguous for both
IPv4 and IPv6 hosts.
In `@grpc_servicer/smg_grpc_servicer/sglang/server.py`:
- Around line 313-318: The metrics_port is being advertised before confirming
the actual port that the metrics sidecar successfully bound to. After awaiting
the start_metrics_sidecar call in the block where enableLanceKnowledgeBase-like
conditions are checked, extract the actual bound port from the returned
metrics_sidecar object and reconcile it with the advertised metrics_port. Ensure
that GetServerInfo publishes only the port that the sidecar actually bound to,
not the requested port, to prevent advertising unreachable endpoints when the
bind fails or a different port is used.
In `@grpc_servicer/smg_grpc_servicer/tokenspeed/server.py`:
- Around line 134-139: The metrics sidecar is started via the
start_metrics_sidecar function but the servicer.metrics_port is never updated
with the actual port that the sidecar is running on, which can cause service
discovery to advertise an incorrect or dead metrics endpoint. After the await
start_metrics_sidecar call completes and assigns the result to metrics_sidecar,
extract the actual metrics port from the metrics_sidecar object and assign it to
servicer.metrics_port to ensure the advertised metadata reflects the real
sidecar state.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 76079758-704f-44fc-ba1a-aeb2a2e2bc10
📒 Files selected for processing (9)
grpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/metrics.pygrpc_servicer/smg_grpc_servicer/sglang/server.pygrpc_servicer/smg_grpc_servicer/sglang/servicer.pygrpc_servicer/smg_grpc_servicer/tokenspeed/__main__.pygrpc_servicer/smg_grpc_servicer/tokenspeed/server.pygrpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pygrpc_servicer/tests/test_metrics_sidecar.pygrpc_servicer/tests/test_servicer_metrics_advertise.py
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
3e9c6d3 to
5ddf6fc
Compare
|
Addressed review feedback (new commit on top of a rebase onto Rebase: Rebased onto Fixes:
Tests ( |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5ddf6fc7ef
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@grpc_servicer/smg_grpc_servicer/metrics.py`:
- Around line 170-179: The timeout is applied individually to each readline()
call within the while loop that reads headers, allowing a slow client to send
one header line just under the _READ_TIMEOUT limit repeatedly and keep the
handler open indefinitely. To fix this, introduce a cumulative timeout for the
entire header-read phase: record the start time before the while loop, and for
each asyncio.wait_for() call within that loop (starting with the one in the
while True block), calculate the remaining time budget by subtracting the
elapsed time from _READ_TIMEOUT and use this remaining time as the timeout
value. This ensures the timeout exception fires if the client doesn't complete
all headers within the total timeout window.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 12bee1f6-0fcb-4412-aaf8-0dfd7370bb83
📒 Files selected for processing (9)
grpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/metrics.pygrpc_servicer/smg_grpc_servicer/sglang/server.pygrpc_servicer/smg_grpc_servicer/sglang/servicer.pygrpc_servicer/smg_grpc_servicer/tokenspeed/__main__.pygrpc_servicer/smg_grpc_servicer/tokenspeed/server.pygrpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pygrpc_servicer/tests/test_metrics_sidecar.pygrpc_servicer/tests/test_servicer_metrics_advertise.py
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Round-2 review fixes pushed (36ed6e0):
Resolved 8 addressed review threads (whole-phase timeout, read-timeout/EOF, writer.wait_closed, collect() try-except, tokenspeed waiting count, advertise-only-bound-port x2, gateway-allowlist). The gateway-allowlist thread was replied to: it is handled in the W1 PR (#1779) — Verification: |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 36ed6e0a37
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
The Rust gateway scrapes per-worker engine metrics over HTTP, but the gRPC servicers (sglang, tokenspeed) had no HTTP endpoint to scrape. Add a shared, dependency-light (stdlib asyncio + prometheus_client) best-effort /metrics sidecar started alongside serve_grpc; any bind/serve failure is logged and swallowed so the gRPC server keeps serving. - metrics.py: the asyncio sidecar (one whole-head-phase read deadline so a Slowloris client can't pin a handler open), a SchedulerLoadCollector that re-exposes the GetLoads scheduler snapshot for engines without a native registry (tokenspeed), and metrics_url/metrics_server_args helpers. - Advertise the sidecar to the gateway via GetServerInfo.server_args: metrics_port always, plus metrics_url only for a routable bind host. Wildcard and loopback binds aren't reachable across the network, so only the port is advertised and the gateway pairs it with the discovered worker address. - Advertise only the actually-bound port, set after a successful start; a failed bind advertises nothing. - sglang/tokenspeed: --metrics-port (or SMG_METRICS_PORT) wiring. Part of #1773. Closes #1777. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
36ed6e0 to
c26f066
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c26f0660e4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| } | ||
| ) | ||
| # Wildcard host: only the port is advertised (gateway knows the real address). | ||
| servicer.server_args = SimpleNamespace(host="0.0.0.0") |
There was a problem hiding this comment.
Use a dataclass ServerArgs stub for SGLang test
When sglang is installed, this test invokes SGLangSchedulerServicer.GetServerInfo, whose first statement serializes self.server_args with dataclasses.asdict(...); the SimpleNamespace assigned here is not a dataclass, so the test raises TypeError: asdict() should be called on dataclass instances before it ever checks the metrics keys. Use a real ServerArgs instance/dataclass stub here, or add a production fallback before enabling this coverage.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py`:
- Line 579: Update GetServerInfo in servicer.py so it can serialize server_args
when it is not a dataclass. The current dataclasses.asdict(self.server_args)
path should be guarded with a fallback to a __dict__-based conversion, matching
the behavior already used by TokenSpeed, and then continue merging in
_metrics_server_args() so the metrics advertisement path works with
SimpleNamespace and other non-dataclass wrappers.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 35ac3e42-b0a1-432d-8b0a-1d85a45bf4eb
📒 Files selected for processing (9)
grpc_servicer/pyproject.tomlgrpc_servicer/smg_grpc_servicer/metrics.pygrpc_servicer/smg_grpc_servicer/sglang/server.pygrpc_servicer/smg_grpc_servicer/sglang/servicer.pygrpc_servicer/smg_grpc_servicer/tokenspeed/__main__.pygrpc_servicer/smg_grpc_servicer/tokenspeed/server.pygrpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pygrpc_servicer/tests/test_metrics_sidecar.pygrpc_servicer/tests/test_servicer_metrics_advertise.py
| return str(obj) | ||
|
|
||
| serializable_args = make_serializable(server_args_dict) | ||
| serializable_args.update(self._metrics_server_args()) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Handle non-dataclass server_args in GetServerInfo.
Line 579 extends this RPC for metrics discovery, but dataclasses.asdict(self.server_args) above still raises on the SimpleNamespace stub used by grpc_servicer/tests/test_servicer_metrics_advertise.py. TokenSpeed already falls back to __dict__; mirroring that here keeps the new advertisement path callable in tests and any non-dataclass wrappers.
Proposed fix
- server_args_dict = dataclasses.asdict(self.server_args)
+ if dataclasses.is_dataclass(self.server_args) and not isinstance(self.server_args, type):
+ server_args_dict = dataclasses.asdict(self.server_args)
+ else:
+ server_args_dict = dict(getattr(self.server_args, "__dict__", {}))
server_args_struct = Struct()🤖 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 `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py` at line 579, Update
GetServerInfo in servicer.py so it can serialize server_args when it is not a
dataclass. The current dataclasses.asdict(self.server_args) path should be
guarded with a fallback to a __dict__-based conversion, matching the behavior
already used by TokenSpeed, and then continue merging in _metrics_server_args()
so the metrics advertisement path works with SimpleNamespace and other
non-dataclass wrappers.
|
Hi @slin1237, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
Part of #1773. Closes #1777.
What
SMG's gRPC servicers exposed no HTTP
/metrics. This adds a shared, dependency-light, best-effort sidecar so they're scrape-able in gRPC mode (the producer side of "enable that for other engines"):smg_grpc_servicer/metrics.py: stdlib-asyncio+prometheus_client/metricsserver; bind/serve failures are logged and swallowed so the gRPC server keeps serving.SchedulerLoadCollectorre-exposes theGetLoadssnapshot for engines without a native registry.serve_grpcfor sglang and tokenspeed (mlx excluded — no PD). Port via--metrics-port/SMG_METRICS_PORT.metrics_port(always, when enabled) +metrics_url(only when the host is routable; omitted for0.0.0.0) inGetServerInfo.server_args, so the gateway (W1, [Metrics] W1: gRPC engine-metrics scraping (/engine_metrics across transports) #1774) auto-discovers it.Tests
20 engine-free sidecar/collector tests (incl. a real HTTP round-trip asserting valid exposition + content-type, and a best-effort bind-failure path) + 3 import-gated servicer tests (verified via a minimal tokenspeed stub).
python -m pytest grpc_servicer/tests→ 91 passed / 5 skipped;ruff check/format --checkclean. No Rust changes (cargo not run).Cross-workstream notes (action on the W1 side, #1774)
metrics_portandmetrics_url. W1'sSGLANG_GRPC_KEYSallowlistsmetrics_url(✓ end-to-end for routable hosts) but notmetrics_port, and doesn't extendTOKENSPEED_GRPC_KEYS. Follow-ups on the W1 side: allowlistmetrics_portand add the keys toTOKENSPEED_GRPC_KEYS.vllm/server.pyhere (vLLM is bootstrapped by upstreamvllm serve --grpc, which serves its own/metrics), andvllm_engineGetServerInfohas noserver_argsStruct (advertising would need a proto field — a Rust workstream).token_usagein the sidecar is0.0(needs an async scheduler round-trip unsafe from a sync Prometheuscollect()); full usage remains available via theGetLoadsRPC.Summary by CodeRabbit
Release Notes
New Features
/metricsHTTP sidecar for gRPC servicers, configurable via--metrics-port/metrics_port(disabled whenNone).GetServerInfonow advertisesmetrics_portand a reachablemetrics_urlwhen metrics are enabled; wildcard-like hosts suppress the URL.Dependency Updates
prometheus-clientas a required runtime dependency.Tests