[Router][Bugfix] KV-aware routing: tokenize chat-completions bodies through the chat template - #1045
[Router][Bugfix] KV-aware routing: tokenize chat-completions bodies through the chat template#1045tyler2cr wants to merge 9 commits into
Conversation
…hrough the chat template KvawareRouter and LoadAwareRouter tokenized only request_json['prompt'], so chat-completions bodies (a 'messages' array, no 'prompt' key) tokenized as the empty string: the LMCache lookup matched nothing and every chat request silently fell back to session/hash-ring routing. Resolves the 'TODO (Yuhan): Handle chat completions' in routing_logic.py. Chat bodies now tokenize through the model's chat template (add_generation_prompt=True), matching the token ids the engine actually caches KV for; the remote /tokenize fallback sends vLLM's TokenizeChatRequest form for the same reason. Prompt-form requests are byte-identical to before. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: tyler <tcr@enfuse.io>
![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 introduces support for chat-completion tokenization in the KV-aware and load-aware routers by applying the model's chat template to the messages array, preventing silent fallback degradation. Feedback highlights two main improvements: first, running the synchronous requests.post call in route_request within an executor to avoid blocking the event loop, and second, using part.get('text') or '' in _normalize_chat_messages to prevent a potential TypeError if a text part is explicitly set to null.
| data = _tokenize_request_payload(endpoints[0].model_names[0], request_json) | ||
| body = requests.post( | ||
| remote_url, headers=headers, json=data, timeout=10 | ||
| ).json() |
There was a problem hiding this comment.
The synchronous requests.post call inside the async def route_request method blocks the event loop during the remote /tokenize fallback. This can severely degrade router throughput and latency under load.
To prevent blocking the event loop, run the synchronous HTTP request in an executor using asyncio.get_running_loop().run_in_executor, similar to how it is done in LoadAwareRouter.tokenize_prompt.
| data = _tokenize_request_payload(endpoints[0].model_names[0], request_json) | |
| body = requests.post( | |
| remote_url, headers=headers, json=data, timeout=10 | |
| ).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 | |
| ), | |
| ) | |
| body = response.json() |
There was a problem hiding this comment.
Addressed in 2574697 — the /tokenize fallback now runs via loop.run_in_executor, matching the pattern in LoadAwareRouter.tokenize_prompt.
| if isinstance(content, list): | ||
| text_content = " ".join( | ||
| part.get("text", "") | ||
| for part in content | ||
| if isinstance(part, dict) and part.get("type") == "text" | ||
| ) | ||
| message = {**message, "content": text_content} |
There was a problem hiding this comment.
If a multimodal message contains a part where "text" is explicitly set to null (e.g., {"type": "text", "text": null}), part.get("text", "") will return None instead of "". This will cause " ".join(...) to raise a TypeError: sequence item: expected str instance, NoneType found.
Using part.get("text") or "" safely defaults to an empty string if the value is None or missing.
| if isinstance(content, list): | |
| text_content = " ".join( | |
| part.get("text", "") | |
| for part in content | |
| if isinstance(part, dict) and part.get("type") == "text" | |
| ) | |
| message = {**message, "content": text_content} | |
| if isinstance(content, list): | |
| text_content = " ".join( | |
| part.get("text") or "" | |
| for part in content | |
| if isinstance(part, dict) and part.get("type") == "text" | |
| ) | |
| message = {**message, "content": text_content} |
There was a problem hiding this comment.
Addressed in 2574697 — switched to part.get("text") or "" and added a regression test for an explicit {"type": "text", "text": null} part.
|
Parking this while we validate end-to-end on real hardware (2x vLLM engines + LMCache controller + this router build, kvaware mode, live chat-completions traffic) rather than unit tests alone. Will reopen with the e2e results. Branch stays available at tyler2cr/production-stack:router-kvaware-chat-completions if anyone wants to pick it up meanwhile. |
…erate null text parts
- run the synchronous requests.post fallback in an executor so it cannot
block the router event loop (matches LoadAwareRouter.tokenize_prompt)
- an explicit {"type": "text", "text": null} part no longer raises
TypeError in message normalization (.get default only covers a missing
key); regression test added
Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds chat-completion tokenization support to the KV-aware and load-aware routers, resolving an issue where chat requests would silently degrade to session/QPS fallback. It introduces helper functions to normalize chat messages, extract token IDs, and construct the payload for remote tokenization, and runs the synchronous fallback request in an executor to prevent blocking the event loop. Additionally, a comprehensive test suite is added. The reviewer suggests improving robustness by catching exceptions during the remote tokenization fallback to gracefully fall back to session/QPS routing instead of failing the entire request.
| except Exception: | ||
| # Remote /tokenize fallback (let errors bubble up to keep behavior simple) | ||
| # Remote /tokenize fallback (let errors bubble up to keep behavior | ||
| # simple). requests is synchronous - run it in an executor so the | ||
| # fallback does not block the router's event loop. | ||
| 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"] | ||
| 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 | ||
| ), | ||
| ) | ||
| token_ids = response.json()["tokens"] |
There was a problem hiding this comment.
If the remote /tokenize fallback fails (e.g., due to a network timeout, connection error, or a non-200 response from the engine), the exception will bubble up and cause the entire request to fail with a 500 Internal Server Error. Since the router can still successfully route requests using the fallback session/QPS routing logic even without token IDs, we should catch any exceptions during the remote tokenization fallback and gracefully fall back to session/QPS routing instead of failing the request.
A similar robustness improvement should also be applied to LoadAwareRouter.tokenize_prompt.
except Exception as local_err:
# Remote /tokenize fallback. requests is synchronous - run it in an executor so the
# fallback does not block the router's event loop.
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
),
)
token_ids = response.json()["tokens"]
except Exception as remote_err:
logger.error(
f"Tokenization failed completely (local: {local_err}, remote: {remote_err}). "
"Falling back to session/QPS routing."
)
session_id = self.extract_session_id(request, request_json)
logger.debug(f"Fallback to using session id: {session_id}")
self._update_hash_ring(endpoints)
if session_id is None:
return self._qps_routing(endpoints, request_stats)
return self.hash_ring.get_node(session_id)There was a problem hiding this comment.
Addressed in ec2becc — the remote /tokenize fallback (now also checked with raise_for_status) catches failures, logs a warning, and routes via the existing session/QPS fallback instead of failing the request. Applied to both KvawareRouter.route_request and LoadAwareRouter.tokenize_prompt (returns None on failure; its caller falls back), with regression tests for both paths.
… instead of failing the request If the remote /tokenize fallback fails (timeout, connection error, non-2xx - now checked via raise_for_status), the kvaware and loadaware routers log a warning and route via their existing session/QPS fallback rather than surfacing a 500. Applied to KvawareRouter.route_request and LoadAwareRouter.tokenize_prompt (which now returns None on failure; its caller falls back). Regression tests added for both paths. Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for chat-completion tokenization in the KV-aware and load-aware routers, ensuring that chat requests are tokenized through the model's chat template so that token IDs align with the engine-side KV cache. It also updates the remote /tokenize fallback to run asynchronously in an executor to prevent blocking the event loop, and gracefully falls back to session/QPS routing on failure. The review feedback correctly points out that loading the tokenizer via AutoTokenizer.from_pretrained is a synchronous, blocking I/O operation that should also be run in an executor to keep the event loop responsive.
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ) |
There was a problem hiding this comment.
Loading the tokenizer using AutoTokenizer.from_pretrained is a synchronous, blocking I/O operation (which may also involve network requests if the model is not cached locally). Calling it directly in an async def function blocks the event loop, preventing the router from processing other concurrent requests. It should be run in an executor to keep the event loop responsive.
| if self.tokenizer is None: | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| endpoints[0].model_names[0] | |
| ) | |
| if self.tokenizer is None: | |
| loop = asyncio.get_running_loop() | |
| self.tokenizer = await loop.run_in_executor( | |
| None, | |
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]) | |
| ) |
There was a problem hiding this comment.
Addressed in 00fcf4d — the tokenizer load now runs via loop.run_in_executor in both routers.
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ) |
There was a problem hiding this comment.
Loading the tokenizer using AutoTokenizer.from_pretrained is a synchronous, blocking I/O operation. Calling it directly in an async def function blocks the event loop, preventing the router from processing other concurrent requests. It should be run in an executor to keep the event loop responsive.
| if self.tokenizer is None: | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| endpoints[0].model_names[0] | |
| ) | |
| if self.tokenizer is None: | |
| loop = asyncio.get_running_loop() | |
| self.tokenizer = await loop.run_in_executor( | |
| None, | |
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]) | |
| ) |
There was a problem hiding this comment.
Addressed in 00fcf4d — the tokenizer load now runs via loop.run_in_executor in both routers.
…is blocking I/O AutoTokenizer.from_pretrained can hit the network (hub download on first use) and blocks the event loop when called inside the async routing path; run it in an executor in both KvawareRouter.route_request and LoadAwareRouter.tokenize_prompt. Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces chat-completion tokenization support for the KV-aware and Load-aware routers, allowing them to correctly route requests based on token-id prefixes after applying chat templates. It also offloads blocking I/O operations (like loading tokenizers and calling remote fallback APIs) to an executor to prevent blocking the asyncio event loop. The feedback highlights a potential race condition during concurrent tokenizer initialization on startup, which can be resolved using an asyncio.Lock. Additionally, it is recommended to refine the chat message normalization logic to filter out empty or null text parts before joining them, preventing unexpected spaces that could alter tokenization results.
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| # from_pretrained is blocking I/O (possibly a hub download on | ||
| # first use) - keep it off the event loop too. | ||
| loop = asyncio.get_running_loop() | ||
| self.tokenizer = await loop.run_in_executor( | ||
| None, | ||
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]), | ||
| ) |
There was a problem hiding this comment.
Under high concurrent load on startup, multiple requests arriving simultaneously when self.tokenizer is None will all trigger AutoTokenizer.from_pretrained concurrently in separate threads. This race condition can cause CPU/memory spikes, disk I/O bottlenecks, or rate-limiting from Hugging Face Hub.
Using an asyncio.Lock with a double-checked locking pattern ensures that only the first request loads the tokenizer, while subsequent concurrent requests wait and reuse the loaded tokenizer.
| if self.tokenizer is None: | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| endpoints[0].model_names[0] | |
| # from_pretrained is blocking I/O (possibly a hub download on | |
| # first use) - keep it off the event loop too. | |
| loop = asyncio.get_running_loop() | |
| self.tokenizer = await loop.run_in_executor( | |
| None, | |
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]), | |
| ) | |
| if self.tokenizer is None: | |
| if not hasattr(self, "_tokenizer_lock"): | |
| self._tokenizer_lock = asyncio.Lock() | |
| async with self._tokenizer_lock: | |
| if self.tokenizer is None: | |
| # from_pretrained is blocking I/O (possibly a hub download on | |
| # first use) - keep it off the event loop too. | |
| loop = asyncio.get_running_loop() | |
| self.tokenizer = await loop.run_in_executor( | |
| None, | |
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]), | |
| ) |
There was a problem hiding this comment.
Applied in 113973c (double-checked asyncio.Lock, lazily created; no await between the check and the assignment, so its own creation can't race on one event loop).
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| # from_pretrained is blocking I/O (possibly a hub download on | ||
| # first use) - keep it off the event loop too. | ||
| loop = asyncio.get_running_loop() | ||
| self.tokenizer = await loop.run_in_executor( | ||
| None, | ||
| lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]), | ||
| ) |
There was a problem hiding this comment.
Similar to KvawareRouter.route_request, LoadAwareRouter.tokenize_prompt suffers from a race condition where concurrent requests on startup can trigger multiple redundant AutoTokenizer.from_pretrained calls.
Applying an asyncio.Lock with a double-checked locking pattern ensures thread-safe, single-execution initialization of the tokenizer.
if self.tokenizer is None:
if not hasattr(self, "_tokenizer_lock"):
self._tokenizer_lock = asyncio.Lock()
async with self._tokenizer_lock:
if self.tokenizer is None:
# from_pretrained is blocking I/O (possibly a hub download on
# first use) - keep it off the event loop too.
loop = asyncio.get_running_loop()
self.tokenizer = await loop.run_in_executor(
None,
lambda: AutoTokenizer.from_pretrained(endpoints[0].model_names[0]),
)There was a problem hiding this comment.
Applied in 113973c — same single-flight pattern as in KvawareRouter.route_request.
| if isinstance(content, list): | ||
| text_content = " ".join( | ||
| part.get("text") or "" | ||
| for part in content | ||
| if isinstance(part, dict) and part.get("type") == "text" | ||
| ) | ||
| message = {**message, "content": text_content} |
There was a problem hiding this comment.
Joining empty or null text parts with " " can introduce unexpected leading, trailing, or double spaces (e.g., " hello" or "hello world"). In many tokenizers, leading or multiple spaces significantly alter the generated token IDs, which can lead to cache misses or routing mismatches.
Filtering out non-string, null, or empty text parts completely before joining ensures a clean, space-aligned text content.
| if isinstance(content, list): | |
| text_content = " ".join( | |
| part.get("text") or "" | |
| for part in content | |
| if isinstance(part, dict) and part.get("type") == "text" | |
| ) | |
| message = {**message, "content": text_content} | |
| if isinstance(content, list): | |
| 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} |
There was a problem hiding this comment.
Applied in 113973c, with one correction to the rationale: stray spaces here don't cause cache misses or routing mismatches — normalization is deterministic per message content, so the router's hashes stay self-consistent across requests, and for multimodal messages the router-side token ids can never exactly match the engine's anyway (image parts become vision tokens engine-side; this normalization is a best-effort prefix approximation for routing). Applying it regardless because filtering null/empty parts yields cleaner text, closer to what a text-only chat template would render, at zero cost.
| def test_null_text_part_becomes_an_empty_string(): | ||
| # {"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. | ||
| messages = [ | ||
| { | ||
| "role": "user", | ||
| "content": [ | ||
| {"type": "text", "text": None}, | ||
| {"type": "text", "text": "hello"}, | ||
| ], | ||
| } | ||
| ] | ||
| normalized = _normalize_chat_messages(messages) | ||
| assert normalized == [{"role": "user", "content": " hello"}] |
There was a problem hiding this comment.
Updating the test assertion to expect "hello" instead of " hello", aligning with the robust normalization logic that filters out null or empty text parts to prevent tokenization discrepancies.
| def test_null_text_part_becomes_an_empty_string(): | |
| # {"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. | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": None}, | |
| {"type": "text", "text": "hello"}, | |
| ], | |
| } | |
| ] | |
| normalized = _normalize_chat_messages(messages) | |
| assert normalized == [{"role": "user", "content": " hello"}] | |
| def test_null_text_part_becomes_an_empty_string(): | |
| # {"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. | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "text", "text": None}, | |
| {"type": "text", "text": "hello"}, | |
| ], | |
| } | |
| ] | |
| normalized = _normalize_chat_messages(messages) | |
| assert normalized == [{"role": "user", "content": "hello"}] |
There was a problem hiding this comment.
Applied in 113973c together with the normalization change — the test now also covers an explicit empty-string part and asserts "hello" with no stray space.
…t parts in normalization - double-checked asyncio.Lock around AutoTokenizer.from_pretrained in both routers: concurrent cold-start requests no longer each load the tokenizer (benign race, but duplicated work and possible hub rate-limiting) - message normalization joins only non-empty string text parts, so null or empty parts cannot inject stray spaces; test updated accordingly Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces chat-completion tokenization support for KV-aware and load-aware routers, including chat message normalization and a fallback mechanism to the remote /tokenize API. It also optimizes tokenizer loading by using a double-checked lock and offloading blocking operations to an executor. A comprehensive suite of unit tests has been added to verify these changes. The review feedback suggests extracting the duplicated tokenizer loading and locking logic into a shared helper function _ensure_tokenizer to simplify the code and improve maintainability.
| 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", "")} |
There was a problem hiding this comment.
To avoid duplicating the tokenizer loading and double-checked locking logic between KvawareRouter.route_request and LoadAwareRouter.tokenize_prompt, we can extract it into a shared helper function _ensure_tokenizer at the module level.
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", "")}
async def _ensure_tokenizer(router, model_name: str) -> None:
"""Ensure the tokenizer is loaded on the router instance, using a double-checked lock."""
if router.tokenizer is None:
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()
router.tokenizer = await loop.run_in_executor(
None,
lambda: AutoTokenizer.from_pretrained(model_name),
)There was a problem hiding this comment.
Applied in a7fd571, with one signature change: the helper takes endpoints and resolves model_names[0] lazily inside the load branch. The suggested call form — _ensure_tokenizer(self, endpoints[0].model_names[0]) — evaluates model_names eagerly even when the tokenizer is already loaded, which broke an existing loadaware test whose endpoint stub doesn't answer model_names (the pre-refactor code only touched it when actually loading).
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ) | ||
| token_ids = self.tokenizer.encode(request_json.get("prompt", "")) | ||
| # 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; there is no await between the | ||
| # hasattr check and the assignment, so coroutines on one | ||
| # event loop cannot race it. | ||
| if not hasattr(self, "_tokenizer_lock"): | ||
| self._tokenizer_lock = asyncio.Lock() | ||
| async with self._tokenizer_lock: | ||
| if self.tokenizer is None: | ||
| # from_pretrained is blocking I/O (possibly a hub | ||
| # download on first use) - keep it off the event loop. | ||
| loop = asyncio.get_running_loop() | ||
| self.tokenizer = await loop.run_in_executor( | ||
| None, | ||
| lambda: AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ), | ||
| ) | ||
| token_ids = _extract_token_ids(self.tokenizer, request_json) |
There was a problem hiding this comment.
We can simplify this block by using the newly extracted _ensure_tokenizer helper function, which eliminates duplicate code and improves readability.
await _ensure_tokenizer(self, endpoints[0].model_names[0])
token_ids = _extract_token_ids(self.tokenizer, request_json)There was a problem hiding this comment.
Applied in a7fd571 (call form is await _ensure_tokenizer(self, endpoints) — see the thread on the helper for why the model name resolves lazily).
| if self.tokenizer is None: | ||
| self.tokenizer = AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ) | ||
| return self.tokenizer.encode(request_json.get("prompt", "")) | ||
| # 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; there is no await between the | ||
| # hasattr check and the assignment, so coroutines on one | ||
| # event loop cannot race it. | ||
| if not hasattr(self, "_tokenizer_lock"): | ||
| self._tokenizer_lock = asyncio.Lock() | ||
| async with self._tokenizer_lock: | ||
| if self.tokenizer is None: | ||
| # from_pretrained is blocking I/O (possibly a hub | ||
| # download on first use) - keep it off the event loop. | ||
| loop = asyncio.get_running_loop() | ||
| self.tokenizer = await loop.run_in_executor( | ||
| None, | ||
| lambda: AutoTokenizer.from_pretrained( | ||
| endpoints[0].model_names[0] | ||
| ), | ||
| ) | ||
| return _extract_token_ids(self.tokenizer, request_json) |
There was a problem hiding this comment.
We can simplify this block by using the newly extracted _ensure_tokenizer helper function, which eliminates duplicate code and improves readability.
await _ensure_tokenizer(self, endpoints[0].model_names[0])
return _extract_token_ids(self.tokenizer, request_json)Dedupes the tokenizer-load-and-lock logic between KvawareRouter and LoadAwareRouter. Deviates from the suggested signature in one respect: the helper takes the endpoint list and resolves the model name lazily inside the load branch - the suggested _ensure_tokenizer(self, endpoints[0].model_names[0]) evaluates model_names eagerly even when the tokenizer is already loaded, which broke an existing loadaware test whose endpoint stub answers no model_names. Signed-off-by: tyler <tcr@enfuse.io>
…dy from a7fd571 Deployments that serve under a vLLM --served-model-name (not a hub id) can never load the tokenizer locally - without a negative cache every request pays a doomed hub lookup before reaching the remote /tokenize fallback. A failed load is now remembered per model name and later requests go straight to the remote path. Also fixes a bug shipped in a7fd571: the helper's body still bound the old model_name parameter while callers passed the endpoint list, so a genuine cold load would call from_pretrained(endpoints) and always fail into the remote path. Unit tests never exercised a successful cold load through the helper - two tests added (successful cold load receives the model NAME and is cached; failed load is attempted exactly once, not per request). Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for chat-completion tokenization in the KV-aware and load-aware routers. It introduces helper functions to normalize chat messages, extract token IDs using the model's chat template, and handle remote /tokenize fallbacks gracefully without blocking the event loop. Comprehensive unit tests are also added to verify these behaviors. Feedback is provided to add a check for empty endpoints in KvawareRouter.route_request to prevent a potential IndexError when no backends are available.
| token_ids = None | ||
| # Local-first tokenization, fall back to remote "/tokenize" API on failure | ||
| # TODO (Yuhan): Handle chat completions | ||
| try: |
There was a problem hiding this comment.
To prevent a potential IndexError when endpoints is empty (e.g., during a temporary service discovery lag or when all backends are unhealthy), we should check if endpoints is empty at the beginning of route_request, similar to the check in LoadAwareRouter.route_request.
| token_ids = None | |
| # Local-first tokenization, fall back to remote "/tokenize" API on failure | |
| # TODO (Yuhan): Handle chat completions | |
| try: | |
| if not endpoints: | |
| raise HTTPException( | |
| status_code=503, detail="No backend endpoints available" | |
| ) | |
| token_ids = None | |
| # Local-first tokenization, fall back to remote "/tokenize" API on failure | |
| try: |
There was a problem hiding this comment.
Valid catch, but pre-existing behavior rather than something this PR introduces — the unpatched KvawareRouter.route_request performs the same unguarded endpoints[0] access, so the IndexError exists on main today. Keeping this PR scoped to the chat-tokenization fix; happy to bring the empty-endpoints guard (mirroring LoadAwareRouter's 503) as a small follow-up PR if maintainers want it.
…enize=True returns Encoding objects on current transformers Found by fresh end-to-end validation with a hub-known model (the served-name deployments validated earlier always took the remote /tokenize path, so the local chat path was never exercised on real transformers): on transformers 5.x, apply_chat_template(..., tokenize=True) returns Encoding objects rather than a flat id list, so the KV lookup received garbage and kvaware silently fell back to QPS routing for chat traffic - the exact failure this patch exists to fix. Rendering the template to text and encoding it with add_special_tokens=False (the rendered text already carries the special tokens) produces ids byte-identical to the engine's /tokenize output (verified against vLLM v0.22.0 + transformers 5.9.0). Test fakes now assert the router never uses the tokenize=True form and never re-adds special tokens. Signed-off-by: tyler <tcr@enfuse.io>
…red template Signed-off-by: tyler <tcr@enfuse.io>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for chat-completion tokenization in KV-aware and load-aware routers, including local chat-template application, multimodal content normalization, and asynchronous remote /tokenize fallbacks, along with a comprehensive test suite. The review feedback suggests adding defensive checks to prevent potential IndexError exceptions when resolving the tokenizer, and refactoring the duplicated tokenization and fallback logic in KvawareRouter and LoadAwareRouter into a shared helper function.
| if router.tokenizer is None: | ||
| model_name = endpoints[0].model_names[0] |
There was a problem hiding this comment.
Add a defensive check to ensure endpoints is not empty and contains at least one model name before accessing endpoints[0].model_names[0]. This prevents potential IndexError exceptions if the router is invoked with an empty endpoint list.
if router.tokenizer is None:
if not endpoints or not endpoints[0].model_names:
raise ValueError("No endpoints or model names available to resolve tokenizer")
model_name = endpoints[0].model_names[0]There was a problem hiding this comment.
Pre-existing behavior class rather than something this PR introduces — the unguarded endpoints[0] access exists on main today (same as the empty-endpoints case flagged in an earlier round). Keeping this PR scoped to the chat-tokenization fix; happy to bring a defensive-guard follow-up PR (together with the empty-endpoints 503 mirror of LoadAwareRouter) if maintainers want it.
| 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", "")} |
There was a problem hiding this comment.
Extract the duplicated local-first tokenization and remote /tokenize fallback logic into a shared module-level helper function _tokenize_prompt. This avoids significant code duplication between KvawareRouter.route_request and LoadAwareRouter.tokenize_prompt.
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", "")}
async def _tokenize_prompt(
router,
endpoints: List[EndpointInfo],
request_json: Dict
) -> Optional[List[int]]:
"""Local-first tokenization with the remote `/tokenize` fallback."""
try:
await _ensure_tokenizer(router, endpoints)
return _extract_token_ids(router.tokenizer, request_json)
except Exception:
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 NoneThere was a problem hiding this comment.
Reasonable refactor — deliberately deferring it to keep this diff reviewable at its current size after six review rounds. The semantics of the two sites differ slightly (KvawareRouter degrades to session/QPS on total failure inside route_request; LoadAwareRouter returns None for its caller to handle), so the shared helper deserves its own focused change. Happy to follow up post-merge if maintainers agree.
| 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 |
There was a problem hiding this comment.
Use the shared _tokenize_prompt helper function to perform local-first tokenization with remote fallback, eliminating duplicated code.
token_ids = await _tokenize_prompt(self, endpoints, request_json)There was a problem hiding this comment.
See the reply on the helper suggestion above — deferred as a focused follow-up to keep this diff stable for review.
| 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 |
There was a problem hiding this comment.
Use the shared _tokenize_prompt helper function to perform local-first tokenization with remote fallback, eliminating duplicated code.
return await _tokenize_prompt(self, endpoints, request_json)There was a problem hiding this comment.
Same as above — deferred to a focused follow-up.
2 participants
Fixes KV-aware/load-aware routing silently skipping chat-completions traffic: only
prompt-form bodies were tokenized, so KV-aware placement never happened formessagesbodies. Validated end-to-end on real hardware on BOTH tokenization paths — the remote/tokenizepath at3f498ff(2× H200, model behind a vLLM served-name alias) and the local chat-template path atc0c4a42(1× H100, hub-known model; that run caught and fixed a real transformers-5.x bug pre-merge, see Validation below). Commits above3f498ff: hardening from five gemini-code-assist review rounds, a negative cache for failed tokenizer loads (downstream-deployment feedback), and the transformers fix with its test.Problem
KvawareRouter.route_requestandLoadAwareRouter.tokenize_promptonly readrequest_json["prompt"](the/v1/completionsshape). For/v1/chat/completionsbodies —messages: [...]— tokenization fails and therouter silently falls back to session/QPS routing: requests still succeed, so
nothing surfaces above debug level, but KV-aware placement never happens for
chat traffic. (Tutorial 17's test flow uses
/v1/completions, which is whythis wasn't visible.)
Fix
Tokenize chat bodies through the model's chat template so the router's token
ids match what the engine actually caches:
apply_chat_template(messages, add_generation_prompt=True, tokenize=False)) andencode(text, add_special_tokens=False)— NOTtokenize=True, whose return type variesacross transformers versions (
Encodingobjects on 5.x) and silentlybreaks the KV lookup; ids verified byte-identical to the engine's
/tokenizeoutput on vLLM v0.22.0 + transformers 5.9.0/tokenizeendpoint with aTokenizeChatRequest(supported since v0.5.3)Prompt-form requests take exactly the old path. 16 unit tests added; the
full router test suite passes (236/236).
All gemini-code-assist review rounds are addressed (
2574697,ec2becc,00fcf4d,113973c,a7fd571,d34db2c, andf6e243f/c0c4a42from the second hardware pass): blocking I/O (remote
/tokenize,tokenizer load) runs in executors, a total tokenization failure degrades to
the existing session/QPS fallback instead of a 500, tokenizer init is
single-flight behind an
asyncio.Lock, and normalization drops null/emptytext parts, with the tokenizer logic shared via an
_ensure_tokenizerhelper.d34db2cadditionally negative-caches failed tokenizer loads: deployments serving
under a vLLM
--served-model-name(not a hub id) can never load locally,and previously paid a doomed hub lookup on every request before the remote
/tokenizefallback — now the failure is remembered per model name.End-to-end validation — both tokenization paths, real hardware
Remote-path leg (2x H200, gemma-4-31B FP8 under a vLLM
--served-model-name, vLLM v0.22.0 engines + lmcache 0.4.5, 2026-08-22) and local-path leg (1x H100, hub-known Qwen2.5-1.5B via the repro kit below, transformers 5.9.0, 2026-08-24). The local leg caught a real bug pre-merge —apply_chat_template(..., tokenize=True)returnsEncodingobjects on transformers 5.x, so the lookup silently missed — fixed inf6e243fby rendering the template to text and encoding it (add_special_tokens=False); ids verified byte-identical to the engine's/tokenizeoutput. At headc0c4a42: all 4 chat requests pin to the caching engine (4 kvaware decisions vs 1 unpatched), prompt-form unregressed.Raw evidence — leg 1: remote /tokenize path (2× H200, Gemma-4-31B FP8 served as a vLLM alias, 2026-08-22)
Probe: 4 chat requests sharing a salted ~19.4k-token prefix; per-request engine
attribution via
vllm:prompt_tokens_totaldeltas (the delta counts submittedprompt tokens, so which engine moved is the attribution; wall-time shows
cache behavior). Served model name is not a hub id, so the router always used
the remote
/tokenizefallback.BEFORE (merge-base
58a0935) — requests alternate; the same prefix isfull-prefilled on BOTH engines; zero kvaware decisions logged for chat:
AFTER (patched) — requests 2–4 pin to the engine that served request 1:
Raw evidence — leg 2: local chat-template path (1× H100, hub-known Qwen2.5-1.5B via the repro kit, transformers 5.9.0, 2026-08-24)
This leg is what caught the
Encoding-vs-ids bug: at the pre-fix head theAFTER run still showed the BEFORE signature (chat alternating A,B,A,B) because
apply_chat_template(..., tokenize=True)returnedEncodingobjects and thelookup keyed on garbage. Diagnosis inside the router container:
After
f6e243f(render template to text,encode(text, add_special_tokens=False))the same comparison prints
MATCH, and the probe at headc0c4a42:BEFORE leg on the same box (merge-base): chat alternated engine0/engine1 with
the full 5,126-token prefix prefilled on both, zero chat kvaware decisions —
same signature as the H200 leg.
Reproduce it yourself (minutes, any 2-GPU box — or 1 GPU split)
A self-contained kit lives on this fork:
https://github.com/tyler2cr/production-stack/tree/kvaware-chat-repro/repro
(compose + Dockerfile + a short probe; small ungated model
Qwen/Qwen2.5-1.5B-Instruct; no dependence on our infra).The kit's compose/Dockerfile comments codify every wiring precondition we
tripped on (single shared image so the vLLM-rooted NONE_HASH matches,
matched lmcache versions, PYTHONHASHSEED, worker heartbeat env, host
networking, distinct worker ports) — see the notes below.
Notes from standing the e2e up (possible follow-up issues)
These are pre-existing operational sharp edges we hit while validating —
happy to file separately:
QueryInstMsg) keys engines by IP alone —two engines on one IP (e.g. one multi-GPU host) are indistinguishable,
and a kv-followed request to the unmapped instance raises an unhandled
KeyError→ 500 (observed live). The durable fix is that a mapping missshould fall back to session/QPS routing rather than error; including the
worker port in identity helps generic multi-engine hosts (topologies with
one cache owner per IP, e.g. a per-node cache-server DaemonSet, are
unaffected by the identity half).
[lmcache]extra pinned lmcache 0.3.11, whose controller rejects a0.4.5 worker's
RegisterMsgas an unknown type — router/engine lmcacheversions must match exactly. Now filed as [CI/Build][Router] Enable lmcache 0.5.4 for the kvaware router stack #1060 (pins bumped +
lockstep-versioning rationale documented at each pin site).
NONE_HASH=0vs the engines'vLLM-derived root, so every kvaware lookup silently misses. Documented in
[CI/Build][Router] Enable lmcache 0.5.4 for the kvaware router stack #1060's
Dockerfile.kvawarechange; tutorial 17's text still doesn'tmention it.
PYTHONHASHSEEDmust be set identically on router and engines for thebuiltin hash (lmcache warns, tutorial doesn't mention).
reaps silent workers after 30 s — the KV index silently empties shortly
after startup. Still the default at lmcache 0.5.4 (called out in [CI/Build][Router] Enable lmcache 0.5.4 for the kvaware router stack #1060);
the root fix is an lmcache-side defaults change (confirmed biting a real
deployment).
LMCacheConnectorV1does not subclassSupportsHMA, so configuring itmakes vLLM silently disable the hybrid KV cache manager — on
sliding-window models (Gemma-family and others) this inflates per-token
KV roughly 10× with only a startup log line as warning. lmcache's newer
LMCacheMPConnector(KVConnectorBase_V1, SupportsHMA, backed by themultiprocess cache server) is the HMA-capable path.
🤖 Generated with Claude Code