feat(embed): add vLLM gRPC embedding support, clean up proto

[CatherineSue]

Description

Problem

The gRPC embedding pipeline was hardcoded to SGLang only. vLLM embedding requests via gRPC would panic. Additionally, the embedding proto and protocol types had accumulated SGLang-specific leaky abstractions (log_metrics, EmbedError, EmbedComplete oneof) that blocked clean multi-backend support.

Solution

Add vLLM embedding support end-to-end (Rust gateway + Python servicer) and clean up the proto/protocol layers to be backend-agnostic.

Changes

• Remove log_metrics from embedding/classify protocols — SGLang-specific Prometheus metrics gate that leaked into shared protocols crate. vLLM/TensorRT-LLM don't use it.
• Add vLLM embedding to Rust gRPC pipeline — embed() + build_embed_request() on VllmEngineClient, Vllm variants on ProtoEmbedRequest/ProtoEmbedComplete, backend-aware dispatch in EmbeddingRequestBuildingStage (TensorRT-LLM returns unsupported error).
• Implement vLLM Embed RPC in Python servicer — calls engine.encode() with PoolingParams, returns flat EmbedResponse.
• Flatten SGLang EmbedResponse proto — remove oneof { EmbedComplete, EmbedError }, Embedding, EmbedError message types. Errors use tonic::Status (same as chat pipeline). Both backends now share the same flat EmbedResponse shape.
• Remove cached_tokens from embed proto — SGLang hardcoded it to 0, vLLM doesn't have it, OpenAI /v1/embeddings spec doesn't return it, gateway never reads it.
• Remove ProtoEmbedResponse wrapper — client.embed() returns ProtoEmbedComplete directly, no oneof unwrapping.
• Enable embedding e2e tests for vLLM — add "vllm" to @pytest.mark.engine on test_basic.py and test_correctness.py.
• Classify gets vLLM support for free — shares stages 1-6 with embedding.

Test Plan

• cargo check passes
• cargo clippy --all-targets --all-features -- -D warnings passes (via pre-commit)
• cargo +nightly fmt passes (via pre-commit)
• ruff passes on Python changes (via pre-commit)
• e2e: pytest e2e_test/embeddings/test_basic.py -v -k "vllm and grpc" with embedding model
• e2e: pytest e2e_test/embeddings/test_correctness.py -v -k "vllm and grpc" with embedding model

Summary by CodeRabbit

• New Features

  • vLLM engine now supports embedding requests alongside SGLang.

• Changes

  • Embedding responses simplified to a single flat success shape — removed batching, wrapper/error variants, request_id, and cached_tokens.
  • Removed the log_metrics option from embedding and classification requests.
  • Backend dispatch for embed requests broadened to include vLLM; embed error handling and status mapping standardized.

• Tests

  • Embedding tests updated to run against both SGLang and vLLM engines.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Removed log_metrics from request types, flattened SGLang EmbedResponse to a single success-only layout, added vLLM embedding support and end-to-end embed RPC paths across clients, servicers, and the model gateway, and adjusted pipeline/request/response wrappers to handle backend-specific flat embed responses.

Changes

Cohort / File(s)	Summary
Protobuf & Protocol Types crates/grpc_client/proto/sglang_scheduler.proto, crates/protocols/src/classify.rs, crates/protocols/src/embedding.rs	Removed log_metrics from requests; flattened SGLang EmbedResponse (deleted EmbedComplete/EmbedError/Embedding), reserved old tags, and inlined embedding, prompt_tokens, and embedding_dim.
gRPC Clients crates/grpc_client/src/sglang_scheduler.rs, crates/grpc_client/src/vllm_engine.rs	SGLang client build_embed_request no longer accepts log_metrics; added VLLM build_embed_request and an async embed helper that injects trace metadata and calls tonic.
Servicers grpc_servicer/smg_grpc_servicer/sglang/servicer.py, grpc_servicer/smg_grpc_servicer/vllm/servicer.py	SGLang servicer adapted to flat response shape and early-error abort; vLLM servicer implements Embed using engine.encode(), validates tokenized, consumes streaming outputs, and maps errors to gRPC statuses.
Model Gateway — Request/Response Wrapping model_gateway/src/routers/grpc/proto_wrapper.rs, model_gateway/src/routers/grpc/client.rs	Added Vllm variant to ProtoEmbedRequest; replaced layered response/error enums with ProtoEmbedComplete carrying backend EmbedResponse variants; updated client dispatch to return ProtoEmbedComplete.
Model Gateway — Pipeline & Execution model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs, model_gateway/src/routers/grpc/common/stages/request_execution.rs	Request-building now selects backend-specific build_embed_request (SGLang
Tests & E2E e2e_test/embeddings/test_basic.py, e2e_test/embeddings/test_correctness.py, model_gateway/tests/spec/embedding.rs	Expanded pytest engine markers to include "vllm"; removed log_metrics usage in unit tests and assertions.

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant ModelGateway
  participant BackendSelector
  participant SglangServicer
  participant VllmServicer
  participant VllmEngine

  Client->>ModelGateway: Send EmbedRequest
  ModelGateway->>BackendSelector: Select backend (SGLang | VLLM)
  alt SGLang
    BackendSelector->>SglangServicer: Forward EmbedRequest
    SglangServicer-->>BackendSelector: EmbedResponse { embedding, prompt_tokens, embedding_dim }
  else VLLM
    BackendSelector->>VllmServicer: Forward EmbedRequest
    VllmServicer->>VllmEngine: engine.encode(TokensPrompt, PoolingParams)
    VllmEngine-->>VllmServicer: streaming outputs -> final output
    VllmServicer-->>BackendSelector: EmbedResponse { embedding, prompt_tokens, embedding_dim }
  end
  BackendSelector-->>ModelGateway: ProtoEmbedComplete (Sglang | Vllm)
  ModelGateway-->>Client: Return completion (embedding, prompt_tokens, embedding_dim)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Possibly related PRs

• fix(grpc): return tonic::Status directly from grpc_client RPC methods #647 — Changes to grpc_client embedding APIs; overlaps on client method/signature updates.
• refactor(grpc): use EngineClient interface instead of AsyncLLM in vLLM servicer #949 — Related vLLM servicer work; overlapping edits in vllm servicer implementation.
• refactor: remove dead embedding/classify code from regular PreparationStage #668 — Model gateway embedding pipeline adjustments; overlapping request-building and routing changes.

Suggested labels

protocols

Suggested reviewers

• key4ng
• njhill
• slin1237
• XinyueZhang369

Poem

🐰 I nibbled proto fields today,
log_metrics hopped gently away.
vLLM learned to encode and send,
flat embeddings race to the end. ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and accurately summarizes the two main changes: adding vLLM gRPC embedding support and cleaning up the proto by removing SGLang-specific abstractions like log_metrics.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch changsu/remove-embedding-log-metrics

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request implements embedding support for the vLLM backend and simplifies the gRPC EmbedResponse structure by flattening it across both SGLang and vLLM. It removes the log_metrics field from embedding and classification requests and updates the model gateway to handle the unified response format. A high-severity issue was identified in the vLLM servicer where calling .tolist() on a standard Python list would cause a runtime error.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 208-212: The handler constructs a TokensPrompt by dereferencing
request.tokenized without checking whether the field is present; add a guard
using request.HasField("tokenized") before accessing request.tokenized in the
vLLM request handler (the block that creates prompt: TokensPrompt and reads
tokenized.input_ids and tokenized.original_text) and if the field is missing
abort the RPC with an INVALID_ARGUMENT grpc error (or raise the appropriate gRPC
exception) so malformed calls return INVALID_ARGUMENT instead of producing an
empty submessage or INTERNAL errors.

In `@model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs`:
- Around line 87-95: The current EmbeddingRequestBuildingStage::execute branch
for GrpcClient::Trtllm returns error::internal_error which incorrectly surfaces
unsupported TensorRT-LLM as a server failure; change this to a client-facing
error (e.g., error::bad_request or error::not_implemented) and update the logged
message accordingly so callers and metrics can distinguish
misrouting/configuration from real outages—locate the GrpcClient::Trtllm match
arm and replace the error::internal_error(...) return with
error::bad_request(...) or error::not_implemented(...), keeping the explanatory
string ("TensorRT-LLM embedding is not yet supported via gRPC") and the existing
error log context in EmbeddingRequestBuildingStage::execute.

🪄 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: 5b1a9f71-cd7d-417a-a31d-08c6effe7be7

📥 Commits

Reviewing files that changed from the base of the PR and between a484b61 and 77b4f0a.

📒 Files selected for processing (14)

• crates/grpc_client/proto/sglang_scheduler.proto
• crates/grpc_client/src/sglang_scheduler.rs
• crates/grpc_client/src/vllm_engine.rs
• crates/protocols/src/classify.rs
• crates/protocols/src/embedding.rs
• e2e_test/embeddings/test_basic.py
• e2e_test/embeddings/test_correctness.py
• grpc_servicer/smg_grpc_servicer/sglang/servicer.py
• grpc_servicer/smg_grpc_servicer/vllm/servicer.py
• model_gateway/src/routers/grpc/client.rs
• model_gateway/src/routers/grpc/common/stages/request_execution.rs
• model_gateway/src/routers/grpc/proto_wrapper.rs
• model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs
• model_gateway/tests/spec/embedding.rs

💤 Files with no reviewable changes (4)

• crates/protocols/src/classify.rs
• crates/protocols/src/embedding.rs
• model_gateway/tests/spec/embedding.rs
• crates/grpc_client/src/sglang_scheduler.rs

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs (1)

60-99: ⚠️ Potential issue | 🟡 Minor

Restore the Embedding/Classify guard in this stage.

This block now always emits ProtoRequest::Embed at Line 99 once a client is present. If any unsupported RequestType is ever routed here, the _ => format!("embed-...") fallback turns a wrong-pipeline invariant violation into a silent embed RPC dispatch instead of failing fast.

Suggested fix

         let request_id = match &ctx.input.request_type {
             RequestType::Embedding(_) => format!("embed-{}", Uuid::now_v7()),
             RequestType::Classify(_) => format!("classify-{}", Uuid::now_v7()),
-            _ => format!("embed-{}", Uuid::now_v7()), // fallback
+            _ => {
+                error!(
+                    function = "EmbeddingRequestBuildingStage::execute",
+                    "Unsupported request type for embedding request building"
+                );
+                return Err(error::internal_error(
+                    "wrong_pipeline",
+                    "Embedding request building only supports embedding/classify requests",
+                ));
+            }
         };

Based on learnings, GrpcClient::embed intentionally keeps invariant-failure handling in the embedding path to catch client/request mismatches rather than silently coercing them.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs`
around lines 60 - 99, The request-building stage currently coerces any
RequestType into an embed RPC; restore the original guard by matching
ctx.input.request_type (the RequestType enum) and only proceed when it's
RequestType::Embedding(_) or RequestType::Classify(_), generating request_id
with the appropriate prefix for those variants; for any other variant return a
fast failure (e.g., from EmbeddingRequestBuildingStage::execute) with an
appropriate not_implemented/invalid_request error instead of using the `_ =>
format!("embed-...")` fallback so we never set ctx.state.proto_request =
Some(ProtoRequest::Embed(...)) for unsupported request types.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Outside diff comments:
In `@model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs`:
- Around line 60-99: The request-building stage currently coerces any
RequestType into an embed RPC; restore the original guard by matching
ctx.input.request_type (the RequestType enum) and only proceed when it's
RequestType::Embedding(_) or RequestType::Classify(_), generating request_id
with the appropriate prefix for those variants; for any other variant return a
fast failure (e.g., from EmbeddingRequestBuildingStage::execute) with an
appropriate not_implemented/invalid_request error instead of using the `_ =>
format!("embed-...")` fallback so we never set ctx.state.proto_request =
Some(ProtoRequest::Embed(...)) for unsupported request types.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 15531094-7528-48aa-a6df-8404d135b13b

📥 Commits

Reviewing files that changed from the base of the PR and between 77b4f0a and e922641.

📒 Files selected for processing (2)

• grpc_servicer/smg_grpc_servicer/vllm/servicer.py
• model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@grpc_servicer/smg_grpc_servicer/sglang/servicer.py`:
- Around line 281-286: The Embed method assumes result contains "embedding" and
may KeyError when the scheduler returns an error payload; update the handler
(the Embed code path that currently builds sglang_scheduler_pb2.EmbedResponse)
to first check if "error" in result, and if so call
abort_code_from_output(result) and abort the RPC with that status and the
scheduler error message (use result["error"] and any meta_info/finish_reason),
otherwise extract embedding and return sglang_scheduler_pb2.EmbedResponse with
prompt_tokens and embedding_dim as before; follow the same error-check pattern
used by Generate to preserve the scheduler's original status and message.

In `@model_gateway/src/routers/grpc/common/stages/request_execution.rs`:
- Around line 200-208: The embed RPC path (client.embed(proto_request).await in
execute_single_embed) never records circuit-breaker outcomes, so embed-only
failures won't be seen by breaker logic; update the error path (inside the
map_err closure) to call workers.record_outcome(...) with the failing worker id
and a failure Outcome before returning the http error, and after a successful
embed (before returning Ok(ExecutionResult::Embedding { ... })) call
workers.record_outcome(...) with the worker id and a success Outcome; ensure you
use the same worker identity and Outcome enum used elsewhere in this module so
breaker/selection logic is updated consistently.

🪄 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: f8a86320-a197-4b1f-b0a7-97a73c8a67c9

📥 Commits

Reviewing files that changed from the base of the PR and between 6a253e3 and 86b7ae8.

📒 Files selected for processing (14)

• crates/grpc_client/proto/sglang_scheduler.proto
• crates/grpc_client/src/sglang_scheduler.rs
• crates/grpc_client/src/vllm_engine.rs
• crates/protocols/src/classify.rs
• crates/protocols/src/embedding.rs
• e2e_test/embeddings/test_basic.py
• e2e_test/embeddings/test_correctness.py
• grpc_servicer/smg_grpc_servicer/sglang/servicer.py
• grpc_servicer/smg_grpc_servicer/vllm/servicer.py
• model_gateway/src/routers/grpc/client.rs
• model_gateway/src/routers/grpc/common/stages/request_execution.rs
• model_gateway/src/routers/grpc/proto_wrapper.rs
• model_gateway/src/routers/grpc/regular/stages/embedding/request_building.rs
• model_gateway/tests/spec/embedding.rs

💤 Files with no reviewable changes (4)

• model_gateway/tests/spec/embedding.rs
• crates/protocols/src/classify.rs
• crates/protocols/src/embedding.rs
• crates/grpc_client/src/sglang_scheduler.rs

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 208-214: The abort branches should use the servicer's normal
logging and control flow: replace the inline context.abort branch for missing
tokenized (the request.HasField("tokenized") check) by raising ValueError for
the missing-tokenized case instead of calling context.abort directly, remove the
unreachable trailing return, and for the unfinished-encode case log an error via
the servicer's logger before calling
context.abort(grpc.StatusCode.INTERNAL,...); apply the same pattern to the
similar branch around the request/unfinished-encode check later (the block
referenced at 231-236) so server-side logs are emitted before aborting and
control stops via context.abort.

In `@model_gateway/src/routers/grpc/common/stages/request_execution.rs`:
- Line 136: The grpc span currently named "grpc_generate" is being used for
ProtoRequest::Embed (via execute_single_embed), which mislabels embed traffic;
update the tracing/span creation in the request_execution stage to distinguish
embed vs generate—either create separate spans (e.g., "grpc_embed" for
execute_single_embed and keep "grpc_generate" for generate paths) or add a
request_type tag/field to the existing span, ensuring you reference the
ProtoRequest::Embed match arm and the execute_single_embed call when making the
change.

🪄 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: 3621caef-98e3-419f-aba0-287d5e9cb915

📥 Commits

Reviewing files that changed from the base of the PR and between 86b7ae8 and ee526bc.

📒 Files selected for processing (4)

• grpc_servicer/smg_grpc_servicer/sglang/server.py
• grpc_servicer/smg_grpc_servicer/sglang/servicer.py
• grpc_servicer/smg_grpc_servicer/vllm/servicer.py
• model_gateway/src/routers/grpc/common/stages/request_execution.rs