diff --git a/src/tests/test_kvaware_chat_tokenization.py b/src/tests/test_kvaware_chat_tokenization.py new file mode 100644 index 000000000..02524f862 --- /dev/null +++ b/src/tests/test_kvaware_chat_tokenization.py @@ -0,0 +1,387 @@ +"""Unit tests for chat-completion tokenization in the KV-aware routers. + +`kvaware` and `loadaware` place a request by asking the LMCache controller +which engine already holds KV for the request's token-id prefix. vLLM +engines cache KV for the token ids *after* chat-template application, so a +chat-completions body (a "messages" array, no "prompt" key) must be +tokenized through `apply_chat_template` - the old +`encode(request_json.get("prompt", ""))` tokenized the empty string, the +lookup matched nothing, and every chat request silently degraded to the +session/QPS fallback. + +As in `test_loadaware_router.py`, the routers are built with `__new__` and +only the attributes the tokenize path reads, so no LMCache controller (and +no network) is needed. +""" + +from typing import Any, Dict + +import pytest +from uhashring import HashRing + +import vllm_router.routers.routing_logic as routing_logic +from vllm_router.routers.routing_logic import ( + KvawareRouter, + LoadAwareRouter, + _extract_token_ids, + _normalize_chat_messages, + _tokenize_request_payload, +) + + +@pytest.fixture(autouse=True) +def lookup_msg_stub(monkeypatch): + """`LookupMsg`/`QueryInstMsg` come from the optional lmcache dependency; + stub them when absent so the routing tests run without the lmcache + extra.""" + + class _Msg: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + for name in ("LookupMsg", "QueryInstMsg"): + if not hasattr(routing_logic, name): + monkeypatch.setattr(routing_logic, name, _Msg, raising=False) + + +URL_A = "http://10.0.0.1:8000" +URL_B = "http://10.0.0.2:8000" +INST_A = "instance-a" +LOCAL = "LocalCPUBackend" +MODEL = "test-model" +CHAT_IDS = [101, 102, 103, 104, 105] +MESSAGES = [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "Hello there"}, +] + + +class EndpointInfo: + def __init__(self, url: str): + self.url = url + self.model_names = [MODEL] + + +class LookupRet: + def __init__(self, layout_info: Dict[str, Any]): + self.layout_info = layout_info + + +def endpoints(*urls): + return [EndpointInfo(url=url) for url in urls] + + +class ChatTokenizer: + """Records calls; template ids are disjoint from encode ids so a test + can tell which path produced them.""" + + def __init__(self): + self.chat_template_calls = [] + self.encode_calls = [] + + RENDERED = "" + + def apply_chat_template( + self, messages, add_generation_prompt=False, tokenize=False + ): + self.chat_template_calls.append( + { + "messages": messages, + "add_generation_prompt": add_generation_prompt, + "tokenize": tokenize, + } + ) + # tokenize=True return types vary across transformers versions + # (ids vs Encoding objects) - the router must NOT use that form. + assert tokenize is False, "router must render text, not tokenize=True" + return self.RENDERED + + def encode(self, prompt, add_special_tokens=True): + self.encode_calls.append( + {"prompt": prompt, "add_special_tokens": add_special_tokens} + ) + if prompt == self.RENDERED: + # the rendered template must be encoded WITHOUT re-adding + # special tokens (the template already carries them) + assert add_special_tokens is False + return list(CHAT_IDS) + return [1] * len(prompt) + + +class TemplatelessTokenizer(ChatTokenizer): + """A tokenizer with no chat template, as `apply_chat_template` raises on + base models.""" + + def apply_chat_template(self, *args, **kwargs): + raise ValueError("no chat template defined") + + +# --- local tokenization ------------------------------------------------------- + + +def test_messages_tokenize_through_the_chat_template(): + tokenizer = ChatTokenizer() + ids = _extract_token_ids(tokenizer, {"messages": MESSAGES}) + assert ids == CHAT_IDS + call = tokenizer.chat_template_calls[0] + assert call["add_generation_prompt"] is True + # rendered to text, then encoded without re-adding special tokens - the + # tokenize=True return type varies across transformers versions + assert call["tokenize"] is False + # exactly one encode: the rendered template text, no special re-adding + assert tokenizer.encode_calls == [ + {"prompt": ChatTokenizer.RENDERED, "add_special_tokens": False} + ] + + +def test_prompt_requests_keep_the_plain_encode_path(): + tokenizer = ChatTokenizer() + ids = _extract_token_ids(tokenizer, {"prompt": "hello"}) + assert ids == [1] * len("hello") + assert tokenizer.chat_template_calls == [] + + +def test_multimodal_content_parts_are_flattened_to_their_text(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe"}, + {"type": "image_url", "image_url": {"url": "data:image/png;..."}}, + {"type": "text", "text": "this image"}, + ], + } + ] + normalized = _normalize_chat_messages(messages) + assert normalized == [{"role": "user", "content": "describe this image"}] + # The request body itself is never mutated. + assert isinstance(messages[0]["content"], list) + + +def test_null_and_empty_text_parts_are_dropped(): + # {"type": "text", "text": null} is valid JSON a client can send; .get + # with a default only covers a MISSING key, so an explicit null must not + # reach " ".join as None - and null/empty parts are dropped entirely so + # they cannot inject stray spaces into the normalized text. + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": None}, + {"type": "text", "text": ""}, + {"type": "text", "text": "hello"}, + ], + } + ] + normalized = _normalize_chat_messages(messages) + assert normalized == [{"role": "user", "content": "hello"}] + + +def test_none_content_becomes_an_empty_string(): + messages = [{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]}] + normalized = _normalize_chat_messages(messages) + assert normalized[0]["content"] == "" + assert normalized[0]["tool_calls"] == [{"id": "1"}] + + +def test_string_content_messages_pass_through_untouched(): + assert _normalize_chat_messages(MESSAGES) == MESSAGES + + +# --- the remote /tokenize fallback payload ------------------------------------ + + +def test_chat_bodies_use_the_tokenize_chat_request_form(): + payload = _tokenize_request_payload(MODEL, {"messages": MESSAGES}) + assert payload == { + "model": MODEL, + "messages": MESSAGES, + "add_generation_prompt": True, + } + + +def test_prompt_bodies_keep_the_completion_form(): + assert _tokenize_request_payload(MODEL, {"prompt": "hello"}) == { + "model": MODEL, + "prompt": "hello", + } + + +# --- through the routers ------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_kvaware_routes_chat_requests_via_the_kv_lookup_path(): + """The regression this file exists for: a messages-form body must reach + the controller as non-empty, template-aligned token ids and route to the + KV holder - not tokenize as "" and fall back to session/QPS.""" + router = KvawareRouter.__new__(KvawareRouter) + router.tokenizer = ChatTokenizer() + router.threshold = 2000 + router.instance_id_to_ip = {INST_A: URL_A} + router.session_key = None + router.hash_ring = HashRing() + seen = {} + + async def query_manager(msg): + seen["tokens"] = msg.tokens + return LookupRet({INST_A: (LOCAL, len(CHAT_IDS))}) + + router.query_manager = query_manager + url = await router.route_request( + endpoints(URL_A, URL_B), {}, {}, None, {"messages": MESSAGES} + ) + assert seen["tokens"] == CHAT_IDS + assert url == URL_A + + +@pytest.mark.asyncio +async def test_loadaware_tokenizes_chat_requests_through_the_template(): + router = LoadAwareRouter.__new__(LoadAwareRouter) + router.tokenizer = ChatTokenizer() + ids = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids == CHAT_IDS + + +@pytest.mark.asyncio +async def test_prompt_requests_are_unchanged_by_the_chat_support(): + router = LoadAwareRouter.__new__(LoadAwareRouter) + tokenizer = ChatTokenizer() + router.tokenizer = tokenizer + ids = await router.tokenize_prompt(endpoints(URL_A), {"prompt": "hello"}) + assert ids == [1] * len("hello") + assert tokenizer.chat_template_calls == [] + + +@pytest.mark.asyncio +async def test_remote_tokenize_fallback_sends_the_messages_for_chat(monkeypatch): + """A tokenizer without a chat template falls back to the engine's + /tokenize with the original messages (vLLM's TokenizeChatRequest), not + {"prompt": ""}.""" + router = LoadAwareRouter.__new__(LoadAwareRouter) + router.tokenizer = TemplatelessTokenizer() + captured = {} + + class Response: + @staticmethod + def raise_for_status(): + pass + + @staticmethod + def json(): + return {"count": len(CHAT_IDS), "tokens": CHAT_IDS} + + def fake_post(url, headers=None, json=None, timeout=None): + captured["url"] = url + captured["json"] = json + return Response() + + monkeypatch.setattr(routing_logic.requests, "post", fake_post) + ids = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids == CHAT_IDS + assert captured["url"] == URL_A + "/tokenize" + assert captured["json"]["messages"] == MESSAGES + assert captured["json"]["add_generation_prompt"] is True + assert "prompt" not in captured["json"] + + +# --- tokenization failure degrades to fallback routing, never a 500 ----------- + + +@pytest.mark.asyncio +async def test_kvaware_falls_back_to_qps_when_tokenization_fails(monkeypatch): + """A dead /tokenize endpoint (plus no usable local template) must not + fail the request - the router still has session/QPS routing.""" + router = KvawareRouter.__new__(KvawareRouter) + router.tokenizer = TemplatelessTokenizer() + router.threshold = 2000 + router.instance_id_to_ip = {} + router.session_key = None + router.hash_ring = HashRing() + lookups = [] + + async def query_manager(msg): + lookups.append(msg) + + router.query_manager = query_manager + + def dead_post(url, headers=None, json=None, timeout=None): + raise ConnectionError("engine unreachable") + + monkeypatch.setattr(routing_logic.requests, "post", dead_post) + url = await router.route_request( + endpoints(URL_A, URL_B), {}, {}, None, {"messages": MESSAGES} + ) + assert url == URL_A # QPS routing with no stats picks the first endpoint + assert lookups == [] # no KV lookup without token ids + + +@pytest.mark.asyncio +async def test_loadaware_tokenize_returns_none_when_both_paths_fail(monkeypatch): + router = LoadAwareRouter.__new__(LoadAwareRouter) + router.tokenizer = TemplatelessTokenizer() + + def dead_post(url, headers=None, json=None, timeout=None): + raise ConnectionError("engine unreachable") + + monkeypatch.setattr(routing_logic.requests, "post", dead_post) + ids = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids is None + + +@pytest.mark.asyncio +async def test_failed_tokenizer_load_is_not_retried_per_request(monkeypatch): + """A served-model alias (vLLM --served-model-name) is not a hub id, so + the local load always fails - the failure must be cached, not re-paid as + a hub lookup on every request before the remote /tokenize fallback.""" + router = LoadAwareRouter.__new__(LoadAwareRouter) + router.tokenizer = None + load_attempts = [] + + class FailingAuto: + @staticmethod + def from_pretrained(name): + load_attempts.append(name) + raise OSError("not a local folder and not a valid repo id") + + # transformers may be absent in the test env (the module import is + # guarded), so set the module attribute itself with raising=False + monkeypatch.setattr(routing_logic, "AutoTokenizer", FailingAuto, raising=False) + + class Response: + @staticmethod + def raise_for_status(): + pass + + @staticmethod + def json(): + return {"count": len(CHAT_IDS), "tokens": CHAT_IDS} + + monkeypatch.setattr(routing_logic.requests, "post", lambda *a, **k: Response()) + for _ in range(3): + ids = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids == CHAT_IDS + assert load_attempts == [MODEL] # exactly one hub attempt, not three + + +@pytest.mark.asyncio +async def test_cold_tokenizer_load_receives_the_model_name(monkeypatch): + """_ensure_tokenizer must resolve endpoints -> model name at load time; + a successful cold load caches the tokenizer for subsequent requests.""" + router = LoadAwareRouter.__new__(LoadAwareRouter) + router.tokenizer = None + loaded = [] + + class FakeAuto: + @staticmethod + def from_pretrained(name): + loaded.append(name) + return ChatTokenizer() + + monkeypatch.setattr(routing_logic, "AutoTokenizer", FakeAuto, raising=False) + ids = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids == CHAT_IDS + assert loaded == [MODEL] # the NAME string, not the endpoint list + ids2 = await router.tokenize_prompt(endpoints(URL_A), {"messages": MESSAGES}) + assert ids2 == CHAT_IDS + assert loaded == [MODEL] # cached - no second load diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py index b25e9184a..fa67a5d53 100644 --- a/src/vllm_router/routers/routing_logic.py +++ b/src/vllm_router/routers/routing_logic.py @@ -90,6 +90,123 @@ def _loadaware_beta(override: Optional[float]) -> float: return DEFAULT_LOADAWARE_BETA +def _normalize_chat_messages(messages: List[Dict]) -> List[Dict]: + """Text-only view of an OpenAI chat ``messages`` array for local + chat-template application. + + Multimodal content parts (a list of ``{"type": ...}`` dicts) are + flattened to their text parts - mirroring ``PrefixAwareRouter`` - because + plain HF chat templates expect string content. ``None`` content (e.g. + assistant tool-call turns) becomes ``""``. Messages that already carry + string content pass through untouched, so template-relevant fields + (``role``, ``name``, ``tool_calls``, ...) are preserved. The input is + never mutated. + """ + normalized = [] + for message in messages: + content = message.get("content") + if isinstance(content, list): + # Join only non-empty string parts: empty/null parts would inject + # stray spaces into the normalized text, drifting it away from + # what a text-only chat template would render. + text_content = " ".join( + part["text"] + for part in content + if isinstance(part, dict) + and part.get("type") == "text" + and isinstance(part.get("text"), str) + and part["text"] + ) + message = {**message, "content": text_content} + elif content is None: + message = {**message, "content": ""} + normalized.append(message) + return normalized + + +def _extract_token_ids(tokenizer, request_json: Dict) -> List[int]: + """Token ids as the serving engine would see them for this request body. + + Chat-completion bodies (``messages``) are tokenized through the model's + chat template with ``add_generation_prompt=True``: vLLM engines cache KV + for the token ids *after* template application, so encoding the raw + message text would never line up with the engine-side prefix. The + template is rendered to TEXT and then encoded (``add_special_tokens= + False`` - the rendered template already carries its special tokens): + ``apply_chat_template(..., tokenize=True)``'s return type varies across + transformers versions (plain ids vs ``Encoding`` objects), and feeding + the non-id form to the KV lookup silently never matches. Completion + bodies keep the plain ``encode`` path, unchanged. May raise (e.g. the + tokenizer defines no chat template) - callers fall back to the engine's + ``/tokenize`` API. + """ + if "messages" in request_json: + text = tokenizer.apply_chat_template( + _normalize_chat_messages(request_json["messages"]), + add_generation_prompt=True, + tokenize=False, + ) + return tokenizer.encode(text, add_special_tokens=False) + return tokenizer.encode(request_json.get("prompt", "")) + + +async def _ensure_tokenizer(router, endpoints: List[EndpointInfo]): + """Load the router's tokenizer once and return it. + + Double-checked lock: concurrent cold-start requests would otherwise each + load the tokenizer (benign but redundant - duplicated disk/CPU work and + possible hub rate-limiting). The lock is created lazily and there is no + await between the check and the assignment, so its creation cannot race + on one event loop. ``from_pretrained`` is blocking I/O (possibly a hub + download on first use), so it runs in an executor. + + The model name is resolved from the endpoint list only when a load + actually happens, and a failed load is remembered per name: a served-model + alias (e.g. a vLLM ``--served-model-name``) is not a real tokenizer id and + can never load, so without the negative cache every request would pay a + doomed hub lookup before reaching the remote ``/tokenize`` fallback. + """ + if router.tokenizer is None: + model_name = endpoints[0].model_names[0] + if model_name in getattr(router, "_tokenizer_load_failures", ()): + raise ValueError( + f"tokenizer load for '{model_name}' already failed; " + f"using the remote /tokenize fallback" + ) + if not hasattr(router, "_tokenizer_lock"): + router._tokenizer_lock = asyncio.Lock() + async with router._tokenizer_lock: + if router.tokenizer is None: + loop = asyncio.get_running_loop() + try: + router.tokenizer = await loop.run_in_executor( + None, lambda: AutoTokenizer.from_pretrained(model_name) + ) + except Exception: + if not hasattr(router, "_tokenizer_load_failures"): + router._tokenizer_load_failures = set() + router._tokenizer_load_failures.add(model_name) + raise + return router.tokenizer + + +def _tokenize_request_payload(model: str, request_json: Dict) -> Dict: + """Request body for the engine's ``/tokenize`` fallback. + + Chat bodies map to vLLM's ``TokenizeChatRequest`` (the engine applies its + own chat template, multimodal content included), completion bodies to + ``TokenizeCompletionRequest``. ``add_generation_prompt=True`` is vLLM's + chat-completions default; sent explicitly to pin the alignment. + """ + if "messages" in request_json: + return { + "model": model, + "messages": request_json["messages"], + "add_generation_prompt": True, + } + return {"model": model, "prompt": request_json.get("prompt", "")} + + class RoutingInterface(metaclass=SingletonABCMeta): def _qps_routing( self, endpoints: List[EndpointInfo], request_stats: Dict[str, RequestStats] @@ -387,36 +504,49 @@ async def route_request( """ token_ids = None # Local-first tokenization, fall back to remote "/tokenize" API on failure - # TODO (Yuhan): Handle chat completions try: - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained( - endpoints[0].model_names[0] - ) - token_ids = self.tokenizer.encode(request_json.get("prompt", "")) + await _ensure_tokenizer(self, endpoints) + token_ids = _extract_token_ids(self.tokenizer, request_json) 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"] + # Remote /tokenize fallback. requests is synchronous - run it in + # an executor so the fallback does not block the router's event + # loop. A failure here (engine timeout, connection error, non-2xx) + # must not fail the request: the session/QPS fallback below routes + # fine without token ids. + try: + remote_url = endpoints[0].url + "/tokenize" + headers = {"Content-Type": "application/json"} + data = _tokenize_request_payload( + endpoints[0].model_names[0], request_json + ) + loop = asyncio.get_running_loop() + response = await loop.run_in_executor( + None, + lambda: requests.post( + remote_url, headers=headers, json=data, timeout=10 + ), + ) + response.raise_for_status() + token_ids = response.json()["tokens"] + except Exception as e: + logger.warning( + f"Tokenization failed locally and via remote /tokenize " + f"({e}); falling back to session/QPS routing" + ) + token_ids = None - event_id = "Lookup" + str(uuid.uuid4()) - msg = LookupMsg(tokens=token_ids, event_id=event_id) - instance_id = await self.query_manager(msg) + instance_id = None matched_tokens = math.inf - logger.debug(f"Lookup return message: {instance_id}") - if len(list(instance_id.layout_info.keys())) > 0: - matched_instance_id = list(instance_id.layout_info.keys())[ - 0 - ] # Get the first key - matched_tokens = instance_id.layout_info[matched_instance_id][1] + if token_ids is not None: + event_id = "Lookup" + str(uuid.uuid4()) + msg = LookupMsg(tokens=token_ids, event_id=event_id) + instance_id = await self.query_manager(msg) + logger.debug(f"Lookup return message: {instance_id}") + if len(list(instance_id.layout_info.keys())) > 0: + matched_instance_id = list(instance_id.layout_info.keys())[ + 0 + ] # Get the first key + matched_tokens = instance_id.layout_info[matched_instance_id][1] if ( instance_id is None @@ -646,33 +776,41 @@ async def query_endpoint(endpoint: EndpointInfo) -> None: async def tokenize_prompt( self, endpoints: List[EndpointInfo], request_json: Dict - ) -> List[int]: + ) -> Optional[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. + Chat-completion bodies go through the chat template so the token ids + match what the engine caches KV for (see `_extract_token_ids`). The + remote fallback is a blocking HTTP call, so it runs in an executor + rather than on the event loop. Returns ``None`` when both paths fail + (the caller then routes via `fallback_url` instead of erroring the + request). """ try: - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained( - endpoints[0].model_names[0] - ) - return self.tokenizer.encode(request_json.get("prompt", "")) + await _ensure_tokenizer(self, endpoints) + return _extract_token_ids(self.tokenizer, request_json) 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"] + try: + remote_url = endpoints[0].url + "/tokenize" + headers = {"Content-Type": "application/json"} + data = _tokenize_request_payload( + endpoints[0].model_names[0], request_json + ) + loop = asyncio.get_running_loop() + response = await loop.run_in_executor( + None, + lambda: requests.post( + remote_url, headers=headers, json=data, timeout=10 + ), + ) + response.raise_for_status() + return response.json()["tokens"] + except Exception as e: + logger.warning( + f"Tokenization failed locally and via remote /tokenize " + f"({e}); falling back to session/QPS routing" + ) + return None def fallback_url( self, @@ -722,6 +860,8 @@ async def route_request( ) token_ids = await self.tokenize_prompt(endpoints, request_json) + if token_ids is None: + return self.fallback_url(endpoints, request_stats, request, request_json) event_id = "Lookup" + str(uuid.uuid4()) msg = LookupMsg(tokens=token_ids, event_id=event_id)