refactor(multimodal): split grpc multimodal.rs into a module directory - #1890
Conversation
`multimodal.rs` had grown to ~2.65k lines -- the largest flat file under routers/grpc (harmony/ and regular/ are already directories). Split it by pipeline phase into a `multimodal/` module directory. This is a pure move with no behavior change. - detect: modality detection + content extraction (chat + messages) - config: model config-file registry + per-router component bundle - process: fetch -> preprocess -> expand tokens -> build intermediate - assemble: intermediate -> backend-specific MultimodalData + per-item split - serialize: tensor byte/dtype serialization (f32 -> bf16/f16) - transport: SHM-vs-inline resolution + /dev/shm namespace verification Tests move next to the code they cover. The crate-facing API is unchanged: the same symbols are re-exported from mod.rs, and internal helpers are now enforced-private (pub(super)) instead of implicitly module-scoped. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Note Currently processing new changes in this PR. This may take a few minutes, please wait... ⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe monolithic ChangesgRPC multimodal module split
Estimated code review effort: 4 (Complex) | ~75 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Code Review
This pull request refactors the multimodal processing integration by splitting the monolithic multimodal.rs file into a dedicated module structure under model_gateway/src/routers/grpc/multimodal/, dividing the logic into separate files for assembly, configuration, detection, processing, serialization, and transport. The review feedback suggests updating internal documentation links across several of these new files (mod.rs, assemble.rs, and process.rs) to use absolute paths starting with crate:: instead of relative paths, in accordance with internal documentation guidelines.
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.
| //! - [`detect`]: find modalities and extract content parts from chat/messages. | ||
| //! - [`config`]: model config-file registry and per-router component bundle. | ||
| //! - [`process`]: fetch media → preprocess → expand placeholder tokens → | ||
| //! build the lightweight [`MultimodalIntermediate`]. | ||
| //! - [`assemble`]: turn the intermediate into backend-specific `MultimodalData` | ||
| //! once the target backend is known (after worker selection). | ||
| //! - [`serialize`]: tensor byte/dtype serialization used by assembly. | ||
| //! - [`transport`]: SHM-vs-inline transport resolution and `/dev/shm` | ||
| //! namespace verification. |
There was a problem hiding this comment.
According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths. This ensures that the links remain valid if modules are moved or refactored in the future.
| //! - [`detect`]: find modalities and extract content parts from chat/messages. | |
| //! - [`config`]: model config-file registry and per-router component bundle. | |
| //! - [`process`]: fetch media → preprocess → expand placeholder tokens → | |
| //! build the lightweight [`MultimodalIntermediate`]. | |
| //! - [`assemble`]: turn the intermediate into backend-specific `MultimodalData` | |
| //! once the target backend is known (after worker selection). | |
| //! - [`serialize`]: tensor byte/dtype serialization used by assembly. | |
| //! - [`transport`]: SHM-vs-inline transport resolution and `/dev/shm` | |
| //! namespace verification. | |
| //! - [`crate::routers::grpc::multimodal::detect`]: find modalities and extract content parts from chat/messages. | |
| //! - [`crate::routers::grpc::multimodal::config`]: model config-file registry and per-router component bundle. | |
| //! - [`crate::routers::grpc::multimodal::process`]: fetch media → preprocess → expand placeholder tokens → | |
| //! build the lightweight [`crate::routers::grpc::multimodal::MultimodalIntermediate`]. | |
| //! - [`crate::routers::grpc::multimodal::assemble`]: turn the intermediate into backend-specific `MultimodalData` | |
| //! once the target backend is known (after worker selection). | |
| //! - [`crate::routers::grpc::multimodal::serialize`]: tensor byte/dtype serialization used by assembly. | |
| //! - [`crate::routers::grpc::multimodal::transport`]: SHM-vs-inline transport resolution and `/dev/shm` | |
| //! namespace verification. |
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.
| @@ -0,0 +1,782 @@ | |||
| //! Assembly: convert a [`MultimodalIntermediate`] into backend-specific | |||
There was a problem hiding this comment.
According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths. Please update [MultimodalIntermediate] to use its absolute path.
| //! Assembly: convert a [`MultimodalIntermediate`] into backend-specific | |
| //! Assembly: convert a [`crate::routers::grpc::multimodal::MultimodalIntermediate`] into backend-specific |
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.
| //! Multimodal processing core: fetch media → preprocess pixels → expand | ||
| //! placeholder tokens → build the lightweight [`MultimodalIntermediate`]. | ||
| //! | ||
| //! The chat and Messages API pipelines share `process_multimodal_parts`; only | ||
| //! the content extraction differs (see [`super::detect`]). |
There was a problem hiding this comment.
According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths (such as super::). Please update [MultimodalIntermediate] and [super::detect] to use their absolute paths.
| //! Multimodal processing core: fetch media → preprocess pixels → expand | |
| //! placeholder tokens → build the lightweight [`MultimodalIntermediate`]. | |
| //! | |
| //! The chat and Messages API pipelines share `process_multimodal_parts`; only | |
| //! the content extraction differs (see [`super::detect`]). | |
| //! Multimodal processing core: fetch media → preprocess pixels → expand | |
| //! placeholder tokens → build the lightweight [`crate::routers::grpc::multimodal::MultimodalIntermediate`]. | |
| //! | |
| //! The chat and Messages API pipelines share `process_multimodal_parts`; only | |
| //! the content extraction differs (see [`crate::routers::grpc::multimodal::detect`]). |
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.
There was a problem hiding this comment.
Clean refactoring — verified that all external call sites resolve to the re-exported symbols in mod.rs, visibility modifiers are correct (pub(super) for internal helpers, pub(crate) for the module contract), and no logic was changed. The module split follows the pipeline's natural phase boundaries nicely.
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 `@model_gateway/src/routers/grpc/multimodal/config.rs`:
- Around line 59-117: The get_or_load flow in MultimodalConfig can trigger
duplicate cold loads for the same tokenizer_id because the cache is checked
before any in-flight load is coordinated. Update the get_or_load path to use a
single-flight or per-tokenizer lock around the resolve_model_config_dir and
config.json/preprocessor_config.json loading so concurrent misses share one
load, then insert the resulting Arc<MultimodalModelConfig> once for that
tokenizer_id.
In `@model_gateway/src/routers/grpc/multimodal/mod.rs`:
- Around line 42-47: The hot-path toggle lookup in log_mm_timing_enabled() is
re-reading and re-parsing SMG_LOG_MM_TIMING on every call from
process_multimodal_parts, which is unnecessary overhead. Cache the parsed
boolean once using a one-time initializer (for example, a OnceLock or
equivalent) inside log_mm_timing_enabled(), and have subsequent calls return the
cached value instead of touching std::env::var repeatedly.
In `@model_gateway/src/routers/grpc/multimodal/process.rs`:
- Around line 427-497: expand_tokens currently warns only when replacements
outnumber placeholder tokens, but it silently leaves extra placeholder_id tokens
in the output once replacement_idx is exhausted. Update the expand_tokens logic
in process.rs to detect and warn on the extra-placeholder case when token_ids
still contain placeholder_id after all PromptReplacement entries are consumed,
while preserving the unchanged token flow. Add a test around expand_tokens
covering the mismatch path so the warning behavior is exercised and the
resulting ExpandedTokens are verified.
In `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 149-162: The SHM decision in
transport::worker_matches_shm_namespace is too broad because it returns based on
encode_assignments alone, but shm_enabled is also used for
model_specific_tensors downstream. Update the logic so global SHM is only
enabled when prefill, decode, and all encode workers share the same /dev/shm
namespace, or split the policy into separate flags for encoder_input and
model_specific_tensors. Keep the fix localized in worker_matches_shm_namespace
and the caller path that sets shm_enabled.
🪄 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: 3df68237-1a73-4ed3-9007-b5cfdde5810e
📒 Files selected for processing (8)
model_gateway/src/routers/grpc/multimodal.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/multimodal/config.rsmodel_gateway/src/routers/grpc/multimodal/detect.rsmodel_gateway/src/routers/grpc/multimodal/mod.rsmodel_gateway/src/routers/grpc/multimodal/process.rsmodel_gateway/src/routers/grpc/multimodal/serialize.rsmodel_gateway/src/routers/grpc/multimodal/transport.rs
💤 Files with no reviewable changes (1)
- model_gateway/src/routers/grpc/multimodal.rs
| pub(crate) async fn get_or_load( | ||
| &self, | ||
| tokenizer_id: &str, | ||
| tokenizer_source: &str, | ||
| ) -> Result<Arc<MultimodalModelConfig>> { | ||
| if let Some(cached) = self.get(tokenizer_id) { | ||
| debug!(%tokenizer_id, "multimodal config cache hit"); | ||
| return Ok(cached); | ||
| } | ||
|
|
||
| debug!( | ||
| %tokenizer_id, | ||
| %tokenizer_source, | ||
| "multimodal config cache miss, loading" | ||
| ); | ||
|
|
||
| let base_dir = llm_multimodal::hub::resolve_model_config_dir(tokenizer_source) | ||
| .await | ||
| .with_context(|| { | ||
| format!("Failed to resolve model config directory for '{tokenizer_source}'") | ||
| })?; | ||
|
|
||
| let config_path = base_dir.join("config.json"); | ||
| let config: serde_json::Value = std::fs::read_to_string(&config_path) | ||
| .with_context(|| format!("Failed to read config.json at {}", config_path.display())) | ||
| .and_then(|s| { | ||
| serde_json::from_str(&s).with_context(|| { | ||
| format!("Failed to parse config.json at {}", config_path.display()) | ||
| }) | ||
| })?; | ||
|
|
||
| // preprocessor_config.json is optional — each vision processor supplies | ||
| // its own model-specific defaults, so missing/unparsable files fall | ||
| // back to `PreProcessorConfig::default()`. This matches the bundle | ||
| // preload path in `try_load_multimodal_config`. | ||
| let pp_config_path = base_dir.join("preprocessor_config.json"); | ||
| let preprocessor_config = | ||
| load_preprocessor_config_file(&pp_config_path, "preprocessor_config.json") | ||
| .unwrap_or_else(|| { | ||
| debug!( | ||
| path = %pp_config_path.display(), | ||
| "No preprocessor_config.json found; using PreProcessorConfig defaults" | ||
| ); | ||
| PreProcessorConfig::default() | ||
| }); | ||
| let video_preprocessor_config = load_video_preprocessor_config(&base_dir); | ||
|
|
||
| let model_config = Arc::new(MultimodalModelConfig { | ||
| config, | ||
| preprocessor_config, | ||
| video_preprocessor_config, | ||
| }); | ||
|
|
||
| self.configs | ||
| .insert(tokenizer_id.to_string(), model_config.clone()); | ||
|
|
||
| debug!(%tokenizer_id, "multimodal config loaded and cached"); | ||
| Ok(model_config) | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map relevant files and nearby symbols first.
git ls-files 'model_gateway/src/routers/grpc/multimodal/*' 'llm_multimodal*' | sed -n '1,200p'
echo '---'
ast-grep outline model_gateway/src/routers/grpc/multimodal/config.rs --view expanded
echo '---'
rg -n "resolve_model_config_dir|get_or_load|OnceCell|SingleFlight|preprocessor_config" model_gateway/src llm_multimodal -SRepository: lightseekorg/smg
Length of output: 9819
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# If the outline is too limited, inspect the implementation of resolve_model_config_dir and the surrounding cache code.
fd -a 'config.rs' model_gateway llm_multimodal . || true
echo '---'
rg -n "pub(crate) async fn get_or_load|struct .*Multimodal|DashMap|resolve_model_config_dir" model_gateway/src -S
echo '---'
rg -n "fn resolve_model_config_dir|async fn resolve_model_config_dir|resolve_model_config_dir" llm_multimodal -SRepository: lightseekorg/smg
Length of output: 11941
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for any code path that indicates network fetching in the model config directory resolver.
rg -n "hub::resolve_model_config_dir|resolve_model_config_dir\(" . -S
echo '---'
rg -n "hf hub|huggingface|network fetch|download|cache miss|snapshot_download|Model::load|hf_hub" . -SRepository: lightseekorg/smg
Length of output: 24414
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read only the minimal source slices needed to determine behavior.
echo '--- config.rs ---'
sed -n '1,220p' model_gateway/src/routers/grpc/multimodal/config.rs
echo '--- resolver ---'
for f in $(git ls-files 'llm_multimodal/**' 'llm-multimodal/**' 2>/dev/null); do
case "$f" in
*hub*|*config*|*resolve*) echo "### $f"; sed -n '1,260p' "$f"; echo '---';;
esac
doneRepository: lightseekorg/smg
Length of output: 8122
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' crates/multimodal/src/hub.rsRepository: lightseekorg/smg
Length of output: 2088
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check whether the resolver has any internal deduplication / caching or only plain hf_hub downloads.
rg -n "download|snapshot|cache|OnceCell|Mutex|Semaphore|resolve_model_config_dir" crates/multimodal/src -SRepository: lightseekorg/smg
Length of output: 1223
Concurrent cold misses can stampede this loader. resolve_model_config_dir downloads config.json and optional processor files for HuggingFace model IDs, so multiple requests for the same uncached tokenizer_id can all perform the same network I/O before the first insert lands. A single-flight cache would collapse the duplicate loads.
🤖 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 `@model_gateway/src/routers/grpc/multimodal/config.rs` around lines 59 - 117,
The get_or_load flow in MultimodalConfig can trigger duplicate cold loads for
the same tokenizer_id because the cache is checked before any in-flight load is
coordinated. Update the get_or_load path to use a single-flight or per-tokenizer
lock around the resolve_model_config_dir and
config.json/preprocessor_config.json loading so concurrent misses share one
load, then insert the resulting Arc<MultimodalModelConfig> once for that
tokenizer_id.
| if !skip_pixel_values { | ||
| if let Some(encode_assignments) = encode_assignments { | ||
| // EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM | ||
| // is safe only if every encode worker assigned in this request shares | ||
| // the gateway's /dev/shm. A mixed local/remote fan-out must fall back | ||
| // to inline/RDMA rather than giving a remote worker an unreadable SHM handle. | ||
| return encode_assignments | ||
| .iter() | ||
| .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local)); | ||
| } | ||
| } | ||
| worker_matches_shm_namespace(prefill, local) | ||
| && worker_matches_shm_namespace(decode, local) | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Don’t enable global SHM based only on encode workers.
shm_enabled is later used for more than encoder_input: downstream conversion also uses it for model_specific_tensors. Returning after only the encode-worker check can produce SHM-backed tensors that remote prefill/decode workers cannot read. Either require prefill/decode to share /dev/shm too, or split the policy into separate encoder/model-specific SHM flags.
Safer localized fix
- return encode_assignments
- .iter()
- .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
+ let encode_workers_share = encode_assignments
+ .iter()
+ .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
+ return encode_workers_share
+ && worker_matches_shm_namespace(prefill, local)
+ && worker_matches_shm_namespace(decode, local);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if !skip_pixel_values { | |
| if let Some(encode_assignments) = encode_assignments { | |
| // EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM | |
| // is safe only if every encode worker assigned in this request shares | |
| // the gateway's /dev/shm. A mixed local/remote fan-out must fall back | |
| // to inline/RDMA rather than giving a remote worker an unreadable SHM handle. | |
| return encode_assignments | |
| .iter() | |
| .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local)); | |
| } | |
| } | |
| worker_matches_shm_namespace(prefill, local) | |
| && worker_matches_shm_namespace(decode, local) | |
| } | |
| if !skip_pixel_values { | |
| if let Some(encode_assignments) = encode_assignments { | |
| // EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM | |
| // is safe only if every encode worker assigned in this request shares | |
| // the gateway's /dev/shm. A mixed local/remote fan-out must fall back | |
| // to inline/RDMA rather than giving a remote worker an unreadable SHM handle. | |
| let encode_workers_share = encode_assignments | |
| .iter() | |
| .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local)); | |
| return encode_workers_share | |
| && worker_matches_shm_namespace(prefill, local) | |
| && worker_matches_shm_namespace(decode, local); | |
| } | |
| } | |
| worker_matches_shm_namespace(prefill, local) | |
| && worker_matches_shm_namespace(decode, local) | |
| } |
🤖 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 `@model_gateway/src/routers/grpc/multimodal/transport.rs` around lines 149 -
162, The SHM decision in transport::worker_matches_shm_namespace is too broad
because it returns based on encode_assignments alone, but shm_enabled is also
used for model_specific_tensors downstream. Update the logic so global SHM is
only enabled when prefill, decode, and all encode workers share the same
/dev/shm namespace, or split the policy into separate flags for encoder_input
and model_specific_tensors. Keep the fix localized in
worker_matches_shm_namespace and the caller path that sets shm_enabled.
Two review fixes on top of the module split, kept as a separate commit so the split itself stays a pure, reviewable move: - log_mm_timing_enabled: read SMG_LOG_MM_TIMING once via OnceLock instead of re-parsing the env on every multimodal request. - expand_tokens: warn when the token sequence has more placeholder tokens than PromptReplacements (the excess were previously left unexpanded silently). Token output is unchanged; adds a regression test. Both address CodeRabbit review comments. The single-flight concern in MultimodalConfigRegistry::get_or_load is deferred to a dedicated PR (it needs real async coordination + concurrency tests, out of scope for this refactor). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Thanks for the reviews. Dispositions:
For context: this PR is a pure module split (commit 1) plus the two small review fixes above (commit 2), kept separate so the move stays reviewable as a no-op via |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Description
Problem
model_gateway/src/routers/grpc/multimodal.rshad grown to ~2.65k lines — thelargest flat file under
routers/grpc(its neighborsharmony/andregular/are already module directories). Everything from content detection to tensor
serialization to SHM transport resolution lived in one file, and the module's
real public surface was buried among ~60 functions with ad-hoc visibility.
Solution
Split
multimodal.rsinto amultimodal/module directory along the pipeline'sown phase boundaries. This is a pure move — no behavior change.
detectconfigprocessassembleMultimodalData+ per-item splitserializetransport/dev/shmnamespace verificationThe crate-facing API is unchanged: the same symbols are re-exported from
mod.rs, so everymultimodal::Xcall site keeps working. Internal helpers arenow enforced-private (
pub(super)) instead of implicitly module-scoped, so there-export block in
mod.rsis now the explicit module contract.Changes
git mv multimodal.rs multimodal/mod.rs;mod.rsnow holds the module docs,the shared intermediate types (
MultimodalOutput,MultimodalIntermediate,PrecomputedMultimodalIntermediate), and thepub(crate)re-export surface.Reviewing: view with
git diff --color-moved=zebra— the vast majority isverbatim moves. The only non-move edits are per-file
useblocks andvisibility bumps (private →
pub(super)/pub(crate)).Test Plan
cargo check -p smg --lib --tests— cleancargo clippy -p smg --lib --tests -- -D warnings— cleancargo +nightly fmt— clean (only the new files touched)cargo test -p smg --lib multimodal— 24 passed, 0 failed (the sametests as before, now under
assemble/detect/config/process/serialize/transport)No functional or config surface changed, so no bindings/e2e impact.
Checklist
cargo +nightly fmtcleancargo clippy -- -D warningscleanSummary by CodeRabbit