[Router][Bugfix] Keep the KV-aware /tokenize fallback off the event loop

src/tests/test_kvaware_router.py

@@ -1,5 +1,7 @@
 """Unit tests for the KV-cache-aware routing logic."""
 
+import asyncio
+import time
 from types import SimpleNamespace
 
 import pytest
@@ -220,3 +222,64 @@ async def query_manager(msg):
     )
 
     assert selected == url_b
+
+
+@pytest.mark.asyncio
+async def test_kvaware_remote_tokenize_fallback_does_not_block_event_loop(
+    monkeypatch,
+):
+    """The remote /tokenize fallback is a blocking HTTP call; it must run off
+    the event loop so other requests keep being served."""
+
+    def no_local_tokenizer(_model_name):
+        raise OSError("tokenizer not available locally")
+
+    def slow_post(*_args, **_kwargs):
+        time.sleep(0.3)
+        return SimpleNamespace(json=lambda: {"tokens": [1, 2, 3]})
+
+    monkeypatch.setattr(
+        routing_logic,
+        "AutoTokenizer",
+        SimpleNamespace(from_pretrained=no_local_tokenizer),
+        raising=False,
+    )
+    monkeypatch.setattr(routing_logic.requests, "post", slow_post)
+
+    url = "http://10.0.0.1:8000"
+    router = KvawareRouter.__new__(KvawareRouter)
+    router.tokenizers = {}
+    router.threshold = 0
+    router.session_key = "x-session-id"
+    router.hash_ring = routing_logic.HashRing()
+    router.instance_id_to_ip = {}
+    lookup_tokens = []
+
+    async def query_manager(msg):
+        lookup_tokens.append(msg.tokens)
+        return SimpleNamespace(layout_info={})
+
+    router.query_manager = query_manager
+
+    ticks = 0
+
+    async def ticker():
+        nonlocal ticks
+        while True:
+            await asyncio.sleep(0.01)
+            ticks += 1
+
+    ticker_task = asyncio.create_task(ticker())
+    selected = await router.route_request(
+        [EndpointInfo(url)],
+        {},
+        {url: SimpleNamespace(qps=0)},
+        SimpleNamespace(headers={}),
+        {"prompt": "test"},
+    )
+    ticker_task.cancel()
+
+    assert selected == url
+    assert lookup_tokens == [[1, 2, 3]]
+    # the event loop kept running while the fallback request was in flight
+    assert ticks >= 10

src/vllm_router/routers/routing_logic.py

@@ -367,6 +367,33 @@ def close(self):
                 pass
             self.lmcache_cluster_monitor_task = None
 
+    async def tokenize_prompt(
+        self, endpoints: List[EndpointInfo], request_json: Dict
+    ) -> List[int]:
+        """Local-first tokenization with the remote `/tokenize` fallback.
+
+        The remote fallback is a blocking HTTP call, so it runs in an
+        executor rather than on the event loop.
+        """
+        try:
+            tokenizer = self._get_tokenizer(endpoints)
+            return tokenizer.encode(request_json.get("prompt", ""))
+        except Exception:
+            remote_url = endpoints[0].url + "/tokenize"
+            headers = {"Content-Type": "application/json"}
+            data = {
+                "model": endpoints[0].model_names[0],
+                "prompt": request_json.get("prompt", ""),
+            }
+            loop = asyncio.get_running_loop()
+            response = await loop.run_in_executor(
+                None,
+                lambda: requests.post(
+                    remote_url, headers=headers, json=data, timeout=10
+                ),
+            )
+            return response.json()["tokens"]
+
     async def route_request(
         self,
         endpoints: List[EndpointInfo],
@@ -391,24 +418,8 @@ async def route_request(
             request_json (Dict): The request body (needed for finding the
             longest prefix match)
         """
-        token_ids = None
-        # Local-first tokenization, fall back to remote "/tokenize" API on failure
         # TODO (Yuhan): Handle chat completions
-        try:
-            tokenizer = self._get_tokenizer(endpoints)
-            token_ids = tokenizer.encode(request_json.get("prompt", ""))
-        except Exception:
-            # Remote /tokenize fallback (let errors bubble up to keep behavior simple)
-            remote_url = endpoints[0].url + "/tokenize"
-            headers = {"Content-Type": "application/json"}
-            data = {
-                "model": endpoints[0].model_names[0],
-                "prompt": request_json.get("prompt", ""),
-            }
-            body = requests.post(
-                remote_url, headers=headers, json=data, timeout=10
-            ).json()
-            token_ids = body["tokens"]
+        token_ids = await self.tokenize_prompt(endpoints, request_json)
 
         event_id = "Lookup" + str(uuid.uuid4())
         msg = LookupMsg(tokens=token_ids, event_id=event_id)
@@ -662,33 +673,6 @@ async def query_endpoint(endpoint: EndpointInfo) -> None:
         await asyncio.gather(*(query_endpoint(e) for e in endpoints))
         logger.info(f"Instance id to ip mapping: {self.instance_id_to_ip}")
 
-    async def tokenize_prompt(
-        self, endpoints: List[EndpointInfo], request_json: Dict
-    ) -> List[int]:
-        """Local-first tokenization with the remote `/tokenize` fallback.
-
-        The remote fallback is a blocking HTTP call, so it runs in an
-        executor rather than on the event loop.
-        """
-        try:
-            tokenizer = self._get_tokenizer(endpoints)
-            return tokenizer.encode(request_json.get("prompt", ""))
-        except Exception:
-            remote_url = endpoints[0].url + "/tokenize"
-            headers = {"Content-Type": "application/json"}
-            data = {
-                "model": endpoints[0].model_names[0],
-                "prompt": request_json.get("prompt", ""),
-            }
-            loop = asyncio.get_running_loop()
-            response = await loop.run_in_executor(
-                None,
-                lambda: requests.post(
-                    remote_url, headers=headers, json=data, timeout=10
-                ),
-            )
-            return response.json()["tokens"]
-
     def fallback_url(
         self,
         endpoints: List[EndpointInfo],