feat(grpc): support /get_server_info for TokenSpeed gRPC routers#1574
feat(grpc): support /get_server_info for TokenSpeed gRPC routers#1574YzXiao101 wants to merge 1 commit into
Conversation
github-actions
Bot
added
documentation
📝 WalkthroughWalkthroughThis PR implements a complete get_server_info endpoint across gRPC routers: ServerInfo protobuf-to-JSON conversion with backend-aware filtering, handler logic in GrpcRouter and GrpcPDRouter with worker selection and error mapping, and updated admin API docs showing the new response shape. ChangesServer Info Endpoint Implementation
Sequence Diagram(s)sequenceDiagram
participant Client
participant GrpcRouter
participant WorkerPool
participant GrpcClient
participant ServerInfo
Client->>GrpcRouter: GET /server_info
GrpcRouter->>WorkerPool: select_first_worker
alt Worker Available
WorkerPool-->>GrpcRouter: worker
GrpcRouter->>GrpcClient: acquire client
GrpcClient->>ServerInfo: get_server_info()
ServerInfo-->>GrpcClient: proto response
GrpcClient-->>GrpcRouter: response
GrpcRouter->>GrpcRouter: to_public_json()
GrpcRouter-->>Client: 200 JSON
else No Worker
GrpcRouter-->>Client: 503 Service Unavailable
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 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. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
|
Hi @YzXiao101, the DCO sign-off check has failed. All commits must include a To fix existing commits: # Sign off the last N commits (replace N with the number of unsigned commits)
git rebase HEAD~N --signoff
git push --force-with-leaseTo sign off future commits automatically:
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the get_server_info endpoint for gRPC routers (GrpcRouter and GrpcPDRouter), allowing them to fetch and return server information from backend workers. It introduces a to_public_json method to format and filter the gRPC server arguments into a curated JSON response, and updates the API documentation and tests accordingly. The review feedback suggests a performance and idiomatic improvement in struct_to_json_map by removing redundant sorting of BTreeMap fields and using .collect() instead of manual loop insertion.
| fn struct_to_json_map(struct_value: &prost_types::Struct) -> serde_json::Map<String, Value> { | ||
| let mut entries: Vec<_> = struct_value.fields.iter().collect(); | ||
| entries.sort_by(|(left, _), (right, _)| left.cmp(right)); | ||
|
|
||
| let mut map = serde_json::Map::new(); | ||
| for (key, value) in entries { | ||
| map.insert(key.clone(), prost_value_to_json(value)); | ||
| } | ||
| map | ||
| } |
There was a problem hiding this comment.
The struct_value.fields field is a BTreeMap<String, Value> in prost_types::Struct, which means its keys are already sorted. Collecting the entries into a Vec and sorting them again is redundant and introduces unnecessary heap allocations and CPU overhead. You can iterate over struct_value.fields directly. Additionally, prefer using .collect() on the iterator to populate the map instead of manually looping to insert elements, as it is more idiomatic and concise.
fn struct_to_json_map(struct_value: &prost_types::Struct) -> serde_json::Map<String, Value> {
struct_value
.fields
.iter()
.map(|(key, value)| (key.clone(), prost_value_to_json(value)))
.collect()
}References
- Prefer using
.collect()on an iterator to populate a collection instead of manually looping to insert elements, as it automatically pre-sizes the collection and is more idiomatic and concise.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 `@docs/reference/api/admin.md`:
- Around line 393-394: Add a blank line between the "**Response:** `200 OK`"
line and the opening JSON code fence (```) so the code fence is separated from
the preceding bold line; this resolves markdownlint MD031 for the block that
starts with the JSON fence and ensures the snippet under the "**Response:** `200
OK`" heading renders and lints correctly.
🪄 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: da2410a0-9930-4df3-ab44-07dde805fb8f
📒 Files selected for processing (4)
docs/reference/api/admin.mdmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/pd_router.rsmodel_gateway/src/routers/grpc/router.rs
Signed-off-by: YzXiao101 <yzxiao101@outlook.com>
YzXiao101
force-pushed
the
feat/grpc-get-server-info
branch
from
a466a0a to
53594a5
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 53594a5570
ℹ️ 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".
| fn to_public_json_sglang_keeps_pb_shape_and_filters_server_args() { | ||
| let server_info = ServerInfo::Sglang(Box::new(sglang_proto::GetServerInfoResponse { | ||
| server_args: Some(prost_types::Struct { | ||
| fields: BTreeMap::from([ |
There was a problem hiding this comment.
In prost-types 0.14, prost_types::Struct::fields is a HashMap<String, Value> by default, but these new tests initialize it with BTreeMap::from(...). When the test module is compiled, this assignment won't type-check, so cargo test for this crate fails before running the new coverage. Use HashMap here (and in the other new Struct initializers) or collect into the field's actual map type.
Useful? React with 👍 / 👎.
|
This is not the right implementation. Get server info should not be a new point. It's already abstracted in worker |
slin1237
left a comment
slin1237
left a comment
There was a problem hiding this comment.
Router shouldn't be changed for this
Worker is the right abstraction. Not router
@slin1237 thx for your feedback.
|
3 participants
Description
/get_server_infowas available on the SMG side, but TokenSpeed gRPC routers did not implement the bridge, so the request could not return backend server info.This tiny PR just fills that gap. It wires
/get_server_infothrough the gRPC router/client path and returns the TokenSpeed payload in SMG's admin response.Changes
get_server_infoto the regular gRPC routerget_server_infoto the PD gRPC routerserver_args/scheduler_infofrom protobufStructto JSONserver_argssubsetValidation
Unit tests:
Struct -> JSONconversionSetup:
Result:
NOTE: I have identified and fixed the JSON formatting issue where integer-like values in server_args and scheduler_info were being emitted as floats (for example, 8192.0), and I am running the full SMG + TokenSpeed validation again.
Summary by CodeRabbit
New Features
Documentation