feat(grpc_servicer): implement GetTokenizer RPC for vLLM backend - #1142
Conversation
Signed-off-by: Chang Su <chang.s.su@oracle.com>
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdds a shared tokenizer-bundle utility and implements async server-streaming Changes
Sequence DiagramsequenceDiagram
participant Client
participant Server as gRPC Server
participant FS as File System
participant ZIP as Zip Encoder
participant Hash as SHA-256
Client->>Server: GetTokenizer(request)
activate Server
Server->>FS: Resolve tokenizer path (model_config.tokenizer or snapshot)
FS-->>Server: tokenizer_dir / files
Server->>ZIP: build_tokenizer_zip(tokenizer_dir)
activate ZIP
ZIP->>FS: Read TOKENIZER_FILES + TOKENIZER_GLOBS
FS-->>ZIP: file contents
ZIP->>ZIP: Create in-memory ZIP bytes
ZIP-->>Server: ZIP bytes
deactivate ZIP
Server->>Hash: Compute SHA-256(ZIP bytes)
Hash-->>Server: hex digest
loop Stream ZIP chunks
Server->>Client: GetTokenizerChunk(data=chunk, sha256="") -- intermediate
end
Server->>Client: GetTokenizerChunk(data=final_chunk, sha256=digest) -- final
deactivate Server
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Clean implementation. The GetTokenizer RPC correctly mirrors the SGLang version with appropriate vLLM adaptations. Proto types match the service definition, error handling follows existing patterns in the file, and the streaming/chunking logic is correct.
Summary: 0 🔴 Important · 1 🟡 Nit · 0 🟣 Pre-existing
The one nit is about the full duplication of tokenizer-bundle constants and _build_tokenizer_zip across the two servicers — worth extracting to a shared module to prevent future drift, but not blocking.
![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 the current code and only fix it if needed.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/vllm/servicer.py`:
- Around line 448-452: The loop that adds tokenizer files to the ZIP uses
tokenizer_dir.glob(pattern) which can yield nondeterministic ordering; to
produce stable ZIP byte order and SHA fingerprints, sort the glob matches before
iterating: for each pattern in _TOKENIZER_GLOBS, replace direct iteration over
tokenizer_dir.glob(pattern) with iterating over a sorted list of matches (e.g.,
sorted(tokenizer_dir.glob(pattern))), then keep the existing checks and
zf.write(match, match.name) / added.add(match.name) logic intact so behavior
doesn't change aside from deterministic ordering.
🪄 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: 73b3217b-7f0e-41c1-bd43-0253cbefb7bd
📒 Files selected for processing (1)
grpc_servicer/smg_grpc_servicer/vllm/servicer.py
![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 adds a GetTokenizer gRPC endpoint to the VllmEngineServicer to stream tokenizer artifacts as a ZIP bundle. The implementation includes logic for file discovery, in-memory compression, and SHA-256 fingerprinting. Feedback suggests offloading the synchronous ZIP creation to a thread pool to avoid blocking the asynchronous event loop and using a more specific gRPC status code when tokenizer files are not found.
Move _TOKENIZER_FILES, _TOKENIZER_GLOBS, _TOKENIZER_CHUNK_SIZE, and _build_tokenizer_zip into smg_grpc_servicer/tokenizer_bundle.py so both sglang and vllm servicers import from one source of truth. Signed-off-by: Chang Su <chang.s.su@oracle.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/tokenizer_bundle.py`:
- Around line 34-53: The build_tokenizer_zip function returns an io.BytesIO
(buf) still positioned at EOF; ensure the returned stream is readable by
rewinding it before returning (call seek(0) on buf), i.e., in
build_tokenizer_zip after closing the zipfile context but before returning buf
so callers of build_tokenizer_zip (which creates buf and writes TOKENIZER_FILES
/ TOKENIZER_GLOBS into it) receive a stream with its cursor at the start.
🪄 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: c87f901d-c169-4992-b3a7-a3537e265300
📒 Files selected for processing (3)
grpc_servicer/smg_grpc_servicer/sglang/servicer.pygrpc_servicer/smg_grpc_servicer/tokenizer_bundle.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.py
Add buf.seek(0) before returning so future callers that use read() instead of getbuffer() get the full archive. Signed-off-by: Chang Su <chang.s.su@oracle.com>
![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: deb2b8df3b
ℹ️ 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".
| tokenizer_path = self.engine.model_config.tokenizer | ||
| if not tokenizer_path: | ||
| await context.abort( | ||
| grpc.StatusCode.FAILED_PRECONDITION, | ||
| "Tokenizer path is not configured on this server.", | ||
| ) | ||
| tokenizer_dir = Path(tokenizer_path) |
There was a problem hiding this comment.
GetTokenizer assumes self.engine.model_config.tokenizer is a filesystem directory and immediately wraps it with Path(...), but in vLLM the default tokenizer value is often the HF model ID string (e.g. meta-llama/...) when --tokenizer is not explicitly set. In that common configuration, build_tokenizer_zip searches a non-existent local path and aborts with INTERNAL, so the new RPC still fails for standard vllm serve <hf-id> --grpc deployments and the gateway cannot fetch tokenizer artifacts in offline/separate-node setups.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The comment is partially correct: model_config.tokenizer is indeed the raw string (HF ID or local path) — vLLM does not resolve it to a local directory in ModelConfig.
However, this is handled gracefully:
-
When vLLM is started with a local path (e.g.,
--model /raid/models/...),model_config.tokenizeris that local path andGetTokenizerworks correctly (verified in testing). -
When vLLM is started with an HF model ID (e.g.,
--model meta-llama/Llama-3.2-1B-Instruct),build_tokenizer_zipwill fail withFileNotFoundError, the handler returnsINTERNAL, and SMG falls through to try the next worker or load from HF directly viaLoadTokenizerStep(which handles HF IDs natively).
The INTERNAL error path is already exercised — tokenizer_registration.rs:278-285 logs it and continues. No crash, no stuck state.
For production IGW/K8s, models are typically mounted from PVCs at local paths. The HF ID case is an edge case where GetTokenizer gracefully degrades.
model_config.tokenizer can be a HuggingFace model ID instead of a local path when vLLM is started with --model meta-llama/... without an explicit --tokenizer. Use snapshot_download(local_files_only=True) to resolve the ID to the HF cache directory. Signed-off-by: Chang Su <chang.s.su@oracle.com>
|
|
||
| tokenizer_dir = Path(snapshot_download(tokenizer_path, local_files_only=True)) | ||
| except Exception: | ||
| pass # Fall through to build_tokenizer_zip which will raise |
There was a problem hiding this comment.
🟡 Nit: Silently swallowing the snapshot_download exception loses diagnostic context. When the HF resolution fails and build_tokenizer_zip subsequently raises on the non-existent directory (line 401), the logged traceback won't show why the resolution failed (e.g., model not in cache, corrupted snapshot, permission error).
A logger.debug here would make this much easier to diagnose without changing control flow:
| pass # Fall through to build_tokenizer_zip which will raise | |
| except Exception: | |
| logger.debug("HF cache lookup failed for %r, will try raw path", tokenizer_path, exc_info=True) |
2 participants
Description
Problem
The vLLM gRPC servicer does not implement the
GetTokenizerRPC, causing the gateway to receiveUNIMPLEMENTEDwhen it attempts to fetch tokenizer artifacts from a vLLM worker. In IGW/K8s deployments where gateway and workers run on separate nodes, the gateway cannot load the tokenizer.Solution
Implement the server-side
GetTokenizerhandler in the vLLM gRPC servicer, following the same pattern as the sglang implementation (#1136). The handler:model_config.tokenizer(vLLM's equivalent of sglang'sserver_args.tokenizer_path)crates/tokenizer/src/hub.rs:is_tokenizer_file())GetTokenizerChunkmessages with SHA-256 fingerprint on the final chunkChanges
GetTokenizerasync generator method toVllmEngineServicer_build_tokenizer_zipstatic method for ZIP archive creation_TOKENIZER_FILESallowlist and_TOKENIZER_GLOBSpatterns (same as sglang)hashlib,io,zipfile,Path,AsyncIterator,common_pb2Test Plan
Setup
Test 1: Direct gRPC call to verify bundle integrity
Test script
vLLM server logs:
Test output:
Test 2: Full end-to-end pipeline (cross-machine, vLLM)
Verified the complete
worker → GetTokenizer → gatewayflow using SSH port-forwarding to simulate the production IGW/K8s scenario.# Local machine: start SMG (no --model-path, no local model files) cargo run --bin smg -- --host 0.0.0.0 --port 3002 --prometheus-port 9322 \ --worker-urls grpc://127.0.0.1:8080 --log-level infoSMG gateway logs (local machine — no model files on disk)
Verified model is serving after GetTokenizer load:
Verified inference works end-to-end:
Note on tokenizer lifecycle
Tokenizer files are not persisted on disk after loading. The flow: stream ZIP → extract to OS temp dir → load into memory → delete temp dir. The tokenizer lives only as
Arc<dyn Tokenizer>in theTokenizerRegistry.Learnings applied from sglang PR review (#1136)
_resolve_tokenizer_dir()— runtime already resolves the pathreturnaftercontext.abort()— it raisesgetbuffer()instead ofgetvalue()for zero-copy streaminglogger.exception()for error logging (matching vLLM file conventions)FAILED_PRECONDITIONChecklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Improvements