Add RenderChatCompletion and RenderCompletion gRPC RPCs via vLLM OpenAIServingRender (#428)

Author: sagearc

Makefile

@@ -263,13 +263,14 @@ uds-tokenizer-install-deps: detect-python ## Set up venv and install UDS tokeniz
 		echo "Virtual environment already exists"; \
 	fi
 	@echo "Installing dependencies..."
-	@$(UDS_TOKENIZER_VENV_BIN)/pip install "$(UDS_TOKENIZER_DIR)[test]"
+	@$(UDS_TOKENIZER_VENV_BIN)/pip install "$(UDS_TOKENIZER_DIR)[test]" --extra-index-url https://download.pytorch.org/whl/cpu
 
 .PHONY: uds-tokenizer-service-test
 uds-tokenizer-service-test: uds-tokenizer-install-deps ## Run UDS tokenizer integration tests (starts server automatically)
 	@printf "\033[33;1m==== Running UDS tokenizer integration tests ====\033[0m\n"
 	@$(UDS_TOKENIZER_VENV_BIN)/python -m pytest \
 		$(UDS_TOKENIZER_DIR)/tests/test_integration.py \
+		$(UDS_TOKENIZER_DIR)/tests/test_renderer.py \
 		-v --timeout=60
 
 .PHONY: bench

api/tokenizerpb/tokenizer.proto

@@ -110,6 +110,56 @@ message InitializeTokenizerResponse {
   string error_message = 2;   // Error message if initialization failed
 }
 
+// PlaceholderRange represents a range of placeholder tokens for a multimodal item
+message PlaceholderRange {
+  int32 offset = 1; // Start index of the placeholder tokens in the prompt
+  int32 length = 2; // Number of placeholder tokens
+}
+
+// StringList holds a list of strings (used as map values)
+message StringList {
+  repeated string values = 1;
+}
+
+// PlaceholderRangeList holds a list of PlaceholderRange (used as map values)
+message PlaceholderRangeList {
+  repeated PlaceholderRange ranges = 1;
+}
+
+// MultiModalFeatures contains multimodal metadata produced by the render step
+message MultiModalFeatures {
+  map<string, StringList> mm_hashes = 1;                  // Per-modality item hashes
+  map<string, PlaceholderRangeList> mm_placeholders = 2;  // Per-modality placeholder ranges
+}
+
+// RenderChatCompletionRequest wraps an OpenAI chat completion request for rendering
+message RenderChatCompletionRequest {
+  string request_json = 1; // JSON-serialized OpenAI ChatCompletionRequest
+  string model_name = 2;   // Model name to use for renderer selection
+}
+
+// RenderChatCompletionResponse contains the rendered output from OpenAIServingRender
+message RenderChatCompletionResponse {
+  string request_id = 1;            // Request ID from the render response
+  repeated uint32 token_ids = 2;    // Token IDs for the rendered input
+  MultiModalFeatures features = 3;  // Multimodal features (populated for MM inputs)
+  bool success = 4;                 // Whether the request was successful
+  string error_message = 5;         // Error message if the request failed
+}
+
+// RenderCompletionRequest wraps an OpenAI completion request for rendering
+message RenderCompletionRequest {
+  string request_json = 1; // JSON-serialized OpenAI CompletionRequest
+  string model_name = 2;   // Model name to use for renderer selection
+}
+
+// RenderCompletionResponse contains the rendered output for each prompt in the completion request
+message RenderCompletionResponse {
+  repeated RenderChatCompletionResponse items = 1; // One item per prompt in the request
+  bool success = 2;                                // Whether the request was successful
+  string error_message = 3;                        // Error message if the request failed
+}
+
 // TokenizationService defines the gRPC service for tokenization
 service TokenizationService {
   // Tokenize converts a text input to token IDs
@@ -120,4 +170,12 @@ service TokenizationService {
 
   // InitializeTokenizer initializes the tokenizer for a specific model
   rpc InitializeTokenizer(InitializeTokenizerRequest) returns (InitializeTokenizerResponse);
+
+  // RenderChatCompletion renders an OpenAI chat completion request via OpenAIServingRender,
+  // returning token IDs and multimodal features as a GenerateRequest-compatible response
+  rpc RenderChatCompletion(RenderChatCompletionRequest) returns (RenderChatCompletionResponse);
+
+  // RenderCompletion renders an OpenAI completion request via OpenAIServingRender,
+  // returning one rendered item per prompt in the request
+  rpc RenderCompletion(RenderCompletionRequest) returns (RenderCompletionResponse);
 }

services/uds_tokenizer/Dockerfile

@@ -23,7 +23,7 @@ WORKDIR /app
 
 # Copy project metadata and install dependencies
 COPY pyproject.toml /app/pyproject.toml
-RUN pip install --no-cache-dir .
+RUN pip install --no-cache-dir . --extra-index-url https://download.pytorch.org/whl/cpu
 
 # Runtime stage
 FROM --platform=$TARGETPLATFORM python:3.12-slim
@@ -60,6 +60,9 @@ RUN mkdir -p /.modelscope && chown -R 65532:65532 /.modelscope
 # Create and set permissions for Hugging Face cache directory
 RUN mkdir -p /.cache && chown -R 65532:65532 /.cache
 
+# Create non-root user so getpwuid() works (required by torch/vllm)
+RUN useradd -u 65532 -m nonroot
+
 # Switch to non-root user
 USER 65532:65532
 

services/uds_tokenizer/pyproject.toml

@@ -4,19 +4,22 @@ version = "0.1.0"
 description = "UDS Tokenizer Service - gRPC tokenization over Unix Domain Socket"
 requires-python = ">=3.12"
 dependencies = [
-    "pydantic==2.11.7",
-    "shortuuid==1.0.13",
-    "transformers==4.53.0",
-    "safetensors==0.5.3",
-    "Jinja2==3.1.6",
-    "modelscope",
-    "huggingface-hub",
-    "aiohttp==3.9.5",
-    "protobuf==6.31.1",
-    "tiktoken>=0.7.0",
-    "grpcio==1.76.0",
-    "grpcio-tools==1.76.0",
-    "grpcio-reflection==1.76.0",
+    # vLLM 0.18.0rc1 CPU wheels — Linux-only, arch-conditional via PEP 508 markers
+    "vllm @ https://wheels.vllm.ai/262ddd0d81a1e4687e209f988d6ea32616e736fa/vllm-0.18.0rc1%2Bcpu-cp38-abi3-manylinux_2_35_x86_64.whl ; platform_system == 'Linux' and platform_machine == 'x86_64'",
+    "vllm @ https://wheels.vllm.ai/262ddd0d81a1e4687e209f988d6ea32616e736fa/vllm-0.18.0rc1%2Bcpu-cp38-abi3-manylinux_2_35_aarch64.whl ; platform_system == 'Linux' and platform_machine == 'aarch64'",
+    "pydantic>=2.12.0",
+    "shortuuid>=1.0.13",
+    "transformers>=4.57.0,<5",
+    "safetensors>=0.7.0",
+    "Jinja2>=3.1.6",
+    "modelscope>=1.35.0",
+    "huggingface-hub>=0.36.0",
+    "aiohttp>=3.13.0",
+    "protobuf>=6.31.1",
+    "tiktoken>=0.12.0",
+    "grpcio>=1.78.0",
+    "grpcio-tools>=1.78.0",
+    "grpcio-reflection>=1.78.0",
 ]
 
 [project.optional-dependencies]

services/uds_tokenizer/run_grpc_server.py

@@ -27,6 +27,7 @@
 
 from aiohttp import web
 from tokenizer_service.tokenizer import TokenizerService
+from tokenizer_service.renderer import RendererService
 from tokenizer_grpc_service import create_grpc_server
 from utils.thread_pool_utils import get_thread_pool
 
@@ -55,6 +56,7 @@
 probe_started_event = threading.Event()  # Event to signal when probe server has started
 current_config = None
 tokenizer_service = None
+renderer_service = None
 tokenizer_ready = False
 shutdown_event = threading.Event()  # Event to signal shutdown
 
@@ -72,10 +74,11 @@ def _signal_handler(signum, frame):
 
 def initialize_tokenizer():
     """Initialize the tokenizer service without pre-loading a specific model"""
-    global tokenizer_service, current_config, tokenizer_ready
+    global tokenizer_service, renderer_service, current_config, tokenizer_ready
     try:
         # Initialize tokenizer service without pre-loading any model
         tokenizer_service = TokenizerService()  # Empty constructor
+        renderer_service = RendererService()
         tokenizer_ready = True
         logging.info("Tokenizer service initialized successfully")
     except Exception as e:
@@ -181,7 +184,7 @@ def run_server():
 
     thread_pool = get_thread_pool()
     grpc_server = create_grpc_server(
-        tokenizer_service, UDS_SOCKET_PATH, thread_pool, GRPC_PORT
+        tokenizer_service, UDS_SOCKET_PATH, thread_pool, renderer_service, GRPC_PORT
     )
     grpc_server.start()
     logging.info(

services/uds_tokenizer/tests/conftest.py

@@ -23,6 +23,7 @@
 
 import tokenizerpb.tokenizer_pb2_grpc as tokenizer_pb2_grpc
 from tokenizer_service.tokenizer import TokenizerService
+from tokenizer_service.renderer import RendererService
 from tokenizer_grpc_service import create_grpc_server
 from utils.thread_pool_utils import get_thread_pool
 
@@ -49,24 +50,29 @@ def uds_socket_path() -> Iterator[str]:
 
 
 @pytest.fixture(scope="session")
-def tokenizer_service(uds_socket_path: str) -> Iterator[TokenizerService]:
-    """Provide the TokenizerService instance used by the gRPC server."""
-    service = TokenizerService()
+def grpc_server(uds_socket_path: str) -> Iterator[None]:
+    """Start and stop the gRPC server for the test session."""
+    tokenizer_service = TokenizerService()
+    renderer_service = RendererService()
+
     thread_pool = get_thread_pool()
-    server = create_grpc_server(service, uds_socket_path, thread_pool)
+    server = create_grpc_server(
+        tokenizer_service,
+        uds_socket_path,
+        thread_pool,
+        renderer_service=renderer_service,
+    )
     server.start()
 
-    yield service
+    yield
 
     # Graceful shutdown with matching timeout
     stop_future = server.stop(grace=5)
     stop_future.wait(timeout=5)
 
 
 @pytest.fixture(scope="session")
-def grpc_channel(
-    tokenizer_service: TokenizerService, uds_socket_path: str
-) -> Iterator[grpc.Channel]:
+def grpc_channel(grpc_server: None, uds_socket_path: str) -> Iterator[grpc.Channel]:
     """Create a gRPC channel connected to the test server.
 
     Uses wait_for_ready to automatically retry connection until server is ready.

services/uds_tokenizer/tests/test_integration.py

@@ -26,9 +26,9 @@
 
 import grpc
 import pytest
+from transformers import AutoTokenizer
 
 import tokenizerpb.tokenizer_pb2 as tokenizer_pb2
-from tokenizer_service.tokenizer import TokenizerService
 
 
 # ---------------------------------------------------------------------------
@@ -99,9 +99,7 @@ def test_tokenize_simple_text(self, grpc_stub, test_model):
         assert resp.success
         assert len(resp.input_ids) > 0
 
-    def test_tokenize_returns_offset_pairs(
-        self, grpc_stub, test_model, tokenizer_service: TokenizerService
-    ):
+    def test_tokenize_returns_offset_pairs(self, grpc_stub, test_model):
         """Tokenize returns offset_pairs alongside token IDs."""
         grpc_stub.InitializeTokenizer(
             tokenizer_pb2.InitializeTokenizerRequest(model_name=test_model)
@@ -118,16 +116,14 @@ def test_tokenize_returns_offset_pairs(
         assert len(resp.offset_pairs) == 2 * len(resp.input_ids)
 
         # Verify token count matches tokenizer
-        tokenizer, _ = tokenizer_service.get_tokenizer_for_model(test_model)
+        tokenizer = AutoTokenizer.from_pretrained(test_model)
         expected_tokens = tokenizer.encode("Hello world", add_special_tokens=True)
         assert list(resp.input_ids) == expected_tokens
 
-    def test_tokenize_without_special_tokens(
-        self, grpc_stub, tokenizer_service: TokenizerService
-    ):
+    def test_tokenize_without_special_tokens(self, grpc_stub):
         """Tokenize with add_special_tokens=False omits special tokens."""
 
-        model_name = "google-bert/bert-base-uncased"
+        model_name = "openai/gpt-oss-120b"
 
         grpc_stub.InitializeTokenizer(
             tokenizer_pb2.InitializeTokenizerRequest(model_name=model_name)
@@ -147,19 +143,15 @@ def test_tokenize_without_special_tokens(
             )
         )
         assert with_special.success and without_special.success
-        # With special tokens should produce > tokens as without.
-        assert len(with_special.input_ids) > len(without_special.input_ids)
 
-        # Verify special tokens using actual tokenizer
-        tokenizer, _ = tokenizer_service.get_tokenizer_for_model(model_name)
-
-        # BERT adds [CLS] at start and [SEP] at end
-        assert with_special.input_ids[0] == tokenizer.cls_token_id
-        assert with_special.input_ids[-1] == tokenizer.sep_token_id
-
-        # Without special tokens should not have [CLS] or [SEP]
-        assert without_special.input_ids[0] != tokenizer.cls_token_id
-        assert without_special.input_ids[-1] != tokenizer.sep_token_id
+        # Verify both match the underlying tokenizer's behavior
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        assert list(with_special.input_ids) == tokenizer.encode(
+            "test", add_special_tokens=True
+        )
+        assert list(without_special.input_ids) == tokenizer.encode(
+            "test", add_special_tokens=False
+        )
 
     def test_tokenize_empty_input(self, grpc_stub, test_model):
         grpc_stub.InitializeTokenizer(
@@ -191,9 +183,7 @@ def test_tokenize_long_input(self, grpc_stub, test_model):
         assert resp.success
         assert len(resp.input_ids) > 100  # Should have many tokens.
 
-    def test_tokenize_special_characters(
-        self, grpc_stub, test_model, tokenizer_service: TokenizerService
-    ):
+    def test_tokenize_special_characters(self, grpc_stub, test_model):
         """Tokenize handles special / unicode characters."""
         grpc_stub.InitializeTokenizer(
             tokenizer_pb2.InitializeTokenizerRequest(model_name=test_model)
@@ -210,8 +200,7 @@ def test_tokenize_special_characters(
         assert len(resp.input_ids) > 0
 
         # Verify tokenization matches actual tokenizer
-        tokenizer, _ = tokenizer_service.get_tokenizer_for_model(test_model)
-
+        tokenizer = AutoTokenizer.from_pretrained(test_model)
         expected_tokens = tokenizer.encode(test_input, add_special_tokens=True)
         assert list(resp.input_ids) == expected_tokens