feat(grpc_servicer): implement GetTokenizer RPC for vLLM backend

grpc_servicer/smg_grpc_servicer/sglang/servicer.py

@@ -8,11 +8,9 @@
 import asyncio
 import dataclasses
 import hashlib
-import io
 import logging
 import os
 import time
-import zipfile
 from collections.abc import AsyncIterator
 from datetime import datetime, timezone
 from pathlib import Path
@@ -55,29 +53,11 @@
 from smg_grpc_servicer.sglang.health_servicer import SGLangHealthServicer
 from smg_grpc_servicer.sglang.request_manager import GrpcRequestManager
 from smg_grpc_servicer.sglang.utils import abort_code_from_output
+from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip
 
 logger = logging.getLogger(__name__)
 HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20))
 
-# Tokenizer bundle streaming constants (aligned with Rust grpc_client limits)
-_TOKENIZER_CHUNK_SIZE = 64 * 1024  # 64 KB per gRPC chunk
-# Files to include in the tokenizer ZIP bundle.
-# Aligned with crates/tokenizer/src/hub.rs:is_tokenizer_file() plus model config files.
-_TOKENIZER_FILES = [
-    "tokenizer.json",
-    "tokenizer_config.json",
-    "config.json",
-    "generation_config.json",
-    "special_tokens_map.json",
-    "vocab.json",
-    "merges.txt",
-    "tokenizer.model",  # SentencePiece
-    "tiktoken.model",  # tiktoken
-    "chat_template.json",
-]
-# Glob patterns for additional tokenizer-related files
-_TOKENIZER_GLOBS = ["*.tiktoken", "*.jinja", "*.model"]
-
 
 def _convert_loads_to_protobuf(
     result: GetLoadsReqOutput,
@@ -596,7 +576,7 @@ async def GetTokenizer(
 
         # Build ZIP archive in memory
         try:
-            zip_buffer = self._build_tokenizer_zip(tokenizer_dir)
+            zip_buffer = build_tokenizer_zip(tokenizer_dir)
         except Exception as e:
             logger.error(f"Failed to build tokenizer ZIP: {e}\n{get_exception_traceback()}")
             await context.abort(grpc.StatusCode.INTERNAL, str(e))
@@ -614,36 +594,14 @@ async def GetTokenizer(
         offset = 0
         total = len(zip_data)
         while offset < total:
-            end = min(offset + _TOKENIZER_CHUNK_SIZE, total)
+            end = min(offset + CHUNK_SIZE, total)
             is_last = end == total
             yield common_pb2.GetTokenizerChunk(
                 data=bytes(zip_data[offset:end]),
                 sha256=sha256 if is_last else "",
             )
             offset = end
 
-    @staticmethod
-    def _build_tokenizer_zip(tokenizer_dir: Path) -> io.BytesIO:
-        """Create an in-memory ZIP archive of tokenizer files from a directory."""
-        buf = io.BytesIO()
-        added: set[str] = set()
-        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
-            # Exact-name files
-            for name in _TOKENIZER_FILES:
-                filepath = tokenizer_dir / name
-                if filepath.is_file():
-                    zf.write(filepath, name)
-                    added.add(name)
-            # Glob patterns (*.tiktoken, *.jinja, *.model)
-            for pattern in _TOKENIZER_GLOBS:
-                for match in tokenizer_dir.glob(pattern):
-                    if match.is_file() and match.name not in added:
-                        zf.write(match, match.name)
-                        added.add(match.name)
-        if not added:
-            raise FileNotFoundError(f"No tokenizer files found in {tokenizer_dir}")
-        return buf
-
     async def SubscribeKvEvents(
         self,
         request: common_pb2.SubscribeKvEventsRequest,

grpc_servicer/smg_grpc_servicer/tokenizer_bundle.py

@@ -0,0 +1,54 @@
+"""
+Shared tokenizer bundle utilities for gRPC servicers.
+
+Builds and streams tokenizer artifacts as ZIP bundles over gRPC.
+Used by both SGLang and vLLM servicers.
+"""
+
+import io
+import zipfile
+from pathlib import Path
+
+# Streaming chunk size (aligned with Rust grpc_client limits)
+CHUNK_SIZE = 64 * 1024  # 64 KB per gRPC chunk
+
+# Files to include in the tokenizer ZIP bundle.
+# Aligned with crates/tokenizer/src/hub.rs:is_tokenizer_file() plus model config files.
+TOKENIZER_FILES = [
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "config.json",
+    "generation_config.json",
+    "special_tokens_map.json",
+    "vocab.json",
+    "merges.txt",
+    "tokenizer.model",  # SentencePiece
+    "tiktoken.model",  # tiktoken
+    "chat_template.json",
+]
+
+# Glob patterns for additional tokenizer-related files
+TOKENIZER_GLOBS = ["*.tiktoken", "*.jinja", "*.model"]
+
+
+def build_tokenizer_zip(tokenizer_dir: Path) -> io.BytesIO:
+    """Create an in-memory ZIP archive of tokenizer files from a directory."""
+    buf = io.BytesIO()
+    added: set[str] = set()
+    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+        # Exact-name files
+        for name in TOKENIZER_FILES:
+            filepath = tokenizer_dir / name
+            if filepath.is_file():
+                zf.write(filepath, name)
+                added.add(name)
+        # Glob patterns (*.tiktoken, *.jinja, *.model)
+        for pattern in TOKENIZER_GLOBS:
+            for match in tokenizer_dir.glob(pattern):
+                if match.is_file() and match.name not in added:
+                    zf.write(match, match.name)
+                    added.add(match.name)
+    if not added:
+        raise FileNotFoundError(f"No tokenizer files found in {tokenizer_dir}")
+    buf.seek(0)
+    return buf

grpc_servicer/smg_grpc_servicer/vllm/servicer.py

@@ -5,13 +5,16 @@
 Implements the VllmEngine gRPC service on top of vLLM's EngineClient.
 """
 
+import hashlib
 import itertools
 import time
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, AsyncIterator
+from pathlib import Path
 
 import grpc
 import torch
 from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
+from smg_grpc_proto.generated import common_pb2
 from transformers import BatchFeature
 from vllm import PoolingParams, SamplingParams, TokensPrompt
 from vllm.engine.protocol import EngineClient
@@ -27,6 +30,8 @@
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams
 
+from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip
+
 logger = init_logger(__name__)
 
 # Proto dtype string → torch dtype
@@ -49,13 +54,14 @@ class VllmEngineServicer(vllm_engine_pb2_grpc.VllmEngineServicer):
     """
     gRPC servicer implementing the VllmEngine service.
 
-    Handles 6 RPCs:
+    Handles 7 RPCs:
     - Generate: Streaming text generation
     - Embed: Embeddings
     - HealthCheck: Health probe
     - Abort: Cancel requests out-of-band
     - GetModelInfo: Model metadata
     - GetServerInfo: Server state
+    - GetTokenizer: Stream tokenizer artifacts
     """
 
     def __init__(self, async_llm: EngineClient, start_time: float):
@@ -358,6 +364,65 @@ async def GetServerInfo(
             kv_role=kv_role,
         )
 
+    async def GetTokenizer(
+        self,
+        request: common_pb2.GetTokenizerRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> AsyncIterator[common_pb2.GetTokenizerChunk]:
+        """Stream tokenizer artifacts as a ZIP bundle.
+
+        Resolves the tokenizer directory from model_config, zips all relevant
+        tokenizer files, and streams them as GetTokenizerChunk messages.
+        The final chunk carries the SHA-256 fingerprint of the full archive.
+        """
+        logger.info("Receive GetTokenizer request")
+
+        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)
+
+        # model_config.tokenizer may be an HF model ID (e.g. "meta-llama/...")
+        # rather than a local path. Resolve it to the HF cache directory.
+        if not tokenizer_dir.is_dir():
+            try:
+                from huggingface_hub import snapshot_download
+
+                tokenizer_dir = Path(snapshot_download(tokenizer_path, local_files_only=True))
+            except Exception:
+                pass  # Fall through to build_tokenizer_zip which will raise
+
+        # Build ZIP archive in memory
+        try:
+            zip_buffer = build_tokenizer_zip(tokenizer_dir)
+        except Exception as e:
+            logger.exception("Failed to build tokenizer ZIP")
+            await context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+        zip_data = zip_buffer.getbuffer()
+        sha256 = hashlib.sha256(zip_data).hexdigest()
+
+        logger.info(
+            "Streaming tokenizer bundle: %d bytes, sha256=%s",
+            len(zip_data),
+            sha256,
+        )
+
+        # Stream chunks; SHA-256 only on the final chunk
+        offset = 0
+        total = len(zip_data)
+        while offset < total:
+            end = min(offset + CHUNK_SIZE, total)
+            is_last = end == total
+            yield common_pb2.GetTokenizerChunk(
+                data=bytes(zip_data[offset:end]),
+                sha256=sha256 if is_last else "",
+            )
+            offset = end
+
     # ========== Helper methods ==========
 
     def _build_preprocessed_mm_inputs(