refactor(grpc): centralize backend×modality capability check - #1914
Conversation
github-actions
Bot
added
grpc
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 9 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughThe PR adds centralized runtime/modality capability validation for multimodal gRPC requests, applies it during worker selection and assembly, and adds runtime derivation helpers plus capability-matrix tests. ChangesMultimodal capability validation
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant WorkerSelectionStage
participant CapabilityMatrix
participant GrpcClient
participant MultimodalAssembly
WorkerSelectionStage->>GrpcClient: derive runtime_type()
WorkerSelectionStage->>CapabilityMatrix: validate runtime and modalities
CapabilityMatrix-->>WorkerSelectionStage: allow or return bad_request
MultimodalAssembly->>GrpcClient: derive concrete runtime
MultimodalAssembly->>CapabilityMatrix: revalidate intermediate batches
CapabilityMatrix-->>MultimodalAssembly: allow assembly or return error
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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 centralizes the backend-by-modality capability mapping into a single source of truth in capability.rs and introduces early validation during worker selection to reject unsupported backend/modality combinations before media preprocessing. It also refactors the assembly code to use a unified into_single_batch helper and adds defense-in-depth checks. The review feedback suggests using absolute paths (starting with crate::) for internal documentation links in capability.rs and assemble.rs. Additionally, it points out that Modality::ImageEmbeds is only supported by TokenSpeed and should be separated from Modality::Image in the capability matrix, with corresponding updates to the unit tests.
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.
| /// `Modality::ImageEmbeds` is treated as image (TokenSpeed folds it into an | ||
| /// image batch during assembly); it cannot reach the early check because a | ||
| /// [`super::MediaBatch`] only ever yields Image/Video/Audio. |
There was a problem hiding this comment.
According to the general rules, we should use absolute paths starting with crate:: for internal documentation links instead of relative paths (such as super::) to ensure links remain valid after refactoring or module moves.
Please update [super::MediaBatch] to use the absolute path [crate::routers::grpc::multimodal::MediaBatch].
| /// `Modality::ImageEmbeds` is treated as image (TokenSpeed folds it into an | |
| /// image batch during assembly); it cannot reach the early check because a | |
| /// [`super::MediaBatch`] only ever yields Image/Video/Audio. | |
| /// `Modality::ImageEmbeds` is treated as image (TokenSpeed folds it into an | |
| /// image batch during assembly); it cannot reach the early check because a | |
| /// [`crate::routers::grpc::multimodal::MediaBatch`] only ever yields Image/Video/Audio. |
References
- Use absolute paths starting with
crate::for internal documentation links instead of relative paths (e.g.,super::) to ensure links remain valid after refactoring or module moves.
| Modality::Image | Modality::ImageEmbeds => matches!( | ||
| runtime, | ||
| RuntimeType::Sglang | RuntimeType::Vllm | RuntimeType::Trtllm | RuntimeType::TokenSpeed | ||
| ), |
There was a problem hiding this comment.
The capability matrix currently groups Modality::ImageEmbeds with Modality::Image, indicating that SGLang, vLLM, and TRT-LLM support image embeddings. However, these backends only support raw images (Modality::Image) and do not support pre-computed image embeddings (Modality::ImageEmbeds), which is a TokenSpeed-only disaggregated EPD feature.
To prevent potential correctness issues or silent failures if ImageEmbeds is ever routed to these backends, they should be separated in the capability matrix so that only TokenSpeed supports ImageEmbeds.
| Modality::Image | Modality::ImageEmbeds => matches!( | |
| runtime, | |
| RuntimeType::Sglang | RuntimeType::Vllm | RuntimeType::Trtllm | RuntimeType::TokenSpeed | |
| ), | |
| Modality::Image => matches!( | |
| runtime, | |
| RuntimeType::Sglang | RuntimeType::Vllm | RuntimeType::Trtllm | RuntimeType::TokenSpeed | |
| ), | |
| Modality::ImageEmbeds => matches!(runtime, RuntimeType::TokenSpeed), |
| // ImageEmbeds tracks image support. | ||
| (Sglang, ImageEmbeds, true), | ||
| (TokenSpeed, ImageEmbeds, true), | ||
| (Mlx, ImageEmbeds, false), |
There was a problem hiding this comment.
If Modality::ImageEmbeds is restricted to TokenSpeed (as it is not supported by SGLang, vLLM, or TRT-LLM), the test cases should be updated to assert that ImageEmbeds is unsupported on non-TokenSpeed runtimes.
| // ImageEmbeds tracks image support. | |
| (Sglang, ImageEmbeds, true), | |
| (TokenSpeed, ImageEmbeds, true), | |
| (Mlx, ImageEmbeds, false), | |
| // ImageEmbeds is only supported by TokenSpeed. | |
| (Sglang, ImageEmbeds, false), | |
| (TokenSpeed, ImageEmbeds, true), | |
| (Mlx, ImageEmbeds, false), |
| /// Defense-in-depth capability assertion mirroring the early worker-selection | ||
| /// check, keyed on the concrete backend client. See | ||
| /// [`super::capability::runtime_supports_modality`]. |
There was a problem hiding this comment.
According to the general rules, we should use absolute paths starting with crate:: for internal documentation links instead of relative paths (such as super::) to ensure links remain valid after refactoring or module moves.
Please update [super::capability::runtime_supports_modality] to use the absolute path [crate::routers::grpc::multimodal::capability::runtime_supports_modality].
| /// Defense-in-depth capability assertion mirroring the early worker-selection | |
| /// check, keyed on the concrete backend client. See | |
| /// [`super::capability::runtime_supports_modality`]. | |
| /// Defense-in-depth capability assertion mirroring the early worker-selection | |
| /// check, keyed on the concrete backend client. See | |
| /// [`crate::routers::grpc::multimodal::capability::runtime_supports_modality`]. |
References
- Use absolute paths starting with
crate::for internal documentation links instead of relative paths (e.g.,super::) to ensure links remain valid after refactoring or module moves.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean, behavior-preserving refactor. The centralized capability matrix in capability.rs is correct and well-tested, the early rejection at worker selection is properly placed (after runtime is known, before media fetch/preprocess), and the defense-in-depth assertion in assembly mirrors the same matrix. The selection_runtime helper correctly picks the prefill leg for disaggregated mode. No issues found.
0 🔴 Important · 0 🟡 Nit · 0 🟣 Pre-existing
![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 `@model_gateway/src/routers/grpc/multimodal/assemble.rs`:
- Around line 163-176: Update ensure_client_supports_intermediate to collect the
intermediate batch modalities and delegate validation to
capability::ensure_backend_supports_modalities using the client’s runtime type,
removing the duplicated loop and anyhow::ensure! logic. Replace the
runtime_supports_modality import with ensure_backend_supports_modalities and
preserve the existing Result<()> behavior.
🪄 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: 81654097-218f-4d1d-9ce2-c5687f846e5e
📒 Files selected for processing (5)
model_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/common/stages/worker_selection.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/multimodal/capability.rsmodel_gateway/src/routers/grpc/multimodal/mod.rs
The 'which engine supports which modality' truth was implicit and
duplicated across assembly: per-backend into_single_image_batch /
into_single_vision_batch checks plus ad-hoc bail!s with divergent
messages, all firing deep in assembly after media was already
fetched and preprocessed.
Introduce a single capability matrix (runtime_supports_modality) in a
new multimodal::capability module and enforce it once, early, at worker
selection -- so an unsupported (backend, modality) request fails fast
with one consistent message ("backend {X} does not support {modality}
inputs") before any fetch/preprocess. Assembly re-asserts against the
same matrix as defense in depth; the merged into_single_batch helper now
carries only the structural single-batch constraint.
Behavior-preserving: exactly the same (engine, modality) combinations
are rejected (Image: SGLang/vLLM/TRT-LLM/TokenSpeed; Video:
vLLM/TokenSpeed; Audio: TokenSpeed; MLX: none), just earlier and with a
uniform message.
Add GrpcClient::runtime_type() so assembly can key the shared matrix on
the concrete client, and unit tests for the full capability matrix plus
the early-rejection path.
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
0747d71 to
9d29d9f
Compare
|
Addressed all review comments (pushed):
|
![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: 9d29d9fd02
ℹ️ 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".
| if let Some(intermediate) = multimodal_intermediate(prep) { | ||
| if let Err(err) = multimodal::ensure_backend_supports_modalities( | ||
| selection_runtime(&workers), | ||
| intermediate, |
There was a problem hiding this comment.
This check still runs after the expensive media path: in the chat/messages pipelines worker selection follows preparation, and multimodal_intermediate is only populated after process_multimodal_plan has fetched and preprocessed the media. For unsupported runtime/modality pairs such as audio routed to SGLang, the request will still download/preprocess the audio before being rejected here, so the intended fail-fast behavior is not achieved; the check needs to use the detected MediaPlan/modalities before full multimodal processing or otherwise run before preparation fetches media.
Useful? React with 👍 / 👎.
| #[test] | ||
| fn non_grpc_runtimes_support_no_modality() { | ||
| for runtime in [RuntimeType::Unspecified, RuntimeType::External] { | ||
| for modality in [Modality::Image, Modality::Video, Modality::Audio] { |
There was a problem hiding this comment.
🟡 Nit: ImageEmbeds is missing from the modality list here. The main capability_matrix_matches_backend_support test covers (Unspecified/External, ImageEmbeds) implicitly (those runtimes aren't listed), but this dedicated test for non-gRPC runtimes should be exhaustive for all modalities — especially since the main test doesn't exercise Unspecified/External at all.
| for modality in [Modality::Image, Modality::Video, Modality::Audio] { | |
| for modality in [Modality::Image, Modality::ImageEmbeds, Modality::Video, Modality::Audio] { |
1 participant
Description
Problem
Which engine supports which modality (audio = TokenSpeed-only; video = vLLM/TokenSpeed; image = all-but-MLX) was enforced by scattered
bail!s deep inassemble.rs, with inconsistent messages, after media was already fetched and preprocessed.Solution
Add a single capability matrix and check it early, at worker selection, before any media is fetched — one consistent
multimodal_not_supportederror. Keep a defense-in-depth assertion in assembly.Changes
multimodal/capability.rs(new):runtime_supports_modality(runtime, modality)+ensure_backend_supports_modalities(...), with a full backend×modality unit test.grpc/client.rs:GrpcClient::runtime_type().common/stages/worker_selection.rs: early rejection using the request'sMultimodalIntermediate+ selected runtime.multimodal/assemble.rs: mergeinto_single_image_batch/into_single_vision_batch→into_single_batch; centralize the per-modality check.Test Plan
cargo test -p smg --lib routers::grpcpasses. Behavior-preserving: the same (backend, modality) combinations are rejected as before — just earlier and with one consistent message (verified against each priorbail!).Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspasses