diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index 7a3065b52213..0bb65cd4272a 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -14,12 +14,30 @@ whole prompt. The new prompt may also rewrite the last tokens of the cached one (those after the seam described below), as a chat template does when it replaces the previous turn's generation prefix with the response. +Agentic clients also rewrite the end of the history between turns, for example +by trimming a tool result or replacing the last message, so the new prompt is +no longer an extension of the cached one. The cache records a checkpoint every +256 tokens of a cached prompt. When a new prompt diverges from a cached one, the +cache reuses the cached ids up to the checkpoint one interval before the last +checkpoint inside their longest common prefix, and tokenizes only the rest. The +same is tried when a new prompt extends a much shorter cached prompt, such as +another conversation with the same opening, and the later split is used. + Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore -backs off a fixed number of tokens from the end of the cached prefix, +backs off a fixed number of tokens from the end of the cached prefix (for a +diverging prompt, at least one checkpoint interval before the divergence), re-tokenizes from there, and accepts the splice only if the re-tokenized ids -equal the cached ids over that seam. Otherwise the prompt is tokenized in full. -The output is always identical to tokenizing the whole prompt. +equal the cached ids over that seam (for a diverging prompt, over the whole +interval up to the next checkpoint). Otherwise the prompt is tokenized in full. Every split is moved to +the first token of a character, so a character tokenized into several tokens is +never split. The output is always identical to tokenizing the whole prompt. + +Some tokenizers alter every text they tokenize; for example, a `tokenizer.json` +whose normalizer prepends `▁` (as for Phi-3 and Mixtral-8x7B-v0.1) adds a space +to every tail. Such a shift can pass the check at the seam, so the cache probes +the tokenizer when it is created and stays disabled, with a warning, if a tail +can be shifted. ## Enabling the cache @@ -45,10 +63,11 @@ The cache size can be tuned with environment variables: | `TLLM_PREFIX_TOKEN_CACHE_ENTRIES` | 512 | Maximum number of cached prompts. Eviction is least-recently-used. | | `TLLM_PREFIX_TOKEN_CACHE_MAX_CHARS` | 67108864 | Maximum total characters of cached prompt text. Cached ids are stored as int32, about one byte per character of English text, so this bounds host memory to roughly twice this many bytes. | | `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` | 4096 | Prompts shorter than this are tokenized normally and never cached. | +| `TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION` | 0 | Set to `1` to tokenize each conversation in one serving frontend; see [Multiple serving frontends](#multiple-serving-frontends). | -Each value must be a positive integer. A malformed value raises at server -startup so that misconfiguration fails loudly rather than silently disabling a -feature you turned on. +The sizing values must be positive integers, and the routing switch must be `0` +or `1`. A malformed value raises at server startup so that misconfiguration +fails loudly rather than silently disabling a feature you turned on. ## Scope and safety @@ -62,7 +81,10 @@ feature you turned on. `add_special_tokens=False`, so they benefit. `/v1/completions` defaults to `add_special_tokens=True` and is not accelerated. - A prompt that reuses a cached entry replaces that entry, so a conversation - costs one entry regardless of how many turns it has. Lookup is bucketed by a + costs one entry regardless of how many turns it has. A prompt that diverges + from a cached entry before its seam is added next to it instead: the + divergence may be a branch, such as a subagent started from a shared history, + and the original conversation can still continue from the old entry. Lookup is bucketed by a hash of the first `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` characters, and each entry keeps a short probe of the characters before its seam so that conversations sharing a long system preamble are told apart without @@ -73,9 +95,69 @@ feature you turned on. - An unexpected exception inside the cache disables it, with a warning, and the request is tokenized normally. A cache problem never fails a request. - The cache logs its configuration at startup and, every 1000 cache-eligible - requests, an `INFO` line with the hit rate, the number of re-synchronization - fallbacks, and its occupancy, so a cache that is enabled but never hitting is - visible in the server log. + requests (routed ones included), an `INFO` line with the hit rate (and how many hits were on + diverging prompts), the number of re-synchronization fallbacks, how many + prompts were routed to another frontend and how many were tokenized locally + because routing failed (see below), and its occupancy, so a cache that is + enabled but never hitting is visible in the server log. + +## Multiple serving frontends + +`trtllm-serve` starts 8 frontend processes by default. When you enable the +cache for multi-turn traffic, use a single serving frontend by passing +`--num_serve_frontends 1` (or setting `num_serve_frontends: 1` in the +`--config` YAML file). Multiple frontends are available only with the classic +IPC executor, so this section applies only there. + +The cache lives in the memory of one frontend process and is not shared +between processes. With `num_serve_frontends: K`, the K frontend processes +share the serving port through `SO_REUSEPORT`, and the kernel assigns each new +TCP connection to one of them without looking at the request. A client that +opens a new connection per request, as many benchmark clients do, therefore +sends each turn of a conversation to an arbitrary frontend. That frontend +usually has not seen the conversation's previous turn, so the prompt is +tokenized in full. + +On an agentic multi-turn trace with prompts of about 160k tokens on average +(GLM-5.3 and its tokenizer on GB300), at one session, the p90 time spent in the +frontend was about 165 ms with 8 frontends and 44 ms with one. Decode throughput +per user was the same from 1 to 32 concurrent sessions, the range measured, +because at that concurrency the frontend is not the bottleneck. + +Multiple frontends help when a single Python frontend cannot keep up, typically +at high request rates with short prompts. That is not the workload this cache +targets. + +A client that can pick the frontend itself avoids the problem: with +`serve_frontend_affinity_port_stride` set, each frontend also listens on its own +port, and a disaggregated server sends every context request that carries a +conversation id to the port of the frontend that conversation maps to. + +Otherwise, if you need several frontends, set +`TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION=1` to tokenize each conversation +in one frontend. A chat request's conversation id (the same id attention DP +routes by: from `conversation_params.conversation_id` or headers such as +`X-Correlation-ID`, or, on a disaggregated context server, the sub-agent parent +id forwarded by the disaggregated server) is mapped to an owner frontend, the +same one the affinity ports use, and the frontend the request arrived at sends +the prompt text to the owner over a local Unix socket and gets back the token +ids; the request itself stays where it arrived. Routing helps only when the id +is stable across a conversation's turns: with an id that is fresh on every +request, as many gateways set `X-Correlation-ID`, it only adds a hop per turn. + +A request without a conversation id is tokenized locally. So is one whose owner +is not serving, does not answer within 5 seconds, or cannot tokenize the +prompt; these are counted as route failures in the periodic log line. An owner +that times out is logged with a warning and not asked again for 30 seconds. +Each frontend keeps one connection to every other frontend, shared by its +tokenizing threads, so with K frontends each one holds 2 × (K − 1) extra Unix +sockets and K − 1 extra threads. + +The switch is off by default. When it is set but cannot take effect, with a +single frontend or with the cache disabled, the server logs a warning; when it +takes effect, each frontend logs an `INFO` line. On the trace above with 8 +frontends, it brought the p90 frontend time at one session from about 165 ms to +about 58 ms. ## Measured effect @@ -97,3 +179,8 @@ The splice was also checked against real `gpt2`, Llama-3, and Qwen3 tokenizers on 300 randomly grown multi-turn prompts each, with deltas deliberately starting mid-word: 300/300 identical ids, 280 cache hits, and 0 re-synchronization fallbacks per tokenizer. + +Splicing diverging prompts was checked on 300 multi-turn prompts whose history +is rewritten between turns: all 300 matched tokenizing the whole prompt, 99.3% +were cache hits, and tokenization took about a tenth of the time of +tokenizing every prompt in full. diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2fcd075d71df..ca0095fef847 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -805,7 +805,9 @@ def launch_server( multi_frontend_serving = MultiFrontendServing( launcher_uds=attached_launcher_uds, is_launcher=False, - launcher_pid=os.getppid()) + launcher_pid=os.getppid(), + frontend_id=int(os.environ["TLLM_EXECUTOR_FRONTEND_ID"]), + num_frontends=multi_frontend.num_frontends) if enable_runtime_control_endpoints and multi_frontend.num_frontends > 1: raise click.BadParameter( "Runtime control endpoints require num_serve_frontends=1 because " @@ -906,7 +908,9 @@ def launch_server( stop_watching_frontends = _watch_attached_frontends( frontend_children) multi_frontend_serving = MultiFrontendServing( - launcher_uds=launcher_uds_path, is_launcher=True) + launcher_uds=launcher_uds_path, + is_launcher=True, + num_frontends=multi_frontend.num_frontends) sockets.append(launcher_uds_sock) server = OpenAIServer( diff --git a/tensorrt_llm/inputs/prefix_token_cache.py b/tensorrt_llm/inputs/prefix_token_cache.py index fe5ebf128dfe..dc2ac841db98 100644 --- a/tensorrt_llm/inputs/prefix_token_cache.py +++ b/tensorrt_llm/inputs/prefix_token_cache.py @@ -10,13 +10,22 @@ of the cached prompt, so the new prompt may also rewrite that end, as a chat template does when the previous turn's generation prefix is replaced by the response. +Agentic clients also rewrite the end of the history between turns (a trimmed +tool result, a replaced last message), so a prompt that diverges from a cached +one before its re-tokenization point is spliced at a checkpoint of the cached +prompt: the last one that has another checkpoint after it inside their longest +common prefix. Tokenizing a tail in isolation is not generally identical to tokenizing the whole prompt, because BPE merges can straddle the split. The cache therefore -backs off ``overlap`` tokens from the end of the cached prefix, re-tokenizes -from there, and accepts the splice only if the first ``resync`` re-tokenized ids -equal the cached ids over the same span. Otherwise the prompt is tokenized in -full. +backs off ``overlap`` tokens from the end of the cached prefix (or, for a +diverging prompt, at least ``checkpoint_tokens`` tokens from the divergence), +re-tokenizes from there, and accepts the splice only if the first ``resync`` +re-tokenized ids equal the cached ids over the same span (for a diverging +prompt, over the whole span up to the next checkpoint, which must also start a +re-tokenized token). Otherwise the prompt is tokenized in full. Tokenizers that alter +every text they are called on, such as those whose normalizer prepends a +``▁``, are probed for once when the cache is created and get no cache. Off by default; enable with the ``enable_tokenization_cache`` LLM argument. Each ``DefaultInputProcessor`` owns its own cache, so cached ids are never @@ -26,19 +35,25 @@ from __future__ import annotations +import bisect +import contextlib import os import threading from array import array from collections import OrderedDict +from contextvars import ContextVar from dataclasses import dataclass -from typing import ClassVar, Mapping, Protocol, Sequence +from typing import Callable, ClassVar, Iterator, Mapping, NamedTuple, Protocol, Sequence from ..logger import logger __all__ = [ + "OffsetTokenizer", "PrefixTokenCache", "PrefixTokenCacheConfig", + "RouteFailed", "create_prefix_token_cache", + "tokenization_routing_key", ] # Log the hit rate every this many cache-eligible requests, so an operator can @@ -50,6 +65,30 @@ # candidate is only compared in full when the incoming text has that probe at the # same position. PROBE_CHARS = 256 +# Block size, in characters, of the coarse scan for the first differing +# character of two prompts. +_COMPARE_BLOCK_CHARS = 4096 + +# The conversation the prompt being tokenized belongs to, if the caller knows it. +_ROUTING_KEY: ContextVar[str | None] = ContextVar("tokenization_routing_key", default=None) + + +@contextlib.contextmanager +def tokenization_routing_key(key: str | None) -> Iterator[None]: + """Tag the prompts tokenized in this block with conversation ``key``. + + A cache with a ``route`` hands such prompts to the process that caches that + conversation. + """ + token = _ROUTING_KEY.set(key) + try: + yield + finally: + _ROUTING_KEY.reset(token) + + +class RouteFailed(Exception): + """``PrefixTokenCache.route`` could not get the ids from the owning process.""" class OffsetTokenizer(Protocol): @@ -64,8 +103,8 @@ def __call__( class PrefixTokenCacheConfig: """Tunables for :class:`PrefixTokenCache`. - The sizing fields have env overrides (see ``ENV_VARS``); ``overlap`` and - ``resync`` are correctness internals and do not. + The sizing fields have env overrides (see ``ENV_VARS``); ``overlap``, + ``resync`` and ``checkpoint_tokens`` are correctness internals and do not. """ max_entries: int = 512 @@ -81,6 +120,10 @@ class PrefixTokenCacheConfig: overlap: int = 64 # Re-tokenized ids that must equal the cached ids for the splice to be used. resync: int = 32 + # Spacing, in tokens, of the split points a diverging prompt can be spliced + # at. A split is used only when the next checkpoint still lies inside the + # common prefix, so it must exceed ``resync``. + checkpoint_tokens: int = 256 ENV_VARS: ClassVar[dict[str, str]] = { "max_entries": "TLLM_PREFIX_TOKEN_CACHE_ENTRIES", @@ -88,6 +131,12 @@ class PrefixTokenCacheConfig: "min_chars": "TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", } + def __post_init__(self) -> None: + if self.checkpoint_tokens <= self.resync: + raise ValueError( + f"checkpoint_tokens ({self.checkpoint_tokens}) must exceed resync ({self.resync})" + ) + @classmethod def from_env(cls) -> PrefixTokenCacheConfig: """Build a config from the environment. @@ -111,22 +160,42 @@ class _Entry: text: str ids: array # int32 # Every prompt that keeps ``text[:split_char]`` is re-tokenized from the same - # place: token ``split_token``, which starts at character ``split_char``. + # place: token ``split_token``, the first token of character ``split_char``. split_token: int split_char: int probe: str # the PROBE_CHARS characters of ``text`` before ``split_char`` + # Where a prompt that diverges from ``text`` can be re-tokenized from: the + # first token of the character holding every ``checkpoint_tokens``-th token + # (0, C, 2C, ...), and the start character of that token. + checkpoint_token_indices: array # int32 + checkpoint_chars: array # int32 + + +class _Match(NamedTuple): + eid: int + entry: _Entry + split_token: int + split_char: int + extends: bool # the prompt keeps ``entry.text[: entry.split_char]`` + # For a diverging prompt, the token and character of the checkpoint after + # the split: the re-tokenized tail must match the cached ids up to there. + next_checkpoint: tuple[int, int] | None = None class PrefixTokenCache: """Splice a cached prefix's token ids with a freshly tokenized tail. Thread-safe. Eviction is least-recently-used, bounded by both entry count - and total cached characters. When a prompt reuses a cached entry, the new - prompt replaces that entry rather than sitting alongside it, so a - conversation costs one entry however many turns it has. + and total cached characters. A prompt that extends a cached entry replaces + that entry rather than sitting alongside it, so a conversation costs one + entry however many turns it has. A prompt that diverges from a cached entry + is added next to it, since the divergence may be a branch of a history that + continues. An unexpected error disables the cache, with a warning, and every prompt is - tokenized in full from then on: a cache bug must never fail a request. + tokenized in full from then on: a cache bug must never fail a request. A + prompt the tokenizer itself rejects fails as it would without the cache and + leaves the cache enabled. """ def __init__(self, config: PrefixTokenCacheConfig) -> None: @@ -143,27 +212,87 @@ def __init__(self, config: PrefixTokenCacheConfig) -> None: self.hits = 0 self.misses = 0 self.resync_failures = 0 + # Hits on a cached prompt the new one diverges from (also in ``hits``). + self.divergent_hits = 0 + # Cache-eligible prompts, routed or not; paces the periodic log line. + self.requests = 0 + # Prompts another process tokenized through the route. + self.routed = 0 + # Prompts tokenized here because the route failed. + self.route_failures = 0 + self._route: Callable[[str, str], list[int] | None] | None = None + + def set_route(self, route: Callable[[str, str], list[int] | None] | None) -> None: + """Hand prompts tagged with a routing key to ``route``, or stop if None. + + ``route(key, text)`` returns the ids of ``text`` as tokenized by the + process that caches conversation ``key``, or None to tokenize it here, + and raises :class:`RouteFailed` if that process did not answer; the + prompt is then tokenized here too. + """ + self._route = route def encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: - """Return the token ids of ``text`` without special tokens.""" + """Return the token ids of ``text`` without special tokens. + + Raises whatever the tokenizer raises for a prompt it cannot tokenize. + """ if self.disabled or len(text) < self._config.min_chars: return self._tokenize(tokenizer, text)[0] + try: + ids = self._routed(text) + if ids is None: + ids = self._encode_or_disable(tokenizer, text) + return ids + finally: + self._count_request() + + def _routed(self, text: str) -> list[int] | None: + """Ids from the process that owns ``text``'s conversation, if any.""" + route, key = self._route, _ROUTING_KEY.get() + if route is None or key is None: + return None + try: + ids = route(key, text) + except RouteFailed as e: + logger.debug(f"Tokenizing a prompt locally: {e}") + with self._lock: + self.route_failures += 1 + return None + if ids is not None: + with self._lock: + self.routed += 1 + return ids + + def _encode_or_disable(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: try: return self._encode(tokenizer, text) except Exception as e: + # Raises, leaving the cache enabled, if the tokenizer rejects the + # prompt: that is the request's error, not the cache's. + ids = self._tokenize(tokenizer, text)[0] self.disabled = True logger.warning( f"Disabling the prefix token cache after an error: {e!r} ({self.summary()})" ) - return self._tokenize(tokenizer, text)[0] + return ids + + def _count_request(self) -> None: + with self._lock: + self.requests += 1 + report = self.requests % LOG_EVERY_REQUESTS == 0 + summary = self.summary() if report else None + if summary is not None: + logger.info(f"Prefix token cache: {summary}") def summary(self) -> str: """One-line hit-rate and occupancy summary. Caller need not hold the lock.""" total = self.hits + self.misses rate = self.hits / total if total else 0.0 return ( - f"{self.hits}/{total} hits ({rate:.1%}), " - f"{self.resync_failures} resync fallbacks, " + f"{self.hits}/{total} hits ({rate:.1%}, {self.divergent_hits} divergent), " + f"{self.resync_failures} resync fallbacks, {self.routed} routed, " + f"{self.route_failures} route failures, " f"{len(self._entries)} entries, {self._total_chars} cached chars" ) @@ -171,11 +300,13 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: found = self._lookup(text) ids = None if found is not None: - eid, entry = found - tail_ids, tail_offsets = self._tokenize(tokenizer, text[entry.split_char :]) - if self._resynced(entry, tail_ids): - ids = entry.ids[: entry.split_token].tolist() + list(tail_ids) - base_token, base_char = entry.split_token, entry.split_char + entry = found.entry + tail_ids, tail_offsets = self._tokenize(tokenizer, text[found.split_char :]) + if self._resynced(entry, found.split_token, tail_ids) and self._aligned( + found, tail_ids, tail_offsets + ): + ids = entry.ids[: found.split_token].tolist() + list(tail_ids) + base_token, base_char = found.split_token, found.split_char if ids is None: tail_ids, tail_offsets = self._tokenize(tokenizer, text) ids, base_token, base_char = tail_ids, 0, 0 @@ -183,6 +314,9 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: # Where a prompt extending this one will be re-tokenized from. split_token = len(ids) - self._config.overlap tail_index = split_token - base_token + if tail_index >= 0: + tail_index = _first_of_char(tail_offsets, tail_index) + split_token = base_token + tail_index with self._lock: if found is None: @@ -192,14 +326,24 @@ def _encode(self, tokenizer: OffsetTokenizer, text: str) -> list[int]: self.resync_failures += 1 else: self.hits += 1 - if found is not None: - self._remove(found[0]) + self.divergent_hits += not found.extends + # A prompt that reuses an entry supersedes it. A diverging prompt + # may be a branch (a subagent off a shared history), so the entry + # it was spliced from stays for the other branch. + if found is not None and found.extends: + self._remove(found.eid) if split_token > 0 and tail_index >= 0: - self._insert(text, ids, split_token, base_char + tail_offsets[tail_index][0]) - report = (self.hits + self.misses) % LOG_EVERY_REQUESTS == 0 - summary = self.summary() if report else None - if summary is not None: - logger.info(f"Prefix token cache: {summary}") + checkpoint_token_indices, checkpoint_chars = self._checkpoints( + found, base_token, base_char, tail_offsets, len(ids) + ) + self._insert( + text, + ids, + split_token, + base_char + tail_offsets[tail_index][0], + checkpoint_token_indices, + checkpoint_chars, + ) return ids @staticmethod @@ -209,41 +353,128 @@ def _tokenize( enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) return enc["input_ids"], enc["offset_mapping"] - def _resynced(self, entry: _Entry, tail_ids: Sequence[int]) -> bool: + def _checkpoints( + self, + found: _Match | None, + base_token: int, + base_char: int, + tail_offsets: Sequence[tuple[int, int]], + num_ids: int, + ) -> tuple[array, array]: + """Checkpoint tokens and chars of ``base_token`` cached ids plus a tail.""" + step = self._config.checkpoint_tokens + # Spliced ids before base_token are the cached entry's, at the same chars. + num_kept = -(-base_token // step) if base_token else 0 + tokens = array("i", found.entry.checkpoint_token_indices[:num_kept] if num_kept else ()) + chars = array("i", found.entry.checkpoint_chars[:num_kept] if num_kept else ()) + for index in range(num_kept * step, num_ids, step): + tail_index = _first_of_char(tail_offsets, index - base_token) + tokens.append(base_token + tail_index) + chars.append(base_char + tail_offsets[tail_index][0]) + return tokens, chars + + def _resynced(self, entry: _Entry, split_token: int, tail_ids: Sequence[int]) -> bool: """Whether the tokenizer re-synchronized with the cached ids at the seam.""" - cached_tail = entry.ids[entry.split_token :] + cached_tail = entry.ids[split_token:] span = min(self._config.resync, len(tail_ids), len(cached_tail)) return span > 0 and list(tail_ids[:span]) == cached_tail[:span].tolist() + @staticmethod + def _aligned( + found: _Match, tail_ids: Sequence[int], tail_offsets: Sequence[tuple[int, int]] + ) -> bool: + """Whether a diverging prompt's tail matches the cached ids up to the next checkpoint. + + The resync span can lie inside a run of equal tokens, such as a long + run of spaces, where a tail shifted by a character still agrees with + the cached ids; the shift shows where the run ends, or as tokens that + start one character off. The next checkpoint starts inside the common + prefix, so an unshifted tail has the cached ids up to it and starts a + token there. + """ + if found.next_checkpoint is None: + return True + token, char = found.next_checkpoint + index = token - found.split_token + return ( + index < len(tail_ids) + and tail_offsets[index][0] == char - found.split_char + and list(tail_ids[:index]) == found.entry.ids[found.split_token : token].tolist() + ) + def _bucket_key(self, text: str) -> int: return hash(text[: self._config.min_chars]) - def _lookup(self, text: str) -> tuple[int, _Entry] | None: - """Cached entry with the longest re-tokenization prefix ``text`` keeps, promoted to MRU.""" + def _lookup(self, text: str) -> _Match | None: + """Best cached entry to splice ``text`` from, promoted to MRU if ``text`` extends it. + + The entry with the longest re-tokenization prefix that ``text`` keeps + wins. Without one, or when it leaves more than two checkpoint intervals + to re-tokenize, the entries ``text`` diverges from are also tried, each + split at the checkpoint one interval before its last checkpoint inside + the common prefix, and the latest split wins. + """ with self._lock: candidates = list(self._buckets.get(self._bucket_key(text), {}).items()) # Entries are immutable, so the comparisons run outside the lock. The # probe check rejects other conversations in O(PROBE_CHARS); only a true # prefix pays the full compare over the shared preamble. - best: tuple[int, _Entry] | None = None + best: _Match | None = None for eid, entry in candidates: n = entry.split_char if ( n <= len(text) - and (best is None or n > best[1].split_char) + and (best is None or n > best.split_char) and text[n - len(entry.probe) : n] == entry.probe and text.startswith(entry.text[:n]) ): - best = (eid, entry) - if best is not None: + best = _Match(eid, entry, entry.split_token, entry.split_char, True) + if best is None or self._leaves_long_tail(best, text): + for eid, entry in candidates: + match = self._divergent_match(eid, entry, text) + if match is not None and (best is None or match.split_token > best.split_token): + best = match + # A diverging prompt is not promoted: when the history was rewritten + # rather than branched, the entry it diverged from is dead. + if best is not None and best.extends: with self._lock: - if best[0] in self._entries: - self._entries.move_to_end(best[0]) + if best.eid in self._entries: + self._entries.move_to_end(best.eid) # If it was evicted meanwhile the entry is still a valid prefix to # splice from; _remove() on it is a no-op. return best - def _insert(self, text: str, ids: Sequence[int], split_token: int, split_char: int) -> None: + def _leaves_long_tail(self, match: _Match, text: str) -> bool: + """Whether ``text`` extends ``match`` by more than two checkpoint intervals.""" + entry = match.entry + chars_per_token = entry.split_char / max(entry.split_token, 1) + tail_tokens = (len(text) - match.split_char) / max(chars_per_token, 1e-6) + return tail_tokens > 2 * self._config.checkpoint_tokens + + def _divergent_match(self, eid: int, entry: _Entry, text: str) -> _Match | None: + common = _common_prefix_len(entry.text, text) + # The checkpoint after the split must start inside the common prefix. + index = bisect.bisect_left(entry.checkpoint_chars, common) - 2 + if index < 1: + return None + return _Match( + eid, + entry, + entry.checkpoint_token_indices[index], + entry.checkpoint_chars[index], + False, + (entry.checkpoint_token_indices[index + 1], entry.checkpoint_chars[index + 1]), + ) + + def _insert( + self, + text: str, + ids: Sequence[int], + split_token: int, + split_char: int, + checkpoint_token_indices: array, + checkpoint_chars: array, + ) -> None: """Caller holds the lock.""" key = self._bucket_key(text) bucket = self._buckets.setdefault(key, {}) @@ -251,7 +482,15 @@ def _insert(self, text: str, ids: Sequence[int], split_token: int, split_char: i return # concurrent misses on the same prompt eid, self._next_id = self._next_id, self._next_id + 1 probe = text[max(split_char - PROBE_CHARS, 0) : split_char] - entry = _Entry(text, array("i", ids), split_token, split_char, probe) + entry = _Entry( + text, + array("i", ids), + split_token, + split_char, + probe, + checkpoint_token_indices, + checkpoint_chars, + ) self._entries[eid] = entry bucket[eid] = entry self._total_chars += len(text) @@ -274,10 +513,59 @@ def _remove(self, eid: int) -> None: del self._buckets[key] +def _common_prefix_len(a: str, b: str) -> int: + """Length of the longest common prefix, compared in blocks.""" + n, i, block = min(len(a), len(b)), 0, _COMPARE_BLOCK_CHARS + while i < n and a[i : i + block] == b[i : i + block]: + i += block + # The first mismatch, if any, lies in the block at i; halve it to find it. + lo, hi = min(i, n), min(i + block, n) + while lo < hi: + mid = (lo + hi + 1) // 2 + if a[lo:mid] == b[lo:mid]: + lo = mid + else: + hi = mid - 1 + return lo + + +def _first_of_char(offsets: Sequence[tuple[int, int]], index: int) -> int: + """First token of the character token ``index`` starts in. + + A character tokenized into several tokens (a byte-fallback character, an + expanded ligature) reports the same start for each; a split must not + separate them, or the tail would repeat the tokens of that character. + """ + while index > 0 and offsets[index - 1][0] == offsets[index][0]: + index -= 1 + return index + + +def _splices_wrongly(tokenizer: OffsetTokenizer, config: PrefixTokenCacheConfig) -> bool: + """Whether a tail the cache would accept at the seam can differ from full tokenization. + + Tokenizers that alter every text they are called on, such as one whose + normalizer prepends ``▁``, shift a tail by a character. Inside a run of + equal tokens the shift passes the resync check and only shows at the end + of the run, so the probe splits inside a long run of spaces. + """ + text = "x" + " " * 4096 + "y" + ids, offsets = PrefixTokenCache._tokenize(tokenizer, text) + split = next((i for i, (start, _) in enumerate(offsets) if start >= 1024), None) + if split is None: + return False + split = _first_of_char(offsets, split) + tail_ids = PrefixTokenCache._tokenize(tokenizer, text[offsets[split][0] :])[0] + span = min(config.resync, len(tail_ids), len(ids) - split) + resynced = list(tail_ids[:span]) == list(ids[split : split + span]) + return resynced and list(ids[:split]) + list(tail_ids) != list(ids) + + def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: """Return a cache for ``tokenizer`` if the tokenizer can support one. - Returns None when there is no tokenizer or it cannot report offsets. + Returns None when there is no tokenizer, it cannot report offsets, or a + spliced tail can differ from tokenizing the whole prompt. Raises ``ValueError`` for a malformed env override. """ if tokenizer is None: @@ -290,5 +578,20 @@ def create_prefix_token_cache(tokenizer: object) -> PrefixTokenCache | None: ) return None config = PrefixTokenCacheConfig.from_env() + try: + unsafe = _splices_wrongly(tokenizer, config) + except Exception as e: + logger.warning( + f"Could not probe the tokenizer for the prefix token cache ({e!r}); disabling it." + ) + return None + if unsafe: + logger.warning( + "enable_tokenization_cache is set but the tokenizer changes the start " + "of every text it tokenizes (for example, its normalizer prepends a " + "space marker), so spliced token ids could differ from tokenizing " + "the whole prompt; disabling the prefix token cache." + ) + return None logger.info(f"Prefix token cache enabled: {config}") return PrefixTokenCache(config) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 71cfa6cefd23..83965750fc91 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -52,7 +52,8 @@ validate_mm_inputs) # yapf: enable from .multimodal_data import MEDIA_MODALITIES, serialize_item -from .prefix_token_cache import PrefixTokenCache, create_prefix_token_cache +from .prefix_token_cache import (OffsetTokenizer, PrefixTokenCache, + create_prefix_token_cache) from .tokenization import tokenize_prompt N = TypeVar("N", bound=Type[nn.Module]) @@ -140,6 +141,13 @@ def __init__(self, self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) + def get_prefix_token_cache( + self) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """This processor's prefix-tokenization cache and the tokenizer it caches, if enabled.""" + if self._prefix_token_cache is None: + return None + return self._prefix_token_cache, self.tokenizer + def __call__( self, inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: @@ -153,6 +161,18 @@ def __call__( return token_ids, None +def prefix_token_cache_of( + input_processor: object +) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """The prefix-tokenization cache of ``input_processor`` and its tokenizer. + + None when the processor has no cache: it is disabled, or the processor + does not support one. + """ + get_cache = getattr(input_processor, "get_prefix_token_cache", None) + return None if get_cache is None else get_cache() + + class BaseMultimodalInputProcessor(ABC): """ Base class for multimodal input processors with default implementations @@ -272,6 +292,13 @@ def _init_tokenization_cache(self, enable_tokenization_cache: bool, self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) + def get_prefix_token_cache( + self) -> Optional[Tuple[PrefixTokenCache, OffsetTokenizer]]: + """This processor's prefix-tokenization cache and the tokenizer it caches, if enabled.""" + if self._prefix_token_cache is None: + return None + return self._prefix_token_cache, self._prefix_token_cache_tokenizer + def _encode_with_tokenization_cache( self, prompt: str, diff --git a/tensorrt_llm/serve/multi_frontend.py b/tensorrt_llm/serve/multi_frontend.py index ea7aa6bd4cea..49766455a570 100644 --- a/tensorrt_llm/serve/multi_frontend.py +++ b/tensorrt_llm/serve/multi_frontend.py @@ -26,6 +26,7 @@ import math import os import tempfile +import zlib from dataclasses import dataclass from typing import Callable, Optional @@ -38,6 +39,9 @@ LAUNCHER_UDS_NAME = "launcher.sock" LAUNCHER_UDS_ENV = "TLLM_SERVE_LAUNCHER_UDS" CLOCK_OFFSET_NAME = "clock_offset" +# Every frontend serves its prefix-tokenization cache on this Unix socket when +# tokenization is routed by conversation. +_TOKENIZATION_UDS_NAME = "tokenize-{}.sock" # Launcher-owned routes. Attached frontends register them ahead of the local # handlers, which they shadow since FastAPI matches in registration order. @@ -85,12 +89,30 @@ class MultiFrontendServing: launcher_uds: str is_launcher: bool launcher_pid: Optional[int] = None # attached frontends: their parent + frontend_id: int = 0 # the launcher is frontend 0 + num_frontends: int = 1 @property def is_attached(self) -> bool: return not self.is_launcher +def frontend_of_conversation(conversation_id: str, num_frontends: int) -> int: + """The frontend that serves conversation ``conversation_id``, the same in every process. + + A disaggregated server sends the conversation's context requests to this + frontend's affinity port, and with tokenization routing this frontend + tokenizes its prompts, so both must use this mapping. + """ + # Lone surrogates, which JSON can carry, must not fail the request. + return zlib.crc32(conversation_id.encode(errors="surrogatepass")) % num_frontends + + +def tokenization_uds_path(ipc_dir: str, frontend_id: int) -> str: + """The Unix socket frontend ``frontend_id`` serves its tokenization cache on.""" + return os.path.join(ipc_dir, _TOKENIZATION_UDS_NAME.format(frontend_id)) + + class LauncherForwarder: """Forward HTTP requests from an attached frontend to the launcher.""" diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 1bf03ddebbef..542550e186a0 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -18,7 +18,6 @@ import json import os import traceback -import zlib from abc import ABC, abstractmethod from contextlib import aclosing from typing import ( @@ -49,6 +48,7 @@ build_subagent_affinity_headers, request_requires_internal_disagg_auth, ) +from tensorrt_llm.serve.multi_frontend import frontend_of_conversation from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -309,7 +309,7 @@ def _url_server(self, server: str, request: UCompletionRequest) -> str: return server num_frontends, stride = self._frontend_affinity host, _, port = server.rpartition(":") - frontend = zlib.crc32(conversation_id.encode()) % num_frontends + frontend = frontend_of_conversation(conversation_id, num_frontends) affinity_server = f"{host}:{int(port) + stride * (frontend + 1)}" refused_at = self._refused_affinity_ports.get(affinity_server) if refused_at is not None: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 233512d9788b..d83c8865aad9 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -23,7 +23,7 @@ from http import HTTPStatus from pathlib import Path from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, - AsyncIterator, Dict, List, Optional, Tuple, Union) + AsyncIterator, Callable, Dict, List, Optional, Tuple, Union) import msgspec import uvicorn @@ -51,13 +51,14 @@ from tensorrt_llm.inputs.media_io import (BaseMediaIO, log_allowed_local_media_path) from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.inputs.prefix_token_cache import tokenization_routing_key from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.inputs.utils import ConversationMessage from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams from tensorrt_llm.llmapi import MultimodalEncoder, SchedulingParams, tracing from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) -from tensorrt_llm.llmapi.llm import LLM, RequestOutput +from tensorrt_llm.llmapi.llm import LLM, PreprocessedInputs, RequestOutput from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor @@ -86,7 +87,8 @@ from tensorrt_llm.serve.chat_utils import (load_chat_template, parse_chat_messages_coroutines) from tensorrt_llm.serve.cluster_storage import create_cluster_storage_client -from tensorrt_llm.serve.conversation_id import resolve_request_conversation_id +from tensorrt_llm.serve.conversation_id import (get_request_conversation_id, + resolve_request_conversation_id) from tensorrt_llm.serve.disagg_auth import ( request_requires_internal_disagg_auth, validate_internal_disagg_request, validate_subagent_affinity) @@ -145,6 +147,8 @@ from tensorrt_llm.serve.runtime_control_auth import ( RuntimeControlReplayCache, validate_runtime_control_request) from tensorrt_llm.serve.serving_extensions import load_builtin_extensions +from tensorrt_llm.serve.tokenization_owners import (TokenizationOwners, + start_tokenization_owners) from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers) @@ -164,6 +168,14 @@ from tensorrt_llm.visual_gen import VisualGen +def _preprocess_in_conversation( + routing_id: Optional[str], + preprocess_fn: Callable[[], PreprocessedInputs]) -> PreprocessedInputs: + """Run ``preprocess_fn`` with its prompt tagged as conversation ``routing_id``.""" + with tokenization_routing_key(routing_id): + return preprocess_fn() + + def _is_visual_gen_instance(obj) -> bool: """isinstance(obj, VisualGen) without importing the visual_gen tree. @@ -754,6 +766,10 @@ def __init__( if multi_frontend_serving is not None and multi_frontend_serving.is_attached: self._launcher_forwarder = LauncherForwarder( multi_frontend_serving.launcher_uds) + self._tokenization_owners: Optional[ + TokenizationOwners] = start_tokenization_owners( + multi_frontend_serving, + getattr(generator, "input_processor", None)) self._shared_clock_offset: Optional[SharedClockOffset] = ( SharedClockOffset( os.path.dirname(multi_frontend_serving.launcher_uds)) @@ -866,6 +882,8 @@ async def lifespan(app: FastAPI): pass if self._launcher_forwarder is not None: await self._launcher_forwarder.close() + if self._tokenization_owners is not None: + await asyncio.to_thread(self._tokenization_owners.close) if self.embedding_batcher is not None: await self.embedding_batcher.shutdown() logger.info("Stopped encode dynamic batcher") @@ -2481,11 +2499,19 @@ async def chat_stream_generator( generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: + preprocess = functools.partial(preprocess_fn, prompt, + sampling_params, + disaggregated_params) + if self._tokenization_owners is not None: + # A disaggregated context worker gets the sub-agent's + # parent id only from the signed header. + preprocess = functools.partial( + _preprocess_in_conversation, + scheduling_params.subagent_affinity_id + or get_request_conversation_id(request), preprocess) loop = asyncio.get_event_loop() generate_inputs = await loop.run_in_executor( - self._input_proc_executor, - functools.partial(preprocess_fn, prompt, sampling_params, - disaggregated_params)) + self._input_proc_executor, preprocess) promise = self.generator.generate_async( inputs=generate_inputs, diff --git a/tensorrt_llm/serve/tokenization_owners.py b/tensorrt_llm/serve/tokenization_owners.py new file mode 100644 index 000000000000..e2c1c4870d03 --- /dev/null +++ b/tensorrt_llm/serve/tokenization_owners.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tokenize each conversation in the frontend that caches it. + +With ``num_serve_frontends > 1`` the kernel hands each connection to any +frontend, so consecutive turns of a conversation land in different processes, +each with its own prefix-tokenization cache. Every frontend therefore serves +its cache on a Unix socket in the multi-frontend ipc dir, and a prompt whose +conversation id maps to another frontend is tokenized there, by the cache that +has seen the conversation's earlier turns. Only the prompt text and the token +ids cross the socket; the request itself stays where it arrived. + +Each frontend keeps one connection to every other frontend, shared by its +tokenizing threads, so a server with K frontends holds 2 * (K - 1) routing +sockets and K - 1 serving threads per frontend. +""" + +import contextlib +import os +import socket +import socketserver +import struct +import threading +import time +from array import array +from typing import Optional + +from tensorrt_llm.inputs.prefix_token_cache import OffsetTokenizer, PrefixTokenCache, RouteFailed +from tensorrt_llm.inputs.registry import prefix_token_cache_of +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.multi_frontend import ( + MultiFrontendServing, + frontend_of_conversation, + tokenization_uds_path, +) + +# Set to 1 to route tokenization by conversation id across serving frontends. +ENABLE_ENV_VAR = "TLLM_PREFIX_TOKEN_CACHE_ROUTE_BY_CONVERSATION" + +_LENGTH = struct.Struct("!q") +# Sent in place of a token count when the owner could not tokenize the prompt. +_FAILED = -1 +# Lone surrogates, which JSON can carry, survive the trip and fail in the +# tokenizer as they would locally. +_TEXT_ERRORS = "surrogatepass" + + +def routing_enabled() -> bool: + """Whether tokenization is routed by conversation id (off unless enabled). + + Raises: + ValueError: if the switch is set to anything but ``0`` or ``1``. + """ + raw = os.environ.get(ENABLE_ENV_VAR) + if raw is None or raw == "0": + return False + if raw == "1": + return True + raise ValueError(f"{ENABLE_ENV_VAR} must be 0 or 1, got {raw!r}") + + +def start_tokenization_owners( + multi_frontend_serving: Optional[MultiFrontendServing], input_processor: object +) -> Optional["TokenizationOwners"]: + """Route this frontend's prompts by conversation if the switch is on and can engage. + + Raises: + ValueError: if the switch is malformed. + """ + if not routing_enabled(): + return None + num_frontends = 1 if multi_frontend_serving is None else multi_frontend_serving.num_frontends + prefix_cache = prefix_token_cache_of(input_processor) + if num_frontends <= 1 or prefix_cache is None: + logger.warning( + f"{ENABLE_ENV_VAR}=1 has no effect: it needs num_serve_frontends > 1 " + f"(got {num_frontends}) and an enabled prefix token cache " + f"(enable_tokenization_cache: {prefix_cache is not None})." + ) + return None + logger.info( + f"Routing tokenization by conversation id: frontend " + f"{multi_frontend_serving.frontend_id} of {num_frontends}." + ) + return TokenizationOwners( + os.path.dirname(multi_frontend_serving.launcher_uds), + multi_frontend_serving.frontend_id, + num_frontends, + *prefix_cache, + ) + + +class _Peer: + """This frontend's connection to one owner, shared by its tokenizing threads.""" + + def __init__(self) -> None: + # Held for a whole exchange, which keeps requests and replies paired. + self.lock = threading.Lock() + self.conn: Optional[socket.socket] = None + # Monotonic time before which the owner, which timed out, is not asked. + self.skip_until = 0.0 + + +class TokenizationOwners: + """Route ``cache``'s prompts to the frontend that owns their conversation. + + Serves this frontend's ``cache`` to the other frontends and installs its + route. A prompt the owner does not tokenize within ``timeout_s``, for + example because it has hung, is tokenized locally, and so is every prompt + it owns for the next ``retry_after_s``. A prompt whose owner is not serving + (not started yet, or dead) or cannot tokenize it is tokenized locally too. + """ + + def __init__( + self, + ipc_dir: str, + frontend_id: int, + num_frontends: int, + cache: PrefixTokenCache, + tokenizer: OffsetTokenizer, + timeout_s: float = 5.0, + retry_after_s: float = 30.0, + ) -> None: + self._ipc_dir = ipc_dir + self._frontend_id = frontend_id + self._num_frontends = num_frontends + self._cache = cache + self._timeout_s = timeout_s + self._retry_after_s = retry_after_s + self._peers = {owner: _Peer() for owner in range(num_frontends) if owner != frontend_id} + + self._server = _Server(tokenization_uds_path(ipc_dir, frontend_id), cache, tokenizer) + threading.Thread( + target=self._server.serve_forever, name="trtllm_tokenize_owner", daemon=True + ).start() + cache.set_route(self._route) + + def close(self, timeout_s: float = 5.0) -> None: + """Stop routing and serving, and close every connection in both directions. + + Blocks for up to about ``timeout_s`` while handlers finish the prompt + they are tokenizing; those still running then are left to exit with + the process. + """ + deadline = time.monotonic() + timeout_s + self._cache.set_route(None) + self._server.shutdown() + self._server.server_close() + self._server.close_connections(deadline) + for peer in self._peers.values(): + conn = peer.conn + if conn is not None: + # Wakes a thread waiting on the owner, which then tokenizes + # its prompt locally. + with contextlib.suppress(OSError): + conn.shutdown(socket.SHUT_RDWR) + conn.close() + + def _route(self, key: str, text: str) -> Optional[list[int]]: + owner = frontend_of_conversation(key, self._num_frontends) + if owner == self._frontend_id: + return None + peer = self._peers[owner] + if time.monotonic() < peer.skip_until: + raise RouteFailed(f"frontend {owner} timed out recently") + if not peer.lock.acquire(timeout=self._timeout_s): + raise RouteFailed(f"frontend {owner} is busy") + try: + return self._request(owner, peer, text) + except RouteFailed: + raise + except TimeoutError as e: + _drop(peer) + peer.skip_until = time.monotonic() + self._retry_after_s + logger.warning( + f"Frontend {owner} did not tokenize a prompt within {self._timeout_s} s; " + f"tokenizing its conversations locally for {self._retry_after_s} s." + ) + raise RouteFailed(f"frontend {owner} timed out") from e + except Exception as e: + # Not serving (connect failed), died mid-exchange, or sent a + # malformed reply: the connection cannot be reused. + _drop(peer) + raise RouteFailed(f"frontend {owner} did not answer: {e!r}") from e + finally: + peer.lock.release() + + def _request(self, owner: int, peer: _Peer, text: str) -> list[int]: + if peer.conn is None: + conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + conn.settimeout(self._timeout_s) + try: + conn.connect(tokenization_uds_path(self._ipc_dir, owner)) + except OSError: + conn.close() + raise + peer.conn = conn + data = text.encode(errors=_TEXT_ERRORS) + peer.conn.sendall(_LENGTH.pack(len(data)) + data) + (count,) = _LENGTH.unpack(_recv_exactly(peer.conn, _LENGTH.size)) + if count == _FAILED: + raise RouteFailed(f"frontend {owner} could not tokenize the prompt") + if count < 0: + raise ValueError(f"invalid token count {count}") + ids = array("i") + ids.frombytes(_recv_exactly(peer.conn, count * ids.itemsize)) + return ids.tolist() + + +def _drop(peer: _Peer) -> None: + """Caller holds ``peer.lock``.""" + if peer.conn is not None: + peer.conn.close() + peer.conn = None + + +class _Server(socketserver.ThreadingUnixStreamServer): + daemon_threads = True + # Every other frontend may connect at once, and a Unix-socket connect to a + # full backlog fails at once instead of waiting. + request_queue_size = socket.SOMAXCONN + + def __init__(self, path: str, cache: PrefixTokenCache, tokenizer: OffsetTokenizer) -> None: + self.cache = cache + self.tokenizer = tokenizer + # Accepted connections and the threads serving them. + self.handlers: dict[socket.socket, threading.Thread] = {} + self.handlers_lock = threading.Lock() + if os.path.exists(path): + os.unlink(path) + super().__init__(path, _Handler) + + def close_connections(self, deadline: float) -> None: + """Wake every handler blocked on its peer and wait for it to exit until ``deadline``.""" + with self.handlers_lock: + handlers = list(self.handlers.items()) + for conn, _ in handlers: + with contextlib.suppress(OSError): + conn.shutdown(socket.SHUT_RDWR) + for _, thread in handlers: + thread.join(max(0.0, deadline - time.monotonic())) + + +class _Handler(socketserver.BaseRequestHandler): + server: _Server + + def setup(self) -> None: + with self.server.handlers_lock: + self.server.handlers[self.request] = threading.current_thread() + + def finish(self) -> None: + with self.server.handlers_lock: + self.server.handlers.pop(self.request, None) + + def handle(self) -> None: + # Ends when the peer closes the connection or close_connections() runs. + with contextlib.suppress(OSError): + while True: + self._answer_one() + + def _answer_one(self) -> None: + conn = self.request + (size,) = _LENGTH.unpack(_recv_exactly(conn, _LENGTH.size)) + text = _recv_exactly(conn, size).decode(errors=_TEXT_ERRORS) + try: + ids = array("i", self.server.cache.encode(self.server.tokenizer, text)) + except Exception as e: + # The asking frontend tokenizes the prompt itself and reports the + # error with its request. + logger.debug(f"Could not tokenize a prompt for another frontend: {e!r}") + conn.sendall(_LENGTH.pack(_FAILED)) + return + conn.sendall(_LENGTH.pack(len(ids)) + ids.tobytes()) + + +def _recv_exactly(conn: socket.socket, size: int) -> bytearray: + buf = bytearray(size) + view = memoryview(buf) + while view: + n = conn.recv_into(view) + if n == 0: + raise ConnectionError("peer closed the connection") + view = view[n:] + return buf diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff302874d5d6..8c796023fd3a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -406,6 +406,7 @@ l0_cpu: - unittest/llmapi/test_serve_context_only_settings.py - unittest/llmapi/test_serve_multi_frontend_helpers.py - unittest/llmapi/test_serve_multi_frontend_route_gate.py + - unittest/llmapi/test_serve_tokenization_owners.py - unittest/llmapi/test_serve_num_frontends_default.py - unittest/llmapi/test_serve_report_addr.py - unittest/llmapi/test_serve_sleep_mode.py diff --git a/tests/unittest/inputs/test_prefix_token_cache.py b/tests/unittest/inputs/test_prefix_token_cache.py index d92cd4783458..fc8550d71e43 100644 --- a/tests/unittest/inputs/test_prefix_token_cache.py +++ b/tests/unittest/inputs/test_prefix_token_cache.py @@ -20,6 +20,7 @@ from tensorrt_llm.inputs.prefix_token_cache import ( PrefixTokenCache, PrefixTokenCacheConfig, + _common_prefix_len, create_prefix_token_cache, ) from tensorrt_llm.inputs.registry import DefaultInputProcessor @@ -75,6 +76,59 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: return {"input_ids": ids, "offset_mapping": offsets} +class _ExpandingTokenizer: + """Char-level, except that "ff" becomes two equal tokens with the same offsets. + + As with a byte-fallback character or an NFKC ligature, a split can land on + the second token of the character, and a long run of them makes the tail + agree with the cached ids however it is misaligned. + """ + + is_fast = True + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + ids, offsets = [], [] + for i, char in enumerate(text): + for _ in range(2 if char == "ff" else 1): + ids.append(_tok_id(char)) + offsets.append((i, i + 1)) + return {"input_ids": ids, "offset_mapping": offsets} + + +class _PrependingTokenizer: + """Prepends a space to every text, as a ``Prepend("▁")`` normalizer does. + + Runs of spaces become tokens of up to four spaces, counted from the start + of the run, so a tail split inside a run is shifted by the prepended space + and still agrees with the cached ids until the end of the run. + """ + + is_fast = True + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + normalized = " " + text + ids, offsets, i = [], [], 0 + while i < len(normalized): + j = i + 1 + if normalized[i] == " ": + while j < len(normalized) and j - i < 4 and normalized[j] == " ": + j += 1 + ids.append(_tok_id(normalized[i:j])) + # The prepended space maps to the start of the text. + offsets.append((max(i - 1, 0), j - 1)) + i = j + return {"input_ids": ids, "offset_mapping": offsets} + + +class _RejectingTokenizer(_MergeTokenizer): + """Rejects a lone surrogate, as HF tokenizers do.""" + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + if "\ud800" in text: + raise TypeError("cannot tokenize a lone surrogate") + return super().__call__(text, **kwargs) + + _MIN_CHARS = 16 # A shared opening longer than min_chars, so growing prompts land in the same # bucket and the prefix lookup can find them. @@ -82,7 +136,7 @@ def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: def _cache(**overrides: int) -> PrefixTokenCache: - config = dict(max_entries=64, overlap=4, resync=2, min_chars=_MIN_CHARS) + config = dict(max_entries=64, overlap=4, resync=2, min_chars=_MIN_CHARS, checkpoint_tokens=8) return PrefixTokenCache(PrefixTokenCacheConfig(**{**config, **overrides})) @@ -97,6 +151,23 @@ def _turns(n: int = 12, preamble: str = _PREAMBLE) -> list[str]: return out +def _assert_checkpoints_are_valid(cache: PrefixTokenCache, tokenizer: Any) -> None: + """Checkpoint i is the first token of the character holding token i * step.""" + step = cache._config.checkpoint_tokens + for entry in cache._entries.values(): + enc = tokenizer(entry.text) + assert entry.ids.tolist() == enc["input_ids"] + starts = [start for start, _ in enc["offset_mapping"]] + tokens = entry.checkpoint_token_indices.tolist() + chars = entry.checkpoint_chars.tolist() + assert len(tokens) == len(chars) == -(-len(starts) // step) + assert tokens == sorted(tokens) and chars == sorted(chars) + for i, (token, char) in enumerate(zip(tokens, chars)): + assert token <= i * step + assert starts[token] == starts[i * step] == char + assert token == 0 or starts[token - 1] != char + + def _conversations(n: int) -> list[str]: """First turns of ``n`` unrelated conversations.""" return [_turns(1, preamble=f"conversation {i:03d}: {_PREAMBLE}")[0] for i in range(n)] @@ -129,6 +200,118 @@ def test_resync_failure_falls_back_to_full_tokenization() -> None: assert cache.resync_failures > 0 +@pytest.mark.parametrize("tokenizer", [_MergeTokenizer(), _UnsyncableTokenizer()]) +def test_rewritten_history_is_identical_to_full_tokenization(tokenizer: Any) -> None: + """A turn that rewrites the end of the history splices from the common prefix.""" + cache = _cache() + turns = _turns(8) + for i, prompt in enumerate(turns): + assert cache.encode(tokenizer, prompt) == tokenizer(prompt)["input_ids"] + # Trim the last message and append a different one, as agents do. + rewritten = prompt[: len(prompt) - 20] + f"rewritten tail {i} there. " + hits, failures = cache.hits, cache.resync_failures + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + if isinstance(tokenizer, _MergeTokenizer): + assert cache.hits == hits + 1 + else: + # Spliced from the history and rejected at the seam. + assert cache.resync_failures == failures + 1 + if isinstance(tokenizer, _MergeTokenizer): + assert cache.divergent_hits >= len(turns) + # The rewritten branch does not evict the history it diverged from. + assert any(entry.text == turns[-1] for entry in cache._entries.values()) + else: + assert cache.divergent_hits == 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) + + +def test_a_deep_edit_splices_from_a_late_checkpoint() -> None: + """A long history edited deep inside is spliced across several compare blocks.""" + tokenizer, cache = _MergeTokenizer(), _cache() + history = _turns(400)[-1][:10_000] + assert cache.encode(tokenizer, history) == tokenizer(history)["input_ids"] + edited = history[:7000] + "#" + history[7001:] + found = cache._lookup(edited) + assert found is not None and not found.extends + # One checkpoint interval (8 tokens of at most 2 chars) before the edit, at most. + assert 7000 - 3 * 2 * 8 <= found.split_char < 7000 + assert cache.encode(tokenizer, edited) == tokenizer(edited)["input_ids"] + assert cache.divergent_hits == 1 and cache.resync_failures == 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) + + +def test_a_rewritten_turn_splices_from_its_history_not_a_shorter_shared_opening() -> None: + tokenizer, cache = _MergeTokenizer(), _cache() + opening = _PREAMBLE + history = _turns(20)[-1] + # Another conversation that so far holds only the shared opening. + for prompt in (history, opening): + assert cache.encode(tokenizer, prompt) == tokenizer(prompt)["input_ids"] + rewritten = history[:-20] + "a rewritten ending. " + # The rewritten turn extends the opening's entry too. + assert any(entry.text == opening for entry in cache._entries.values()) + found = cache._lookup(rewritten) + assert found is not None and found.entry.text == history and not found.extends + divergent_hits = cache.divergent_hits + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + assert cache.divergent_hits == divergent_hits + 1 + + +def test_a_tail_shifted_inside_a_run_of_spaces_is_rejected() -> None: + """A shift the resync span cannot see is caught at the next checkpoint.""" + tokenizer, cache = _PrependingTokenizer(), _cache() + run_end = len(_PREAMBLE) + 400 + history = _PREAMBLE + " " * 400 + "end of the padded tool output. " * 4 + for kept in range(12): + cache = _cache() + assert cache.encode(tokenizer, history) == tokenizer(history)["input_ids"] + # Diverge just after the run, so the split lands inside it. + rewritten = history[: run_end + kept] + "#rewritten" + found = cache._lookup(rewritten) + assert found is not None and run_end - 100 < found.split_char < run_end + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + assert cache.resync_failures == 1 + assert not cache.disabled + + +def test_a_tokenizer_that_shifts_every_tail_gets_no_cache() -> None: + assert create_prefix_token_cache(_PrependingTokenizer()) is None + for tokenizer in (_MergeTokenizer(), _ExpandingTokenizer(), _UnsyncableTokenizer()): + assert create_prefix_token_cache(tokenizer) is not None + + +@pytest.mark.parametrize("mismatch", [0, 1, 4095, 4096, 4097, 6000, 8191, None]) +def test_common_prefix_len_finds_the_first_mismatch(mismatch: int | None) -> None: + a = "".join(chr(97 + (i * 7) % 26) for i in range(8192)) + b = a if mismatch is None else a[:mismatch] + "#" + a[mismatch + 1 :] + expected = len(a) if mismatch is None else mismatch + assert _common_prefix_len(a, b) == _common_prefix_len(b, a) == expected + + +@pytest.mark.parametrize("n", [0, 4095, 4096, 4097, 8192]) +def test_common_prefix_len_of_a_prefix_is_its_length(n: int) -> None: + a = "".join(chr(97 + (i * 7) % 26) for i in range(8192)) + assert _common_prefix_len(a, a[:n]) == _common_prefix_len(a[:n], a) == n + + +def test_a_split_is_moved_to_the_first_token_of_a_multi_token_character() -> None: + tokenizer, cache = _ExpandingTokenizer(), _cache() + text = _PREAMBLE + for i in range(12): + # The odd character shifts every later split onto the other token of a pair. + text += "ff" * 40 + "." * (i % 2) + assert cache.encode(tokenizer, text) == tokenizer(text)["input_ids"] + rewritten = text[:-10] + "rewritten. " + assert cache.encode(tokenizer, rewritten) == tokenizer(rewritten)["input_ids"] + # The splits are usable, not rejected at the seam. + assert cache.resync_failures == 0 + assert cache.divergent_hits > 0 + assert not cache.disabled + _assert_checkpoints_are_valid(cache, tokenizer) + + def test_short_prompts_are_tokenized_but_not_cached() -> None: tokenizer, cache = _MergeTokenizer(), _cache() prompt = "the cab" @@ -160,13 +343,14 @@ def test_conversations_sharing_a_long_preamble_find_their_own_entry() -> None: for prompt in turns.values(): cache.encode(tokenizer, prompt) assert len(cache._buckets) == 1 + hits_before = cache.hits for i, prompt in turns.items(): # The next turn keeps the previous prompt, generation suffix included. extended = prompt + f"reply to {i}.<|user|>\nmore from {i}.{suffix}" found = cache._lookup(extended) assert found is not None and found[1].text == prompt assert cache.encode(tokenizer, extended) == tokenizer(extended)["input_ids"] - assert cache.hits == 16 and cache.resync_failures == 0 + assert cache.hits - hits_before == 16 and cache.resync_failures == 0 def test_rewritten_end_of_cached_prompt_reuses_its_prefix() -> None: @@ -257,6 +441,22 @@ def test_error_disables_cache_and_still_returns_ids() -> None: assert not cache._entries +def test_a_prompt_the_tokenizer_rejects_fails_without_disabling_the_cache() -> None: + tokenizer, cache = _RejectingTokenizer(), _cache() + turns = _turns(3) + cache.encode(tokenizer, turns[0]) + with pytest.raises(TypeError): + cache.encode(tokenizer, turns[1] + "\ud800") + assert not cache.disabled + assert cache.encode(tokenizer, turns[2]) == tokenizer(turns[2])["input_ids"] + assert cache.hits == 1 + + +def test_checkpoints_must_be_wider_than_the_resync_span() -> None: + with pytest.raises(ValueError, match="checkpoint_tokens"): + PrefixTokenCacheConfig(resync=32, checkpoint_tokens=32) + + def test_config_from_env_reads_overrides(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_ENTRIES", "7") monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", "100") diff --git a/tests/unittest/llmapi/test_serve_frontend_affinity.py b/tests/unittest/llmapi/test_serve_frontend_affinity.py index 90569b6b8750..78a3c963ee29 100644 --- a/tests/unittest/llmapi/test_serve_frontend_affinity.py +++ b/tests/unittest/llmapi/test_serve_frontend_affinity.py @@ -106,6 +106,7 @@ def fake_spawn(llm, num_frontends, *, extra_env, report_failure, affinity_socket class FakeServer: def __init__(self, **kwargs): + served["multi_frontend_serving"] = kwargs["multi_frontend_serving"] self.app = Mock() def __call__(self, host, port, sockets): @@ -128,6 +129,13 @@ def __call__(self, host, port, sockets): ) assert served["ports"] == [base, base + stride] + # Tokenization routing needs the launcher's frontend id and the frontend count. + serving = served["multi_frontend_serving"] + assert (serving.is_launcher, serving.frontend_id, serving.num_frontends) == ( + True, + 0, + num_frontends, + ) assert served["uds"] == [str(tmp_path / serve.LAUNCHER_UDS_NAME)] assert handed_off == {1: base + 2 * stride, 2: base + 3 * stride} # Everything the launcher bound is released once it stops serving. diff --git a/tests/unittest/llmapi/test_serve_tokenization_owners.py b/tests/unittest/llmapi/test_serve_tokenization_owners.py new file mode 100644 index 000000000000..b2abe7491719 --- /dev/null +++ b/tests/unittest/llmapi/test_serve_tokenization_owners.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Each conversation is tokenized by the frontend that owns it. + +Two frontends are simulated in one process: each has its own cache and serves it +on its own socket, as two trtllm-serve frontend processes do. +""" + +import contextlib +import os +import socket +import tempfile +import threading +import time +from types import SimpleNamespace +from typing import Any, Iterator +from unittest.mock import patch + +import pytest + +from tensorrt_llm.inputs import prefix_token_cache +from tensorrt_llm.inputs.prefix_token_cache import ( + PrefixTokenCache, + PrefixTokenCacheConfig, + tokenization_routing_key, +) +from tensorrt_llm.serve import tokenization_owners +from tensorrt_llm.serve.multi_frontend import ( + MultiFrontendServing, + frontend_of_conversation, + tokenization_uds_path, +) +from tensorrt_llm.serve.tokenization_owners import ( + ENABLE_ENV_VAR, + TokenizationOwners, + routing_enabled, + start_tokenization_owners, +) + +pytestmark = pytest.mark.cpu_only + + +class _CharTokenizer: + is_fast = True + + def __init__(self) -> None: + self.calls = 0 + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + self.calls += 1 + return { + "input_ids": [ord(c) for c in text], + "offset_mapping": [(i, i + 1) for i in range(len(text))], + } + + +class _RejectingTokenizer(_CharTokenizer): + """Rejects prompts containing ``MARKER``, as HF tokenizers reject a lone surrogate.""" + + MARKER = "" + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + if self.MARKER in text: + self.calls += 1 + raise TypeError("cannot tokenize") + return super().__call__(text, **kwargs) + + +class _SlowTokenizer(_CharTokenizer): + def __init__(self, delay_s: float) -> None: + super().__init__() + self.delay_s = delay_s + self.entered = threading.Event() + self.release = threading.Event() + + def __call__(self, text: str, **kwargs: Any) -> dict[str, Any]: + self.entered.set() + self.release.wait(self.delay_s) + return super().__call__(text, **kwargs) + + +def _cache() -> PrefixTokenCache: + return PrefixTokenCache( + PrefixTokenCacheConfig( + max_entries=64, overlap=4, resync=2, min_chars=64, checkpoint_tokens=8 + ) + ) + + +def _turns(n: int = 6) -> list[str]: + text, out = "system preamble " * 8, [] + for i in range(n): + text += f"turn {i} says something new. " + out.append(text) + return out + + +def _key_owned_by(owner: int, num_frontends: int, skip: int = 0) -> str: + keys = ( + key + for key in (f"conversation-{i}" for i in range(100)) + if frontend_of_conversation(key, num_frontends) == owner + ) + for _ in range(skip): + next(keys) + return next(keys) + + +@pytest.fixture +def ipc_dir() -> Iterator[str]: + # Short, so socket paths stay under the AF_UNIX path limit. + with tempfile.TemporaryDirectory(dir="/tmp") as path: + yield path + + +@contextlib.contextmanager +def _owners(ipc_dir: str, *specs: tuple[int, PrefixTokenCache, Any], **kwargs: Any): + started: list[TokenizationOwners] = [] + try: + for frontend_id, cache, tokenizer in specs: + started.append(TokenizationOwners(ipc_dir, frontend_id, 2, cache, tokenizer, **kwargs)) + yield started + finally: + for owner in started: + owner.close() + + +def test_a_conversation_is_tokenized_by_its_owner_wherever_its_turns_arrive(ipc_dir) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + with _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)): + key = _key_owned_by(1, 2) + turns = _turns() + # Turns alternate between the frontends, as SO_REUSEPORT spreads them. + for i, prompt in enumerate(turns): + with tokenization_routing_key(key): + assert caches[i % 2].encode(tokenizer, prompt) == [ord(c) for c in prompt] + + assert caches[0].routed == len(turns) // 2 + assert caches[0].hits + caches[0].misses == 0 + # The owner saw every turn, so every turn after the first hit. + assert caches[1].hits == len(turns) - 1 + + # A prompt without a conversation stays where it arrived. + caches[0].encode(tokenizer, turns[0]) + assert caches[0].misses == 1 + + +def test_frontends_share_one_connection_per_owner(ipc_dir) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + with _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)) as owners: + prompts = _turns(8) + errors: list[BaseException] = [] + + def worker(thread: int) -> None: + try: + with tokenization_routing_key(_key_owned_by(1, 2, skip=thread)): + for prompt in prompts: + assert caches[0].encode(tokenizer, prompt) == [ord(c) for c in prompt] + except BaseException as exc: # surface in the main thread + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + assert caches[0].routed == 4 * len(prompts) + assert len(owners[1]._server.handlers) == 1 + + +def test_a_prompt_is_tokenized_locally_while_its_owner_is_not_serving(ipc_dir) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + prompt = _turns(1)[0] + with _owners(ipc_dir, (0, cache, tokenizer)) as owners: + with tokenization_routing_key(_key_owned_by(1, 2)): + # Frontend 1 has not started. + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures, cache.misses) == (0, 1, 1) + + # Once it serves, it is asked at once: a refused connect is not a timeout. + owners.append(TokenizationOwners(ipc_dir, 1, 2, _cache(), tokenizer)) + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures) == (1, 1) + assert "1 route failures" in cache.summary() + + +def test_a_prompt_the_owner_cannot_tokenize_disables_no_cache(ipc_dir) -> None: + tokenizer, owner_tokenizer = _RejectingTokenizer(), _RejectingTokenizer() + cache, owner_cache = _cache(), _cache() + with _owners(ipc_dir, (0, cache, tokenizer), (1, owner_cache, owner_tokenizer)) as owners: + turns = _turns(3) + rejected = turns[0] + _RejectingTokenizer.MARKER + with tokenization_routing_key(_key_owned_by(1, 2)): + assert cache.encode(tokenizer, turns[0]) == [ord(c) for c in turns[0]] + connection = owners[0]._peers[1].conn + calls = owner_tokenizer.calls + + with pytest.raises(TypeError): + cache.encode(tokenizer, rejected) + # The owner tried it, with and without the cache, and so did the + # asking frontend. + assert owner_tokenizer.calls == calls + 2 + assert tokenizer.calls == 2 + assert not cache.disabled and not owner_cache.disabled + assert cache.route_failures == 1 + + # The connection survives, and the conversation still hits at its owner. + for prompt in turns[1:]: + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert owners[0]._peers[1].conn is connection + assert cache.routed == 3 + assert owner_cache.hits == 2 + + +def test_a_hung_owner_is_skipped_after_it_times_out(ipc_dir) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + timeout_s = 0.5 + # Frontend 1 accepts connections but never answers. + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as hung: + hung.bind(tokenization_uds_path(ipc_dir, 1)) + hung.listen() + with ( + _owners(ipc_dir, (0, cache, tokenizer), timeout_s=timeout_s), + patch.object(tokenization_owners.logger, "warning") as warning, + ): + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + elapsed = [] + for _ in range(3): + start = time.monotonic() + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + elapsed.append(time.monotonic() - start) + # Only the first prompt waited for the owner. + assert elapsed[0] >= timeout_s + assert max(elapsed[1:]) < timeout_s / 2 + assert (cache.routed, cache.route_failures) == (0, 3) + # Logged once, when the owner started being skipped. + assert warning.call_count == 1 + assert "Frontend 1" in warning.call_args.args[0] + + +def test_close_does_not_wait_long_for_a_handler_in_the_middle_of_a_prompt(ipc_dir) -> None: + tokenizer, slow = _CharTokenizer(), _SlowTokenizer(delay_s=3.0) + cache = _cache() + owner = TokenizationOwners(ipc_dir, 1, 2, _cache(), slow) + with _owners(ipc_dir, (0, cache, tokenizer)): + prompt = _turns(1)[0] + + def ask() -> None: + with tokenization_routing_key(_key_owned_by(1, 2)): + cache.encode(tokenizer, prompt) + + asking = threading.Thread(target=ask, daemon=True) + asking.start() + assert slow.entered.wait(5.0) + handlers = list(owner._server.handlers.values()) + start = time.monotonic() + owner.close(timeout_s=0.2) + # Up to the serve_forever poll interval (0.5 s) plus the shared deadline. + assert time.monotonic() - start < 1.5 + # close() left the handler tokenizing; finish it so no thread outlives the test. + slow.release.set() + for thread in [asking, *handlers]: + thread.join(5.0) + assert not thread.is_alive() + + +def test_routed_prompts_count_toward_the_periodic_log_line(ipc_dir) -> None: + tokenizer = _CharTokenizer() + caches = [_cache(), _cache()] + with ( + _owners(ipc_dir, (0, caches[0], tokenizer), (1, caches[1], tokenizer)), + patch.object(prefix_token_cache, "LOG_EVERY_REQUESTS", 2), + patch.object(prefix_token_cache.logger, "info") as info, + ): + with tokenization_routing_key(_key_owned_by(1, 2)): + for prompt in _turns(4): + caches[0].encode(tokenizer, prompt) + assert caches[0].requests == 4 + # Two lines from frontend 0, whose prompts were all routed, and two from the owner. + lines = [call.args[0] for call in info.call_args_list] + assert sum("4 routed" in line or "2 routed" in line for line in lines) == 2 + assert len(lines) == 4 + + +def test_close_removes_the_route(ipc_dir) -> None: + tokenizer, cache = _CharTokenizer(), _cache() + with _owners(ipc_dir, (0, cache, tokenizer)): + pass + prompt = _turns(1)[0] + with tokenization_routing_key(_key_owned_by(1, 2)): + assert cache.encode(tokenizer, prompt) == [ord(c) for c in prompt] + assert (cache.routed, cache.route_failures, cache.misses) == (0, 0, 1) + + +def test_the_owner_of_a_conversation_is_fixed() -> None: + # Disaggregated servers pick the affinity port with the same mapping, in + # other processes, so it must not change (for example to the seeded hash()). + assert [frontend_of_conversation(f"conversation-{i}", 8) for i in range(4)] == [2, 4, 6, 0] + assert frontend_of_conversation("session-a", 3) == 1 + + +def test_a_conversation_id_with_a_lone_surrogate_has_an_owner() -> None: + # JSON can carry one; it must not fail the request it arrives with. + assert frontend_of_conversation("conversation-\ud800", 2) in (0, 1) + + +@pytest.mark.parametrize("value, enabled", [(None, False), ("0", False), ("1", True)]) +def test_routing_switch_values(monkeypatch: pytest.MonkeyPatch, value: Any, enabled: bool) -> None: + if value is None: + monkeypatch.delenv(ENABLE_ENV_VAR, raising=False) + else: + monkeypatch.setenv(ENABLE_ENV_VAR, value) + assert routing_enabled() is enabled + + +@pytest.mark.parametrize("value", ["true", "on", "1 ", "", "2"]) +def test_a_malformed_routing_switch_raises(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, value) + with pytest.raises(ValueError, match=ENABLE_ENV_VAR): + routing_enabled() + + +def _processor(cache: Any) -> SimpleNamespace: + tokenizer = _CharTokenizer() + return SimpleNamespace( + get_prefix_token_cache=lambda: None if cache is None else (cache, tokenizer) + ) + + +def _serving(ipc_dir: str, num_frontends: int, frontend_id: int = 0) -> MultiFrontendServing: + return MultiFrontendServing( + launcher_uds=os.path.join(ipc_dir, "launcher.sock"), + is_launcher=frontend_id == 0, + frontend_id=frontend_id, + num_frontends=num_frontends, + ) + + +def test_routing_starts_when_switched_on_with_frontends_and_a_cache( + monkeypatch: pytest.MonkeyPatch, ipc_dir: str +) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, "1") + cache = _cache() + with patch.object(tokenization_owners.logger, "info") as info: + owners = start_tokenization_owners(_serving(ipc_dir, 4, frontend_id=2), _processor(cache)) + try: + assert owners is not None + assert os.path.exists(tokenization_uds_path(ipc_dir, 2)) + assert "frontend 2 of 4" in info.call_args.args[0] + finally: + owners.close() + + +@pytest.mark.parametrize("num_frontends, has_cache", [(1, True), (4, False)]) +def test_routing_warns_when_switched_on_but_it_cannot_engage( + monkeypatch: pytest.MonkeyPatch, ipc_dir: str, num_frontends: int, has_cache: bool +) -> None: + monkeypatch.setenv(ENABLE_ENV_VAR, "1") + processor = _processor(_cache() if has_cache else None) + with patch.object(tokenization_owners.logger, "warning") as warning: + assert start_tokenization_owners(_serving(ipc_dir, num_frontends), processor) is None + assert start_tokenization_owners(None, processor) is None + assert warning.call_count == 2 + assert ENABLE_ENV_VAR in warning.call_args.args[0] + assert not os.listdir(ipc_dir) + + +def test_routing_is_off_by_default(monkeypatch: pytest.MonkeyPatch, ipc_dir: str) -> None: + monkeypatch.delenv(ENABLE_ENV_VAR, raising=False) + assert start_tokenization_owners(_serving(ipc_dir, 4), _processor(_cache())) is None + assert not os.listdir(ipc_dir)