From 4b00266aaefaaaa0f8ea95aa75528de6aa399192 Mon Sep 17 00:00:00 2001 From: "Verma, Karan" Date: Wed, 4 Nov 2026 22:34:13 +0000 Subject: [PATCH] fix(metrics): fan the structured token drain out across tokenizer shards At end of run the aggregator drains every pending tokenization. Plain-text inputs already fan out across the ProcessPoolExecutor shards, but structured inputs (MessageInput / PromptInput) were counted one at a time on the in-process thread. apply_chat_template is Jinja and Python and holds the GIL, so that lane runs on a single core. On a large agentic run the drain can take hours, and result_summary.json is not written until it finishes. Move the counting into module-level count_message_tokens / count_prompt_tokens so the in-process lane and the shards share one implementation. During the drain, deal structured items to the shards in chunks of 64, round-robin, so long prompts spread across shards. A failed chunk marks only its own items as errors, and progress is logged every 30 s. The live (mid-run) lane is unchanged. Co-authored-by: Cursor --- .../metrics_aggregator/token_metrics.py | 344 ++++++++++++++---- .../metrics_aggregator/test_token_metrics.py | 120 ++++++ 2 files changed, 384 insertions(+), 80 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 881437303..740b16552 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 @@ -139,6 +139,118 @@ def _normalize_prompt_messages_for_template( # --------------------------------------------------------------------------- _WORKER_TEXT_BACKEND: Any = None +# Full tokenizer + chat-template baseline for the structured (MessageInput / +# PromptInput) drain path. apply_chat_template is Jinja + Python and holds the +# GIL, so the in-process thread pool counts one item at a time; shipping the +# drain's structured items to the shards is what lets it use every core. +_WORKER_TOKENIZER: Any = None +_WORKER_PREFIX_LEN = 0 +_WORKER_BASELINE = 0 +_WORKER_FALLBACK_WARNED: set[str] = set() + +# Structured drain items per shard task: small enough that pickling one task +# stays cheap and progress is reported often, large enough to amortize IPC. +_STRUCTURED_CHUNK_ITEMS = 64 + + +def chat_template_baseline(tok: Any) -> tuple[int, int]: + """``(prefix_len, baseline)`` used to strip chat framing from message counts. + + Baseline = tokens from a [user, empty-assistant] pair minus the [user] + prefix alone, so the assistant frame is subtracted from message counts. + Raises if the template cannot render; callers fall back to ``(0, 0)``. + """ + prefix = tok.apply_chat_template( + [_PREFIX_USER_MSG], + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + with_assistant_tokens = tok.apply_chat_template( + [_PREFIX_USER_MSG, {"role": "assistant", "content": ""}], + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + return len(prefix), len(with_assistant_tokens) - len(prefix) + + +def count_message_tokens( + tok: Any, + prefix_len: int, + baseline: int, + content: str, + reasoning: str | None, + tool_calls: tuple[dict[str, Any], ...] | None, + on_fallback: Callable[[Exception, str], 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. Shared by the in-process lane and the drain shards so both + produce identical numbers. + """ + msg: dict[str, Any] = {"role": "assistant", "content": content or ""} + if reasoning: + msg["reasoning_content"] = reasoning + if tool_calls: + msg["tool_calls"] = _normalize_tool_calls_for_template(tool_calls) + try: + encoded = tok.apply_chat_template( + [_PREFIX_USER_MSG, msg], + tokenize=True, + add_generation_prompt=False, + return_dict=False, + ) + return max(0, len(encoded) - prefix_len - baseline) + except Exception as exc: + on_fallback(exc, "Tool-call OSL/TPOT may diverge.") + tool_calls_json = ( + msgspec.json.encode(list(tool_calls)).decode() if tool_calls else None + ) + parts = [p for p in (content or None, reasoning or None, tool_calls_json) if p] + return len(tok.tokenize("\n".join(parts))) + + +def count_prompt_tokens( + tok: Any, + messages: tuple[dict[str, Any], ...], + tools: tuple[dict[str, Any], ...] | None, + chat_template_kwargs: dict[str, Any] | None, + chat_template: str | None, + tool_choice: str | dict[str, Any] | None, + on_fallback: Callable[[Exception, str], 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, + add_generation_prompt=True, + return_dict=False, + ) + if tools is not None: + 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: + return len(tok.apply_chat_template(prompt_messages, **kwargs)) + except Exception as exc: + on_fallback(exc, "Structured ISL may diverge.") + prompt = {"messages": prompt_messages} + if tools is not None: + prompt["tools"] = list(tools) + return len(tok.tokenize(msgspec.json.encode(prompt).decode())) def load_reference_tokenizer(tokenizer_name: str) -> Any: @@ -182,8 +294,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() - global _WORKER_TEXT_BACKEND - _WORKER_TEXT_BACKEND = load_reference_backend(tokenizer_name) + global _WORKER_TEXT_BACKEND, _WORKER_TOKENIZER, _WORKER_PREFIX_LEN, _WORKER_BASELINE + _WORKER_TOKENIZER = load_reference_tokenizer(tokenizer_name) + _WORKER_TEXT_BACKEND = getattr(_WORKER_TOKENIZER, "backend_tokenizer", None) + try: + _WORKER_PREFIX_LEN, _WORKER_BASELINE = chat_template_baseline(_WORKER_TOKENIZER) + except Exception: + # Same fallback as the parent (BatchTokenizer._load_tokenizer). + _WORKER_PREFIX_LEN, _WORKER_BASELINE = 0, 0 if _WORKER_TEXT_BACKEND is not None: _WORKER_TEXT_BACKEND.encode("warmup", add_special_tokens=False) @@ -203,6 +321,70 @@ def _worker_encode_lengths(texts: list[str]) -> list[int]: return encode_lengths(backend, texts) +def _worker_warn_fallback(exc: Exception, impact: str) -> None: + key = type(exc).__name__ + if key in _WORKER_FALLBACK_WARNED: + return + _WORKER_FALLBACK_WARNED.add(key) + logger.warning( + "apply_chat_template failed in tokenizer shard (%s: %s); falling back " + "to whitespace tokenization. %s", + key, + exc, + impact, + ) + + +def _worker_count_structured( + items: list[MessageInput | PromptInput], +) -> list[int | Exception]: + """Per-item counts for structured inputs; a failing item is returned as + its exception so it does not poison the rest of the chunk.""" + tok = _WORKER_TOKENIZER + if tok is None: + raise RuntimeError("tokenizer worker tokenizer unavailable") + out: list[int | Exception] = [] + for item in items: + try: + match item: + case MessageInput(content, reasoning, tool_calls): + out.append( + count_message_tokens( + tok, + _WORKER_PREFIX_LEN, + _WORKER_BASELINE, + content, + reasoning, + tool_calls, + _worker_warn_fallback, + ) + ) + case PromptInput( + messages, tools, chat_template_kwargs, chat_template, tool_choice + ): + out.append( + count_prompt_tokens( + tok, + messages, + tools, + chat_template_kwargs, + chat_template, + tool_choice, + _worker_warn_fallback, + ) + ) + case _: + out.append( + TypeError( + f"unsupported structured input: {type(item).__name__}" + ) + ) + except Exception as exc: # noqa: BLE001 - isolate this input. + # Exceptions from tokenizers/jinja may not pickle; send a plain one. + out.append(RuntimeError(f"{type(exc).__name__}: {exc}")) + return out + + def _worker_ready(_: int) -> bool: """Warmup probe: returns once the worker's backend is loaded.""" return _WORKER_TEXT_BACKEND is not None @@ -289,26 +471,7 @@ def _load_tokenizer(self) -> None: # 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, - return_dict=False, - ), - ) - self._prefix_len = len(prefix) - 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_tokens) - self._prefix_len + self._prefix_len, self._baseline = chat_template_baseline(tok) except Exception: self._prefix_len = 0 self._baseline = 0 @@ -469,36 +632,16 @@ def _token_count_message( 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: - msg["reasoning_content"] = reasoning - if tool_calls: - 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, - return_dict=False, - ) - return max(0, len(encoded) - self._prefix_len - self._baseline) - except Exception as exc: - 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 - ) - parts = [ - p for p in (content or None, reasoning or None, tool_calls_json) if p - ] - return self._token_count_text("\n".join(parts)) + """Count one assistant output (see :func:`count_message_tokens`).""" + return count_message_tokens( + self._tokenizer, + self._prefix_len, + self._baseline, + content, + reasoning, + tool_calls, + self._warn_template_fallback, + ) def _token_count_prompt( self, @@ -508,37 +651,73 @@ def _token_count_prompt( chat_template: str | None, tool_choice: str | dict[str, Any] | None, ) -> int: - """Count a complete structured input prompt for ISL. + """Count a complete structured prompt (see :func:`count_prompt_tokens`).""" + return count_prompt_tokens( + self._tokenizer, + messages, + tools, + chat_template_kwargs, + chat_template, + tool_choice, + self._warn_template_fallback, + ) + + async def _fan_out_structured( + self, structured: list[tuple[int, MessageInput | PromptInput]] + ) -> list[tuple[int, int | Exception]]: + """Drain-path structured counting across every shard. - 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. + Items are cut into small chunks dealt round-robin to the shards, so one + shard does not end up with all the long agentic prompts. A failed chunk + (e.g. a broken worker) marks only its own items as failed. Progress is + logged because this phase can run for many minutes with no other output. """ - kwargs = dict(chat_template_kwargs or {}) - kwargs.update( - tokenize=True, - add_generation_prompt=True, - return_dict=False, + procs = self._procs + chunks = [ + structured[i : i + _STRUCTURED_CHUNK_ITEMS] + for i in range(0, len(structured), _STRUCTURED_CHUNK_ITEMS) + ] + total = len(structured) + done = 0 + next_log = time.monotonic() + 30.0 + t0 = time.monotonic() + + async def run(k: int, chunk: list[tuple[int, MessageInput | PromptInput]]): + nonlocal done, next_log + ex = procs[k % len(procs)] + try: + counts = await asyncio.wrap_future( + ex.submit(_worker_count_structured, [item for _, item in chunk]) + ) + if len(counts) != len(chunk): + raise RuntimeError( + f"tokenizer shard returned {len(counts)} counts for " + f"{len(chunk)} inputs" + ) + result = [ + (index, c) for (index, _), c in zip(chunk, counts, strict=True) + ] + except Exception as exc: # noqa: BLE001 - isolate this chunk. + result = [(index, exc) for index, _ in chunk] + done += len(chunk) + if time.monotonic() >= next_log: + next_log = time.monotonic() + 30.0 + logger.info( + "Structured tokenization drain: %d/%d (%.0fs)", + done, + total, + time.monotonic() - t0, + ) + return result + + results = await asyncio.gather(*(run(k, c) for k, c in enumerate(chunks))) + logger.info( + "Structured tokenization drain done: %d items on %d shards in %.1fs", + total, + len(procs), + time.monotonic() - t0, ) - if tools is not None: - 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] - prompt_messages, **kwargs - ) - return len(encoded) - except Exception as exc: - 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()) + return list(chain.from_iterable(results)) async def _count_indexed_texts_async( self, @@ -629,6 +808,11 @@ async def count_batch_async( for index, outcome in text_outcomes: outcomes[index] = outcome + if structured and self._procs and not live: + for index, outcome in await self._fan_out_structured(structured): + outcomes[index] = outcome + structured = [] + for index, item in structured: if self._thread is None: outcomes[index] = RuntimeError("BatchTokenizer is closed") 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 ab15bf0d8..772356b36 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 @@ -1047,3 +1047,123 @@ def test_count_sync_counts_plain_text(self): with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: assert tok.count_sync(TextInput("hello world")) == 2 + + +class _InlineProc: + """A shard that runs the submitted worker function in this process.""" + + def __init__(self): + self.calls = 0 + + def submit(self, fn, *args): + self.calls += 1 + fut: Future = Future() + try: + fut.set_result(fn(*args)) + except Exception as exc: + fut.set_exception(exc) + return fut + + def shutdown(self, wait=False, cancel_futures=False): + pass + + +def _structured_items(n: int) -> list[MessageInput | PromptInput]: + tools = ({"type": "function", "function": {"name": "lookup", "parameters": {}}},) + items: list[MessageInput | PromptInput] = [] + for i in range(n): + words = " ".join(["w"] * (i % 7 + 1)) + if i % 2: + items.append(MessageInput(words, "thought" if i % 3 else None, None)) + else: + messages = ({"role": "user", "content": words},) + items.append(PromptInput(messages, tools if i % 4 else None, None, None)) + return items + + +@pytest.fixture +def _worker_tokenizer(monkeypatch): + """Load the shard globals the way the shard initializer does.""" + tok = _FakeTokenizerWithTemplateAndBackend() + prefix_len, baseline = token_metrics_module.chat_template_baseline(tok) + monkeypatch.setattr(token_metrics_module, "_WORKER_TOKENIZER", tok) + monkeypatch.setattr( + token_metrics_module, "_WORKER_TEXT_BACKEND", tok.backend_tokenizer + ) + monkeypatch.setattr(token_metrics_module, "_WORKER_PREFIX_LEN", prefix_len) + monkeypatch.setattr(token_metrics_module, "_WORKER_BASELINE", baseline) + + +@pytest.mark.unit +@pytest.mark.usefixtures("_worker_tokenizer") +class TestStructuredDrainFanOut: + """The drain counts structured inputs on the shards, not one thread.""" + + @pytest.mark.asyncio + async def test_shard_counts_match_in_process_counts(self): + n = token_metrics_module._STRUCTURED_CHUNK_ITEMS * 2 + 5 + items = _structured_items(n) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + expected = await tok.count_batch_async(items, loop) + procs = [_InlineProc(), _InlineProc()] + tok._procs = procs + counts = await tok.count_batch_async(items, loop) + + assert counts == expected + assert all(isinstance(c, int) for c in counts) + # Three chunks dealt round-robin across two shards. + assert [p.calls for p in procs] == [2, 1] + + @pytest.mark.asyncio + async def test_mixed_batch_keeps_input_order(self): + structured = _structured_items(3) + inputs = [ + TokenIdsInput((1, 2, 3)), + structured[0], + TextInput("a b"), + structured[1], + structured[2], + ] + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + expected = await tok.count_batch_async(inputs, loop) + tok._procs = [_InlineProc()] + assert await tok.count_batch_async(inputs, loop) == expected + + @pytest.mark.asyncio + async def test_failed_chunk_marks_only_its_own_items(self): + size = token_metrics_module._STRUCTURED_CHUNK_ITEMS + items = _structured_items(size * 3) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + tok._procs = [_InlineProc(), _BrokenProc()] + counts = await tok.count_batch_async(items, loop) + + broken = counts[size : 2 * size] + assert all(isinstance(c, BrokenProcessPool) for c in broken) + assert all(isinstance(c, int) for c in counts[:size] + counts[2 * size :]) + + @pytest.mark.asyncio + async def test_live_lane_never_uses_shards(self): + items = _structured_items(4) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + tok._procs = [_BrokenProc()] + counts = await tok.count_batch_async(items, loop, live=True) + + assert all(isinstance(c, int) for c in counts) + + def test_worker_returns_item_errors_instead_of_raising(self): + class _Bad: + pass + + out = token_metrics_module._worker_count_structured( + [MessageInput("a b", None, None), _Bad()] # type: ignore[list-item] + ) + assert isinstance(out[0], int) + assert isinstance(out[1], Exception)