From 8526800e7e0728d9cf0fc5cdf4df468cd426c2d1 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Fri, 7 Aug 2026 16:01:21 -0700 Subject: [PATCH 01/11] Count structured chat tokens accurately --- .../metrics_aggregator/metrics_table.py | 21 +- .../metrics_aggregator/token_metrics.py | 192 ++++++++++++++---- src/inference_endpoint/core/types.py | 7 +- .../load_generator/session.py | 30 +-- .../services/metrics_aggregator/conftest.py | 17 ++ .../metrics_aggregator/test_aggregator.py | 57 +++++- .../metrics_aggregator/test_metrics_table.py | 41 ++++ .../metrics_aggregator/test_token_metrics.py | 127 +++++++++++- tests/unit/core/test_record.py | 23 +++ .../unit/load_generator/test_async_session.py | 47 +++++ 10 files changed, 482 insertions(+), 80 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index 536c793b9..4960e304c 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -186,8 +186,8 @@ class TokenTrigger(EmitTrigger): Subclasses implement ``_extract_text()`` to pull the text to tokenize from the event record, and may override ``_extract_message()`` to return - (content, reasoning, tool_calls) for chat-template–aware tokenization when - tool calls are present. ``fire()`` does not tokenize inline — it enqueues + (content, reasoning, tool_calls) for chat-template-aware tokenization of + structured output. ``fire()`` does not tokenize inline — it enqueues the work plus a recorder callback onto the shared ``TokenBatchQueue``, which the aggregator flushes in batches. ``_compute_value()`` can transform the token count before it is recorded. @@ -305,7 +305,7 @@ def __init__(self, registry: MetricsRegistry): class IslTrigger(TokenTrigger): - """ISL from PromptData: ``len(token_ids)`` or the tokenized prompt text.""" + """ISL from token IDs, structured chat messages, or plain prompt text.""" def __init__( self, @@ -319,6 +319,13 @@ def fire(self, ev_rec, row, pre_change): if isinstance(ev_rec.data, PromptData) and ev_rec.data.token_ids is not None: self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) return + if isinstance(ev_rec.data, PromptData) and ev_rec.data.messages is not None: + if self._queue is not None: + self._queue.enqueue_prompt( + (ev_rec.data.messages, ev_rec.data.tools), + self._make_recorder(ev_rec, pre_change), + ) + return # Text path: tokenize raw prompt text — used when token_ids are # unavailable (e.g. OpenAI-compatible endpoints). Enqueued by the base. super().fire(ev_rec, row, pre_change) @@ -349,7 +356,9 @@ def _extract_text(self, ev_rec, row, pre_change): return None def _extract_message(self, ev_rec, row, pre_change): - if isinstance(ev_rec.data, TextModelOutput) and ev_rec.data.tool_calls: + if isinstance(ev_rec.data, TextModelOutput) and ( + ev_rec.data.reasoning or ev_rec.data.tool_calls + ): return ev_rec.data.as_message_parts() return None @@ -391,7 +400,9 @@ def _extract_text(self, ev_rec, row, pre_change): def _extract_message(self, ev_rec, row, pre_change): if pre_change.get(SampleField.RECV_FIRST_NS) is None: return None - if isinstance(ev_rec.data, TextModelOutput) and ev_rec.data.tool_calls: + if isinstance(ev_rec.data, TextModelOutput) and ( + ev_rec.data.reasoning or ev_rec.data.tool_calls + ): return ev_rec.data.as_message_parts_after_first_chunk() return None diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 4d14ead7b..3dca8ec53 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -17,13 +17,15 @@ ``BatchTokenizer`` tokenizes whole batches at once, sharded across worker processes each pinned to a block of ``CORES_PER_WORKER`` cores (a single BPE -rayon pool is memory-bound and saturates ~8 cores). The aggregator buffers +backend pool is memory-bound and saturates ~8 cores). The aggregator buffers per-sample text. The sharded pool is the drain-phase accelerator and is auto-sized (one shard per core block); live mid-run flushes run on a small in-process thread pool (``--tokenizer-workers``, default 2) owned by the -queue's live loop. A tokenizer without a fast (Rust) backend is a startup -error, never a silent slow path. Platforms without CPU affinity (e.g. macOS) -shard unpinned at full speed; only cache/NUMA locality is lost. +queue's live loop. Hugging Face fast tokenizers use their Rust backend; +tokenizers that expose tiktoken through custom Python code use that official +wrapper so its preprocessing semantics are preserved. Platforms without CPU +affinity (e.g. macOS) shard unpinned at full speed; only cache/NUMA locality is +lost. """ from __future__ import annotations @@ -111,6 +113,37 @@ def _normalize_tool_calls_for_template( _WORKER_BACKEND: Any = None +class _PythonTokenizerBackend: + """Use a custom tokenizer's public API instead of bypassing its Python logic.""" + + def __init__(self, tokenizer: Any) -> None: + self.tokenizer = tokenizer + + def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: + return self.tokenizer.encode(text, add_special_tokens=add_special_tokens) + + def encode_lengths(self, texts: list[str]) -> list[int]: + return [len(self.encode(text, add_special_tokens=False)) for text in texts] + + +def _backend_from_tokenizer(tokenizer: Any) -> Any | None: + """Return the tokenizer's supported length-counting path.""" + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is not None: + return backend + + model = getattr(tokenizer, "model", None) + model_type = type(model) + if ( + model_type.__module__ == "tiktoken.core" + and model_type.__name__ == "Encoding" + and callable(getattr(model, "encode", None)) + and callable(getattr(model, "encode_batch", None)) + ): + return _PythonTokenizerBackend(tokenizer) + return None + + def load_reference_tokenizer(tokenizer_name: str) -> Any: """Load the run's reference tokenizer. @@ -123,29 +156,31 @@ def load_reference_tokenizer(tokenizer_name: str) -> Any: def load_reference_backend(tokenizer_name: str) -> Any | None: - """Raw tokenizers backend (fast Rust path) for length counting. + """Supported token-counting path for the reference tokenizer. - Counting through the backend avoids the transformers "sequence longer than - model_max_length" warning the Python wrapper emits, so no ``model_max_length`` - override is needed. ``None`` if the tokenizer has no fast backend. + Hugging Face fast tokenizers use their native backend. Custom tiktoken + tokenizers use their public Python API so model-specific preprocessing is + not reimplemented here. ``None`` if neither path is available. """ - return getattr(load_reference_tokenizer(tokenizer_name), "backend_tokenizer", None) + return _backend_from_tokenizer(load_reference_tokenizer(tokenizer_name)) def _init_worker(tokenizer_name: str, core_set: list[int]) -> None: - """Pin this worker to ``core_set``, then load the raw tokenizers backend. + """Pin this worker to ``core_set``, then load its token-counting path. - Affinity is set before the first encode so the Rust rayon pool sizes itself - to the pinned core count (num_cpus respects sched_getaffinity on Linux). + Affinity is set before the first encode so the Hugging Face rayon pool sizes + itself to the pinned core count (num_cpus respects sched_getaffinity on + Linux). Custom wrappers scale through these process shards. """ # Ctrl-C sends SIGINT to the whole foreground process group; the parent # drives worker shutdown, so a worker dying mid-drain would break the pool # and lose the buffered tokenizations it was counting. signal.signal(signal.SIGINT, signal.SIG_IGN) if core_set: - # Size the rayon pool to the block explicitly: the parent process caps - # its own pool for the live lane, and spawn children inherit that env — - # without the override every shard would run at the live-lane width. + # Size the Hugging Face rayon pool to the block explicitly: the parent + # process caps its own pool for the live lane, and spawn children inherit + # that env — without the override every shard would run at the live-lane + # width. os.environ["RAYON_NUM_THREADS"] = str(len(core_set)) try: os.sched_setaffinity(0, set(core_set)) @@ -164,9 +199,12 @@ def _init_worker(tokenizer_name: str, core_set: list[int]) -> None: def encode_lengths(backend: Any, texts: list[str]) -> list[int]: - """Per-text token counts via the raw tokenizers backend, one rayon call.""" + """Per-text token counts via one bounded backend batch call.""" + if isinstance(backend, _PythonTokenizerBackend): + return backend.encode_lengths(texts) encode_batch = getattr(backend, "encode_batch_fast", None) or backend.encode_batch - return [len(e.ids) for e in encode_batch(texts, add_special_tokens=False)] + encoded = encode_batch(texts, add_special_tokens=False) + return [len(getattr(item, "ids", item)) for item in encoded] def _worker_encode_lengths(texts: list[str]) -> list[int]: @@ -233,10 +271,8 @@ def __init__( n_workers: int = -1, ) -> None: self._tokenizer_name = tokenizer_name - # The live lane runs in-process: cap this process's rayon pool so a - # mid-run batched encode uses ~live_workers cores, not the whole - # machine. Must be set before the first encode initializes the pool; - # setdefault lets an operator-exported RAYON_NUM_THREADS win. + # The live lane runs in-process: cap the Hugging Face rayon pool before + # its first encode. setdefault lets an operator-exported HF cap win. os.environ.setdefault("RAYON_NUM_THREADS", str(max(1, live_workers))) self._fallback_warned: set[str] = set() self._tokenizer: PreTrainedTokenizerBase | None = None @@ -260,25 +296,26 @@ def __init__( def _load_tokenizer(self) -> None: tok = load_reference_tokenizer(self._tokenizer_name) self._tokenizer = tok + self._backend = _backend_from_tokenizer(tok) # Baseline = tokens from a [user, empty-assistant] pair minus the [user] # prefix alone, so the assistant frame is subtracted from message counts. try: prefix = cast( - str, + list[int], tok.apply_chat_template( - [_PREFIX_USER_MSG], tokenize=False, add_generation_prompt=False + [_PREFIX_USER_MSG], tokenize=True, add_generation_prompt=False ), ) - self._prefix_len = len(tok.tokenize(prefix)) + self._prefix_len = len(prefix) with_assistant = cast( - str, + list[int], tok.apply_chat_template( [_PREFIX_USER_MSG, {"role": "assistant", "content": ""}], - tokenize=False, + tokenize=True, add_generation_prompt=False, ), ) - self._baseline = len(tok.tokenize(with_assistant)) - self._prefix_len + self._baseline = len(with_assistant) - self._prefix_len except Exception: self._prefix_len = 0 self._baseline = 0 @@ -295,19 +332,17 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: (``< 0``) fits one shard per ``cores_per_worker`` block of this process's affinity mask (or the online CPU count when the platform has no affinity API — shards then run unpinned), always at least one; - an explicit count is clamped to that capacity. An environment that - cannot shard — no fast Rust backend, a warmup that fails or exceeds - its budget — raises instead of silently degrading to a slow path - that cannot keep up with completions. + an explicit count is clamped to that capacity. An unsupported tokenizer + or a shard warmup that fails or exceeds its budget raises at startup. """ if cores_per_worker <= 0 or n_workers == 0: logger.info("BatchTokenizer: in-process tokenization (explicit)") return - if getattr(self._tokenizer, "backend_tokenizer", None) is None: + if self._backend is None: raise RuntimeError( - f"tokenizer {self._tokenizer_name!r} has no fast (Rust) " - "backend; token metrics require one to keep up with " - "completions. Use a fast tokenizer, or disable token metrics." + f"tokenizer {self._tokenizer_name!r} has no supported " + "token-counting path; use a Hugging Face fast tokenizer or a " + "supported custom tiktoken tokenizer." ) # The full allowed CPU universe (cgroup-clamped) drives the shard block # math. cgroup_clamped_cpus owns the probe-and-restore of this process's @@ -365,7 +400,7 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: def _encode_lengths_inproc(self, texts: list[str]) -> list[int]: tok = self._tokenizer - backend = getattr(tok, "backend_tokenizer", None) + backend = self._backend if backend is not None: return encode_lengths(backend, texts) return [len(tok.tokenize(t)) for t in texts] # type: ignore[union-attr] @@ -424,11 +459,10 @@ def _token_count_message( if tool_calls: msg["tool_calls"] = _normalize_tool_calls_for_template(tool_calls) try: - rendered = tok.apply_chat_template( # type: ignore[union-attr] - [_PREFIX_USER_MSG, msg], tokenize=False, add_generation_prompt=False + encoded = tok.apply_chat_template( # type: ignore[union-attr] + [_PREFIX_USER_MSG, msg], tokenize=True, add_generation_prompt=False ) - full = len(tok.tokenize(rendered)) # type: ignore[union-attr] - return max(0, full - self._prefix_len - self._baseline) + return max(0, len(encoded) - self._prefix_len - self._baseline) except Exception as exc: key = f"{self._tokenizer_name}:{type(exc).__name__}" if key not in self._fallback_warned: @@ -462,6 +496,36 @@ async def token_count_message_async( self._thread, self._token_count_message, content, reasoning, tool_calls ) + def _token_count_prompt( + self, + messages: tuple[dict[str, Any], ...], + tools: tuple[dict[str, Any], ...] | None, + ) -> int: + kwargs: dict[str, Any] = { + "tokenize": True, + "add_generation_prompt": True, + } + if tools is not None: + kwargs["tools"] = list(tools) + encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] + list(messages), **kwargs + ) + return len(encoded) + + async def token_count_prompt_async( + self, + messages: tuple[dict[str, Any], ...], + tools: tuple[dict[str, Any], ...] | None, + loop: asyncio.AbstractEventLoop, + /, + ) -> int: + """Complete chat-prompt token count without blocking the loop.""" + if self._thread is None: + raise RuntimeError("BatchTokenizer is closed") + return await loop.run_in_executor( + self._thread, self._token_count_prompt, messages, tools + ) + def close(self) -> None: """Shut down all workers. Idempotent. @@ -484,9 +548,10 @@ def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: self.close() -# Type alias for the (content, reasoning, tool_calls) tuple a message trigger -# enqueues for chat-template tokenization. +# Message means the structured assistant output used for OSL/TPOT tokenization. MessageParts = tuple[str, str | None, tuple[dict[str, Any], ...] | None] +# Prompt means the complete chat input and tool definitions used for ISL. +PromptParts = tuple[tuple[dict[str, Any], ...], tuple[dict[str, Any], ...] | None] class TokenCounter(Protocol): @@ -519,12 +584,22 @@ async def token_count_message_async( """Chat-template token count for one assistant message.""" raise NotImplementedError + async def token_count_prompt_async( + self, + messages: tuple[dict[str, Any], ...], + tools: tuple[dict[str, Any], ...] | None, + loop: asyncio.AbstractEventLoop, + /, + ) -> int: + """Chat-template token count for one complete input prompt.""" + raise NotImplementedError + class TokenBatchQueue: """Buffers per-sample tokenization work and clears it in batches. - Triggers call ``enqueue_text`` / ``enqueue_message`` at event time with an - ``on_count`` callback that records the resulting metric. The queue owns + Triggers enqueue plain text, assistant messages, or complete chat prompts + at event time with a callback that records the resulting metric. The queue owns its own flush cadence: ``start_live`` begins a periodic flush through the tokenizer's bounded live lane (so live ISL/OSL/TPOT stay current without touching the benchmark's cores), and ``flush_remaining`` drains everything @@ -542,6 +617,7 @@ def __init__( self._loop = loop self._text: list[tuple[str, Callable[[int], None]]] = [] self._msg: list[tuple[MessageParts, Callable[[int], None]]] = [] + self._prompt: list[tuple[PromptParts, Callable[[int], None]]] = [] self._inflight = 0 self._live_task: asyncio.Task | None = None # Serializes flushes so the periodic live flush and the end-of-run @@ -590,6 +666,12 @@ def enqueue_message( self._inflight += 1 self._msg.append((parts, on_count)) + def enqueue_prompt( + self, parts: PromptParts, on_count: Callable[[int], None] + ) -> None: + self._inflight += 1 + self._prompt.append((parts, on_count)) + async def flush_live_once(self) -> None: """One bounded mid-run flush (live lane). @@ -618,7 +700,7 @@ async def _flush(self, live: bool) -> None: in the next flush. Callers use ``flush_live_once`` / ``drain_all``. """ async with self._lock: - if not (self._text or self._msg): + if not (self._text or self._msg or self._prompt): return if live: cap = _LIVE_FLUSH_MAX_ITEMS @@ -626,9 +708,12 @@ async def _flush(self, live: bool) -> None: del self._text[:cap] # in-place: O(cap), not O(backlog). msg_items = self._msg[:cap] del self._msg[:cap] + prompt_items = self._prompt[:cap] + del self._prompt[:cap] else: text_items, self._text = self._text, [] msg_items, self._msg = self._msg, [] + prompt_items, self._prompt = self._prompt, [] # The text and message phases fail independently — they run on # separate executors, so a dead text shard must not drop message # items that would still succeed (and vice versa). The first @@ -643,6 +728,7 @@ async def _flush(self, live: bool) -> None: if live: self._text[:0] = text_items self._msg[:0] = msg_items + self._prompt[:0] = prompt_items raise except Exception as exc: # noqa: BLE001 — isolate phases. failure = exc @@ -676,6 +762,7 @@ async def _flush(self, live: bool) -> None: except asyncio.CancelledError: if live: self._msg[:0] = msg_items[i:] + self._prompt[:0] = prompt_items raise except Exception as exc: # noqa: BLE001 — isolate items. failure = failure or exc @@ -683,6 +770,21 @@ async def _flush(self, live: bool) -> None: self._msg.append(((content, reasoning, tool_calls), on_count)) continue self._record(on_count, count) + for i, ((messages, tools), on_count) in enumerate(prompt_items): + try: + count = await self._tokenizer.token_count_prompt_async( + messages, tools, self._loop + ) + except asyncio.CancelledError: + if live: + self._prompt[:0] = prompt_items[i:] + raise + except Exception as exc: # noqa: BLE001 — isolate items. + failure = failure or exc + if live: + self._prompt.append(((messages, tools), on_count)) + continue + self._record(on_count, count) if failure is not None: raise failure diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index 6e1cbab46..5d3b39f35 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -284,17 +284,22 @@ class PromptData( ): # type: ignore[call-arg] """Prompt input data attached to ISSUED events for ISL computation. - Exactly one of ``text`` or ``token_ids`` should be set: + Exactly one of ``text``, ``token_ids``, or ``messages`` should be set: - ``text``: raw prompt string (OpenAI path) — requires tokenization for ISL. - ``token_ids``: pre-tokenized token ID list (SGLang/Harmonize path) — ISL is len(). + - ``messages``: structured chat history rendered with the model's template. Attributes: text: Raw prompt string. Set when the adapter sends text prompts. token_ids: Pre-computed token IDs. Set when the adapter pre-tokenizes (e.g. SGLang). + messages: Structured messages sent to a chat-completions endpoint. + tools: Tool declarations accompanying ``messages``. """ text: str | None = None token_ids: tuple[int, ...] | None = None + messages: tuple[dict[str, Any], ...] | None = None + tools: tuple[dict[str, Any], ...] | None = None class ErrorData( diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index d0d7ad897..d36782a16 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -263,19 +263,23 @@ def issue( prompt_data: PromptData if isinstance(data, dict): token_ids = data.get("input_tokens") or data.get("token_ids") - # Multimodal datasets store ``prompt`` as a list of OpenAI content - # parts (e.g. [{"type": "text", ...}, {"type": "image_url", ...}]) - # which the HTTP adapter handles directly. `PromptData.text` is only - # meaningful for ISL reporting on text-only prompts. - # Therefore, setting `text=None` for non-string prompts - # means that ISL reporting will be unavailable for multimodal samples. - prompt_text = data.get("prompt") - if prompt_text is None and "messages" in data: - prompt_text = _extract_prompt_text(data["messages"]) - prompt_data = PromptData( - text=prompt_text if isinstance(prompt_text, str) else None, - token_ids=tuple(token_ids) if token_ids is not None else None, - ) + # Prefer the exact representation sent to the endpoint: existing + # token IDs, then structured chat messages, then a plain prompt. + # Non-string standalone prompts (for example multimodal content + # parts) remain unavailable for ISL reporting. + if token_ids is not None: + prompt_data = PromptData(token_ids=tuple(token_ids)) + elif isinstance(data.get("messages"), list | tuple): + tools = data.get("tools") + prompt_data = PromptData( + messages=tuple(data["messages"]), + tools=tuple(tools) if isinstance(tools, list | tuple) else None, + ) + else: + prompt_text = data.get("prompt") + prompt_data = PromptData( + text=prompt_text if isinstance(prompt_text, str) else None + ) else: prompt_data = PromptData() self._publisher.publish( diff --git a/tests/unit/async_utils/services/metrics_aggregator/conftest.py b/tests/unit/async_utils/services/metrics_aggregator/conftest.py index aae7a07ac..0d95825a8 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/conftest.py +++ b/tests/unit/async_utils/services/metrics_aggregator/conftest.py @@ -87,6 +87,23 @@ async def token_count_message_async( combined = (content or "") + " " + (reasoning or "") + " " + tool_calls_str return len(combined.split()) + async def token_count_prompt_async( + self, + messages: tuple[dict, ...], + tools: tuple[dict, ...] | None, + _loop: asyncio.AbstractEventLoop, + ) -> int: + if self._delay: + await asyncio.sleep(self._delay) + parts = [ + str(message.get(key, "")) + for message in messages + for key in ("content", "reasoning_content", "tool_calls") + if message.get(key) + ] + parts.extend(str(tool) for tool in tools or ()) + return len(" ".join(parts).split()) + def close(self) -> None: pass diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py index ca2aa2c99..79b2d1cb3 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -492,9 +492,9 @@ async def test_duplicate_started_logs_error_and_preserves_state( error_records = [ r for r in caplog.records if "Duplicate STARTED" in r.message ] - assert ( - len(error_records) == 1 - ), "duplicate STARTED must log exactly one error" + assert len(error_records) == 1, ( + "duplicate STARTED must log exactly one error" + ) assert "1000" in error_records[0].getMessage() assert "5000" in error_records[0].getMessage() finally: @@ -879,6 +879,45 @@ async def test_isl_text_path_async(self, tmp_path): finally: agg.close() + @pytest.mark.asyncio + async def test_isl_structured_prompt_includes_reasoning_and_tools(self, tmp_path): + loop = asyncio.get_event_loop() + tokenizer = MockBatchTokenizer() + messages = ( + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "reasoning_content": "think carefully", + "content": None, + "tool_calls": ({"function": {"name": "lookup"}},), + }, + {"role": "tool", "content": "tool result"}, + ) + tools = ({"function": {"name": "lookup"}},) + with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: + agg, registry, _ = make_aggregator( + ctx, loop, "agg_isl_structured", tokenizer=tokenizer + ) + try: + await agg.process( + [ + session_event( + SessionEventType.START_PERFORMANCE_TRACKING, ts=0 + ), + sample_event( + SampleEventType.ISSUED, + "s1", + ts=1000, + data=PromptData(messages=messages, tools=tools), + ), + ] + ) + await agg._token_queue.drain_all() + assert snapshot_series_count(registry, MetricSeriesKey.ISL.value) == 1 + assert snapshot_series_total(registry, MetricSeriesKey.ISL.value) > 5 + finally: + agg.close() + @pytest.mark.asyncio async def test_osl_emitted_on_complete(self, tmp_path): """OSL is emitted via async tokenization when COMPLETE carries text.""" @@ -1084,9 +1123,9 @@ async def test_started_arms_the_live_flush_loop(self, tmp_path): assert agg._token_queue is not None assert agg._token_queue._live_task is not None await agg.process([session_event(SessionEventType.ENDED, ts=100)]) - assert ( - agg._token_queue._live_task is None - ), "drain must stop the live loop" + assert agg._token_queue._live_task is None, ( + "drain must stop the live loop" + ) finally: agg.close() @@ -1242,9 +1281,9 @@ async def token_count_message_async(self, *args): ] ) assert agg._token_queue is not None - assert ( - agg._token_queue.pending > 0 - ), "precondition: ISL must be buffered before ENDED" + assert agg._token_queue.pending > 0, ( + "precondition: ISL must be buffered before ENDED" + ) await agg.process([session_event(SessionEventType.ENDED, ts=2000)]) publisher.publish_final.assert_awaited_once() diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py index 016af2a93..ba60f1ffd 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py @@ -354,6 +354,47 @@ async def test_osl_without_tool_calls_uses_text_path(self): assert snapshot_series_count(registry, "osl") == 1 + async def test_osl_with_reasoning_uses_structured_message_path(self): + from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import ( + OslTrigger, + SampleRow, + ) + from inference_endpoint.core.types import TextModelOutput + + class CapturingTokenizer: + def __init__(self): + self.messages = [] + + async def count_texts_async(self, *args, **kwargs): + raise AssertionError("reasoning output must not use flattened text") + + async def token_count_message_async( + self, content, reasoning, tool_calls, _loop + ): + self.messages.append((content, reasoning, tool_calls)) + return 4 + + registry = MetricsRegistry() + registry.register_series("osl", hdr_low=1, hdr_high=100_000) + tokenizer = CapturingTokenizer() + queue = TokenBatchQueue(tokenizer, asyncio.get_running_loop()) + trigger = OslTrigger(registry, queue) + output = TextModelOutput(output="answer", reasoning="private reasoning") + + trigger.fire( + EventRecord( + event_type=SampleEventType.COMPLETE, + timestamp_ns=1000, + sample_uuid="s1", + data=output, + ), + SampleRow(sample_uuid="s1"), + {}, + ) + await queue.drain_all() + + assert tokenizer.messages == [("answer", "private reasoning", None)] + @pytest.mark.unit @pytest.mark.asyncio diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index 3587bc0e1..6bfe93905 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -139,7 +139,7 @@ class _FakeTokenizerWithTemplate(_FakeTokenizer): """Tokenizer that supports apply_chat_template for tool-call testing.""" def apply_chat_template( - self, messages, tokenize=False, add_generation_prompt=False + self, messages, tools=None, tokenize=False, add_generation_prompt=False ): # Simulate 2 wrapper tokens for the template frame. parts = ["WRAPPER", "WRAPPER"] @@ -153,6 +153,10 @@ def apply_chat_template( import msgspec parts.append(msgspec.json.encode(msg["tool_calls"]).decode()) + if tools: + parts.extend(tool["function"]["name"] for tool in tools) + if add_generation_prompt: + parts.append("GENERATION") rendered = " ".join(parts) if tokenize: return list(range(len(rendered.split()))) @@ -161,6 +165,41 @@ def apply_chat_template( @pytest.mark.unit class TestBatchTokenizerMessageTokenization: + @pytest.mark.asyncio + async def test_token_count_prompt_preserves_messages_tools_and_generation_prompt( + self, + ): + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + messages = ( + {"role": "user", "content": "question"}, + { + "role": "assistant", + "reasoning_content": "reasoning", + "content": None, + "tool_calls": ( + { + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + ), + }, + {"role": "tool", "content": "result"}, + ) + tools = ( + { + "type": "function", + "function": {"name": "lookup", "parameters": {}}, + }, + ) + + count = await tok.token_count_prompt_async(messages, tools, loop) + + # Wrapper + question + reasoning + tool call + tool result + + # declared tool + generation prompt. + assert count == 8 + @pytest.mark.asyncio async def test_token_count_message_subtracts_baseline(self): """token_count_message_async returns full_tokens - baseline.""" @@ -236,6 +275,31 @@ def encode_batch(self, texts, add_special_tokens=False): return [_Encoding(len(t.split())) for t in texts] +class Encoding: + """Stand-in that fails if code bypasses the tokenizer wrapper.""" + + __module__ = "tiktoken.core" + + def encode(self, *args, **kwargs): + raise AssertionError("raw tiktoken encode must not be called") + + def encode_batch(self, *args, **kwargs): + raise AssertionError("raw tiktoken encode_batch must not be called") + + +class _FakeTokenizerWithTikTokenModel(_FakeTokenizer): + """Kimi-shaped custom tokenizer backed by the Rust tiktoken core.""" + + def __init__(self, load_delay: float = 0.0): + super().__init__(load_delay) + self.model = Encoding() + self.encode_calls: list[tuple[str, bool]] = [] + + def encode(self, text, *, add_special_tokens=False): + self.encode_calls.append((text, add_special_tokens)) + return text.split() + + @pytest.mark.unit class TestEncodeHelpers: def test_encode_lengths_prefers_fast(self): @@ -263,6 +327,41 @@ def from_pretrained(name, **kwargs): assert captured["name"] == "m" assert captured["kwargs"].get("trust_remote_code") is True + def test_load_reference_backend_uses_kimi_python_wrapper(self, monkeypatch): + """Kimi counting must not bypass its Python tokenizer wrapper.""" + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(name, **kwargs): + assert name == "kimi" + assert kwargs == {"trust_remote_code": True} + return _FakeTokenizerWithTikTokenModel() + + monkeypatch.setattr(token_metrics_module, "AutoTokenizer", _FakeAutoTokenizer) + backend = token_metrics_module.load_reference_backend("kimi") + assert backend is not None + assert encode_lengths(backend, ["a b", "c"]) == [2, 1] + assert backend.tokenizer.encode_calls == [("a b", False), ("c", False)] + + def test_load_reference_backend_rejects_tiktoken_lookalike(self, monkeypatch): + class _LookalikeEncoding: + def encode(self, text, *, allowed_special): + return text.split() + + def encode_batch(self, texts, *, allowed_special): + return [text.split() for text in texts] + + class _FakeTok: + model = _LookalikeEncoding() + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(name, **kwargs): + return _FakeTok() + + monkeypatch.setattr(token_metrics_module, "AutoTokenizer", _FakeAutoTokenizer) + assert token_metrics_module.load_reference_backend("lookalike") is None + def test_worker_encode_lengths_raises_without_backend(self, monkeypatch): monkeypatch.setattr(token_metrics_module, "_WORKER_BACKEND", None) with pytest.raises(RuntimeError, match="backend unavailable"): @@ -343,6 +442,20 @@ def test_no_fast_backend_is_a_startup_error(self, monkeypatch): with pytest.raises(RuntimeError, match="fast"): BatchTokenizer("fake", live_workers=2) + def test_tiktoken_wrapper_is_supported_by_process_shards(self, monkeypatch): + monkeypatch.setattr( + token_metrics_module, "ProcessPoolExecutor", _SpawnlessExecutor + ) + monkeypatch.setattr( + token_metrics_module, "cgroup_clamped_cpus", lambda: list(range(16)) + ) + with patch(_MOCK_TARGET, _FakeTokenizerWithTikTokenModel): + with BatchTokenizer("kimi", live_workers=2) as tok: + assert len(tok._procs) == 2 + assert tok._backend is not None + assert tok._backend.tokenizer is tok._tokenizer + assert encode_lengths(tok._backend, ["a b", "c"]) == [2, 1] + def test_affinity_unavailable_shards_unpinned(self, monkeypatch): """No affinity API (e.g. macOS): shard from the CPU count, unpinned.""" monkeypatch.setattr( @@ -747,18 +860,18 @@ def test_terminate_procs_kills_running_workers(): assert future.running(), "worker task did not start" procs = list((getattr(ex, "_processes", None) or {}).values()) assert procs, "worker did not spawn" - assert not wait_for_process( - [p.sentinel for p in procs], timeout=0 - ), "worker exited before termination" + assert not wait_for_process([p.sentinel for p in procs], timeout=0), ( + "worker exited before termination" + ) _terminate_procs([ex]) for p in procs: # The executor manager may reap the child concurrently; sentinel # readiness observes exit without racing its return-code update. - assert wait_for_process( - [p.sentinel], timeout=5 - ), "worker was not terminated" + assert wait_for_process([p.sentinel], timeout=5), ( + "worker was not terminated" + ) finally: cleanup_procs = procs or list((getattr(ex, "_processes", None) or {}).values()) for p in cleanup_procs: diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index cd9173f6a..dc137dd70 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -172,6 +172,29 @@ def test_sample_event_round_trips_with_prompt_data_token_ids(self): assert decoded.data.token_ids == (101, 202, 303) assert decoded.data.text is None + def test_sample_event_round_trips_with_structured_prompt_data(self): + messages = ( + {"role": "user", "content": "question"}, + { + "role": "assistant", + "reasoning_content": "reasoning", + "tool_calls": [{"function": {"name": "lookup"}}], + }, + ) + tools = ({"function": {"name": "lookup"}},) + record = EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="sample-chat", + data=PromptData(messages=messages, tools=tools), + ) + + _, payload = _codec.encode(record) + decoded = _codec.decode(payload) + + assert isinstance(decoded.data, PromptData) + assert decoded.data.messages == messages + assert decoded.data.tools == tools + def test_error_event_round_trips_with_error_data(self): record = EventRecord( event_type=ErrorEventType.LOADGEN, diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 55fb2403b..f79abff3d 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -163,6 +163,53 @@ def test_issue_builds_query_and_publishes(self): assert issued_events[0].conversation_id == "" assert issued_events[0].turn is None + def test_issue_preserves_structured_chat_input_for_isl(self): + class ChatDataset(FakeDataset): + def load_sample(self, index: int) -> dict: + return { + "messages": [ + {"role": "user", "content": "question"}, + { + "role": "assistant", + "reasoning_content": "reasoning", + "content": None, + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": { + "name": "lookup", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call-1", + "content": "result", + }, + ], + "tools": [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {}}, + } + ], + } + + issuer = FakeIssuer() + issuer._auto_respond = False + publisher = FakePublisher() + phase_issuer = PhaseIssuer(ChatDataset(1), issuer, publisher, lambda: False) + + phase_issuer.issue(0, conversation_id="conv-1", turn=2) + + prompt = publisher.events_of_type(SampleEventType.ISSUED)[0].data + assert prompt.messages == tuple(issuer.issued_queries[0].data["messages"]) + assert prompt.tools == tuple(issuer.issued_queries[0].data["tools"]) + assert prompt.text is None + def test_issue_returns_none_when_stopped(self): dataset = FakeDataset(5) issuer = FakeIssuer() From 1ee7b471b7bc040ee6a8b8e80e2be0ab281426ad Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 14:04:57 -0700 Subject: [PATCH 02/11] Refine structured message tokenization --- pyproject.toml | 1 + .../services/metrics_aggregator/__main__.py | 6 +- .../metrics_aggregator/metrics_table.py | 28 ++- .../metrics_aggregator/token_metrics.py | 217 ++++++++++------- src/inference_endpoint/core/types.py | 2 + .../load_generator/session.py | 83 ++++--- .../services/metrics_aggregator/conftest.py | 1 + .../metrics_aggregator/test_token_metrics.py | 220 ++++++++++-------- tests/unit/core/test_record.py | 8 +- .../unit/load_generator/test_async_session.py | 123 +++++----- uv.lock | 6 +- 11 files changed, 394 insertions(+), 301 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 640b76eb7..fc3d7b6f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,7 @@ dependencies = [ "rich==14.3.3", # Needed for tokenization and OSL reporting "transformers==5.5.0", + "tiktoken==0.13.0", # Required by transformers' apply_chat_template "jinja2==3.1.6", "numpy>=1.26.4", diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py index 3238db0e0..3cff4f247 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py @@ -236,9 +236,9 @@ async def main() -> None: args.tokenizer, live_workers=args.tokenizer_workers ) except RuntimeError as exc: - # Fail-fast contract: a tokenizer environment that cannot shard - # must surface as a clear service-launch failure, not a silent - # slow path that cannot keep up with completions. + # Shard initialization failures must surface as a clear + # service-launch failure. Tokenizers without a plain-text backend + # do not create shards and remain usable for structured messages. raise SystemExit(f"FATAL: {exc}") from exc else: tokenizer_cm = nullcontext() diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index 4960e304c..207be7b65 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -315,17 +315,23 @@ def __init__( super().__init__(MetricSeriesKey.ISL, registry, queue) def fire(self, ev_rec, row, pre_change): - # Sync fast path: any backend that pre-populates token_ids (e.g. SGLang). - if isinstance(ev_rec.data, PromptData) and ev_rec.data.token_ids is not None: - self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) - return - if isinstance(ev_rec.data, PromptData) and ev_rec.data.messages is not None: - if self._queue is not None: - self._queue.enqueue_prompt( - (ev_rec.data.messages, ev_rec.data.tools), - self._make_recorder(ev_rec, pre_change), - ) - return + if isinstance(ev_rec.data, PromptData): + # Sync fast path: any backend that pre-populates token_ids (e.g. SGLang). + if ev_rec.data.token_ids is not None: + self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) + return + # Structured chat path: render the complete prompt with its template. + if ev_rec.data.messages is not None: + if self._queue is not None: + self._queue.enqueue_prompt( + ( + ev_rec.data.messages, + ev_rec.data.tools, + ev_rec.data.chat_template_kwargs, + ), + self._make_recorder(ev_rec, pre_change), + ) + return # Text path: tokenize raw prompt text — used when token_ids are # unavailable (e.g. OpenAI-compatible endpoints). Enqueued by the base. super().fire(ev_rec, row, pre_change) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 3dca8ec53..c950ea1f2 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -21,11 +21,11 @@ per-sample text. The sharded pool is the drain-phase accelerator and is auto-sized (one shard per core block); live mid-run flushes run on a small in-process thread pool (``--tokenizer-workers``, default 2) owned by the -queue's live loop. Hugging Face fast tokenizers use their Rust backend; -tokenizers that expose tiktoken through custom Python code use that official -wrapper so its preprocessing semantics are preserved. Platforms without CPU -affinity (e.g. macOS) shard unpinned at full speed; only cache/NUMA locality is -lost. +queue's live loop. Plain text counting uses a Hugging Face fast tokenizer's +Rust backend. Structured chat counting uses the full tokenizer's +``apply_chat_template`` path and does not require that backend. Platforms +without CPU affinity (e.g. macOS) shard unpinned at full speed; only cache/NUMA +locality is lost. """ from __future__ import annotations @@ -104,44 +104,60 @@ def _normalize_tool_calls_for_template( return normalized +def _text_only_multimodal_messages( + messages: tuple[dict[str, Any], ...] | list[dict[str, Any]], +) -> tuple[list[dict[str, Any]], bool]: + """Replace multimodal content parts with their text for text-only templates.""" + normalized: list[dict[str, Any]] = [] + changed = False + for message in messages: + content = message.get("content") + if not isinstance(content, list): + normalized.append(message) + continue + changed = True + text = " ".join( + part["text"] + for part in content + if isinstance(part, dict) + and part.get("type") == "text" + and isinstance(part.get("text"), str) + ) + normalized.append({**message, "content": text}) + return normalized, changed + + +def _normalize_prompt_messages_for_template( + messages: tuple[dict[str, Any], ...], +) -> list[dict[str, Any]]: + """Normalize historical tool calls without mutating the event payload.""" + normalized: list[dict[str, Any]] = [] + for message in messages: + tool_calls = message.get("tool_calls") + if not isinstance(tool_calls, list | tuple) or not tool_calls: + normalized.append(message) + continue + normalized.append( + { + **message, + "tool_calls": _normalize_tool_calls_for_template(tool_calls), + } + ) + return normalized + + # --------------------------------------------------------------------------- # Process-worker entry points (module-level so ProcessPoolExecutor can pickle # them by name). Each worker holds one raw tokenizers backend, pinned to a # fixed core block. # --------------------------------------------------------------------------- -_WORKER_BACKEND: Any = None - - -class _PythonTokenizerBackend: - """Use a custom tokenizer's public API instead of bypassing its Python logic.""" - - def __init__(self, tokenizer: Any) -> None: - self.tokenizer = tokenizer - - def encode(self, text: str, *, add_special_tokens: bool = False) -> list[int]: - return self.tokenizer.encode(text, add_special_tokens=add_special_tokens) - - def encode_lengths(self, texts: list[str]) -> list[int]: - return [len(self.encode(text, add_special_tokens=False)) for text in texts] +_WORKER_TEXT_BACKEND: Any = None def _backend_from_tokenizer(tokenizer: Any) -> Any | None: - """Return the tokenizer's supported length-counting path.""" - backend = getattr(tokenizer, "backend_tokenizer", None) - if backend is not None: - return backend - - model = getattr(tokenizer, "model", None) - model_type = type(model) - if ( - model_type.__module__ == "tiktoken.core" - and model_type.__name__ == "Encoding" - and callable(getattr(model, "encode", None)) - and callable(getattr(model, "encode_batch", None)) - ): - return _PythonTokenizerBackend(tokenizer) - return None + """Return the optional fast backend used only for plain-text counting.""" + return getattr(tokenizer, "backend_tokenizer", None) def load_reference_tokenizer(tokenizer_name: str) -> Any: @@ -156,12 +172,7 @@ def load_reference_tokenizer(tokenizer_name: str) -> Any: def load_reference_backend(tokenizer_name: str) -> Any | None: - """Supported token-counting path for the reference tokenizer. - - Hugging Face fast tokenizers use their native backend. Custom tiktoken - tokenizers use their public Python API so model-specific preprocessing is - not reimplemented here. ``None`` if neither path is available. - """ + """Load the optional fast backend used only for plain-text counting.""" return _backend_from_tokenizer(load_reference_tokenizer(tokenizer_name)) @@ -170,7 +181,7 @@ def _init_worker(tokenizer_name: str, core_set: list[int]) -> None: Affinity is set before the first encode so the Hugging Face rayon pool sizes itself to the pinned core count (num_cpus respects sched_getaffinity on - Linux). Custom wrappers scale through these process shards. + Linux). """ # Ctrl-C sends SIGINT to the whole foreground process group; the parent # drives worker shutdown, so a worker dying mid-drain would break the pool @@ -189,19 +200,14 @@ def _init_worker(tokenizer_name: str, core_set: list[int]) -> None: # unpinned shards from oversubscribing each other. logger.debug("could not pin tokenizer worker to %s", core_set) transformers_logging.set_verbosity_error() - # Sharding executes the tokenizer's (possibly custom) code in each worker - # process; the name is operator-supplied — same trust boundary as the - # in-process load. - global _WORKER_BACKEND - _WORKER_BACKEND = load_reference_backend(tokenizer_name) - if _WORKER_BACKEND is not None: - _WORKER_BACKEND.encode("warmup", add_special_tokens=False) + global _WORKER_TEXT_BACKEND + _WORKER_TEXT_BACKEND = load_reference_backend(tokenizer_name) + if _WORKER_TEXT_BACKEND is not None: + _WORKER_TEXT_BACKEND.encode("warmup", add_special_tokens=False) def encode_lengths(backend: Any, texts: list[str]) -> list[int]: """Per-text token counts via one bounded backend batch call.""" - if isinstance(backend, _PythonTokenizerBackend): - return backend.encode_lengths(texts) encode_batch = getattr(backend, "encode_batch_fast", None) or backend.encode_batch encoded = encode_batch(texts, add_special_tokens=False) return [len(getattr(item, "ids", item)) for item in encoded] @@ -209,7 +215,7 @@ def encode_lengths(backend: Any, texts: list[str]) -> list[int]: def _worker_encode_lengths(texts: list[str]) -> list[int]: """Per-text token counts for a shard, in one rayon-parallel call.""" - backend = _WORKER_BACKEND + backend = _WORKER_TEXT_BACKEND if backend is None: raise RuntimeError("tokenizer worker backend unavailable") return encode_lengths(backend, texts) @@ -217,7 +223,7 @@ def _worker_encode_lengths(texts: list[str]) -> list[int]: def _worker_ready(_: int) -> bool: """Warmup probe: returns once the worker's backend is loaded.""" - return _WORKER_BACKEND is not None + return _WORKER_TEXT_BACKEND is not None def _terminate_procs(procs: list[ProcessPoolExecutor]) -> None: @@ -296,26 +302,30 @@ def __init__( def _load_tokenizer(self) -> None: tok = load_reference_tokenizer(self._tokenizer_name) self._tokenizer = tok - self._backend = _backend_from_tokenizer(tok) + self._text_backend = _backend_from_tokenizer(tok) # Baseline = tokens from a [user, empty-assistant] pair minus the [user] # prefix alone, so the assistant frame is subtracted from message counts. try: prefix = cast( list[int], tok.apply_chat_template( - [_PREFIX_USER_MSG], tokenize=True, add_generation_prompt=False + [_PREFIX_USER_MSG], + tokenize=True, + add_generation_prompt=False, + return_dict=False, ), ) self._prefix_len = len(prefix) - with_assistant = cast( + with_assistant_tokens = cast( list[int], tok.apply_chat_template( [_PREFIX_USER_MSG, {"role": "assistant", "content": ""}], tokenize=True, add_generation_prompt=False, + return_dict=False, ), ) - self._baseline = len(with_assistant) - self._prefix_len + self._baseline = len(with_assistant_tokens) - self._prefix_len except Exception: self._prefix_len = 0 self._baseline = 0 @@ -332,18 +342,20 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: (``< 0``) fits one shard per ``cores_per_worker`` block of this process's affinity mask (or the online CPU count when the platform has no affinity API — shards then run unpinned), always at least one; - an explicit count is clamped to that capacity. An unsupported tokenizer - or a shard warmup that fails or exceeds its budget raises at startup. + an explicit count is clamped to that capacity. A tokenizer without a + fast text backend skips shard creation; structured chat tokenization + remains available. A shard warmup failure or timeout raises at startup. """ if cores_per_worker <= 0 or n_workers == 0: logger.info("BatchTokenizer: in-process tokenization (explicit)") return - if self._backend is None: - raise RuntimeError( - f"tokenizer {self._tokenizer_name!r} has no supported " - "token-counting path; use a Hugging Face fast tokenizer or a " - "supported custom tiktoken tokenizer." + if self._text_backend is None: + logger.info( + "BatchTokenizer: no plain-text backend for %s; structured " + "chat tokenization remains available", + self._tokenizer_name, ) + return # The full allowed CPU universe (cgroup-clamped) drives the shard block # math. cgroup_clamped_cpus owns the probe-and-restore of this process's # mask, so the aggregator's event loop, publisher, and live tokenizer @@ -399,11 +411,13 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: # -- batched text path -------------------------------------------------- def _encode_lengths_inproc(self, texts: list[str]) -> list[int]: - tok = self._tokenizer - backend = self._backend - if backend is not None: - return encode_lengths(backend, texts) - return [len(tok.tokenize(t)) for t in texts] # type: ignore[union-attr] + backend = self._text_backend + if backend is None: + raise RuntimeError( + f"plain-text tokenization for {self._tokenizer_name!r} requires " + "a supported fast backend" + ) + return encode_lengths(backend, texts) async def count_texts_async( self, @@ -460,7 +474,10 @@ def _token_count_message( msg["tool_calls"] = _normalize_tool_calls_for_template(tool_calls) try: encoded = tok.apply_chat_template( # type: ignore[union-attr] - [_PREFIX_USER_MSG, msg], tokenize=True, add_generation_prompt=False + [_PREFIX_USER_MSG, msg], + tokenize=True, + add_generation_prompt=False, + return_dict=False, ) return max(0, len(encoded) - self._prefix_len - self._baseline) except Exception as exc: @@ -500,16 +517,36 @@ def _token_count_prompt( self, messages: tuple[dict[str, Any], ...], tools: tuple[dict[str, Any], ...] | None, + chat_template_kwargs: dict[str, Any] | None = None, ) -> int: - kwargs: dict[str, Any] = { - "tokenize": True, - "add_generation_prompt": True, - } + kwargs = dict(chat_template_kwargs or {}) + kwargs.update( + tokenize=True, + add_generation_prompt=True, + return_dict=False, + ) if tools is not None: kwargs["tools"] = list(tools) - encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] - list(messages), **kwargs - ) + prompt_messages = _normalize_prompt_messages_for_template(messages) + try: + encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] + prompt_messages, **kwargs + ) + except Exception: + text_only_messages, changed = _text_only_multimodal_messages(messages) + if not changed: + raise + key = f"{self._tokenizer_name}:multimodal-prompt" + if key not in self._fallback_warned: + self._fallback_warned.add(key) + logger.warning( + "Chat template for %s rejected multimodal messages; " + "retrying with text content only", + self._tokenizer_name, + ) + encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] + text_only_messages, **kwargs + ) return len(encoded) async def token_count_prompt_async( @@ -518,12 +555,17 @@ async def token_count_prompt_async( tools: tuple[dict[str, Any], ...] | None, loop: asyncio.AbstractEventLoop, /, + chat_template_kwargs: dict[str, Any] | None = None, ) -> int: """Complete chat-prompt token count without blocking the loop.""" if self._thread is None: raise RuntimeError("BatchTokenizer is closed") return await loop.run_in_executor( - self._thread, self._token_count_prompt, messages, tools + self._thread, + self._token_count_prompt, + messages, + tools, + chat_template_kwargs, ) def close(self) -> None: @@ -551,7 +593,11 @@ def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: # Message means the structured assistant output used for OSL/TPOT tokenization. MessageParts = tuple[str, str | None, tuple[dict[str, Any], ...] | None] # Prompt means the complete chat input and tool definitions used for ISL. -PromptParts = tuple[tuple[dict[str, Any], ...], tuple[dict[str, Any], ...] | None] +PromptParts = tuple[ + tuple[dict[str, Any], ...], + tuple[dict[str, Any], ...] | None, + dict[str, Any] | None, +] class TokenCounter(Protocol): @@ -590,6 +636,7 @@ async def token_count_prompt_async( tools: tuple[dict[str, Any], ...] | None, loop: asyncio.AbstractEventLoop, /, + chat_template_kwargs: dict[str, Any] | None = None, ) -> int: """Chat-template token count for one complete input prompt.""" raise NotImplementedError @@ -770,10 +817,16 @@ async def _flush(self, live: bool) -> None: self._msg.append(((content, reasoning, tool_calls), on_count)) continue self._record(on_count, count) - for i, ((messages, tools), on_count) in enumerate(prompt_items): + for i, ( + (messages, tools, chat_template_kwargs), + on_count, + ) in enumerate(prompt_items): try: count = await self._tokenizer.token_count_prompt_async( - messages, tools, self._loop + messages, + tools, + self._loop, + chat_template_kwargs=chat_template_kwargs, ) except asyncio.CancelledError: if live: @@ -782,7 +835,9 @@ async def _flush(self, live: bool) -> None: except Exception as exc: # noqa: BLE001 — isolate items. failure = failure or exc if live: - self._prompt.append(((messages, tools), on_count)) + self._prompt.append( + ((messages, tools, chat_template_kwargs), on_count) + ) continue self._record(on_count, count) if failure is not None: diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index 5d3b39f35..942331c9d 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -294,12 +294,14 @@ class PromptData( token_ids: Pre-computed token IDs. Set when the adapter pre-tokenizes (e.g. SGLang). messages: Structured messages sent to a chat-completions endpoint. tools: Tool declarations accompanying ``messages``. + chat_template_kwargs: Model-specific arguments used to render ``messages``. """ text: str | None = None token_ids: tuple[int, ...] | None = None messages: tuple[dict[str, Any], ...] | None = None tools: tuple[dict[str, Any], ...] | None = None + chat_template_kwargs: dict[str, Any] | None = None class ErrorData( diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index d36782a16..df10d53c5 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -21,7 +21,6 @@ from __future__ import annotations import asyncio -import json import logging import time import uuid @@ -48,31 +47,6 @@ logger = logging.getLogger(__name__) _SESSION_ID_HEADER = "X-Session-ID" - - -def _extract_prompt_text(messages: list[Any]) -> str | None: - """Join text content from an OpenAI messages list; handles list-form multimodal content.""" - parts: list[str] = [] - for m in messages: - if not isinstance(m, dict): - continue - c = m.get("content") - if isinstance(c, str) and c: - parts.append(c) - elif isinstance(c, list): - parts.extend( - p["text"] - for p in c - if isinstance(p, dict) - and p.get("type") == "text" - and isinstance(p.get("text"), str) - ) - tc = m.get("tool_calls") - if tc: - parts.append(json.dumps(tc, separators=(",", ":"))) - return "\n".join(parts) if parts else None - - # --------------------------------------------------------------------------- # Phase configuration # --------------------------------------------------------------------------- @@ -262,25 +236,58 @@ def issue( ts = time.monotonic_ns() prompt_data: PromptData if isinstance(data, dict): - token_ids = data.get("input_tokens") or data.get("token_ids") - # Prefer the exact representation sent to the endpoint: existing - # token IDs, then structured chat messages, then a plain prompt. - # Non-string standalone prompts (for example multimodal content - # parts) remain unavailable for ISL reporting. - if token_ids is not None: - prompt_data = PromptData(token_ids=tuple(token_ids)) - elif isinstance(data.get("messages"), list | tuple): + input_tokens = data.get("input_tokens") + token_ids = data.get("token_ids") + messages = data.get("messages") + prompt = data.get("prompt") + representations = [ + name + for name, present in ( + ("input_tokens", input_tokens is not None), + ("token_ids", token_ids is not None), + ("messages", isinstance(messages, list | tuple) and bool(messages)), + ("prompt", isinstance(prompt, str)), + ) + if present + ] + if len(representations) > 1: + raise ValueError( + "sample contains multiple prompt representations: " + + ", ".join(representations) + ) + + if input_tokens is not None or token_ids is not None: + selected_token_ids = ( + input_tokens if input_tokens is not None else token_ids + ) + prompt_data = PromptData(token_ids=tuple(selected_token_ids)) + elif isinstance(messages, list | tuple) and messages: tools = data.get("tools") + chat_template_kwargs = data.get("chat_template_kwargs") prompt_data = PromptData( - messages=tuple(data["messages"]), + messages=tuple(messages), tools=tuple(tools) if isinstance(tools, list | tuple) else None, + chat_template_kwargs=( + dict(chat_template_kwargs) + if isinstance(chat_template_kwargs, dict) + else None + ), ) + elif isinstance(prompt, str): + prompt_data = PromptData(text=prompt) else: - prompt_text = data.get("prompt") - prompt_data = PromptData( - text=prompt_text if isinstance(prompt_text, str) else None + logger.warning( + "Sample %s has no supported prompt representation for ISL; " + "expected token IDs, non-empty messages, or a string prompt", + sample_index, ) + prompt_data = PromptData() else: + logger.warning( + "Sample %s has no supported prompt representation for ISL; " + "expected a mapping", + sample_index, + ) prompt_data = PromptData() self._publisher.publish( EventRecord( diff --git a/tests/unit/async_utils/services/metrics_aggregator/conftest.py b/tests/unit/async_utils/services/metrics_aggregator/conftest.py index 0d95825a8..bc6c2e54c 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/conftest.py +++ b/tests/unit/async_utils/services/metrics_aggregator/conftest.py @@ -92,6 +92,7 @@ async def token_count_prompt_async( messages: tuple[dict, ...], tools: tuple[dict, ...] | None, _loop: asyncio.AbstractEventLoop, + chat_template_kwargs: dict | None = None, ) -> int: if self._delay: await asyncio.sleep(self._delay) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index 6bfe93905..73ade78af 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -42,8 +42,8 @@ class _FakeTokenizer: """Deterministic tokenizer that splits on whitespace. - Has no ``backend_tokenizer``, so BatchTokenizer keeps the batch path - in-process (no subprocess shards) and counts via ``tokenize`` per text. + Has no ``backend_tokenizer`` and therefore supports only the structured + chat-template path when a subclass supplies ``apply_chat_template``. """ def __init__(self, load_delay: float = 0.0): @@ -86,7 +86,7 @@ def shutdown(self, wait=False, cancel_futures=False): class TestBatchTokenizer: @pytest.mark.asyncio async def test_count_texts_async(self): - with patch(_MOCK_TARGET, _FakeTokenizer): + with patch(_MOCK_TARGET, _FakeTokenizerWithBackend): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: counts = await tok.count_texts_async(["Hello world foo", "a"], loop) @@ -99,6 +99,16 @@ async def test_count_texts_async_empty(self): with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: assert await tok.count_texts_async([], loop) == [] + @pytest.mark.asyncio + async def test_plain_text_requires_a_text_backend(self): + with patch(_MOCK_TARGET, _FakeTokenizer): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + with pytest.raises( + RuntimeError, match="plain-text tokenization.*backend" + ): + await tok.count_texts_async(["Hello world"], loop) + @pytest.mark.asyncio async def test_count_texts_async_sharded(self): """With shards present, chunks are reassembled in original order.""" @@ -139,7 +149,12 @@ class _FakeTokenizerWithTemplate(_FakeTokenizer): """Tokenizer that supports apply_chat_template for tool-call testing.""" def apply_chat_template( - self, messages, tools=None, tokenize=False, add_generation_prompt=False + self, + messages, + tools=None, + tokenize=False, + add_generation_prompt=False, + return_dict=True, ): # Simulate 2 wrapper tokens for the template frame. parts = ["WRAPPER", "WRAPPER"] @@ -159,12 +174,27 @@ def apply_chat_template( parts.append("GENERATION") rendered = " ".join(parts) if tokenize: - return list(range(len(rendered.split()))) + token_ids = list(range(len(rendered.split()))) + if return_dict: + return { + "input_ids": token_ids, + "attention_mask": [1] * len(token_ids), + } + return token_ids return rendered @pytest.mark.unit class TestBatchTokenizerMessageTokenization: + def test_chat_template_requests_token_ids_not_batch_encoding(self): + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + count = tok._token_count_prompt( + ({"role": "user", "content": "one two three four"},), None + ) + + assert count == 7 + @pytest.mark.asyncio async def test_token_count_prompt_preserves_messages_tools_and_generation_prompt( self, @@ -255,6 +285,74 @@ def apply_chat_template(self, *args, **kwargs): ) assert count > 0 + def test_token_count_prompt_falls_back_to_text_only_multimodal_messages(self): + class _TextOnlyTemplateTokenizer(_FakeTokenizerWithTemplate): + def apply_chat_template(self, messages, **kwargs): + if any( + isinstance(message.get("content"), list) for message in messages + ): + raise TypeError("text-only template") + return super().apply_chat_template(messages, **kwargs) + + with patch(_MOCK_TARGET, _TextOnlyTemplateTokenizer): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + count = tok._token_count_prompt( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this image"}, + {"type": "image_url", "image_url": {"url": "x"}}, + ], + }, + ), + None, + ) + + assert count == 6 + + def test_token_count_prompt_normalizes_tools_and_forwards_template_kwargs(self): + class _RecordingTokenizer(_FakeTokenizerWithTemplate): + last_messages = None + last_kwargs = None + + def apply_chat_template(self, messages, **kwargs): + type(self).last_messages = messages + type(self).last_kwargs = dict(kwargs) + kwargs.pop("enable_thinking", None) + return super().apply_chat_template(messages, **kwargs) + + messages = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "type": "function", + "function": { + "name": "lookup", + "arguments": '{"city": "SF"}', + }, + } + ], + }, + ) + with patch(_MOCK_TARGET, _RecordingTokenizer): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + tok._token_count_prompt( + messages, + None, + {"enable_thinking": False}, + ) + + normalized_call = _RecordingTokenizer.last_messages[0]["tool_calls"][0] + assert normalized_call["function"]["arguments"] == {"city": "SF"} + assert messages[0]["tool_calls"][0]["function"]["arguments"] == ( + '{"city": "SF"}' + ) + assert _RecordingTokenizer.last_kwargs["enable_thinking"] is False + assert _RecordingTokenizer.last_kwargs["return_dict"] is False + class _Encoding: def __init__(self, n: int): @@ -275,31 +373,6 @@ def encode_batch(self, texts, add_special_tokens=False): return [_Encoding(len(t.split())) for t in texts] -class Encoding: - """Stand-in that fails if code bypasses the tokenizer wrapper.""" - - __module__ = "tiktoken.core" - - def encode(self, *args, **kwargs): - raise AssertionError("raw tiktoken encode must not be called") - - def encode_batch(self, *args, **kwargs): - raise AssertionError("raw tiktoken encode_batch must not be called") - - -class _FakeTokenizerWithTikTokenModel(_FakeTokenizer): - """Kimi-shaped custom tokenizer backed by the Rust tiktoken core.""" - - def __init__(self, load_delay: float = 0.0): - super().__init__(load_delay) - self.model = Encoding() - self.encode_calls: list[tuple[str, bool]] = [] - - def encode(self, text, *, add_special_tokens=False): - self.encode_calls.append((text, add_special_tokens)) - return text.split() - - @pytest.mark.unit class TestEncodeHelpers: def test_encode_lengths_prefers_fast(self): @@ -327,48 +400,15 @@ def from_pretrained(name, **kwargs): assert captured["name"] == "m" assert captured["kwargs"].get("trust_remote_code") is True - def test_load_reference_backend_uses_kimi_python_wrapper(self, monkeypatch): - """Kimi counting must not bypass its Python tokenizer wrapper.""" - - class _FakeAutoTokenizer: - @staticmethod - def from_pretrained(name, **kwargs): - assert name == "kimi" - assert kwargs == {"trust_remote_code": True} - return _FakeTokenizerWithTikTokenModel() - - monkeypatch.setattr(token_metrics_module, "AutoTokenizer", _FakeAutoTokenizer) - backend = token_metrics_module.load_reference_backend("kimi") - assert backend is not None - assert encode_lengths(backend, ["a b", "c"]) == [2, 1] - assert backend.tokenizer.encode_calls == [("a b", False), ("c", False)] - - def test_load_reference_backend_rejects_tiktoken_lookalike(self, monkeypatch): - class _LookalikeEncoding: - def encode(self, text, *, allowed_special): - return text.split() - - def encode_batch(self, texts, *, allowed_special): - return [text.split() for text in texts] - - class _FakeTok: - model = _LookalikeEncoding() - - class _FakeAutoTokenizer: - @staticmethod - def from_pretrained(name, **kwargs): - return _FakeTok() - - monkeypatch.setattr(token_metrics_module, "AutoTokenizer", _FakeAutoTokenizer) - assert token_metrics_module.load_reference_backend("lookalike") is None - def test_worker_encode_lengths_raises_without_backend(self, monkeypatch): - monkeypatch.setattr(token_metrics_module, "_WORKER_BACKEND", None) + monkeypatch.setattr(token_metrics_module, "_WORKER_TEXT_BACKEND", None) with pytest.raises(RuntimeError, match="backend unavailable"): _worker_encode_lengths(["a"]) def test_worker_encode_lengths_uses_backend(self, monkeypatch): - monkeypatch.setattr(token_metrics_module, "_WORKER_BACKEND", _FastBackend()) + monkeypatch.setattr( + token_metrics_module, "_WORKER_TEXT_BACKEND", _FastBackend() + ) assert _worker_encode_lengths(["a b", "c d e"]) == [2, 3] @@ -399,8 +439,8 @@ class TestSetupShardsDecisions: clamped / 0 explicit in-process (auto-sized in production — the CLI's --tokenizer-workers maps to the live thread lane, not to shards). - An environment that cannot shard is a startup error — never a silent - in-process fallback. + A fast text backend is optional. If present, shard warmup failures are + startup errors rather than silent in-process fallbacks. """ def _make(self, monkeypatch, cpus, n_workers, executor=_SpawnlessExecutor): @@ -434,27 +474,19 @@ def test_blocks_are_disjoint_consecutive_core_sets(self, monkeypatch): blocks = [set(ex.initargs[1]) for ex in tok._procs] assert blocks == [set(range(0, 8)), set(range(8, 16))] - def test_no_fast_backend_is_a_startup_error(self, monkeypatch): - monkeypatch.setattr( - token_metrics_module, "ProcessPoolExecutor", _SpawnlessExecutor - ) - with patch(_MOCK_TARGET, _FakeTokenizer): # no backend_tokenizer - with pytest.raises(RuntimeError, match="fast"): - BatchTokenizer("fake", live_workers=2) - - def test_tiktoken_wrapper_is_supported_by_process_shards(self, monkeypatch): + def test_structured_tokenization_does_not_require_a_text_backend(self, monkeypatch): monkeypatch.setattr( token_metrics_module, "ProcessPoolExecutor", _SpawnlessExecutor ) - monkeypatch.setattr( - token_metrics_module, "cgroup_clamped_cpus", lambda: list(range(16)) - ) - with patch(_MOCK_TARGET, _FakeTokenizerWithTikTokenModel): - with BatchTokenizer("kimi", live_workers=2) as tok: - assert len(tok._procs) == 2 - assert tok._backend is not None - assert tok._backend.tokenizer is tok._tokenizer - assert encode_lengths(tok._backend, ["a b", "c"]) == [2, 1] + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + with BatchTokenizer("fake", live_workers=2) as tok: + assert tok._procs == [] + assert ( + tok._token_count_prompt( + ({"role": "user", "content": "one two"},), None + ) + == 5 + ) def test_affinity_unavailable_shards_unpinned(self, monkeypatch): """No affinity API (e.g. macOS): shard from the CPU count, unpinned.""" @@ -494,7 +526,7 @@ class TestLiveLane: @pytest.mark.asyncio async def test_live_never_touches_the_shard_pool(self): """Mid-run flushes run in-process; the shards are drain-only.""" - with patch(_MOCK_TARGET, _FakeTokenizer): + with patch(_MOCK_TARGET, _FakeTokenizerWithBackend): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=1) as tok: procs = [_RecordingProc(), _RecordingProc(), _RecordingProc()] @@ -860,18 +892,18 @@ def test_terminate_procs_kills_running_workers(): assert future.running(), "worker task did not start" procs = list((getattr(ex, "_processes", None) or {}).values()) assert procs, "worker did not spawn" - assert not wait_for_process([p.sentinel for p in procs], timeout=0), ( - "worker exited before termination" - ) + assert not wait_for_process( + [p.sentinel for p in procs], timeout=0 + ), "worker exited before termination" _terminate_procs([ex]) for p in procs: # The executor manager may reap the child concurrently; sentinel # readiness observes exit without racing its return-code update. - assert wait_for_process([p.sentinel], timeout=5), ( - "worker was not terminated" - ) + assert wait_for_process( + [p.sentinel], timeout=5 + ), "worker was not terminated" finally: cleanup_procs = procs or list((getattr(ex, "_processes", None) or {}).values()) for p in cleanup_procs: diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index dc137dd70..ad2920312 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -182,10 +182,15 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): }, ) tools = ({"function": {"name": "lookup"}},) + chat_template_kwargs = {"enable_thinking": False} record = EventRecord( event_type=SampleEventType.ISSUED, sample_uuid="sample-chat", - data=PromptData(messages=messages, tools=tools), + data=PromptData( + messages=messages, + tools=tools, + chat_template_kwargs=chat_template_kwargs, + ), ) _, payload = _codec.encode(record) @@ -194,6 +199,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): assert isinstance(decoded.data, PromptData) assert decoded.data.messages == messages assert decoded.data.tools == tools + assert decoded.data.chat_template_kwargs == chat_template_kwargs def test_error_event_round_trips_with_error_data(self): record = EventRecord( diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index f79abff3d..aa619f4e9 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -38,7 +38,6 @@ PhaseResult, PhaseType, SessionResult, - _extract_prompt_text, ) from inference_endpoint.metrics.metric import Throughput @@ -196,6 +195,7 @@ def load_sample(self, index: int) -> dict: "function": {"name": "lookup", "parameters": {}}, } ], + "chat_template_kwargs": {"enable_thinking": False}, } issuer = FakeIssuer() @@ -208,8 +208,59 @@ def load_sample(self, index: int) -> dict: prompt = publisher.events_of_type(SampleEventType.ISSUED)[0].data assert prompt.messages == tuple(issuer.issued_queries[0].data["messages"]) assert prompt.tools == tuple(issuer.issued_queries[0].data["tools"]) + assert prompt.chat_template_kwargs == {"enable_thinking": False} assert prompt.text is None + def test_issue_rejects_multiple_prompt_representations(self): + class AmbiguousDataset(FakeDataset): + def load_sample(self, index: int) -> dict: + return { + "messages": [{"role": "user", "content": "question"}], + "prompt": "question", + } + + issuer = FakeIssuer() + publisher = FakePublisher() + phase_issuer = PhaseIssuer( + AmbiguousDataset(1), issuer, publisher, lambda: False + ) + + with pytest.raises(ValueError, match="multiple prompt representations"): + phase_issuer.issue(0) + + def test_issue_empty_messages_falls_through_to_prompt(self): + class EmptyMessagesDataset(FakeDataset): + def load_sample(self, index: int) -> dict: + return {"messages": [], "prompt": "fallback prompt"} + + issuer = FakeIssuer() + publisher = FakePublisher() + phase_issuer = PhaseIssuer( + EmptyMessagesDataset(1), issuer, publisher, lambda: False + ) + + phase_issuer.issue(0) + + prompt = publisher.events_of_type(SampleEventType.ISSUED)[0].data + assert prompt.messages is None + assert prompt.text == "fallback prompt" + + def test_issue_warns_when_no_supported_prompt_representation(self, caplog): + class UnsupportedPromptDataset(FakeDataset): + def load_sample(self, index: int) -> dict: + return {"prompt": [{"type": "image_url"}]} + + issuer = FakeIssuer() + publisher = FakePublisher() + phase_issuer = PhaseIssuer( + UnsupportedPromptDataset(1), issuer, publisher, lambda: False + ) + + with caplog.at_level("WARNING"): + phase_issuer.issue(0) + + assert "no supported prompt representation" in caplog.text + def test_issue_returns_none_when_stopped(self): dataset = FakeDataset(5) issuer = FakeIssuer() @@ -1132,76 +1183,6 @@ def test_perf_results_filter(self): assert sr.perf_results[0].name == "perf1" -@pytest.mark.unit -class TestExtractPromptText: - def test_string_content_extracted(self): - messages = [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi"}, - ] - assert _extract_prompt_text(messages) == "Hello\nHi" - - def test_multimodal_list_content_text_parts_extracted(self): - messages = [ - { - "role": "user", - "content": [ - {"type": "text", "text": "Describe this image"}, - {"type": "image_url"}, - ], - } - ] - assert _extract_prompt_text(messages) == "Describe this image" - - def test_mixed_string_and_list_content(self): - messages = [ - {"role": "system", "content": "You are helpful"}, - { - "role": "user", - "content": [ - {"type": "text", "text": "What is this?"}, - {"type": "image_url"}, - ], - }, - ] - assert _extract_prompt_text(messages) == "You are helpful\nWhat is this?" - - def test_none_content_skipped(self): - messages = [ - {"role": "assistant", "content": None}, - {"role": "user", "content": "Hello"}, - ] - assert _extract_prompt_text(messages) == "Hello" - - def test_list_content_with_no_text_parts_returns_none(self): - messages = [{"role": "user", "content": [{"type": "image_url"}]}] - assert _extract_prompt_text(messages) is None - - def test_non_dict_messages_skipped(self): - messages = ["not a dict", {"role": "user", "content": "Valid"}] - assert _extract_prompt_text(messages) == "Valid" - - def test_tool_calls_included(self): - messages = [ - {"role": "user", "content": "What's the weather?"}, - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "c1", - "type": "function", - "function": {"name": "get_weather", "arguments": "{}"}, - } - ], - }, - ] - result = _extract_prompt_text(messages) - assert result is not None - assert "What's the weather?" in result - assert "get_weather" in result - - @pytest.mark.unit class TestBenchmarkSessionHandleResponse: """Direct invocation of BenchmarkSession._handle_response (no session.run).""" diff --git a/uv.lock b/uv.lock index 44e64f5af..188b30b07 100644 --- a/uv.lock +++ b/uv.lock @@ -1328,6 +1328,7 @@ dependencies = [ { name = "pyzmq", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, { name = "rich", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, { name = "sentencepiece", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, + { name = "tiktoken", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, { name = "transformers", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, { name = "typing-extensions", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, { name = "urllib3", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, @@ -1430,6 +1431,7 @@ requires-dist = [ { name = "sphinx-autodoc-typehints", marker = "extra == 'dev'", specifier = "==3.9.11" }, { name = "sphinx-rtd-theme", marker = "extra == 'dev'", specifier = "==3.1.0" }, { name = "sqlalchemy", marker = "extra == 'sql'", specifier = "==2.0.48" }, + { name = "tiktoken", specifier = "==0.13.0" }, { name = "transformers", specifier = "==5.5.0" }, { name = "typing-extensions", specifier = "==4.15.0" }, { name = "urllib3", specifier = "==2.7.0" }, @@ -3672,8 +3674,8 @@ name = "tiktoken" version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "regex", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "requests", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "regex", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, + { name = "requests", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'darwin' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-dev') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-performance') or (sys_platform == 'linux' and extra == 'extra-18-inference-endpoint-bfcl' and extra == 'extra-18-inference-endpoint-test')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e4/e5/5f3cb2159769d0f4324c0e9e87f9de3c4b1cd45848a96b2eb3566ad5ca77/tiktoken-0.13.0.tar.gz", hash = "sha256:c9435714c3a84c2319499de9a300c0e604449dd0799ff246458b3bb6a7f433c1", size = 38986, upload-time = "2026-05-15T04:51:27.153Z" } wheels = [ From 485bc5a57645093e63d023154b2ffb93fe53ac0d Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 16:48:34 -0700 Subject: [PATCH 03/11] Simplify tokenizer backend loading --- .../services/metrics_aggregator/token_metrics.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index c950ea1f2..937486bae 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -155,11 +155,6 @@ def _normalize_prompt_messages_for_template( _WORKER_TEXT_BACKEND: Any = None -def _backend_from_tokenizer(tokenizer: Any) -> Any | None: - """Return the optional fast backend used only for plain-text counting.""" - return getattr(tokenizer, "backend_tokenizer", None) - - def load_reference_tokenizer(tokenizer_name: str) -> Any: """Load the run's reference tokenizer. @@ -173,7 +168,8 @@ def load_reference_tokenizer(tokenizer_name: str) -> Any: def load_reference_backend(tokenizer_name: str) -> Any | None: """Load the optional fast backend used only for plain-text counting.""" - return _backend_from_tokenizer(load_reference_tokenizer(tokenizer_name)) + tokenizer = load_reference_tokenizer(tokenizer_name) + return getattr(tokenizer, "backend_tokenizer", None) def _init_worker(tokenizer_name: str, core_set: list[int]) -> None: @@ -302,7 +298,7 @@ def __init__( def _load_tokenizer(self) -> None: tok = load_reference_tokenizer(self._tokenizer_name) self._tokenizer = tok - self._text_backend = _backend_from_tokenizer(tok) + self._text_backend = getattr(tok, "backend_tokenizer", None) # Baseline = tokens from a [user, empty-assistant] pair minus the [user] # prefix alone, so the assistant frame is subtracted from message counts. try: From 31b064686788400fb4ea8b65e147e4236f1ac956 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 18:37:02 -0700 Subject: [PATCH 04/11] Clarify token counting input routing --- .../metrics_aggregator/metrics_table.py | 136 +++----- .../metrics_aggregator/token_metrics.py | 302 ++++++++---------- .../metrics_aggregator/tokenization.py | 51 +++ .../services/metrics_aggregator/conftest.py | 79 +++-- .../metrics_aggregator/test_aggregator.py | 13 +- .../metrics_aggregator/test_metrics_table.py | 17 +- .../metrics_aggregator/test_token_metrics.py | 210 +++++++----- 7 files changed, 406 insertions(+), 402 deletions(-) create mode 100644 src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index 207be7b65..5fa263810 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -25,6 +25,13 @@ from typing import TYPE_CHECKING, Any import msgspec +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, + PromptInput, + TextInput, + TokenIdsInput, + TokenizationInput, +) from inference_endpoint.core.record import SampleEventType, SessionEventType from inference_endpoint.core.types import PromptData, TextModelOutput @@ -33,7 +40,6 @@ MetricsRegistry, ) from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( - MessageParts, TokenBatchQueue, ) from inference_endpoint.core.record import EventRecord @@ -184,13 +190,9 @@ def fire(self, ev_rec, row, pre_change): class TokenTrigger(EmitTrigger): """Base for triggers whose metric needs tokenization. - Subclasses implement ``_extract_text()`` to pull the text to tokenize from - the event record, and may override ``_extract_message()`` to return - (content, reasoning, tool_calls) for chat-template-aware tokenization of - structured output. ``fire()`` does not tokenize inline — it enqueues - the work plus a recorder callback onto the shared ``TokenBatchQueue``, which - the aggregator flushes in batches. ``_compute_value()`` can transform the - token count before it is recorded. + Subclasses return one explicit tokenization input. ``fire()`` dispatches + already-tokenized IDs synchronously and sends every other variant through + the shared ``TokenBatchQueue``. """ def __init__( @@ -205,22 +207,12 @@ def __init__( self._queue = queue @abstractmethod - def _extract_text( + def _extract_tokenization_input( self, ev_rec: EventRecord, row: SampleRow, pre_change: dict[str, Any] - ) -> str | None: - """Return the text to tokenize, or None to skip.""" + ) -> TokenizationInput | None: + """Return the operation-specific tokenization input, or None to skip.""" raise NotImplementedError() - def _extract_message( - self, ev_rec: EventRecord, row: SampleRow, pre_change: dict[str, Any] - ) -> MessageParts | None: - """Return (content, reasoning, tool_calls) for message-aware tokenization. - - When non-None, the message (chat-template) path is used instead of the - plain-text path. Default returns None (use text path). - """ - return None - def _compute_value( self, token_count: int, ev_rec: EventRecord, pre_change: dict[str, Any] ) -> int | float | None: @@ -241,21 +233,14 @@ def record(count: int) -> None: return record def fire(self, ev_rec, row, pre_change): - if self._queue is None: + item = self._extract_tokenization_input(ev_rec, row, pre_change) + if item is None: return - message_parts = self._extract_message(ev_rec, row, pre_change) - if message_parts is not None: - self._queue.enqueue_message( - message_parts, self._make_recorder(ev_rec, pre_change) - ) - return - text = self._extract_text(ev_rec, row, pre_change) - if not text: - # Empty output (no text and no tool calls) is not an anomaly: - # there is nothing to tokenize, so we record no token-count - # sample rather than a spurious 0 that would skew the series. - return - self._queue.enqueue_text(text, self._make_recorder(ev_rec, pre_change)) + if isinstance(item, TokenIdsInput): + self.registry.record(self.metric_name, len(item.token_ids)) + elif isinstance(item, TextInput | MessageInput | PromptInput): + if self._queue is not None: + self._queue.enqueue(item, self._make_recorder(ev_rec, pre_change)) # --------------------------------------------------------------------------- @@ -314,31 +299,19 @@ def __init__( ): super().__init__(MetricSeriesKey.ISL, registry, queue) - def fire(self, ev_rec, row, pre_change): - if isinstance(ev_rec.data, PromptData): - # Sync fast path: any backend that pre-populates token_ids (e.g. SGLang). - if ev_rec.data.token_ids is not None: - self.registry.record(self.metric_name, len(ev_rec.data.token_ids)) - return - # Structured chat path: render the complete prompt with its template. - if ev_rec.data.messages is not None: - if self._queue is not None: - self._queue.enqueue_prompt( - ( - ev_rec.data.messages, - ev_rec.data.tools, - ev_rec.data.chat_template_kwargs, - ), - self._make_recorder(ev_rec, pre_change), - ) - return - # Text path: tokenize raw prompt text — used when token_ids are - # unavailable (e.g. OpenAI-compatible endpoints). Enqueued by the base. - super().fire(ev_rec, row, pre_change) - - def _extract_text(self, ev_rec, row, pre_change): - if isinstance(ev_rec.data, PromptData) and ev_rec.data.text is not None: - return ev_rec.data.text + def _extract_tokenization_input(self, ev_rec, row, pre_change): + if not isinstance(ev_rec.data, PromptData): + return None + if ev_rec.data.token_ids is not None: + return TokenIdsInput(ev_rec.data.token_ids) + if ev_rec.data.text: + return TextInput(ev_rec.data.text) + if ev_rec.data.messages is not None: + return PromptInput( + ev_rec.data.messages, + ev_rec.data.tools, + ev_rec.data.chat_template_kwargs, + ) return None @@ -352,20 +325,14 @@ def __init__( ): super().__init__(MetricSeriesKey.OSL, registry, queue) - def _extract_text(self, ev_rec, row, pre_change): - if isinstance(ev_rec.data, TextModelOutput): - if ev_rec.data.tool_calls: - # Delegate to _extract_message for chat-template tokenization. - return None - text = str(ev_rec.data) - return text if text else None - return None - - def _extract_message(self, ev_rec, row, pre_change): - if isinstance(ev_rec.data, TextModelOutput) and ( - ev_rec.data.reasoning or ev_rec.data.tool_calls - ): - return ev_rec.data.as_message_parts() + def _extract_tokenization_input(self, ev_rec, row, pre_change): + if not isinstance(ev_rec.data, TextModelOutput): + return None + if ev_rec.data.reasoning or ev_rec.data.tool_calls: + return MessageInput(*ev_rec.data.as_message_parts()) + text = str(ev_rec.data) + if text: + return TextInput(text) return None @@ -393,23 +360,16 @@ def __init__( dtype=float, ) - def _extract_text(self, ev_rec, row, pre_change): + def _extract_tokenization_input(self, ev_rec, row, pre_change): if pre_change.get(SampleField.RECV_FIRST_NS) is None: return None - if isinstance(ev_rec.data, TextModelOutput): - if ev_rec.data.tool_calls: - # Delegate to _extract_message for chat-template tokenization. - return None - return ev_rec.data.text_after_first_chunk() or None - return None - - def _extract_message(self, ev_rec, row, pre_change): - if pre_change.get(SampleField.RECV_FIRST_NS) is None: + if not isinstance(ev_rec.data, TextModelOutput): return None - if isinstance(ev_rec.data, TextModelOutput) and ( - ev_rec.data.reasoning or ev_rec.data.tool_calls - ): - return ev_rec.data.as_message_parts_after_first_chunk() + if ev_rec.data.reasoning or ev_rec.data.tool_calls: + return MessageInput(*ev_rec.data.as_message_parts_after_first_chunk()) + text = ev_rec.data.text_after_first_chunk() + if text: + return TextInput(text) return None def _compute_value(self, token_count, ev_rec, pre_change): diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 937486bae..abe818d1e 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -43,6 +43,13 @@ from typing import TYPE_CHECKING, Any, Protocol, cast import msgspec +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, + PromptInput, + TextInput, + TokenIdsInput, + TokenizationInput, +) from inference_endpoint.endpoint_client.cpu_affinity import ( cgroup_clamped_cpus, ) @@ -259,9 +266,8 @@ def _even_chunks(items: list[str], n: int) -> list[list[str]]: class BatchTokenizer: """Counts tokens for batches of text, sharded across pinned CPU cores. - ``count_texts_async`` tokenizes a whole list in one sharded call. The - chat-template ``token_count_message_async`` path runs on a small in-process - thread — rare (tool calls) relative to the batched OSL/ISL/TPOT flush. + ``count_batch_async`` explicitly routes token IDs, text, assistant + messages, and complete prompts to their corresponding counting path. """ def __init__( @@ -415,7 +421,7 @@ def _encode_lengths_inproc(self, texts: list[str]) -> list[int]: ) return encode_lengths(backend, texts) - async def count_texts_async( + async def _count_texts_async( self, texts: list[str], loop: asyncio.AbstractEventLoop, @@ -494,21 +500,6 @@ def _token_count_message( ] return self._token_count_text("\n".join(parts)) - async def token_count_message_async( - self, - content: str, - reasoning: str | None, - tool_calls: tuple[dict[str, Any], ...] | None, - loop: asyncio.AbstractEventLoop, - /, - ) -> int: - """Chat-template message token count without blocking the loop.""" - if self._thread is None: - raise RuntimeError("BatchTokenizer is closed") - return await loop.run_in_executor( - self._thread, self._token_count_message, content, reasoning, tool_calls - ) - def _token_count_prompt( self, messages: tuple[dict[str, Any], ...], @@ -545,24 +536,72 @@ def _token_count_prompt( ) return len(encoded) - async def token_count_prompt_async( + async def count_batch_async( self, - messages: tuple[dict[str, Any], ...], - tools: tuple[dict[str, Any], ...] | None, + inputs: list[TokenizationInput], loop: asyncio.AbstractEventLoop, /, - chat_template_kwargs: dict[str, Any] | None = None, - ) -> int: - """Complete chat-prompt token count without blocking the loop.""" - if self._thread is None: - raise RuntimeError("BatchTokenizer is closed") - return await loop.run_in_executor( - self._thread, - self._token_count_prompt, - messages, - tools, - chat_template_kwargs, - ) + *, + live: bool = False, + ) -> list[int | Exception]: + """Count a mixed batch while preserving input order.""" + outcomes: list[int | Exception | None] = [None] * len(inputs) + text_indices: list[int] = [] + texts: list[str] = [] + structured: list[tuple[int, MessageInput | PromptInput]] = [] + + for index, item in enumerate(inputs): + if isinstance(item, TokenIdsInput): + outcomes[index] = len(item.token_ids) + elif isinstance(item, TextInput): + text_indices.append(index) + texts.append(item.text) + elif isinstance(item, MessageInput): + structured.append((index, item)) + elif isinstance(item, PromptInput): + structured.append((index, item)) + + if texts: + try: + counts = await self._count_texts_async(texts, loop, live=live) + if len(counts) != len(texts): + raise RuntimeError( + f"tokenizer returned {len(counts)} counts for " + f"{len(texts)} texts" + ) + for index, count in zip(text_indices, counts, strict=True): + outcomes[index] = count + except Exception as exc: # noqa: BLE001 - isolate this input kind. + for index in text_indices: + outcomes[index] = exc + + for index, item in structured: + if self._thread is None: + outcomes[index] = RuntimeError("BatchTokenizer is closed") + continue + try: + if isinstance(item, MessageInput): + outcomes[index] = await loop.run_in_executor( + self._thread, + self._token_count_message, + item.content, + item.reasoning, + item.tool_calls, + ) + elif isinstance(item, PromptInput): + outcomes[index] = await loop.run_in_executor( + self._thread, + self._token_count_prompt, + item.messages, + item.tools, + item.chat_template_kwargs, + ) + except Exception as exc: # noqa: BLE001 - isolate this input. + outcomes[index] = exc + + if any(outcome is None for outcome in outcomes): + raise AssertionError("unhandled TokenizationInput variant") + return cast(list[int | Exception], outcomes) def close(self) -> None: """Shut down all workers. Idempotent. @@ -586,16 +625,6 @@ def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: self.close() -# Message means the structured assistant output used for OSL/TPOT tokenization. -MessageParts = tuple[str, str | None, tuple[dict[str, Any], ...] | None] -# Prompt means the complete chat input and tool definitions used for ISL. -PromptParts = tuple[ - tuple[dict[str, Any], ...], - tuple[dict[str, Any], ...] | None, - dict[str, Any] | None, -] - - class TokenCounter(Protocol): """The async tokenization surface ``TokenBatchQueue`` depends on. @@ -604,38 +633,16 @@ class TokenCounter(Protocol): tokenizer and test doubles type-check without inheritance. """ - async def count_texts_async( + async def count_batch_async( self, - texts: list[str], + inputs: list[TokenizationInput], loop: asyncio.AbstractEventLoop, /, *, live: bool = False, - ) -> list[int]: - """Per-text token counts (``live=True`` = the bounded mid-run lane).""" - raise NotImplementedError - - async def token_count_message_async( - self, - content: str, - reasoning: str | None, - tool_calls: tuple[dict[str, Any], ...] | None, - loop: asyncio.AbstractEventLoop, - /, - ) -> int: - """Chat-template token count for one assistant message.""" - raise NotImplementedError - - async def token_count_prompt_async( - self, - messages: tuple[dict[str, Any], ...], - tools: tuple[dict[str, Any], ...] | None, - loop: asyncio.AbstractEventLoop, - /, - chat_template_kwargs: dict[str, Any] | None = None, - ) -> int: - """Chat-template token count for one complete input prompt.""" - raise NotImplementedError + ) -> list[int | Exception]: + """Return one count or error per input, in input order.""" + ... class TokenBatchQueue: @@ -658,9 +665,7 @@ def __init__( ) -> None: self._tokenizer = tokenizer self._loop = loop - self._text: list[tuple[str, Callable[[int], None]]] = [] - self._msg: list[tuple[MessageParts, Callable[[int], None]]] = [] - self._prompt: list[tuple[PromptParts, Callable[[int], None]]] = [] + self._items: list[tuple[TokenizationInput, Callable[[int], None]]] = [] self._inflight = 0 self._live_task: asyncio.Task | None = None # Serializes flushes so the periodic live flush and the end-of-run @@ -699,21 +704,9 @@ def pending(self) -> int: """Enqueued items not yet tokenized-and-recorded.""" return self._inflight - def enqueue_text(self, text: str, on_count: Callable[[int], None]) -> None: - self._inflight += 1 - self._text.append((text, on_count)) - - def enqueue_message( - self, parts: MessageParts, on_count: Callable[[int], None] - ) -> None: + def enqueue(self, item: TokenizationInput, on_count: Callable[[int], None]) -> None: self._inflight += 1 - self._msg.append((parts, on_count)) - - def enqueue_prompt( - self, parts: PromptParts, on_count: Callable[[int], None] - ) -> None: - self._inflight += 1 - self._prompt.append((parts, on_count)) + self._items.append((item, on_count)) async def flush_live_once(self) -> None: """One bounded mid-run flush (live lane). @@ -743,99 +736,56 @@ async def _flush(self, live: bool) -> None: in the next flush. Callers use ``flush_live_once`` / ``drain_all``. """ async with self._lock: - if not (self._text or self._msg or self._prompt): + if not self._items: return if live: - cap = _LIVE_FLUSH_MAX_ITEMS - text_items = self._text[:cap] - del self._text[:cap] # in-place: O(cap), not O(backlog). - msg_items = self._msg[:cap] - del self._msg[:cap] - prompt_items = self._prompt[:cap] - del self._prompt[:cap] + selected: list[tuple[TokenizationInput, Callable[[int], None]]] = [] + remaining: list[tuple[TokenizationInput, Callable[[int], None]]] = [] + selected_by_type: dict[type, int] = {} + for queued in self._items: + item_type = type(queued[0]) + count = selected_by_type.get(item_type, 0) + if count < _LIVE_FLUSH_MAX_ITEMS: + selected.append(queued) + selected_by_type[item_type] = count + 1 + else: + remaining.append(queued) + items = selected + self._items = remaining else: - text_items, self._text = self._text, [] - msg_items, self._msg = self._msg, [] - prompt_items, self._prompt = self._prompt, [] - # The text and message phases fail independently — they run on - # separate executors, so a dead text shard must not drop message - # items that would still succeed (and vice versa). The first - # failure is re-raised after both phases so callers still see it. + items, self._items = self._items, [] + + try: + outcomes = await self._tokenizer.count_batch_async( + [item for item, _ in items], self._loop, live=live + ) + except asyncio.CancelledError: + if live: + self._items[:0] = items + raise + except Exception: + if live: + self._items[:0] = items + raise + + if len(outcomes) != len(items): + if live: + self._items[:0] = items + raise RuntimeError( + f"tokenizer returned {len(outcomes)} outcomes for " + f"{len(items)} inputs" + ) + failure: Exception | None = None - if text_items: - try: - counts = await self._tokenizer.count_texts_async( - [t for t, _ in text_items], self._loop, live=live - ) - except asyncio.CancelledError: - if live: - self._text[:0] = text_items - self._msg[:0] = msg_items - self._prompt[:0] = prompt_items - raise - except Exception as exc: # noqa: BLE001 — isolate phases. - failure = exc - if live: - # A live hiccup must not lose samples: give the items - # back so the end-of-run drain (full pool) retries. - # Drain failures are terminal and stay pending-only. - self._text[:0] = text_items + retry: list[tuple[TokenizationInput, Callable[[int], None]]] = [] + for queued, outcome in zip(items, outcomes, strict=True): + if isinstance(outcome, Exception): + failure = failure or outcome + retry.append(queued) else: - if len(counts) != len(text_items): - # Tokenizer contract violation (wrong-length result). - # Treat it like any other text-phase failure so the - # message phase still runs: re-queue live items, leave - # drain items pending (terminal), re-raise after both. - failure = RuntimeError( - f"tokenizer returned {len(counts)} counts for " - f"{len(text_items)} texts" - ) - if live: - self._text[:0] = text_items - else: - for (_, on_count), count in zip( - text_items, counts, strict=False - ): - self._record(on_count, count) - for i, ((content, reasoning, tool_calls), on_count) in enumerate(msg_items): - try: - count = await self._tokenizer.token_count_message_async( - content, reasoning, tool_calls, self._loop - ) - except asyncio.CancelledError: - if live: - self._msg[:0] = msg_items[i:] - self._prompt[:0] = prompt_items - raise - except Exception as exc: # noqa: BLE001 — isolate items. - failure = failure or exc - if live: - self._msg.append(((content, reasoning, tool_calls), on_count)) - continue - self._record(on_count, count) - for i, ( - (messages, tools, chat_template_kwargs), - on_count, - ) in enumerate(prompt_items): - try: - count = await self._tokenizer.token_count_prompt_async( - messages, - tools, - self._loop, - chat_template_kwargs=chat_template_kwargs, - ) - except asyncio.CancelledError: - if live: - self._prompt[:0] = prompt_items[i:] - raise - except Exception as exc: # noqa: BLE001 — isolate items. - failure = failure or exc - if live: - self._prompt.append( - ((messages, tools, chat_template_kwargs), on_count) - ) - continue - self._record(on_count, count) + self._record(queued[1], outcome) + if live and retry: + self._items[:0] = retry if failure is not None: raise failure diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py new file mode 100644 index 000000000..98614892c --- /dev/null +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Internal token-counting inputs. + +Each variant names the tokenization operation its payload requires. Keeping +this distinction explicit prevents structured messages and complete prompts +from accidentally falling through the plain-text tokenizer path. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypeAlias + + +@dataclass(frozen=True, slots=True) +class TokenIdsInput: + """Already-tokenized input; counting is simply ``len(token_ids)``.""" + + token_ids: tuple[int, ...] + + +@dataclass(frozen=True, slots=True) +class TextInput: + """Unstructured text counted by the tokenizer's text backend.""" + + text: str + + +@dataclass(frozen=True, slots=True) +class MessageInput: + """One structured assistant output rendered by the chat template.""" + + content: str + reasoning: str | None + tool_calls: tuple[dict[str, Any], ...] | None + + +@dataclass(frozen=True, slots=True) +class PromptInput: + """A complete structured chat prompt rendered by the chat template.""" + + messages: tuple[dict[str, Any], ...] + tools: tuple[dict[str, Any], ...] | None + chat_template_kwargs: dict[str, Any] | None + + +TokenizationInput: TypeAlias = ( # noqa: UP040 - mypy version lacks PEP 695. + TokenIdsInput | TextInput | MessageInput | PromptInput +) diff --git a/tests/unit/async_utils/services/metrics_aggregator/conftest.py b/tests/unit/async_utils/services/metrics_aggregator/conftest.py index bc6c2e54c..a88656f81 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/conftest.py +++ b/tests/unit/async_utils/services/metrics_aggregator/conftest.py @@ -40,6 +40,12 @@ SeriesStat, SessionState, ) +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, + PromptInput, + TextInput, + TokenIdsInput, +) from inference_endpoint.async_utils.transport.zmq.context import ManagedZMQContext from inference_endpoint.core.record import ( EventRecord, @@ -50,7 +56,7 @@ # --------------------------------------------------------------------------- # Mock BatchTokenizer — whitespace token counts; matches the BatchTokenizer -# surface the TokenBatchQueue calls (count_texts_async + message path). +# surface the TokenBatchQueue calls (one mixed-batch method). # --------------------------------------------------------------------------- @@ -60,50 +66,41 @@ class MockBatchTokenizer: def __init__(self, delay: float = 0.0) -> None: self._delay = delay - async def count_texts_async( - self, - texts: list[str], - _loop: asyncio.AbstractEventLoop, - live: bool = False, - ) -> list[int]: + async def count_batch_async(self, inputs, _loop, live=False): if self._delay: await asyncio.sleep(self._delay) - return [len(t.split()) for t in texts] - - async def token_count_message_async( - self, - content: str, - reasoning: str | None, - tool_calls, - _loop: asyncio.AbstractEventLoop, - ) -> int: import msgspec - if self._delay: - await asyncio.sleep(self._delay) - tool_calls_str = ( - msgspec.json.encode(list(tool_calls)).decode() if tool_calls else "" - ) - combined = (content or "") + " " + (reasoning or "") + " " + tool_calls_str - return len(combined.split()) - - async def token_count_prompt_async( - self, - messages: tuple[dict, ...], - tools: tuple[dict, ...] | None, - _loop: asyncio.AbstractEventLoop, - chat_template_kwargs: dict | None = None, - ) -> int: - if self._delay: - await asyncio.sleep(self._delay) - parts = [ - str(message.get(key, "")) - for message in messages - for key in ("content", "reasoning_content", "tool_calls") - if message.get(key) - ] - parts.extend(str(tool) for tool in tools or ()) - return len(" ".join(parts).split()) + outcomes = [] + for item in inputs: + if isinstance(item, TokenIdsInput): + outcomes.append(len(item.token_ids)) + elif isinstance(item, TextInput): + outcomes.append(len(item.text.split())) + elif isinstance(item, MessageInput): + tool_calls = ( + msgspec.json.encode(list(item.tool_calls)).decode() + if item.tool_calls + else "" + ) + combined = ( + (item.content or "") + + " " + + (item.reasoning or "") + + " " + + tool_calls + ) + outcomes.append(len(combined.split())) + elif isinstance(item, PromptInput): + parts = [ + str(message.get(key, "")) + for message in item.messages + for key in ("content", "reasoning_content", "tool_calls") + if message.get(key) + ] + parts.extend(str(tool) for tool in item.tools or ()) + outcomes.append(len(" ".join(parts).split())) + return outcomes def close(self) -> None: pass diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py index 79b2d1cb3..1046e601b 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -1203,10 +1203,7 @@ async def test_drain_failure_reports_pending_and_finalizes(self, tmp_path): loop = asyncio.get_event_loop() class FailingBatchTokenizer: - async def count_texts_async(self, texts, _loop, live=False): - raise RuntimeError("tokenizer backend died") - - async def token_count_message_async(self, *args): + async def count_batch_async(self, inputs, _loop, live=False): raise RuntimeError("tokenizer backend died") with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: @@ -1250,13 +1247,9 @@ async def test_drain_timeout_reports_pending_count(self, tmp_path): loop = asyncio.get_event_loop() class BlockingBatchTokenizer: - async def count_texts_async(self, texts, _loop, live=False): + async def count_batch_async(self, inputs, _loop, live=False): await asyncio.sleep(10.0) # exceeds drain timeout - return [0] * len(texts) - - async def token_count_message_async(self, *args): - await asyncio.sleep(10.0) - return 0 + return [0] * len(inputs) with ManagedZMQContext.scoped(socket_dir=str(tmp_path)) as ctx: agg, _, publisher = make_aggregator( diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py index ba60f1ffd..87df387c4 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py @@ -37,6 +37,9 @@ from inference_endpoint.async_utils.services.metrics_aggregator.token_metrics import ( TokenBatchQueue, ) +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, +) from inference_endpoint.core.record import ( EventRecord, SampleEventType, @@ -365,14 +368,12 @@ class CapturingTokenizer: def __init__(self): self.messages = [] - async def count_texts_async(self, *args, **kwargs): - raise AssertionError("reasoning output must not use flattened text") - - async def token_count_message_async( - self, content, reasoning, tool_calls, _loop - ): - self.messages.append((content, reasoning, tool_calls)) - return 4 + async def count_batch_async(self, inputs, _loop, live=False): + assert all(isinstance(item, MessageInput) for item in inputs) + self.messages.extend( + (item.content, item.reasoning, item.tool_calls) for item in inputs + ) + return [4] * len(inputs) registry = MetricsRegistry() registry.register_series("osl", hdr_low=1, hdr_high=100_000) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index 73ade78af..ef93b82ce 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -35,10 +35,24 @@ _worker_encode_lengths, encode_lengths, ) +from inference_endpoint.async_utils.services.metrics_aggregator.tokenization import ( + MessageInput, + PromptInput, + TextInput, + TokenIdsInput, +) _MOCK_TARGET = "inference_endpoint.async_utils.services.metrics_aggregator.token_metrics.AutoTokenizer" +@pytest.mark.unit +def test_tokenization_inputs_make_the_four_paths_explicit(): + assert TokenIdsInput((1, 2)).token_ids == (1, 2) + assert TextInput("hello").text == "hello" + assert MessageInput("answer", "thought", None).reasoning == "thought" + assert PromptInput(({"role": "user", "content": "hi"},), None, None).messages + + class _FakeTokenizer: """Deterministic tokenizer that splits on whitespace. @@ -89,7 +103,7 @@ async def test_count_texts_async(self): with patch(_MOCK_TARGET, _FakeTokenizerWithBackend): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: - counts = await tok.count_texts_async(["Hello world foo", "a"], loop) + counts = await tok._count_texts_async(["Hello world foo", "a"], loop) assert counts == [3, 1] @pytest.mark.asyncio @@ -97,7 +111,7 @@ async def test_count_texts_async_empty(self): with patch(_MOCK_TARGET, _FakeTokenizer): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: - assert await tok.count_texts_async([], loop) == [] + assert await tok._count_texts_async([], loop) == [] @pytest.mark.asyncio async def test_plain_text_requires_a_text_backend(self): @@ -107,7 +121,7 @@ async def test_plain_text_requires_a_text_backend(self): with pytest.raises( RuntimeError, match="plain-text tokenization.*backend" ): - await tok.count_texts_async(["Hello world"], loop) + await tok._count_texts_async(["Hello world"], loop) @pytest.mark.asyncio async def test_count_texts_async_sharded(self): @@ -116,7 +130,7 @@ async def test_count_texts_async_sharded(self): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: tok._procs = [_FakeProc(), _FakeProc()] - counts = await tok.count_texts_async(["a", "b b", "c c c", "d"], loop) + counts = await tok._count_texts_async(["a", "b b", "c c c", "d"], loop) assert counts == [1, 2, 3, 1] @pytest.mark.asyncio @@ -127,7 +141,7 @@ async def test_count_texts_async_shard_failure_propagates(self): with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: tok._procs = [_BrokenProc()] with pytest.raises(BrokenProcessPool): - await tok.count_texts_async(["a b"], loop) + await tok._count_texts_async(["a b"], loop) def test_close_is_idempotent(self): with patch(_MOCK_TARGET, _FakeTokenizer): @@ -142,7 +156,7 @@ async def test_use_after_close_raises(self): tok = BatchTokenizer("fake", n_workers=0, live_workers=2) tok.close() with pytest.raises(RuntimeError, match="closed"): - await tok.count_texts_async(["hello"], loop) + await tok._count_texts_async(["hello"], loop) class _FakeTokenizerWithTemplate(_FakeTokenizer): @@ -224,7 +238,11 @@ async def test_token_count_prompt_preserves_messages_tools_and_generation_prompt }, ) - count = await tok.token_count_prompt_async(messages, tools, loop) + count = ( + await tok.count_batch_async( + [PromptInput(messages, tools, None)], loop + ) + )[0] # Wrapper + question + reasoning + tool call + tool result + # declared tool + generation prompt. @@ -232,14 +250,16 @@ async def test_token_count_prompt_preserves_messages_tools_and_generation_prompt @pytest.mark.asyncio async def test_token_count_message_subtracts_baseline(self): - """token_count_message_async returns full_tokens - baseline.""" + """Structured message counting returns full tokens minus baseline.""" with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: # "hello world" -> 2 content + 2 wrapper = 4; baseline = 0, prefix = 2 - count = await tok.token_count_message_async( - "hello world", None, None, loop - ) + count = ( + await tok.count_batch_async( + [MessageInput("hello world", None, None)], loop + ) + )[0] assert count == 2 @pytest.mark.asyncio @@ -255,10 +275,16 @@ async def test_token_count_message_includes_tool_calls(self): "function": {"name": "f", "arguments": "{}"}, }, ) - without = await tok.token_count_message_async("hello", None, None, loop) - with_calls = await tok.token_count_message_async( - "hello", None, tool_calls, loop - ) + without = ( + await tok.count_batch_async( + [MessageInput("hello", None, None)], loop + ) + )[0] + with_calls = ( + await tok.count_batch_async( + [MessageInput("hello", None, tool_calls)], loop + ) + )[0] assert with_calls > without @pytest.mark.asyncio @@ -280,9 +306,11 @@ def apply_chat_template(self, *args, **kwargs): }, ) # Must not raise; falls back to whitespace tokenizer. - count = await tok.token_count_message_async( - "hello world", None, tool_calls, loop - ) + count = ( + await tok.count_batch_async( + [MessageInput("hello world", None, tool_calls)], loop + ) + )[0] assert count > 0 def test_token_count_prompt_falls_back_to_text_only_multimodal_messages(self): @@ -418,6 +446,29 @@ class _FakeTokenizerWithBackend(_FakeTokenizer): backend_tokenizer = _FastBackend() +class _FakeTokenizerWithTemplateAndBackend(_FakeTokenizerWithTemplate): + backend_tokenizer = _FastBackend() + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_count_batch_routes_all_four_inputs_and_preserves_order(): + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + outcomes = await tok.count_batch_async( + [ + TokenIdsInput((1, 2, 3)), + TextInput("plain text"), + MessageInput("answer here", None, None), + PromptInput(({"role": "user", "content": "ask now"},), None, None), + ], + loop, + ) + + assert outcomes == [3, 2, 2, 5] + + class _SpawnlessExecutor: """Stands in for ProcessPoolExecutor: records ctor args, instant warmup.""" @@ -531,7 +582,7 @@ async def test_live_never_touches_the_shard_pool(self): with BatchTokenizer("fake", n_workers=0, live_workers=1) as tok: procs = [_RecordingProc(), _RecordingProc(), _RecordingProc()] tok._procs = procs - counts = await tok.count_texts_async(["a b", "c"], loop, live=True) + counts = await tok._count_texts_async(["a b", "c"], loop, live=True) assert counts == [2, 1] assert all(p.chunks == [] for p in procs) @@ -542,7 +593,7 @@ async def test_drain_uses_every_shard(self): with BatchTokenizer("fake", n_workers=0, live_workers=1) as tok: procs = [_RecordingProc(), _RecordingProc()] tok._procs = procs - await tok.count_texts_async(["a", "b", "c", "d"], loop) + await tok._count_texts_async(["a", "b", "c", "d"], loop) assert all(p.chunks for p in procs) @@ -553,7 +604,7 @@ async def test_start_live_flushes_periodically(self): loop = asyncio.get_running_loop() queue = TokenBatchQueue(_CapturingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("a b c", recorded.append) + queue.enqueue(TextInput("a b c"), recorded.append) queue.start_live(0.01) queue.start_live(0.01) # idempotent await asyncio.sleep(0.05) @@ -563,15 +614,15 @@ async def test_start_live_flushes_periodically(self): async def test_live_loop_survives_tokenizer_failure(self): class _FailingLive(_CapturingTokenizer): - async def count_texts_async(self, texts, _loop, live=False): + async def count_batch_async(self, inputs, _loop, live=False): if live: raise RuntimeError("live lane boom") - return await super().count_texts_async(texts, _loop) + return await super().count_batch_async(inputs, _loop) loop = asyncio.get_running_loop() queue = TokenBatchQueue(_FailingLive(), loop) recorded: list[int] = [] - queue.enqueue_text("a b", recorded.append) + queue.enqueue(TextInput("a b"), recorded.append) queue.start_live(0.01) await asyncio.sleep(0.05) assert recorded == [] @@ -628,7 +679,7 @@ async def test_live_flush_takes_at_most_the_cap(self, monkeypatch): queue = TokenBatchQueue(_CapturingTokenizer(), loop) recorded: list[int] = [] for i in range(5): - queue.enqueue_text(f"t{i}", recorded.append) + queue.enqueue(TextInput(f"t{i}"), recorded.append) await queue.flush_live_once() assert len(recorded) == 3 assert queue.pending == 2 @@ -638,22 +689,22 @@ async def test_live_flush_takes_at_most_the_cap(self, monkeypatch): async def test_live_cancellation_requeues_texts(self): class _Hanging(_CapturingTokenizer): - async def count_texts_async(self, texts, _loop, live=False): + async def count_batch_async(self, inputs, _loop, live=False): if live: await asyncio.sleep(30) - return await super().count_texts_async(texts, _loop) + return await super().count_batch_async(inputs, _loop) loop = asyncio.get_running_loop() queue = TokenBatchQueue(_Hanging(), loop) recorded: list[int] = [] - queue.enqueue_text("a b", recorded.append) + queue.enqueue(TextInput("a b"), recorded.append) task = loop.create_task(queue.flush_live_once()) await asyncio.sleep(0.01) task.cancel() with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(task, timeout=1.0) assert queue.pending == 1 - assert len(queue._text) == 1, "cancelled live flush must give items back" + assert len(queue._items) == 1, "cancelled live flush must give items back" assert await queue.flush_remaining(timeout=1.0) == 0 assert recorded == [2] @@ -661,40 +712,39 @@ async def test_live_cancellation_requeues_messages_too(self): """A cancel landing in the text encode must give back BOTH kinds.""" class _Hanging(_CapturingTokenizer): - async def count_texts_async(self, texts, _loop, live=False): + async def count_batch_async(self, inputs, _loop, live=False): if live: await asyncio.sleep(30) - return await super().count_texts_async(texts, _loop) + return await super().count_batch_async(inputs, _loop) loop = asyncio.get_running_loop() queue = TokenBatchQueue(_Hanging(), loop) recorded: list[int] = [] - queue.enqueue_text("a b", recorded.append) - queue.enqueue_message(("hello world", None, None), recorded.append) + queue.enqueue(TextInput("a b"), recorded.append) + queue.enqueue(MessageInput("hello world", None, None), recorded.append) task = loop.create_task(queue.flush_live_once()) await asyncio.sleep(0.01) task.cancel() with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(task, timeout=1.0) assert queue.pending == 2 - assert len(queue._text) == 1 - assert len(queue._msg) == 1, "detached messages must be re-queued" + assert len(queue._items) == 2, "all detached items must be re-queued" assert await queue.flush_remaining(timeout=1.0) == 0 assert sorted(recorded) == [2, 2] async def test_live_message_failure_requeues_message(self): class _MsgFailing(_CapturingTokenizer): - async def token_count_message_async(self, *args): - raise RuntimeError("template boom") + async def count_batch_async(self, inputs, _loop, live=False): + return [RuntimeError("template boom") for _ in inputs] loop = asyncio.get_running_loop() queue = TokenBatchQueue(_MsgFailing(), loop) recorded: list[int] = [] - queue.enqueue_message(("hello world", None, None), recorded.append) + queue.enqueue(MessageInput("hello world", None, None), recorded.append) with pytest.raises(RuntimeError, match="template boom"): await queue.flush_live_once() assert queue.pending == 1 - assert len(queue._msg) == 1, "failed live message must be re-queued" + assert len(queue._items) == 1, "failed live message must be re-queued" @pytest.mark.unit @@ -721,12 +771,22 @@ def test_preserves_order_and_bounds_chunk_count(self): class _CapturingTokenizer: """Minimal tokenizer stub for queue tests: whitespace counts, no procs.""" - async def count_texts_async(self, texts, _loop, live=False): - return [len(t.split()) for t in texts] - - async def token_count_message_async(self, content, reasoning, tool_calls, _loop): - parts = [p for p in (content, reasoning) if p] - return len(" ".join(parts).split()) + (len(tool_calls) if tool_calls else 0) + async def count_batch_async(self, inputs, _loop, live=False): + outcomes = [] + for item in inputs: + if isinstance(item, TokenIdsInput): + outcomes.append(len(item.token_ids)) + elif isinstance(item, TextInput): + outcomes.append(len(item.text.split())) + elif isinstance(item, MessageInput): + parts = [p for p in (item.content, item.reasoning) if p] + outcomes.append( + len(" ".join(parts).split()) + + (len(item.tool_calls) if item.tool_calls else 0) + ) + elif isinstance(item, PromptInput): + outcomes.append(len(item.messages)) + return outcomes @pytest.mark.unit @@ -736,8 +796,8 @@ async def test_flush_records_text_via_callback(self): loop = asyncio.get_running_loop() queue = TokenBatchQueue(_CapturingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("a b c", recorded.append) - queue.enqueue_text("d e", recorded.append) + queue.enqueue(TextInput("a b c"), recorded.append) + queue.enqueue(TextInput("d e"), recorded.append) assert queue.pending == 2 await queue.drain_all() assert sorted(recorded) == [2, 3] @@ -747,7 +807,7 @@ async def test_flush_records_message_via_callback(self): loop = asyncio.get_running_loop() queue = TokenBatchQueue(_CapturingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_message(("hello world", None, None), recorded.append) + queue.enqueue(MessageInput("hello world", None, None), recorded.append) await queue.drain_all() assert recorded == [2] @@ -757,18 +817,16 @@ async def test_flush_wrong_length_text_result_isolated(self): message still records, and the failure is raised after both phases.""" class _WrongLength: - async def count_texts_async(self, texts, _loop, live=False): - return [1] # too few regardless of input length - - async def token_count_message_async(self, *args): - return 7 + async def count_batch_async(self, inputs, _loop, live=False): + error = RuntimeError("tokenizer returned 1 counts for 2 texts") + return [error, error, 7] loop = asyncio.get_running_loop() queue = TokenBatchQueue(_WrongLength(), loop) recorded: list[int] = [] - queue.enqueue_text("a b", recorded.append) - queue.enqueue_text("c d", recorded.append) - queue.enqueue_message(("hi", None, None), recorded.append) + queue.enqueue(TextInput("a b"), recorded.append) + queue.enqueue(TextInput("c d"), recorded.append) + queue.enqueue(MessageInput("hi", None, None), recorded.append) assert queue.pending == 3 with pytest.raises(RuntimeError, match="counts for"): @@ -787,7 +845,7 @@ async def test_flush_remaining_clean_returns_zero(self): loop = asyncio.get_running_loop() queue = TokenBatchQueue(_CapturingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("a b", recorded.append) + queue.enqueue(TextInput("a b"), recorded.append) assert await queue.flush_remaining(timeout=5.0) == 0 assert recorded == [2] @@ -795,17 +853,14 @@ async def test_flush_remaining_timeout_reports_pending(self): """A tokenizer slower than the budget leaves items pending.""" class _BlockingTokenizer: - async def count_texts_async(self, texts, _loop, live=False): + async def count_batch_async(self, inputs, _loop, live=False): await asyncio.sleep(10.0) - return [0] * len(texts) - - async def token_count_message_async(self, *args): - return 0 + return [0] * len(inputs) loop = asyncio.get_running_loop() queue = TokenBatchQueue(_BlockingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("never counted", recorded.append) + queue.enqueue(TextInput("never counted"), recorded.append) n_pending = await queue.flush_remaining(timeout=0.05) assert n_pending == 1 assert recorded == [] @@ -814,16 +869,13 @@ async def test_flush_remaining_failure_reports_pending(self): """A tokenizer error leaves items pending and never raises.""" class _FailingTokenizer: - async def count_texts_async(self, texts, _loop, live=False): - raise RuntimeError("tokenizer boom") - - async def token_count_message_async(self, *args): + async def count_batch_async(self, inputs, _loop, live=False): raise RuntimeError("tokenizer boom") loop = asyncio.get_running_loop() queue = TokenBatchQueue(_FailingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("x y", recorded.append) + queue.enqueue(TextInput("x y"), recorded.append) assert await queue.flush_remaining(timeout=5.0) == 1 assert recorded == [] @@ -831,19 +883,19 @@ async def test_flush_text_failure_does_not_drop_message_items(self): """The message phase runs (and records) even when the text batch fails.""" class _TextFailingTokenizer: - async def count_texts_async(self, texts, _loop, live=False): - raise RuntimeError("text shard died") - - async def token_count_message_async( - self, content, reasoning, tool_calls, _loop - ): - return len(content.split()) + async def count_batch_async(self, inputs, _loop, live=False): + return [ + RuntimeError("text shard died") + if isinstance(item, TextInput) + else len(item.content.split()) + for item in inputs + ] loop = asyncio.get_running_loop() queue = TokenBatchQueue(_TextFailingTokenizer(), loop) recorded: list[int] = [] - queue.enqueue_text("never counted", recorded.append) - queue.enqueue_message(("hello world", None, None), recorded.append) + queue.enqueue(TextInput("never counted"), recorded.append) + queue.enqueue(MessageInput("hello world", None, None), recorded.append) with pytest.raises(RuntimeError, match="text shard died"): await queue.drain_all() assert recorded == [2], "message item must survive the text failure" @@ -858,8 +910,8 @@ async def test_flush_recorder_failure_does_not_poison_batch(self): def bad_recorder(count: int) -> None: raise ValueError("recorder bug") - queue.enqueue_text("a b", bad_recorder) - queue.enqueue_text("c d e", recorded.append) + queue.enqueue(TextInput("a b"), bad_recorder) + queue.enqueue(TextInput("c d e"), recorded.append) await queue.drain_all() assert recorded == [3] assert queue.pending == 0, "a raising recorder still counts as recorded" From 31dd129180208c036cfc859b4d357c97d01bddc8 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 21:28:53 -0700 Subject: [PATCH 05/11] Handle structured prompt tokenization failures --- .../metrics_aggregator/token_metrics.py | 46 +++++-------------- src/inference_endpoint/core/types.py | 4 ++ .../metrics_aggregator/test_token_metrics.py | 4 +- 3 files changed, 17 insertions(+), 37 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index abe818d1e..33f3272bf 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -111,29 +111,6 @@ def _normalize_tool_calls_for_template( return normalized -def _text_only_multimodal_messages( - messages: tuple[dict[str, Any], ...] | list[dict[str, Any]], -) -> tuple[list[dict[str, Any]], bool]: - """Replace multimodal content parts with their text for text-only templates.""" - normalized: list[dict[str, Any]] = [] - changed = False - for message in messages: - content = message.get("content") - if not isinstance(content, list): - normalized.append(message) - continue - changed = True - text = " ".join( - part["text"] - for part in content - if isinstance(part, dict) - and part.get("type") == "text" - and isinstance(part.get("text"), str) - ) - normalized.append({**message, "content": text}) - return normalized, changed - - def _normalize_prompt_messages_for_template( messages: tuple[dict[str, Any], ...], ) -> list[dict[str, Any]]: @@ -519,22 +496,21 @@ def _token_count_prompt( encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] prompt_messages, **kwargs ) - except Exception: - text_only_messages, changed = _text_only_multimodal_messages(messages) - if not changed: - raise - key = f"{self._tokenizer_name}:multimodal-prompt" + return len(encoded) + except Exception as exc: + key = f"{self._tokenizer_name}:{type(exc).__name__}" if key not in self._fallback_warned: self._fallback_warned.add(key) - logger.warning( - "Chat template for %s rejected multimodal messages; " - "retrying with text content only", + logger.exception( + "apply_chat_template failed for %s (%s); falling back to " + "whitespace tokenization. Structured ISL may diverge.", self._tokenizer_name, + type(exc).__name__, ) - encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] - text_only_messages, **kwargs - ) - return len(encoded) + prompt = {"messages": prompt_messages} + if tools is not None: + prompt["tools"] = list(tools) + return self._token_count_text(msgspec.json.encode(prompt).decode()) async def count_batch_async( self, diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index 942331c9d..b64635cb4 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -284,6 +284,10 @@ class PromptData( ): # type: ignore[call-arg] """Prompt input data attached to ISSUED events for ISL computation. + AT-RISK (gc=False): Has mutable container fields ``messages``, ``tools``, + and ``chat_template_kwargs``. Any change that introduces cyclic references + must be audited; if cycles become possible, remove ``gc=False``. + Exactly one of ``text``, ``token_ids``, or ``messages`` should be set: - ``text``: raw prompt string (OpenAI path) — requires tokenization for ISL. - ``token_ids``: pre-tokenized token ID list (SGLang/Harmonize path) — ISL is len(). diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index ef93b82ce..a01e48d65 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -313,7 +313,7 @@ def apply_chat_template(self, *args, **kwargs): )[0] assert count > 0 - def test_token_count_prompt_falls_back_to_text_only_multimodal_messages(self): + def test_token_count_prompt_falls_back_on_template_error(self): class _TextOnlyTemplateTokenizer(_FakeTokenizerWithTemplate): def apply_chat_template(self, messages, **kwargs): if any( @@ -337,7 +337,7 @@ def apply_chat_template(self, messages, **kwargs): None, ) - assert count == 6 + assert count > 0 def test_token_count_prompt_normalizes_tools_and_forwards_template_kwargs(self): class _RecordingTokenizer(_FakeTokenizerWithTemplate): From 584547c99def4c59f7dda9f4cea80af80b4bebe0 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 21:35:04 -0700 Subject: [PATCH 06/11] Tighten tokenizer backend handling --- .../async_utils/services/metrics_aggregator/token_metrics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 33f3272bf..75f18856b 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -190,7 +190,7 @@ def encode_lengths(backend: Any, texts: list[str]) -> list[int]: """Per-text token counts via one bounded backend batch call.""" encode_batch = getattr(backend, "encode_batch_fast", None) or backend.encode_batch encoded = encode_batch(texts, add_special_tokens=False) - return [len(getattr(item, "ids", item)) for item in encoded] + return [len(item.ids) for item in encoded] def _worker_encode_lengths(texts: list[str]) -> list[int]: @@ -261,6 +261,7 @@ def __init__( os.environ.setdefault("RAYON_NUM_THREADS", str(max(1, live_workers))) self._fallback_warned: set[str] = set() self._tokenizer: PreTrainedTokenizerBase | None = None + self._text_backend: Any | None = None self._prefix_len = 0 self._baseline = 0 # In-process threads: the live token-metric lane plus the From 77e9a6e4d69ec9d8d8affa949c5adb422c3fd2d6 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 22:03:29 -0700 Subject: [PATCH 07/11] Support tokenizer wrapper text fallback --- .../metrics_aggregator/metrics_table.py | 1 + .../metrics_aggregator/token_metrics.py | 20 ++++++++------ .../metrics_aggregator/tokenization.py | 1 + src/inference_endpoint/core/types.py | 2 ++ .../load_generator/session.py | 5 ++++ .../metrics_aggregator/test_token_metrics.py | 27 +++++++++++++------ tests/unit/core/test_record.py | 3 +++ .../unit/load_generator/test_async_session.py | 2 ++ 8 files changed, 45 insertions(+), 16 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index 5fa263810..901a5fae1 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -311,6 +311,7 @@ def _extract_tokenization_input(self, ev_rec, row, pre_change): ev_rec.data.messages, ev_rec.data.tools, ev_rec.data.chat_template_kwargs, + ev_rec.data.chat_template, ) return None diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index 75f18856b..fb38d7119 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -331,8 +331,8 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: return if self._text_backend is None: logger.info( - "BatchTokenizer: no plain-text backend for %s; structured " - "chat tokenization remains available", + "BatchTokenizer: no fast backend for %s; using the tokenizer " + "wrapper for in-process plain-text tokenization", self._tokenizer_name, ) return @@ -392,12 +392,12 @@ def _setup_shards(self, cores_per_worker: int, n_workers: int) -> None: def _encode_lengths_inproc(self, texts: list[str]) -> list[int]: backend = self._text_backend - if backend is None: - raise RuntimeError( - f"plain-text tokenization for {self._tokenizer_name!r} requires " - "a supported fast backend" - ) - return encode_lengths(backend, texts) + if backend is not None: + return encode_lengths(backend, texts) + tokenizer = self._tokenizer + if tokenizer is None: + raise RuntimeError("BatchTokenizer is closed") + return [len(tokenizer.encode(text, add_special_tokens=False)) for text in texts] async def _count_texts_async( self, @@ -483,6 +483,7 @@ def _token_count_prompt( messages: tuple[dict[str, Any], ...], tools: tuple[dict[str, Any], ...] | None, chat_template_kwargs: dict[str, Any] | None = None, + chat_template: str | None = None, ) -> int: kwargs = dict(chat_template_kwargs or {}) kwargs.update( @@ -492,6 +493,8 @@ def _token_count_prompt( ) if tools is not None: kwargs["tools"] = list(tools) + if chat_template is not None: + kwargs["chat_template"] = chat_template prompt_messages = _normalize_prompt_messages_for_template(messages) try: encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] @@ -572,6 +575,7 @@ async def count_batch_async( item.messages, item.tools, item.chat_template_kwargs, + item.chat_template, ) except Exception as exc: # noqa: BLE001 - isolate this input. outcomes[index] = exc diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py index 98614892c..032e188f5 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py @@ -44,6 +44,7 @@ class PromptInput: messages: tuple[dict[str, Any], ...] tools: tuple[dict[str, Any], ...] | None chat_template_kwargs: dict[str, Any] | None + chat_template: str | None TokenizationInput: TypeAlias = ( # noqa: UP040 - mypy version lacks PEP 695. diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index b64635cb4..6d25bf5c5 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -299,6 +299,7 @@ class PromptData( messages: Structured messages sent to a chat-completions endpoint. tools: Tool declarations accompanying ``messages``. chat_template_kwargs: Model-specific arguments used to render ``messages``. + chat_template: Per-request template used to render ``messages``. """ text: str | None = None @@ -306,6 +307,7 @@ class PromptData( messages: tuple[dict[str, Any], ...] | None = None tools: tuple[dict[str, Any], ...] | None = None chat_template_kwargs: dict[str, Any] | None = None + chat_template: str | None = None class ErrorData( diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index df10d53c5..30128875a 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -272,6 +272,11 @@ def issue( if isinstance(chat_template_kwargs, dict) else None ), + chat_template=( + data["chat_template"] + if isinstance(data.get("chat_template"), str) + else None + ), ) elif isinstance(prompt, str): prompt_data = PromptData(text=prompt) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index a01e48d65..d5f190270 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -50,7 +50,7 @@ def test_tokenization_inputs_make_the_four_paths_explicit(): assert TokenIdsInput((1, 2)).token_ids == (1, 2) assert TextInput("hello").text == "hello" assert MessageInput("answer", "thought", None).reasoning == "thought" - assert PromptInput(({"role": "user", "content": "hi"},), None, None).messages + assert PromptInput(({"role": "user", "content": "hi"},), None, None, None).messages class _FakeTokenizer: @@ -66,6 +66,10 @@ def __init__(self, load_delay: float = 0.0): def tokenize(self, text: str) -> list[str]: return text.split() + def encode(self, text: str, *, add_special_tokens: bool) -> list[int]: + assert add_special_tokens is False + return list(range(len(text.split()))) + @classmethod def from_pretrained(cls, name: str, **kwargs: object) -> "_FakeTokenizer": assert kwargs == {"trust_remote_code": True} @@ -114,14 +118,14 @@ async def test_count_texts_async_empty(self): assert await tok._count_texts_async([], loop) == [] @pytest.mark.asyncio - async def test_plain_text_requires_a_text_backend(self): + async def test_plain_text_falls_back_to_tokenizer_encode_without_backend(self): with patch(_MOCK_TARGET, _FakeTokenizer): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: - with pytest.raises( - RuntimeError, match="plain-text tokenization.*backend" - ): - await tok._count_texts_async(["Hello world"], loop) + counts = await tok._count_texts_async( + ["Hello world", "one two three"], loop + ) + assert counts == [2, 3] @pytest.mark.asyncio async def test_count_texts_async_sharded(self): @@ -240,7 +244,7 @@ async def test_token_count_prompt_preserves_messages_tools_and_generation_prompt count = ( await tok.count_batch_async( - [PromptInput(messages, tools, None)], loop + [PromptInput(messages, tools, None, None)], loop ) )[0] @@ -371,6 +375,7 @@ def apply_chat_template(self, messages, **kwargs): messages, None, {"enable_thinking": False}, + "custom template", ) normalized_call = _RecordingTokenizer.last_messages[0]["tool_calls"][0] @@ -379,6 +384,7 @@ def apply_chat_template(self, messages, **kwargs): '{"city": "SF"}' ) assert _RecordingTokenizer.last_kwargs["enable_thinking"] is False + assert _RecordingTokenizer.last_kwargs["chat_template"] == "custom template" assert _RecordingTokenizer.last_kwargs["return_dict"] is False @@ -461,7 +467,12 @@ async def test_count_batch_routes_all_four_inputs_and_preserves_order(): TokenIdsInput((1, 2, 3)), TextInput("plain text"), MessageInput("answer here", None, None), - PromptInput(({"role": "user", "content": "ask now"},), None, None), + PromptInput( + ({"role": "user", "content": "ask now"},), + None, + None, + None, + ), ], loop, ) diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index ad2920312..c23e06fea 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -183,6 +183,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): ) tools = ({"function": {"name": "lookup"}},) chat_template_kwargs = {"enable_thinking": False} + chat_template = "custom template" record = EventRecord( event_type=SampleEventType.ISSUED, sample_uuid="sample-chat", @@ -190,6 +191,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): messages=messages, tools=tools, chat_template_kwargs=chat_template_kwargs, + chat_template=chat_template, ), ) @@ -200,6 +202,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): assert decoded.data.messages == messages assert decoded.data.tools == tools assert decoded.data.chat_template_kwargs == chat_template_kwargs + assert decoded.data.chat_template == chat_template def test_error_event_round_trips_with_error_data(self): record = EventRecord( diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index aa619f4e9..14ddc8026 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -195,6 +195,7 @@ def load_sample(self, index: int) -> dict: "function": {"name": "lookup", "parameters": {}}, } ], + "chat_template": "custom template", "chat_template_kwargs": {"enable_thinking": False}, } @@ -208,6 +209,7 @@ def load_sample(self, index: int) -> dict: prompt = publisher.events_of_type(SampleEventType.ISSUED)[0].data assert prompt.messages == tuple(issuer.issued_queries[0].data["messages"]) assert prompt.tools == tuple(issuer.issued_queries[0].data["tools"]) + assert prompt.chat_template == "custom template" assert prompt.chat_template_kwargs == {"enable_thinking": False} assert prompt.text is None From 190055fa6399353183d4d3956f858634083dcdbb Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Tue, 11 Aug 2026 22:22:11 -0700 Subject: [PATCH 08/11] Refine prompt representation handling --- .../load_generator/session.py | 58 +++++++++---------- .../unit/load_generator/test_async_session.py | 55 +++++++++++++++--- 2 files changed, 76 insertions(+), 37 deletions(-) diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index 30128875a..e37c300ef 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -155,6 +155,7 @@ class PhaseIssuer: "_issuer", "_on_inflight_drained", "_performance_tracking_stopped", + "_prompt_warning_reasons", "_publisher", "_stop_check", "uuid_to_index", @@ -183,6 +184,14 @@ def __init__( self.inflight: int = 0 self.issued_count: int = 0 self._performance_tracking_stopped = False + self._prompt_warning_reasons: set[str] = set() + + def _warn_prompt_once(self, reason: str, message: str) -> None: + """Warn once per phase when ISL cannot be derived from a sample.""" + if reason in self._prompt_warning_reasons: + return + self._prompt_warning_reasons.add(reason) + logger.warning(message) def mark_inflight_complete(self) -> None: self.inflight -= 1 @@ -240,27 +249,13 @@ def issue( token_ids = data.get("token_ids") messages = data.get("messages") prompt = data.get("prompt") - representations = [ - name - for name, present in ( - ("input_tokens", input_tokens is not None), - ("token_ids", token_ids is not None), - ("messages", isinstance(messages, list | tuple) and bool(messages)), - ("prompt", isinstance(prompt, str)), - ) - if present - ] - if len(representations) > 1: - raise ValueError( - "sample contains multiple prompt representations: " - + ", ".join(representations) - ) + if input_tokens is not None and token_ids is not None: + raise ValueError("sample contains both input_tokens and token_ids") - if input_tokens is not None or token_ids is not None: - selected_token_ids = ( - input_tokens if input_tokens is not None else token_ids - ) - prompt_data = PromptData(token_ids=tuple(selected_token_ids)) + if input_tokens is not None: + prompt_data = PromptData(token_ids=tuple(input_tokens)) + elif token_ids is not None: + prompt_data = PromptData(token_ids=tuple(token_ids)) elif isinstance(messages, list | tuple) and messages: tools = data.get("tools") chat_template_kwargs = data.get("chat_template_kwargs") @@ -281,17 +276,22 @@ def issue( elif isinstance(prompt, str): prompt_data = PromptData(text=prompt) else: - logger.warning( - "Sample %s has no supported prompt representation for ISL; " - "expected token IDs, non-empty messages, or a string prompt", - sample_index, - ) + if isinstance(prompt, list): + self._warn_prompt_once( + "list_prompt", + "List-form prompts are issued normally, but ISL is unavailable", + ) + else: + self._warn_prompt_once( + "unsupported_mapping", + "Samples without token IDs, non-empty messages, or a string " + "prompt are issued normally, but ISL is unavailable", + ) prompt_data = PromptData() else: - logger.warning( - "Sample %s has no supported prompt representation for ISL; " - "expected a mapping", - sample_index, + self._warn_prompt_once( + "non_mapping", + "Non-mapping samples are issued normally, but ISL is unavailable", ) prompt_data = PromptData() self._publisher.publish( diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 14ddc8026..b4903b224 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -213,21 +213,34 @@ def load_sample(self, index: int) -> dict: assert prompt.chat_template_kwargs == {"enable_thinking": False} assert prompt.text is None - def test_issue_rejects_multiple_prompt_representations(self): - class AmbiguousDataset(FakeDataset): + def test_issue_messages_take_precedence_over_prompt(self): + class MessagesDataset(FakeDataset): def load_sample(self, index: int) -> dict: return { "messages": [{"role": "user", "content": "question"}], - "prompt": "question", + "prompt": "generic fallback", } issuer = FakeIssuer() publisher = FakePublisher() + phase_issuer = PhaseIssuer(MessagesDataset(1), issuer, publisher, lambda: False) + + phase_issuer.issue(0) + + prompt = publisher.events_of_type(SampleEventType.ISSUED)[0].data + assert prompt.messages == ({"role": "user", "content": "question"},) + assert prompt.text is None + + def test_issue_rejects_both_token_id_fields(self): + class ConflictingTokensDataset(FakeDataset): + def load_sample(self, index: int) -> dict: + return {"input_tokens": [1, 2], "token_ids": [3, 4]} + phase_issuer = PhaseIssuer( - AmbiguousDataset(1), issuer, publisher, lambda: False + ConflictingTokensDataset(1), FakeIssuer(), FakePublisher(), lambda: False ) - with pytest.raises(ValueError, match="multiple prompt representations"): + with pytest.raises(ValueError, match="both input_tokens and token_ids"): phase_issuer.issue(0) def test_issue_empty_messages_falls_through_to_prompt(self): @@ -247,7 +260,7 @@ def load_sample(self, index: int) -> dict: assert prompt.messages is None assert prompt.text == "fallback prompt" - def test_issue_warns_when_no_supported_prompt_representation(self, caplog): + def test_issue_warns_once_for_list_prompt_without_isl_representation(self, caplog): class UnsupportedPromptDataset(FakeDataset): def load_sample(self, index: int) -> dict: return {"prompt": [{"type": "image_url"}]} @@ -255,13 +268,39 @@ def load_sample(self, index: int) -> dict: issuer = FakeIssuer() publisher = FakePublisher() phase_issuer = PhaseIssuer( - UnsupportedPromptDataset(1), issuer, publisher, lambda: False + UnsupportedPromptDataset(2), issuer, publisher, lambda: False + ) + + with caplog.at_level("WARNING"): + phase_issuer.issue(0) + phase_issuer.issue(1) + + warnings = [ + record + for record in caplog.records + if "List-form prompts are issued normally" in record.message + ] + assert len(warnings) == 1 + + def test_issue_warns_once_for_non_mapping_sample(self, caplog): + class NonMappingDataset(FakeDataset): + def load_sample(self, index: int): + return ["prompt"] + + phase_issuer = PhaseIssuer( + NonMappingDataset(2), FakeIssuer(), FakePublisher(), lambda: False ) with caplog.at_level("WARNING"): phase_issuer.issue(0) + phase_issuer.issue(1) - assert "no supported prompt representation" in caplog.text + warnings = [ + record + for record in caplog.records + if "Non-mapping samples are issued normally" in record.message + ] + assert len(warnings) == 1 def test_issue_returns_none_when_stopped(self): dataset = FakeDataset(5) From 193f111d106dc2c840fc82a494f32808a8e6d631 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Wed, 12 Aug 2026 13:21:59 -0700 Subject: [PATCH 09/11] Address prompt warning and formatting feedback --- src/inference_endpoint/load_generator/session.py | 7 +------ .../metrics_aggregator/test_token_metrics.py | 12 ++++++------ tests/unit/load_generator/test_async_session.py | 9 ++------- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index e37c300ef..aa70fd80e 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -276,12 +276,7 @@ def issue( elif isinstance(prompt, str): prompt_data = PromptData(text=prompt) else: - if isinstance(prompt, list): - self._warn_prompt_once( - "list_prompt", - "List-form prompts are issued normally, but ISL is unavailable", - ) - else: + if not isinstance(prompt, list): self._warn_prompt_once( "unsupported_mapping", "Samples without token IDs, non-empty messages, or a string " diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index d5f190270..82f2f5db8 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -955,18 +955,18 @@ def test_terminate_procs_kills_running_workers(): assert future.running(), "worker task did not start" procs = list((getattr(ex, "_processes", None) or {}).values()) assert procs, "worker did not spawn" - assert not wait_for_process( - [p.sentinel for p in procs], timeout=0 - ), "worker exited before termination" + assert not wait_for_process([p.sentinel for p in procs], timeout=0), ( + "worker exited before termination" + ) _terminate_procs([ex]) for p in procs: # The executor manager may reap the child concurrently; sentinel # readiness observes exit without racing its return-code update. - assert wait_for_process( - [p.sentinel], timeout=5 - ), "worker was not terminated" + assert wait_for_process([p.sentinel], timeout=5), ( + "worker was not terminated" + ) finally: cleanup_procs = procs or list((getattr(ex, "_processes", None) or {}).values()) for p in cleanup_procs: diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index b4903b224..58048f0aa 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -260,7 +260,7 @@ def load_sample(self, index: int) -> dict: assert prompt.messages is None assert prompt.text == "fallback prompt" - def test_issue_warns_once_for_list_prompt_without_isl_representation(self, caplog): + def test_issue_list_prompt_without_isl_representation_does_not_warn(self, caplog): class UnsupportedPromptDataset(FakeDataset): def load_sample(self, index: int) -> dict: return {"prompt": [{"type": "image_url"}]} @@ -275,12 +275,7 @@ def load_sample(self, index: int) -> dict: phase_issuer.issue(0) phase_issuer.issue(1) - warnings = [ - record - for record in caplog.records - if "List-form prompts are issued normally" in record.message - ] - assert len(warnings) == 1 + assert not caplog.records def test_issue_warns_once_for_non_mapping_sample(self, caplog): class NonMappingDataset(FakeDataset): From 128e7454386858ce07c1c384bd34d32b43d3c27f Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Wed, 12 Aug 2026 13:50:33 -0700 Subject: [PATCH 10/11] Apply repository formatter output --- .../metrics_aggregator/test_aggregator.py | 18 +++++++++--------- .../metrics_aggregator/test_token_metrics.py | 12 ++++++------ 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py index 1046e601b..f91daa567 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py @@ -492,9 +492,9 @@ async def test_duplicate_started_logs_error_and_preserves_state( error_records = [ r for r in caplog.records if "Duplicate STARTED" in r.message ] - assert len(error_records) == 1, ( - "duplicate STARTED must log exactly one error" - ) + assert ( + len(error_records) == 1 + ), "duplicate STARTED must log exactly one error" assert "1000" in error_records[0].getMessage() assert "5000" in error_records[0].getMessage() finally: @@ -1123,9 +1123,9 @@ async def test_started_arms_the_live_flush_loop(self, tmp_path): assert agg._token_queue is not None assert agg._token_queue._live_task is not None await agg.process([session_event(SessionEventType.ENDED, ts=100)]) - assert agg._token_queue._live_task is None, ( - "drain must stop the live loop" - ) + assert ( + agg._token_queue._live_task is None + ), "drain must stop the live loop" finally: agg.close() @@ -1274,9 +1274,9 @@ async def count_batch_async(self, inputs, _loop, live=False): ] ) assert agg._token_queue is not None - assert agg._token_queue.pending > 0, ( - "precondition: ISL must be buffered before ENDED" - ) + assert ( + agg._token_queue.pending > 0 + ), "precondition: ISL must be buffered before ENDED" await agg.process([session_event(SessionEventType.ENDED, ts=2000)]) publisher.publish_final.assert_awaited_once() diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index 82f2f5db8..d5f190270 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -955,18 +955,18 @@ def test_terminate_procs_kills_running_workers(): assert future.running(), "worker task did not start" procs = list((getattr(ex, "_processes", None) or {}).values()) assert procs, "worker did not spawn" - assert not wait_for_process([p.sentinel for p in procs], timeout=0), ( - "worker exited before termination" - ) + assert not wait_for_process( + [p.sentinel for p in procs], timeout=0 + ), "worker exited before termination" _terminate_procs([ex]) for p in procs: # The executor manager may reap the child concurrently; sentinel # readiness observes exit without racing its return-code update. - assert wait_for_process([p.sentinel], timeout=5), ( - "worker was not terminated" - ) + assert wait_for_process( + [p.sentinel], timeout=5 + ), "worker was not terminated" finally: cleanup_procs = procs or list((getattr(ex, "_processes", None) or {}).values()) for p in cleanup_procs: From 9549a0571ab5b15c81e5679e7b604db6da28b8a9 Mon Sep 17 00:00:00 2001 From: Harshil Vagadia Date: Thu, 13 Aug 2026 10:20:24 -0700 Subject: [PATCH 11/11] Address structured token counting review feedback --- .../metrics_aggregator/metrics_table.py | 23 ++- .../metrics_aggregator/token_metrics.py | 164 +++++++++++------- .../metrics_aggregator/tokenization.py | 1 + .../commands/benchmark/accuracy.py | 8 +- src/inference_endpoint/core/types.py | 2 + .../load_generator/session.py | 5 + .../metrics_aggregator/test_token_metrics.py | 17 +- tests/unit/core/test_record.py | 3 + .../unit/load_generator/test_async_session.py | 2 + 9 files changed, 154 insertions(+), 71 deletions(-) diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py index 901a5fae1..5e66a787a 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/metrics_table.py @@ -192,7 +192,10 @@ class TokenTrigger(EmitTrigger): Subclasses return one explicit tokenization input. ``fire()`` dispatches already-tokenized IDs synchronously and sends every other variant through - the shared ``TokenBatchQueue``. + the shared ``TokenBatchQueue``. The queue is ``None`` when no tokenizer was + configured or discovered for the run. In that state, pre-tokenized IDs are + still counted synchronously, while text and structured token metrics are + unavailable and therefore skipped. """ def __init__( @@ -290,7 +293,14 @@ def __init__(self, registry: MetricsRegistry): class IslTrigger(TokenTrigger): - """ISL from token IDs, structured chat messages, or plain prompt text.""" + """ISL from one prompt representation, in explicit priority order. + + ``PromptData`` produced by ``PhaseIssuer`` normally has exactly one populated + representation. For defensive handling, token IDs take precedence over + plain text, which takes precedence over structured messages. An explicitly + empty token-ID tuple is therefore counted as zero rather than falling + through to another representation. + """ def __init__( self, @@ -312,12 +322,19 @@ def _extract_tokenization_input(self, ev_rec, row, pre_change): ev_rec.data.tools, ev_rec.data.chat_template_kwargs, ev_rec.data.chat_template, + ev_rec.data.tool_choice, ) return None class OslTrigger(TokenTrigger): - """OSL = token_count(full output text) from COMPLETE event data.""" + """OSL from one complete model output representation. + + An output containing reasoning or tool calls uses ``MessageInput`` so the + chat template renders all structured fields, including ordinary content. + Otherwise, non-empty output content uses ``TextInput``. ``TextModelOutput`` + does not carry server token IDs, so OSL has no token-ID/text precedence. + """ def __init__( self, diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py index fb38d7119..264779ba1 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py @@ -280,6 +280,7 @@ def __init__( # -- setup -------------------------------------------------------------- def _load_tokenizer(self) -> None: + transformers_logging.set_verbosity_error() tok = load_reference_tokenizer(self._tokenizer_name) self._tokenizer = tok self._text_backend = getattr(tok, "backend_tokenizer", None) @@ -440,12 +441,33 @@ async def _fan_out(procs: list[ProcessPoolExecutor], texts: list[str]) -> list[i def _token_count_text(self, text: str) -> int: return len(self._tokenizer.tokenize(text)) # type: ignore[union-attr] + def _warn_template_fallback(self, exc: Exception, impact: str) -> None: + """Log a chat-template fallback once per tokenizer and error type.""" + key = f"{self._tokenizer_name}:{type(exc).__name__}" + if key in self._fallback_warned: + return + self._fallback_warned.add(key) + logger.exception( + "apply_chat_template failed for %s (%s); falling back to " + "whitespace tokenization. %s", + self._tokenizer_name, + type(exc).__name__, + impact, + ) + def _token_count_message( self, content: str, reasoning: str | None, tool_calls: tuple[dict[str, Any], ...] | None, ) -> int: + """Count one assistant output without surrounding chat-template framing. + + Render the structured assistant content, reasoning, and tool calls with + a minimal user prefix, then subtract both that prefix and the empty + assistant frame. The result is the assistant payload count used for + OSL and TPOT. + """ tok = self._tokenizer msg: dict[str, Any] = {"role": "assistant", "content": content or ""} if reasoning: @@ -461,15 +483,7 @@ def _token_count_message( ) return max(0, len(encoded) - self._prefix_len - self._baseline) except Exception as exc: - key = f"{self._tokenizer_name}:{type(exc).__name__}" - if key not in self._fallback_warned: - self._fallback_warned.add(key) - logger.exception( - "apply_chat_template failed for %s (%s); falling back to " - "whitespace tokenization. Tool-call OSL/TPOT may diverge.", - self._tokenizer_name, - type(exc).__name__, - ) + self._warn_template_fallback(exc, "Tool-call OSL/TPOT may diverge.") tool_calls_json = ( msgspec.json.encode(list(tool_calls)).decode() if tool_calls else None ) @@ -482,9 +496,17 @@ def _token_count_prompt( self, messages: tuple[dict[str, Any], ...], tools: tuple[dict[str, Any], ...] | None, - chat_template_kwargs: dict[str, Any] | None = None, - chat_template: str | None = None, + chat_template_kwargs: dict[str, Any] | None, + chat_template: str | None, + tool_choice: str | dict[str, Any] | None, ) -> int: + """Count a complete structured input prompt for ISL. + + Render the full message history and optional tools using the selected + chat template and model-specific keyword arguments. Unlike assistant + output counting, this keeps all conversation framing and appends the + generation prompt because those tokens are part of the server input. + """ kwargs = dict(chat_template_kwargs or {}) kwargs.update( tokenize=True, @@ -495,6 +517,8 @@ def _token_count_prompt( kwargs["tools"] = list(tools) if chat_template is not None: kwargs["chat_template"] = chat_template + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice prompt_messages = _normalize_prompt_messages_for_template(messages) try: encoded = self._tokenizer.apply_chat_template( # type: ignore[union-attr] @@ -502,20 +526,38 @@ def _token_count_prompt( ) return len(encoded) except Exception as exc: - key = f"{self._tokenizer_name}:{type(exc).__name__}" - if key not in self._fallback_warned: - self._fallback_warned.add(key) - logger.exception( - "apply_chat_template failed for %s (%s); falling back to " - "whitespace tokenization. Structured ISL may diverge.", - self._tokenizer_name, - type(exc).__name__, - ) + self._warn_template_fallback(exc, "Structured ISL may diverge.") prompt = {"messages": prompt_messages} if tools is not None: prompt["tools"] = list(tools) return self._token_count_text(msgspec.json.encode(prompt).decode()) + async def _count_indexed_texts_async( + self, + indexed_texts: list[tuple[int, str]], + loop: asyncio.AbstractEventLoop, + *, + live: bool, + ) -> list[tuple[int, int | Exception]]: + """Count one text batch and pair each outcome with its input index.""" + texts = [text for _, text in indexed_texts] + try: + counts = await self._count_texts_async(texts, loop, live=live) + except Exception as exc: # noqa: BLE001 - isolate this input kind. + return [(index, exc) for index, _ in indexed_texts] + + if len(counts) != len(indexed_texts): + length_error = RuntimeError( + f"tokenizer returned {len(counts)} counts for " + f"{len(indexed_texts)} texts" + ) + return [(index, length_error) for index, _ in indexed_texts] + + return [ + (index, count) + for (index, _), count in zip(indexed_texts, counts, strict=True) + ] + async def count_batch_async( self, inputs: list[TokenizationInput], @@ -526,57 +568,55 @@ async def count_batch_async( ) -> list[int | Exception]: """Count a mixed batch while preserving input order.""" outcomes: list[int | Exception | None] = [None] * len(inputs) - text_indices: list[int] = [] - texts: list[str] = [] + indexed_texts: list[tuple[int, str]] = [] structured: list[tuple[int, MessageInput | PromptInput]] = [] for index, item in enumerate(inputs): - if isinstance(item, TokenIdsInput): - outcomes[index] = len(item.token_ids) - elif isinstance(item, TextInput): - text_indices.append(index) - texts.append(item.text) - elif isinstance(item, MessageInput): - structured.append((index, item)) - elif isinstance(item, PromptInput): - structured.append((index, item)) - - if texts: - try: - counts = await self._count_texts_async(texts, loop, live=live) - if len(counts) != len(texts): - raise RuntimeError( - f"tokenizer returned {len(counts)} counts for " - f"{len(texts)} texts" - ) - for index, count in zip(text_indices, counts, strict=True): - outcomes[index] = count - except Exception as exc: # noqa: BLE001 - isolate this input kind. - for index in text_indices: - outcomes[index] = exc + match item: + case TokenIdsInput(token_ids=token_ids): + outcomes[index] = len(token_ids) + case TextInput(text=text): + indexed_texts.append((index, text)) + case MessageInput() | PromptInput(): + structured.append((index, item)) + + if indexed_texts: + text_outcomes = await self._count_indexed_texts_async( + indexed_texts, loop, live=live + ) + for index, outcome in text_outcomes: + outcomes[index] = outcome for index, item in structured: if self._thread is None: outcomes[index] = RuntimeError("BatchTokenizer is closed") continue try: - if isinstance(item, MessageInput): - outcomes[index] = await loop.run_in_executor( - self._thread, - self._token_count_message, - item.content, - item.reasoning, - item.tool_calls, - ) - elif isinstance(item, PromptInput): - outcomes[index] = await loop.run_in_executor( - self._thread, - self._token_count_prompt, - item.messages, - item.tools, - item.chat_template_kwargs, - item.chat_template, - ) + match item: + case MessageInput(content, reasoning, tool_calls): + outcomes[index] = await loop.run_in_executor( + self._thread, + self._token_count_message, + content, + reasoning, + tool_calls, + ) + case PromptInput( + messages, + tools, + chat_template_kwargs, + chat_template, + tool_choice, + ): + outcomes[index] = await loop.run_in_executor( + self._thread, + self._token_count_prompt, + messages, + tools, + chat_template_kwargs, + chat_template, + tool_choice, + ) except Exception as exc: # noqa: BLE001 - isolate this input. outcomes[index] = exc diff --git a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py index 032e188f5..6ba0cf946 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/tokenization.py @@ -45,6 +45,7 @@ class PromptInput: tools: tuple[dict[str, Any], ...] | None chat_template_kwargs: dict[str, Any] | None chat_template: str | None + tool_choice: str | dict[str, Any] | None = None TokenizationInput: TypeAlias = ( # noqa: UP040 - mypy version lacks PEP 695. diff --git a/src/inference_endpoint/commands/benchmark/accuracy.py b/src/inference_endpoint/commands/benchmark/accuracy.py index 713932079..bdfdc9ac4 100644 --- a/src/inference_endpoint/commands/benchmark/accuracy.py +++ b/src/inference_endpoint/commands/benchmark/accuracy.py @@ -103,10 +103,10 @@ def _phase_osl_stats( population in one call would hold every Encoding in memory at once. """ # Skip empty/failed completions (a failed request still logs a COMPLETE - # event with output == ""). The perf-side OslTrigger does the same - # (metrics_table.OslTrigger._extract_text returns None for empty text), so - # accuracy OSL matches its population and a failure isn't counted as a - # 0-token sample that would drag min/avg down. + # event with output == ""). The performance-side OslTrigger follows the + # same rule: _extract_tokenization_input returns None for an empty text-only + # output, so a failure is not counted as a 0-token sample that would drag + # min/avg down. texts = [ uuid_to_text[u] for u in sample_uuids if u in uuid_to_text and uuid_to_text[u] ] diff --git a/src/inference_endpoint/core/types.py b/src/inference_endpoint/core/types.py index 6d25bf5c5..ebed949ac 100644 --- a/src/inference_endpoint/core/types.py +++ b/src/inference_endpoint/core/types.py @@ -300,6 +300,7 @@ class PromptData( tools: Tool declarations accompanying ``messages``. chat_template_kwargs: Model-specific arguments used to render ``messages``. chat_template: Per-request template used to render ``messages``. + tool_choice: Tool-selection mode passed to the chat template. """ text: str | None = None @@ -308,6 +309,7 @@ class PromptData( tools: tuple[dict[str, Any], ...] | None = None chat_template_kwargs: dict[str, Any] | None = None chat_template: str | None = None + tool_choice: str | dict[str, Any] | None = None class ErrorData( diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index aa70fd80e..847cb4cea 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/src/inference_endpoint/load_generator/session.py @@ -272,6 +272,11 @@ def issue( if isinstance(data.get("chat_template"), str) else None ), + tool_choice=( + data["tool_choice"] + if isinstance(data.get("tool_choice"), str | dict) + else None + ), ) elif isinstance(prompt, str): prompt_data = PromptData(text=prompt) diff --git a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py index d5f190270..530435ce0 100644 --- a/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py +++ b/tests/unit/async_utils/services/metrics_aggregator/test_token_metrics.py @@ -208,7 +208,11 @@ def test_chat_template_requests_token_ids_not_batch_encoding(self): with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: count = tok._token_count_prompt( - ({"role": "user", "content": "one two three four"},), None + ({"role": "user", "content": "one two three four"},), + None, + None, + None, + None, ) assert count == 7 @@ -339,6 +343,9 @@ def apply_chat_template(self, messages, **kwargs): }, ), None, + None, + None, + None, ) assert count > 0 @@ -376,6 +383,7 @@ def apply_chat_template(self, messages, **kwargs): None, {"enable_thinking": False}, "custom template", + "auto", ) normalized_call = _RecordingTokenizer.last_messages[0]["tool_calls"][0] @@ -385,6 +393,7 @@ def apply_chat_template(self, messages, **kwargs): ) assert _RecordingTokenizer.last_kwargs["enable_thinking"] is False assert _RecordingTokenizer.last_kwargs["chat_template"] == "custom template" + assert _RecordingTokenizer.last_kwargs["tool_choice"] == "auto" assert _RecordingTokenizer.last_kwargs["return_dict"] is False @@ -545,7 +554,11 @@ def test_structured_tokenization_does_not_require_a_text_backend(self, monkeypat assert tok._procs == [] assert ( tok._token_count_prompt( - ({"role": "user", "content": "one two"},), None + ({"role": "user", "content": "one two"},), + None, + None, + None, + None, ) == 5 ) diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index c23e06fea..dccf1560a 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -184,6 +184,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): tools = ({"function": {"name": "lookup"}},) chat_template_kwargs = {"enable_thinking": False} chat_template = "custom template" + tool_choice = "auto" record = EventRecord( event_type=SampleEventType.ISSUED, sample_uuid="sample-chat", @@ -192,6 +193,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): tools=tools, chat_template_kwargs=chat_template_kwargs, chat_template=chat_template, + tool_choice=tool_choice, ), ) @@ -203,6 +205,7 @@ def test_sample_event_round_trips_with_structured_prompt_data(self): assert decoded.data.tools == tools assert decoded.data.chat_template_kwargs == chat_template_kwargs assert decoded.data.chat_template == chat_template + assert decoded.data.tool_choice == tool_choice def test_error_event_round_trips_with_error_data(self): record = EventRecord( diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 58048f0aa..e0b5ab9c6 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/tests/unit/load_generator/test_async_session.py @@ -197,6 +197,7 @@ def load_sample(self, index: int) -> dict: ], "chat_template": "custom template", "chat_template_kwargs": {"enable_thinking": False}, + "tool_choice": "auto", } issuer = FakeIssuer() @@ -211,6 +212,7 @@ def load_sample(self, index: int) -> dict: assert prompt.tools == tuple(issuer.issued_queries[0].data["tools"]) assert prompt.chat_template == "custom template" assert prompt.chat_template_kwargs == {"enable_thinking": False} + assert prompt.tool_choice == "auto" assert prompt.text is None def test_issue_messages_take_precedence_over_prompt(self):