Skip to content

fix(metrics): fan the structured token drain out across tokenizer shards #561

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
karverma-amd wants to merge 1 commit into from
Closed

fix(metrics): fan the structured token drain out across tokenizer shards #561

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
344 changes: 264 additions & 80 deletions src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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")
Expand Down
Loading
Loading