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 9da76a2e1..e8ef68061 100644 --- a/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py +++ b/src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py @@ -18,6 +18,7 @@ import argparse import asyncio import logging +import os import signal from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext @@ -111,6 +112,19 @@ def _on_sigterm() -> None: return _on_sigterm, pending_tasks +def _cap_live_lane_rayon_pool(tokenizer_workers: int) -> None: + """Cap the Hugging Face rayon pool the live (mid-run) lane encodes on. + + The live lane runs in this process, and rayon sizes its pool once, at the + first encode, so the cap must be in the environment before the tokenizer is + used. ``setdefault`` lets an operator-exported ``RAYON_NUM_THREADS`` win; the + ``max(1, ...)`` keeps ``--tokenizer-workers 0`` from reading as rayon's "use + every CPU". The end-of-run drain is not capped by this: its shard processes + size their own pools in ``_init_worker``. + """ + os.environ.setdefault("RAYON_NUM_THREADS", str(max(1, tokenizer_workers))) + + async def main() -> None: parser = argparse.ArgumentParser( description="Metrics aggregator service - subscribes to EventRecords and computes real-time metrics" @@ -257,6 +271,7 @@ async def main() -> None: # (coalesces to 'object' not 'AbstractContextManager[BatchTokenizer | None]') tokenizer_cm: AbstractContextManager[BatchTokenizer | None] if args.tokenizer: + _cap_live_lane_rayon_pool(args.tokenizer_workers) try: tokenizer_cm = BatchTokenizer( args.tokenizer, live_workers=args.tokenizer_workers 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..36fcf16de 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 @@ -37,7 +37,7 @@ import os import signal import time -from collections.abc import Callable +from collections.abc import Callable, Sequence from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from itertools import chain from typing import TYPE_CHECKING, Any, Protocol, cast @@ -139,6 +139,151 @@ 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 count_structured_item( + tok: Any, + prefix_len: int, + baseline: int, + item: MessageInput | PromptInput, + on_fallback: Callable[[Exception, str], None], +) -> int: + """Count one structured input with the logic that matches its kind. + + The single dispatch shared by the in-process thread and the drain shards, + so a sample gets the same number whichever lane counts it. + """ + match item: + case MessageInput(content, reasoning, tool_calls): + return count_message_tokens( + tok, prefix_len, baseline, content, reasoning, tool_calls, on_fallback + ) + case PromptInput( + messages, tools, chat_template_kwargs, chat_template, tool_choice + ): + return count_prompt_tokens( + tok, + messages, + tools, + chat_template_kwargs, + chat_template, + tool_choice, + on_fallback, + ) + case _: + raise TypeError(f"unsupported structured input: {type(item).__name__}") def load_reference_tokenizer(tokenizer_name: str) -> Any: @@ -182,8 +327,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 +354,46 @@ 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: + out.append( + count_structured_item( + tok, + _WORKER_PREFIX_LEN, + _WORKER_BASELINE, + item, + _worker_warn_fallback, + ) + ) + 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 @@ -258,9 +449,6 @@ def __init__( n_workers: int = -1, ) -> None: self._tokenizer_name = tokenizer_name - # 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 self._text_backend: Any | None = None @@ -289,26 +477,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 @@ -446,9 +615,6 @@ async def _fan_out(procs: list[ProcessPoolExecutor], texts: list[str]) -> list[i # -- sync + chat-template paths (in-process thread) --------------------- - 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__}" @@ -463,82 +629,62 @@ def _warn_template_fallback(self, exc: Exception, impact: str) -> None: 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: - 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)) + async def _fan_out_structured( + self, structured: list[tuple[int, MessageInput | PromptInput]] + ) -> list[tuple[int, int | Exception]]: + """Drain-path structured counting across every shard. - def _token_count_prompt( - self, - 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, - ) -> 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. + 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, @@ -566,39 +712,116 @@ async def _count_indexed_texts_async( for (index, _), count in zip(indexed_texts, counts, strict=True) ] - def count_sync(self, item: TokenizationInput) -> int: - """Count one input on the calling thread. + async def _count_structured_async( + self, + structured: list[tuple[int, MessageInput | PromptInput]], + loop: asyncio.AbstractEventLoop, + *, + live: bool, + ) -> list[tuple[int, int | Exception]]: + """Count structured inputs and pair each outcome with its input index. + + Routed like :meth:`_count_indexed_texts_async`: ``live=True`` (mid-run) + never touches the shard processes and counts on the in-process thread; + the drain fans out across every shard. Both lanes count with the same + ``count_message_tokens`` / ``count_prompt_tokens`` logic, so a sample + gets the same number whichever lane handles it. + """ + if self._procs and not live: + # Draining will call this path for multiprocessing + return await self._fan_out_structured(structured) + # during benchmarking, we will call this path to the in-process thread pool + return await self._count_structured_inproc(structured, loop) + + async def _count_structured_inproc( + self, + structured: list[tuple[int, MessageInput | PromptInput]], + loop: asyncio.AbstractEventLoop, + ) -> list[tuple[int, int | Exception]]: + """Count structured inputs one at a time on the in-process thread. + + The live lane, and the drain when there are no shards (``n_workers=0``). + ``apply_chat_template`` holds the GIL, so each item is its own executor + call, which keeps cancellation points between items. + """ + results: list[tuple[int, int | Exception]] = [] + for index, item in structured: + if self._thread is None: + results.append((index, RuntimeError("BatchTokenizer is closed"))) + continue + try: + count = await loop.run_in_executor( + self._thread, + count_structured_item, + self._tokenizer, + self._prefix_len, + self._baseline, + item, + self._warn_template_fallback, + ) + results.append((index, count)) + except Exception as exc: # noqa: BLE001 - isolate this input. + results.append((index, exc)) + return results + + def count_batch_sync( + self, items: Sequence[TokenizationInput] + ) -> list[int | Exception]: + """Count many inputs on this tokenizer's thread pool, in input order. The post-run counterpart to :meth:`count_batch_async`, for callers that - run after the event loop is gone (full-run OSL at finalize). Dispatch - mirrors ``count_batch_async`` exactly — same text backend, same chat - template, same baseline subtraction — so a turn counted here and a turn - counted on the perf hot path yield the same number. Divergence between - the two lanes would make the windowed and full-run OSL statistics - incomparable (mlcommons/endpoints#500). + run after the event loop is gone (full-run OSL at finalize). Structured + inputs are counted one per task with :func:`count_structured_item`, the + logic the aggregator's live lane and drain shards use, so a turn gets + byte-identical counts here and in the aggregator. The ``live_workers`` + threads overlap the tokenizer work of different turns. Plain text is + encoded in one batch, as the aggregator does. """ - match item: - case TokenIdsInput(token_ids=token_ids): - return len(token_ids) - case TextInput(text=text): - return self._encode_lengths_inproc([text])[0] - case MessageInput(content, reasoning, tool_calls): - return self._token_count_message(content, reasoning, tool_calls) - case PromptInput( - messages, - tools, - chat_template_kwargs, - chat_template, - tool_choice, - ): - return self._token_count_prompt( - messages, - tools, - chat_template_kwargs, - chat_template, - tool_choice, + if self._thread is None: + raise RuntimeError("BatchTokenizer is closed") + outcomes: list[int | Exception | None] = [None] * len(items) + text_slots: list[int] = [] + texts: list[str] = [] + structured_slots: list[int] = [] + structured: list[MessageInput | PromptInput] = [] + for slot, item in enumerate(items): + match item: + case TokenIdsInput(token_ids=token_ids): + outcomes[slot] = len(token_ids) + case TextInput(text=text): + text_slots.append(slot) + texts.append(text) + case MessageInput() | PromptInput(): + structured_slots.append(slot) + structured.append(item) + case _: + outcomes[slot] = TypeError( + f"unsupported tokenization input: {type(item).__name__}" + ) + + def count_one(item: MessageInput | PromptInput) -> int | Exception: + try: + return count_structured_item( + self._tokenizer, + self._prefix_len, + self._baseline, + item, + self._warn_template_fallback, ) - raise TypeError(f"unsupported tokenization input: {type(item).__name__}") + except Exception as exc: # noqa: BLE001 - isolate this input. + return exc + + # Start the structured turns first so the pool works while the calling + # thread encodes the (few) plain texts. + futures = [self._thread.submit(count_one, item) for item in structured] + if texts: + for slot, count in zip( + text_slots, self._encode_lengths_inproc(texts), strict=True + ): + outcomes[slot] = count + for slot, future in zip(structured_slots, futures, strict=True): + outcomes[slot] = future.result() + return cast(list[int | Exception], outcomes) async def count_batch_async( self, @@ -629,38 +852,12 @@ async def count_batch_async( 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: - 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 + if structured: + structured_outcomes = await self._count_structured_async( + structured, loop, live=live + ) + for index, outcome in structured_outcomes: + outcomes[index] = outcome if any(outcome is None for outcome in outcomes): raise AssertionError("unhandled TokenizationInput variant") diff --git a/src/inference_endpoint/commands/benchmark/full_run_osl.py b/src/inference_endpoint/commands/benchmark/full_run_osl.py index 140e350ea..5a081ad14 100644 --- a/src/inference_endpoint/commands/benchmark/full_run_osl.py +++ b/src/inference_endpoint/commands/benchmark/full_run_osl.py @@ -22,7 +22,8 @@ from __future__ import annotations import logging -from collections.abc import Callable, Collection +import os +from collections.abc import Callable, Collection, Sequence from pathlib import Path from typing import Any @@ -36,20 +37,40 @@ TextInput, extract_tokenization_input, ) -from inference_endpoint.core.record import EventRecord, EventType +from inference_endpoint.core.record import ( + EventRecord, + EventType, + SampleEventType, + leading_event_topic, +) from inference_endpoint.core.types import PhaseType, TextModelOutput +from inference_endpoint.endpoint_client.cpu_affinity import ( + UnsupportedPlatformError, + get_all_online_cpus, +) from inference_endpoint.metrics.report import series_metric_dict logger = logging.getLogger(__name__) -# Matched against each line's leading bytes so the 6+ GB of prompt and response -# payload in a large run is never decoded just to classify the record. -_COMPLETE_MARKER = b'"event_type":"sample.complete"' -_PREFIX_BYTES = 256 +# A line in the writer's compact layout is classified from its leading bytes, so +# the 6+ GB of prompt and response payload in a large run is never decoded just +# to find the COMPLETE records. Any other layout is decoded and classified by its +# decoded event_type (see compute_full_run_osl). +_COMPLETE_TOPIC: bytes = SampleEventType.COMPLETE.topic_bytes # type: ignore[attr-defined] # AgenticInferenceInlineScorer's dataset_name; the key its turns live under in # sample_idx_map.json. -TokenCounter = Callable[[MessageInput | TextInput], int] +# Counts a batch of turns, returning a count or the exception for each turn. +BatchTokenCounter = Callable[ + [Sequence[MessageInput | TextInput]], list[int | Exception] +] + +# Turns counted per tokenizer batch at finalize. +_DEFAULT_BATCH_SIZE = 1024 +# Counting threads at finalize. Each turn is counted with the aggregator's own +# per-item logic, whose Python parts hold the GIL: measured throughput levels +# off by about 32 threads and is slightly lower beyond that. +_MAX_COUNT_THREADS = 32 def _in_population( @@ -85,9 +106,10 @@ def _in_population( def compute_full_run_osl( events_path: Path, - count_tokens: TokenCounter, + count_batch: BatchTokenCounter, *, performance_uuids: Collection[str] | None = None, + batch_size: int = _DEFAULT_BATCH_SIZE, ) -> dict[str, Any] | None: """Roll up OSL over every completed turn in ``events_path``. @@ -109,10 +131,43 @@ def compute_full_run_osl( n_undecodable = 0 first_error: str | None = None seen_uuids: set[str] = set() + batch_size = max(1, batch_size) + pending: list[MessageInput | TextInput] = [] + + def count_pending() -> None: + nonlocal n_errors, first_error + n_pending = len(pending) + try: + outcomes = count_batch(pending) + except Exception as e: # noqa: BLE001 - charge the batch, keep the stat + outcomes = [e] * n_pending + if len(outcomes) != n_pending: + # One outcome per turn is the counter's contract. A different number + # cannot be matched back to its turns, so charge the whole batch + # (the block turns partial) rather than average an unknown subset. + mismatch = RuntimeError( + f"token counter returned {len(outcomes)} outcomes for " + f"{n_pending} turns" + ) + outcomes = [mismatch] * n_pending + for outcome in outcomes: + if isinstance(outcome, Exception): + n_errors += 1 + if first_error is None: + first_error = str(outcome) + else: + lengths.append(outcome) + pending.clear() with events_path.open("rb") as events_file: for line in events_file: - if _COMPLETE_MARKER not in line[:_PREFIX_BYTES]: + topic = leading_event_topic(line) + if topic is None: + # Not the compact layout (spaces, reordered fields, ...): decode + # the line and classify it below by its decoded event_type. + if not line.strip(): + continue + elif topic != _COMPLETE_TOPIC: continue # Decode before the population filter so a corrupt line — which we # cannot attribute to the performance population — is skipped rather @@ -122,6 +177,8 @@ def compute_full_run_osl( except msgspec.DecodeError: n_undecodable += 1 continue + if record.event_type != SampleEventType.COMPLETE: + continue if not _in_population(record, performance_uuids): continue seen_uuids.add(record.sample_uuid) @@ -142,11 +199,15 @@ def compute_full_run_osl( if tokenization_input is None: n_empty += 1 continue - lengths.append(count_tokens(tokenization_input)) + pending.append(tokenization_input) + if len(pending) >= batch_size: + count_pending() except Exception as e: # noqa: BLE001 - one bad turn must not void the stat n_errors += 1 if first_error is None: first_error = str(e) + if pending: + count_pending() # A performance turn whose COMPLETE record never reached the log is invisible # to the counters above; reconcile observed UUIDs against the expected @@ -253,10 +314,17 @@ def full_run_osl_for_report( logger.info("Full-run OSL skipped: no performance-phase samples to bound to") return None try: - with BatchTokenizer(tokenizer_name, live_workers=1, n_workers=0) as tokenizer: + try: + threads = len(get_all_online_cpus()) or 1 + except UnsupportedPlatformError: + threads = os.cpu_count() or 1 + threads = min(threads, _MAX_COUNT_THREADS) + with BatchTokenizer( + tokenizer_name, live_workers=threads, n_workers=0 + ) as tokenizer: return compute_full_run_osl( events_path, - tokenizer.count_sync, + tokenizer.count_batch_sync, performance_uuids=performance_uuids, ) except Exception as e: # noqa: BLE001 - optional block; never fail finalize diff --git a/src/inference_endpoint/core/record.py b/src/inference_endpoint/core/record.py index 6dddfe1d5..66fa430d4 100644 --- a/src/inference_endpoint/core/record.py +++ b/src/inference_endpoint/core/record.py @@ -170,6 +170,34 @@ class EventRecord(msgspec.Struct, kw_only=True, frozen=True, gc=False): # type: worker_id: int | msgspec.UnsetType = msgspec.UNSET +EVENT_LINE_LEAD: Final[bytes] = b'{"event_type":"' +"""bytes: How an events.jsonl line begins when it is written compactly. + +``JSONLWriter`` encodes an ``EventRecord`` with ``event_type`` first and without +spaces, so every line it writes starts with these bytes. +""" + + +def leading_event_topic(line: bytes) -> bytes | None: + """The event_type topic that leads a compact events.jsonl line, if it does. + + Lets a reader pick out some event types from a line's first bytes and skip + the rest undecoded; a run's log is hundreds of GB, mostly ISSUED prompts. + Compare the result with a member's ``topic_bytes``, e.g. + ``SampleEventType.COMPLETE.topic_bytes``. + + ``None`` means the line is not in that layout (a space after the colon, + reordered fields, ``event_type`` not first, a blank line). The caller must + then decode the whole line and classify it by its decoded ``event_type``, so + a change in how the log is written cannot silently drop events. + """ + if not line.startswith(EVENT_LINE_LEAD): + return None + start = len(EVENT_LINE_LEAD) + end = line.find(b'"', start, start + TOPIC_FRAME_SIZE) + return None if end == -1 else line[start:end] + + class EventRecordCodec: """MessageCodec[EventRecord] — binds the pub/sub layer to EventRecord wire format. diff --git a/src/inference_endpoint/evaluation/scoring.py b/src/inference_endpoint/evaluation/scoring.py index 790a17bcb..61bec6575 100644 --- a/src/inference_endpoint/evaluation/scoring.py +++ b/src/inference_endpoint/evaluation/scoring.py @@ -49,7 +49,7 @@ _evaluate = None _nltk = None -from ..core.record import EventRecord, EventType, SampleEventType +from ..core.record import EventRecord, EventType, SampleEventType, leading_event_topic from ..core.types import TextModelOutput from ..dataset_manager.agentic_inference_dataset import AgenticInferenceDataset from ..dataset_manager.dataset import Dataset @@ -62,6 +62,75 @@ logger = logging.getLogger(__name__) +# The compact layout JSONLWriter emits puts event_type first, so a line can be +# classified from its leading bytes (leading_event_topic) without decoding its +# payload. Any other layout is decoded whole; see guess_and_decode. +_ISSUED_TOPIC: bytes = SampleEventType.ISSUED.topic_bytes # type: ignore[attr-defined] +_COMPLETE_TOPIC: bytes = SampleEventType.COMPLETE.topic_bytes # type: ignore[attr-defined] +# An ISSUED line's key fields come before "data"; the cut is only attempted when +# "data" starts within this many bytes. +_ISSUED_KEY_WINDOW_BYTES = 1024 +_DATA_FIELD = b',"data":' +_KEY_FIELDS = (b'"sample_uuid":', b'"conversation_id":', b'"turn":') + + +class _IssuedKey(msgspec.Struct): + """The fields of an ISSUED record that the inline scorer reads.""" + + sample_uuid: str = "" + conversation_id: str = "" + turn: int | None = None + + +_EVENT_DECODER = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) +_ISSUED_DECODER = msgspec.json.Decoder(_IssuedKey) + + +def _decode_issued_key(line: bytes) -> _IssuedKey: + """The key fields of an ISSUED line, decoding only what is needed. + + The cut at "data" is an optimization: it applies when the key fields precede + "data" (the writer's field order) and the cut is valid JSON. In every other + case, including a cut that landed inside a nested object, the whole line is + decoded, which always works for a valid line. + """ + end = line.find(_DATA_FIELD, 0, _ISSUED_KEY_WINDOW_BYTES) + if end != -1 and all(field in line[:end] for field in _KEY_FIELDS): + try: + return _ISSUED_DECODER.decode(line[:end] + b"}") + except msgspec.DecodeError: + pass # "data" was not at the top level here: decode the whole line + return _ISSUED_DECODER.decode(line) + + +def guess_and_decode(line: bytes) -> EventRecord | _IssuedKey | None: + """Decode the ISSUED and COMPLETE events of one ``events.jsonl`` line. + + A line in the writer's compact layout (``event_type`` first) is classified + from those leading bytes alone, so other events return ``None`` undecoded and + an ISSUED line, which carries the whole prompt after "data", decodes only its + key fields into an ``_IssuedKey``. Any other layout (spaces, reordered fields, + ``event_type`` past the lead) is decoded whole and classified by its decoded + ``event_type``, so a change in how the log is written cannot silently drop + events. A blank line returns ``None``. + + Raises: + msgspec.DecodeError: the line is not a valid event. + """ + topic = leading_event_topic(line) + if topic is None: + if not line.strip(): + return None + record = _EVENT_DECODER.decode(line) + if record.event_type == SampleEventType.ISSUED: + return _IssuedKey(record.sample_uuid, record.conversation_id, record.turn) + return record if record.event_type == SampleEventType.COMPLETE else None + if topic == _ISSUED_TOPIC: + return _decode_issued_key(line) + if topic == _COMPLETE_TOPIC: + return _EVENT_DECODER.decode(line) + return None + class Scorer(ABC): """Scorers will read in a dataset and outputs from a log and compute an accuracy score. @@ -174,13 +243,30 @@ def _iter_complete(self) -> Iterator[tuple[str, Any]]: if not events_log_path.exists(): raise FileNotFoundError(f"Events log file not found at {events_log_path}") - decoder = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) - with events_log_path.open("r") as f: - for line in f: - stripped = line.strip() - if not stripped: + with events_log_path.open("rb") as f: + for line_no, line in enumerate(f, 1): + topic = leading_event_topic(line) + if topic is None: + # Not the compact layout: decode the line and classify it by + # its decoded event_type rather than dropping it unseen. + if not line.strip(): + continue + try: + record = _EVENT_DECODER.decode(line) + except msgspec.DecodeError as exc: + logger.warning( + "Skipping undecodable event log line %d in %s: %s", + line_no, + events_log_path, + exc, + ) + continue + elif topic == _COMPLETE_TOPIC: + record = _EVENT_DECODER.decode(line) + else: + # Other events (ISSUED lines carry the whole prompt) are + # skipped undecoded. continue - record = decoder.decode(stripped) if record.event_type == SampleEventType.COMPLETE: yield record.sample_uuid, record.data @@ -822,18 +908,14 @@ def _issued_and_completed_model_turns( if not events_path.exists(): raise FileNotFoundError(f"Events log file not found at {events_path}") - decoder = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) uuid_to_key: dict[str, tuple[str, int, int]] = {} completed_by_uuid: dict[str, dict[str, Any] | None] = {} issued_turns: set[tuple[str, int, int]] = set() model_turns: dict[tuple[str, int, int], dict[str, Any] | None] = {} - with events_path.open() as f: + with events_path.open("rb") as f: for line_no, line in enumerate(f, 1): - stripped = line.strip() - if not stripped: - continue try: - record = decoder.decode(stripped) + record = guess_and_decode(line) except msgspec.DecodeError as exc: logger.warning( "Skipping malformed event log line %d in %s: %s", @@ -842,10 +924,7 @@ def _issued_and_completed_model_turns( exc, ) continue - if record.event_type not in ( - SampleEventType.ISSUED, - SampleEventType.COMPLETE, - ): + if record is None: continue if record.turn is None or not record.conversation_id: continue @@ -861,7 +940,7 @@ def _issued_and_completed_model_turns( continue key = (conversation_id, repeat_id, turn) - if record.event_type == SampleEventType.ISSUED: + if isinstance(record, _IssuedKey): uuid_to_key[record.sample_uuid] = key issued_turns.add(key) if record.sample_uuid in completed_by_uuid: 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..f42d0742c 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 @@ -16,15 +16,21 @@ """Tests for BatchTokenizer and TokenBatchQueue.""" import asyncio +import logging import multiprocessing +import os import signal import time from concurrent.futures import Future, ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool from multiprocessing.connection import wait as wait_for_process +from types import SimpleNamespace from unittest.mock import patch import pytest +from inference_endpoint.async_utils.services.metrics_aggregator import ( + __main__ as aggregator_main, +) from inference_endpoint.async_utils.services.metrics_aggregator import ( token_metrics as token_metrics_module, ) @@ -205,15 +211,18 @@ def apply_chat_template( @pytest.mark.unit class TestBatchTokenizerMessageTokenization: - def test_chat_template_requests_token_ids_not_batch_encoding(self): + def test_prompt_counts_the_rendered_template(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, - None, - None, - None, + (count,) = tok.count_batch_sync( + [ + PromptInput( + ({"role": "user", "content": "one two three four"},), + None, + None, + None, + ) + ] ) assert count == 7 @@ -333,23 +342,29 @@ def apply_chat_template(self, 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, - None, - None, - None, + (count,) = tok.count_batch_sync( + [ + PromptInput( + ( + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this image"}, + { + "type": "image_url", + "image_url": {"url": "x"}, + }, + ], + }, + ), + None, + None, + None, + ) + ] ) - assert count > 0 + assert isinstance(count, int) and count > 0 def test_token_count_prompt_normalizes_tools_and_forwards_template_kwargs(self): class _RecordingTokenizer(_FakeTokenizerWithTemplate): @@ -379,12 +394,16 @@ def apply_chat_template(self, messages, **kwargs): ) 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}, - "custom template", - "auto", + tok.count_batch_sync( + [ + PromptInput( + messages, + None, + {"enable_thinking": False}, + "custom template", + "auto", + ) + ] ) normalized_call = _RecordingTokenizer.last_messages[0]["tool_calls"][0] @@ -395,7 +414,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 + assert _RecordingTokenizer.last_kwargs["tokenize"] is True class _Encoding: @@ -570,16 +589,13 @@ def test_structured_tokenization_does_not_require_a_text_backend(self, monkeypat 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, - None, - None, - None, - ) - == 5 - ) + assert tok.count_batch_sync( + [ + PromptInput( + ({"role": "user", "content": "one two"},), None, None, None + ) + ] + ) == [5] def test_affinity_unavailable_shards_unpinned(self, monkeypatch): """No affinity API (e.g. macOS): shard from the CPU count, unpinned.""" @@ -686,17 +702,48 @@ async def test_flush_remaining_stops_live_loop(self): @pytest.mark.unit class TestRayonCaps: - def test_ctor_caps_rayon_to_live_workers(self, monkeypatch): - monkeypatch.delenv("RAYON_NUM_THREADS", raising=False) + @staticmethod + def _unset_rayon_cap(monkeypatch): + """Make RAYON_NUM_THREADS absent for the test and restore it afterwards. + + setenv-then-delenv registers the original state (absent) for undo, so a + variable the code under test writes does not outlive the test. + """ + monkeypatch.setenv("RAYON_NUM_THREADS", "placeholder") + monkeypatch.delenv("RAYON_NUM_THREADS") + + def test_ctor_leaves_the_environment_alone(self, monkeypatch): + """Built in the CLI at finalize, a constructor that wrote the cap would + leak it into children started later (the audit run's aggregator).""" + self._unset_rayon_cap(monkeypatch) + before = dict(os.environ) with patch(_MOCK_TARGET, _FakeTokenizer): with BatchTokenizer("fake", n_workers=0, live_workers=3): - assert token_metrics_module.os.environ["RAYON_NUM_THREADS"] == "3" + assert dict(os.environ) == before + assert "RAYON_NUM_THREADS" not in os.environ - def test_ctor_respects_operator_exported_cap(self, monkeypatch): + def test_ctor_does_not_overwrite_an_exported_cap(self, monkeypatch): monkeypatch.setenv("RAYON_NUM_THREADS", "7") with patch(_MOCK_TARGET, _FakeTokenizer): with BatchTokenizer("fake", n_workers=0, live_workers=3): - assert token_metrics_module.os.environ["RAYON_NUM_THREADS"] == "7" + assert os.environ["RAYON_NUM_THREADS"] == "7" + + def test_aggregator_caps_the_live_lane_to_tokenizer_workers(self, monkeypatch): + self._unset_rayon_cap(monkeypatch) + aggregator_main._cap_live_lane_rayon_pool(3) + assert os.environ["RAYON_NUM_THREADS"] == "3" + + def test_aggregator_cap_never_reads_as_rayon_default(self, monkeypatch): + """--tokenizer-workers 0 must not become RAYON_NUM_THREADS=0, which + rayon treats as "use every CPU".""" + self._unset_rayon_cap(monkeypatch) + aggregator_main._cap_live_lane_rayon_pool(0) + assert os.environ["RAYON_NUM_THREADS"] == "1" + + def test_aggregator_cap_respects_an_operator_exported_cap(self, monkeypatch): + monkeypatch.setenv("RAYON_NUM_THREADS", "7") + aggregator_main._cap_live_lane_rayon_pool(3) + assert os.environ["RAYON_NUM_THREADS"] == "7" def test_init_worker_overrides_inherited_cap_with_block_size(self, monkeypatch): """Spawn children inherit the parent's live cap; each shard must @@ -1012,8 +1059,8 @@ def test_terminate_procs_kills_running_workers(): @pytest.mark.unit -class TestCountSync: - """The finalize-side synchronous counter. +class TestCountBatchSync: + """The finalize-side synchronous batch counter. Full-run OSL is computed after the run, outside the event loop, but it must produce the same number the perf-side async path would for the same @@ -1021,29 +1068,371 @@ class TestCountSync: (mlcommons/endpoints#500). """ + _MESSAGES = [ + MessageInput( + "hello world", + "reasoning here", + ( + { + "id": "c1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + }, + ), + ), + MessageInput("plain answer", "some longer reasoning text", None), + TextInput("hello world"), + PromptInput( + ( + {"role": "user", "content": "find the bug"}, + {"role": "assistant", "content": "looking", "reasoning_content": "hmm"}, + ), + ({"type": "function", "function": {"name": "grep", "parameters": {}}},), + None, + None, + ), + PromptInput(({"role": "user", "content": "just text"},), None, None, None), + ] + @pytest.mark.asyncio - async def test_count_sync_agrees_with_count_batch_async(self): - """Same input, same count, whichever lane counted it.""" - with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + @pytest.mark.parametrize( + "tokenizer_cls", + [ + _FakeTokenizerWithTemplate, + _FakeTokenizerWithTemplateAndBackend, + _FakeTokenizer, + _FakeTokenizerWithBackend, + ], + ids=[ + "chat-template-encode", + "chat-template-backend", + "fallback-encode", + "fallback-backend", + ], + ) + async def test_agrees_with_count_batch_async(self, tokenizer_cls): + """Same inputs, same counts, whichever lane counted them.""" + with patch(_MOCK_TARGET, tokenizer_cls): loop = asyncio.get_running_loop() with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: - message = MessageInput( - "hello world", - "reasoning here", - ( - { - "id": "c1", - "type": "function", - "function": {"name": "f", "arguments": "{}"}, - }, - ), + expected = await tok.count_batch_async(self._MESSAGES, loop) + + assert tok.count_batch_sync(self._MESSAGES) == expected + + def test_counts_plain_text_and_token_ids(self): + """Text takes the text path, not the chat template; token IDs count by length.""" + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + counts = tok.count_batch_sync( + [TextInput("hello world"), TokenIdsInput((1, 2, 3))] ) - expected = (await tok.count_batch_async([message], loop))[0] + assert counts == [2, 3] - assert tok.count_sync(message) == expected + def test_unsupported_input_fails_in_its_slot(self): + """An input it cannot count yields its exception without voiding the batch.""" + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + counts = tok.count_batch_sync([TextInput("a b"), object()]) # type: ignore[list-item] + assert counts[0] == 2 + assert isinstance(counts[1], TypeError) + + @pytest.mark.asyncio + async def test_many_turns_keep_input_order_across_threads(self): + """Interleaved kinds, more turns than threads: counts stay in input order + and equal the async lanes' numbers.""" + turns: list = [] + for i in range(60): + turns.append(MessageInput(" ".join(["w"] * (i % 9 + 1)), None, None)) + turns.append(TextInput(" ".join(["t"] * (i % 5 + 1)))) + turns.append(TokenIdsInput(tuple(range(i % 4)))) + turns.append( + PromptInput( + ({"role": "user", "content": " ".join(["p"] * (i % 7 + 1))},), + None, + None, + None, + ) + ) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=4) as tok: + expected = await tok.count_batch_async(turns, loop) + assert tok.count_batch_sync(turns) == expected - def test_count_sync_counts_plain_text(self): - """Text inputs take the text path, not the chat template.""" + def test_closed_tokenizer_raises(self): with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + tok = BatchTokenizer("fake", n_workers=0, live_workers=2) + tok.close() + with pytest.raises(RuntimeError, match="closed"): + tok.count_batch_sync([MessageInput("a b", None, None)]) + + +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: - assert tok.count_sync(TextInput("hello world")) == 2 + 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) + + @pytest.mark.asyncio + async def test_count_structured_async_routes_by_live_flag(self): + """Drain goes to the shards; live (and no-shard) stays in-process.""" + items = list(enumerate(_structured_items(4))) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + inproc = await tok._count_structured_async(items, loop, live=False) + proc = _InlineProc() + tok._procs = [proc] + + live = await tok._count_structured_async(items, loop, live=True) + assert proc.calls == 0 + drain = await tok._count_structured_async(items, loop, live=False) + assert proc.calls == 1 + + assert [i for i, _ in drain] == [i for i, _ in items] + assert drain == live == inproc + + @pytest.mark.asyncio + async def test_shard_returning_wrong_count_fails_only_its_chunk(self): + class _ShortProc: + """A shard that answers one count short for every chunk.""" + + def submit(self, _fn, chunk): + fut: Future = Future() + fut.set_result([1] * (len(chunk) - 1)) + return fut + + def shutdown(self, wait=False, cancel_futures=False): + pass + + size = token_metrics_module._STRUCTURED_CHUNK_ITEMS + items = _structured_items(size * 2) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + # Chunk 0 goes to the first shard, chunk 1 to the short one. + tok._procs = [_InlineProc(), _ShortProc()] + counts = await tok.count_batch_async(items, loop) + + assert all(isinstance(c, int) for c in counts[:size]) + assert all( + isinstance(c, RuntimeError) and "counts for" in str(c) + for c in counts[size:] + ) + + @pytest.mark.asyncio + async def test_drain_logs_progress_after_the_interval(self, monkeypatch, caplog): + """The drain can run for minutes with no other output, so it logs progress.""" + ticks = iter(range(0, 100_000, 31)) # every clock read advances 31 s + monkeypatch.setattr( + token_metrics_module, + "time", + SimpleNamespace( + monotonic=lambda: next(ticks), + perf_counter=time.perf_counter, + sleep=time.sleep, + ), + ) + caplog.set_level(logging.INFO, logger=token_metrics_module.logger.name) + items = _structured_items(token_metrics_module._STRUCTURED_CHUNK_ITEMS * 2) + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplateAndBackend): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + tok._procs = [_InlineProc()] + await tok.count_batch_async(items, loop) + + messages = [record.getMessage() for record in caplog.records] + assert any(m.startswith("Structured tokenization drain: ") for m in messages) + assert any( + m.startswith("Structured tokenization drain done: ") for m in messages + ) + + def test_worker_without_a_tokenizer_raises(self, monkeypatch): + monkeypatch.setattr(token_metrics_module, "_WORKER_TOKENIZER", None) + with pytest.raises(RuntimeError, match="unavailable"): + token_metrics_module._worker_count_structured( + [MessageInput("a b", None, None)] + ) + + def test_worker_fallback_warning_is_logged_once_per_error_type( + self, monkeypatch, caplog + ): + monkeypatch.setattr(token_metrics_module, "_WORKER_FALLBACK_WARNED", set()) + caplog.set_level(logging.WARNING, logger=token_metrics_module.logger.name) + warn = token_metrics_module._worker_warn_fallback + + warn(ValueError("first"), "impact") + warn(ValueError("second"), "impact") # same type: not logged again + warn(KeyError("other"), "impact") + + shard_warnings = [ + r for r in caplog.records if "tokenizer shard" in r.getMessage() + ] + assert len(shard_warnings) == 2 + + 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) + + +class _PoisonTokenizer(_FakeTokenizerWithTemplate): + """Fails on any text containing POISON, in the template path and in the fallback.""" + + def apply_chat_template(self, messages, **kwargs): + if any("POISON" in (m.get("content") or "") for m in messages): + raise ValueError("poison template") + return super().apply_chat_template(messages, **kwargs) + + def tokenize(self, text): + if "POISON" in text: + raise RuntimeError("poison tokenize") + return super().tokenize(text) + + +@pytest.mark.unit +class TestFailureIsolation: + """One bad turn yields its own exception and leaves the rest of the batch counted.""" + + _TURNS = [ + MessageInput("ok one", None, None), + MessageInput("POISON", None, None), + MessageInput("ok two three", None, None), + ] + + @staticmethod + def _assert_only_the_middle_failed(counts): + assert isinstance(counts[0], int) + assert isinstance(counts[1], RuntimeError) + assert isinstance(counts[2], int) + + @pytest.mark.asyncio + async def test_live_lane_isolates_the_failing_turn(self): + with patch(_MOCK_TARGET, _PoisonTokenizer): + loop = asyncio.get_running_loop() + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + counts = await tok.count_batch_async(self._TURNS, loop, live=True) + self._assert_only_the_middle_failed(counts) + + def test_sync_batch_isolates_the_failing_turn(self): + with patch(_MOCK_TARGET, _PoisonTokenizer): + with BatchTokenizer("fake", n_workers=0, live_workers=2) as tok: + counts = tok.count_batch_sync(self._TURNS) + self._assert_only_the_middle_failed(counts) + + @pytest.mark.asyncio + async def test_closed_tokenizer_marks_structured_turns(self): + with patch(_MOCK_TARGET, _FakeTokenizerWithTemplate): + loop = asyncio.get_running_loop() + tok = BatchTokenizer("fake", n_workers=0, live_workers=2) + tok.close() + counts = await tok.count_batch_async(self._TURNS, loop) + assert all(isinstance(c, RuntimeError) and "closed" in str(c) for c in counts) diff --git a/tests/unit/commands/test_full_run_osl.py b/tests/unit/commands/test_full_run_osl.py index 5bf3e1f74..c826cabb2 100644 --- a/tests/unit/commands/test_full_run_osl.py +++ b/tests/unit/commands/test_full_run_osl.py @@ -11,6 +11,7 @@ from __future__ import annotations +import json import logging from pathlib import Path @@ -20,6 +21,7 @@ MessageInput, TextInput, ) +from inference_endpoint.commands.benchmark import full_run_osl as full_run_osl_module from inference_endpoint.commands.benchmark.full_run_osl import ( compute_full_run_osl, full_run_osl_for_report, @@ -31,7 +33,8 @@ SampleEventType, SessionEventType, ) -from inference_endpoint.core.types import TextModelOutput +from inference_endpoint.core.types import PhaseType, TextModelOutput +from inference_endpoint.endpoint_client.cpu_affinity import UnsupportedPlatformError pytestmark = pytest.mark.unit @@ -48,6 +51,24 @@ def _word_count(tok_input: MessageInput | TextInput) -> int: return len(" ".join(parts).split()) +def _batched(count): + """A batch counter from a per-turn one: each turn's exception lands in its slot.""" + + def count_batch(items: list[MessageInput | TextInput]) -> list[int | Exception]: + outcomes: list[int | Exception] = [] + for item in items: + try: + outcomes.append(count(item)) + except Exception as exc: # noqa: BLE001 - mirror count_batch_sync + outcomes.append(exc) + return outcomes + + return count_batch + + +_word_counts = _batched(_word_count) + + def _write_events(path: Path, records: list[EventRecord]) -> Path: path.write_bytes(b"\n".join(_ENCODER.encode(r) for r in records) + b"\n") return path @@ -88,7 +109,7 @@ def test_counts_turns_completed_after_the_performance_window(tmp_path: Path) -> ], ) - result = compute_full_run_osl(events, _word_count) + result = compute_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 4 @@ -120,8 +141,8 @@ def build(name: str, stop_after: int) -> Path: ], ) - early = compute_full_run_osl(build("early.jsonl", 2), _word_count) - late = compute_full_run_osl(build("late.jsonl", 5), _word_count) + early = compute_full_run_osl(build("early.jsonl", 2), _word_counts) + late = compute_full_run_osl(build("late.jsonl", 5), _word_counts) assert early is not None and late is not None assert early["n_turns_counted"] == late["n_turns_counted"] == 6 @@ -148,7 +169,7 @@ def test_excludes_completions_from_other_phases(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count, performance_uuids={"u1", "u3"}) + result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 2 @@ -171,7 +192,7 @@ def test_counts_empty_completions_separately(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count) + result = compute_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -233,10 +254,10 @@ def windowed_mean(stop_after: int) -> float: return sum(counted) / len(counted) early = compute_full_run_osl( - build("e.jsonl", 2), _word_count, performance_uuids=uuids + build("e.jsonl", 2), _word_counts, performance_uuids=uuids ) late = compute_full_run_osl( - build("l.jsonl", 5), _word_count, performance_uuids=uuids + build("l.jsonl", 5), _word_counts, performance_uuids=uuids ) assert windowed_mean(2) != windowed_mean(5) # the windowed mean moves @@ -261,7 +282,7 @@ def test_data_none_completion_counts_as_empty(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count) + result = compute_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -286,7 +307,7 @@ def flaky(tok_input: MessageInput | TextInput) -> int: raise ValueError("tokenizer blew up") return _word_count(tok_input) - result = compute_full_run_osl(events, flaky) + result = compute_full_run_osl(events, _batched(flaky)) assert result is not None assert result["n_turns_counted"] == 1 @@ -294,6 +315,89 @@ def flaky(tok_input: MessageInput | TextInput) -> int: assert result["output_sequence_lengths"]["total"] == 1 # only "one" counted +def _five_turns(tmp_path: Path) -> Path: + texts = ["a", "b b", "c c c", "d d d d", "e e e e e"] + return _write_events( + tmp_path / "events.jsonl", + [ + _session(SessionEventType.STARTED), + *(_complete(turn, text) for turn, text in enumerate(texts, start=1)), + _session(SessionEventType.ENDED), + ], + ) + + +def test_batches_count_every_turn(tmp_path: Path) -> None: + """Full batches plus a final partial one count the same as a single batch.""" + events = _five_turns(tmp_path) + batch_sizes: list[int] = [] + + def recording(items): + batch_sizes.append(len(items)) + return _word_counts(items) + + batched = compute_full_run_osl(events, recording, batch_size=2) + single = compute_full_run_osl(events, _word_counts, batch_size=1024) + + assert batch_sizes == [2, 2, 1] + assert batched is not None and single is not None + assert batched["n_turns_counted"] == 5 + assert batched["output_sequence_lengths"] == single["output_sequence_lengths"] + + +def test_failed_batch_charges_each_of_its_turns(tmp_path: Path) -> None: + """A batch call that raises charges each of its turns to n_errors; later + batches still count.""" + events = _five_turns(tmp_path) + batch_sizes: list[int] = [] + + def first_batch_fails(items): + batch_sizes.append(len(items)) + if len(batch_sizes) == 1: + raise RuntimeError("encoder crashed") + return _word_counts(items) + + result = compute_full_run_osl(events, first_batch_fails, batch_size=2) + + assert batch_sizes == [2, 2, 1] # the failed batch is not retried + assert result is not None + assert result["n_errors"] == 2 + assert result["n_turns_counted"] == 3 + assert result["output_sequence_lengths"]["total"] == 3 + 4 + 5 + assert result["partial"] is True + + +@pytest.mark.parametrize("n_outcomes", [0, 1, 4], ids=["none", "short", "long"]) +def test_counter_returning_the_wrong_number_of_outcomes_charges_its_batch( + tmp_path: Path, caplog: pytest.LogCaptureFixture, n_outcomes: int +) -> None: + """One outcome per turn is the counter's contract. Any other number cannot be + matched back to its turns, so the batch is charged to n_errors and the block + is partial, not published as complete over an unknown subset.""" + events = _five_turns(tmp_path) + batch_sizes: list[int] = [] + + def wrong_count_on_first_batch(items): + batch_sizes.append(len(items)) + if len(batch_sizes) == 1: + return [1] * n_outcomes + return _word_counts(items) + + with caplog.at_level(logging.WARNING): + result = compute_full_run_osl(events, wrong_count_on_first_batch, batch_size=2) + + assert batch_sizes == [2, 2, 1] + assert result is not None + assert result["n_errors"] == 2 # the first batch's two turns + assert result["n_turns_counted"] == 3 + assert result["output_sequence_lengths"]["total"] == 3 + 4 + 5 + assert result["partial"] is True + assert any( + f"returned {n_outcomes} outcomes for 2 turns" in r.getMessage() + for r in caplog.records + ) + + def test_all_empty_population_returns_a_visible_block(tmp_path: Path) -> None: """A population with records but nothing countable returns a block (not None) so the run is visible rather than looking uncomputed.""" @@ -309,7 +413,7 @@ def test_all_empty_population_returns_a_visible_block(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count) + result = compute_full_run_osl(events, _word_counts) assert result is not None assert result["n_turns_counted"] == 0 @@ -329,7 +433,7 @@ def test_undecodable_complete_line_is_skipped( path.write_bytes(good + b"\n" + corrupt + b"\n") with caplog.at_level(logging.WARNING): - result = compute_full_run_osl(path, _word_count) + result = compute_full_run_osl(path, _word_counts) assert result is not None assert result["n_turns_counted"] == 1 @@ -340,6 +444,60 @@ def test_undecodable_complete_line_is_skipped( assert sum("undecodable COMPLETE line" in m for m in messages) == 1 +def _spaced_line(record: EventRecord) -> bytes: + """The record as valid JSON with json.dumps' default spacing.""" + return json.dumps(json.loads(_ENCODER.encode(record))).encode() + + +def _event_type_last_line(record: EventRecord) -> bytes: + """The record with its fields reversed, so event_type comes last.""" + fields = json.loads(_ENCODER.encode(record)) + return json.dumps(dict(reversed(fields.items())), separators=(",", ":")).encode() + + +def test_complete_lines_in_other_valid_layouts_are_counted(tmp_path: Path) -> None: + """The reader must not depend on the writer's compact layout: spaced and + reordered COMPLETE lines count like compact ones, and blank lines and other + events in any layout are skipped, not charged.""" + path = tmp_path / "events.jsonl" + lines = [ + _ENCODER.encode(_complete(1, "a b")), + _spaced_line(_complete(3, "c d e")), + _event_type_last_line(_complete(5, "f")), + _spaced_line(_session(SessionEventType.ENDED)), + b"", + ] + path.write_bytes(b"\n".join(lines) + b"\n") + + result = compute_full_run_osl( + path, _word_counts, performance_uuids={"u1", "u3", "u5"} + ) + + assert result is not None + assert result["n_turns_counted"] == 3 + assert result["output_sequence_lengths"]["total"] == 6 # 2 + 3 + 1 words + assert result["n_missing"] == 0 + assert result["n_undecodable"] == 0 + assert result["partial"] is False + + +def test_unreadable_line_that_does_not_decode_is_charged_as_undecodable( + tmp_path: Path, +) -> None: + """A line whose event type cannot be read from its lead and that is not a + valid event may have been a COMPLETE turn, so it is charged (and the block + marked partial) rather than skipped unseen.""" + path = tmp_path / "events.jsonl" + path.write_bytes(_ENCODER.encode(_complete(1, "one")) + b"\nnot json\n") + + result = compute_full_run_osl(path, _word_counts) + + assert result is not None + assert result["n_turns_counted"] == 1 + assert result["n_undecodable"] == 1 + assert result["partial"] is True + + def test_all_undecodable_returns_visible_block(tmp_path: Path) -> None: """A stream of only corrupt COMPLETE lines returns a block (not None) so the accuracy field is not mistaken for uncomputed.""" @@ -347,7 +505,7 @@ def test_all_undecodable_returns_visible_block(tmp_path: Path) -> None: path = tmp_path / "events.jsonl" path.write_bytes(corrupt + b"\n" + corrupt + b"\n") - result = compute_full_run_osl(path, _word_count) + result = compute_full_run_osl(path, _word_counts) assert result is not None assert result["n_turns_counted"] == 0 @@ -364,6 +522,138 @@ def test_for_report_fails_closed_without_performance_map(tmp_path: Path) -> None assert full_run_osl_for_report(tmp_path, "any-tokenizer") is None +class _SpyTokenizer: + """Stands in for BatchTokenizer: records how it was built, counts by words.""" + + built: list[dict] = [] + + def __init__(self, tokenizer_name: str, *, live_workers: int, n_workers: int): + type(self).built.append( + { + "tokenizer_name": tokenizer_name, + "live_workers": live_workers, + "n_workers": n_workers, + } + ) + + def __enter__(self) -> _SpyTokenizer: + return self + + def __exit__(self, *exc_info: object) -> None: + return None + + def count_batch_sync(self, items): + return _word_counts(items) + + +def _perf_report(report_dir: Path) -> None: + """A two-turn performance report that full_run_osl_for_report will accept.""" + _write_events( + report_dir / "events.jsonl", + [_complete(1, "a"), _complete(3, "b c")], + ) + (report_dir / "sample_idx_map.json").write_bytes( + msgspec.json.encode({PhaseType.PERFORMANCE.value: {"u1": 0, "u3": 1}}) + ) + + +@pytest.mark.parametrize("n_cpus", [4, 32, 352]) +def test_for_report_caps_counting_threads( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, n_cpus: int +) -> None: + """Finalize counts on at most one thread per usable CPU, capped where adding + threads stops helping, and without tokenizer shards (n_workers=0).""" + _perf_report(tmp_path) + _SpyTokenizer.built = [] + monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr( + full_run_osl_module, "get_all_online_cpus", lambda: list(range(n_cpus)) + ) + + result = full_run_osl_for_report(tmp_path, "tok") + + assert result is not None and result["n_turns_counted"] == 2 + assert _SpyTokenizer.built == [ + { + "tokenizer_name": "tok", + "live_workers": min(n_cpus, full_run_osl_module._MAX_COUNT_THREADS), + "n_workers": 0, + } + ] + + +@pytest.mark.parametrize( + ("cpu_count", "expected"), + [(64, full_run_osl_module._MAX_COUNT_THREADS), (8, 8), (None, 1)], +) +def test_for_report_threads_without_cpu_affinity( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + cpu_count: int | None, + expected: int, +) -> None: + """Where the affinity API is unsupported, fall back to os.cpu_count().""" + _perf_report(tmp_path) + _SpyTokenizer.built = [] + monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + + def _unsupported() -> list[int]: + raise UnsupportedPlatformError("no affinity here") + + monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", _unsupported) + monkeypatch.setattr(full_run_osl_module.os, "cpu_count", lambda: cpu_count) + + assert full_run_osl_for_report(tmp_path, "tok") is not None + assert _SpyTokenizer.built[0]["live_workers"] == expected + + +def test_for_report_empty_cpu_list_counts_on_one_thread( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A probe that finds no CPUs must not ask for zero counting threads.""" + _perf_report(tmp_path) + _SpyTokenizer.built = [] + monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", lambda: []) + + assert full_run_osl_for_report(tmp_path, "tok") is not None + assert _SpyTokenizer.built[0]["live_workers"] == 1 + + +class _FailingTokenizer(_SpyTokenizer): + """A tokenizer that cannot be constructed.""" + + def __init__(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("tokenizer could not be loaded") + + +@pytest.mark.parametrize("failure", ["cpu_probe", "tokenizer"]) +def test_for_report_failure_drops_the_block_with_a_warning( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + failure: str, +) -> None: + """The block is best-effort: whatever fails while counting (the CPU probe, + the tokenizer), finalize gets None and a warning, not an exception.""" + _perf_report(tmp_path) + if failure == "cpu_probe": + + def _broken_probe() -> list[int]: + raise RuntimeError("probe broke") + + monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _SpyTokenizer) + monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", _broken_probe) + else: + monkeypatch.setattr(full_run_osl_module, "BatchTokenizer", _FailingTokenizer) + monkeypatch.setattr(full_run_osl_module, "get_all_online_cpus", lambda: [0, 1]) + + with caplog.at_level(logging.WARNING): + assert full_run_osl_for_report(tmp_path, "tok") is None + + assert any("Full-run OSL skipped" in r.getMessage() for r in caplog.records) + + def test_for_report_skips_without_tokenizer_or_events(tmp_path: Path) -> None: """No tokenizer, or a missing events.jsonl, short-circuits to None.""" assert full_run_osl_for_report(tmp_path, None) is None # no tokenizer @@ -382,7 +672,7 @@ def test_missing_turn_is_detected_via_uuid_reconciliation(tmp_path: Path) -> Non ], ) - result = compute_full_run_osl(events, _word_count, performance_uuids={"u1", "u3"}) + result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 1 @@ -408,7 +698,9 @@ def flaky(tok_input: MessageInput | TextInput) -> int: raise ValueError("tokenizer blew up") return _word_count(tok_input) - result = compute_full_run_osl(events, flaky, performance_uuids={"u1", "u3"}) + result = compute_full_run_osl( + events, _batched(flaky), performance_uuids={"u1", "u3"} + ) assert result is not None assert result["n_turns_counted"] == 1 @@ -428,7 +720,7 @@ def test_clean_run_is_not_partial(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count, performance_uuids={"u1", "u3"}) + result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_turns_counted"] == 2 @@ -449,7 +741,7 @@ def test_empty_turns_do_not_make_block_partial(tmp_path: Path) -> None: ], ) - result = compute_full_run_osl(events, _word_count, performance_uuids={"u1", "u3"}) + result = compute_full_run_osl(events, _word_counts, performance_uuids={"u1", "u3"}) assert result is not None assert result["n_empty"] == 1 diff --git a/tests/unit/core/test_record.py b/tests/unit/core/test_record.py index d94934297..ecade65a6 100644 --- a/tests/unit/core/test_record.py +++ b/tests/unit/core/test_record.py @@ -27,6 +27,7 @@ EventType, SampleEventType, SessionEventType, + leading_event_topic, ) from inference_endpoint.core.types import ( ErrorData, @@ -314,3 +315,35 @@ def test_reraises_non_decode_error(self): # silently swallow decode-path bugs. with pytest.raises(ValueError, match="not a decode error"): _codec.on_decode_error(b"", ValueError("not a decode error")) + + +@pytest.mark.unit +class TestLeadingEventTopic: + """leading_event_topic reads the event type from a compact line's first bytes.""" + + _ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + + @pytest.mark.parametrize( + "event_type", + [*SessionEventType, *SampleEventType, *ErrorEventType], + ids=lambda member: member.topic, + ) + def test_every_event_type_a_writer_emits_is_found_at_the_lead(self, event_type): + line = self._ENCODER.encode(EventRecord(event_type=event_type)) + b"\n" + + assert leading_event_topic(line) == event_type.topic_bytes + + @pytest.mark.parametrize( + "line", + [ + b'{"event_type": "sample.complete"}', # a space after the colon + b'{"sample_uuid":"u","event_type":"sample.complete"}', # not first + b'{"event_type":"sample.complete', # no closing quote + b'{"event_type":"' + b"x" * 100 + b'"}', # longer than any topic + b"not json", + b"\n", + b"", + ], + ) + def test_any_other_layout_has_no_leading_topic(self, line): + assert leading_event_topic(line) is None diff --git a/tests/unit/evaluation/test_scoring.py b/tests/unit/evaluation/test_scoring.py index 42fbb7bee..6c1a89e8d 100644 --- a/tests/unit/evaluation/test_scoring.py +++ b/tests/unit/evaluation/test_scoring.py @@ -16,14 +16,22 @@ """Unit tests for evaluation scoring module.""" import json +import logging from pathlib import Path from unittest.mock import MagicMock import msgspec import pandas as pd import pytest -from inference_endpoint.core.record import EventRecord, EventType, SampleEventType -from inference_endpoint.core.types import TextModelOutput +from inference_endpoint.async_utils.services.event_logger.file_writer import JSONLWriter +from inference_endpoint.core.record import ( + ErrorEventType, + EventRecord, + EventType, + SampleEventType, + SessionEventType, +) +from inference_endpoint.core.types import PromptData, TextModelOutput from inference_endpoint.dataset_manager.agentic_inference_dataset import ( AgenticInferenceDataset, ) @@ -40,8 +48,10 @@ VBenchScorer, _calculate_hierarchical_f1, _create_pred_pad_category, + _IssuedKey, _match_hierarchical_paths, _parse_response_to_category, + guess_and_decode, ) from pydantic import ValidationError @@ -228,6 +238,398 @@ def test_score_requires_sample_index_map_and_events(self, mock_dataset, report_d ) +_LINE_ENCODER = msgspec.json.Encoder(enc_hook=EventType.encode_hook) +# Long enough that a layout putting event_type after it starts it past byte 256. +_LONG_ID = "conv-" + "c" * 300 + + +def _compact(record: EventRecord) -> bytes: + """The line JSONLWriter writes: event_type first, no spaces.""" + return _LINE_ENCODER.encode(record) + b"\n" + + +def _spaced(record: EventRecord) -> bytes: + """The same record as valid JSON with json.dumps' default spacing.""" + return json.dumps(json.loads(_LINE_ENCODER.encode(record))).encode() + b"\n" + + +def _event_type_last(record: EventRecord) -> bytes: + """The same record with its fields reversed, so event_type comes last.""" + fields = json.loads(_LINE_ENCODER.encode(record)) + reordered = dict(reversed(fields.items())) + return json.dumps(reordered, separators=(",", ":")).encode() + b"\n" + + +def _event_type_past_the_lead(record: EventRecord) -> bytes: + """event_type after the (long) conversation_id, beyond the first 256 bytes.""" + fields = json.loads(_LINE_ENCODER.encode(record)) + reordered = {"conversation_id": fields.pop("conversation_id"), **fields} + return json.dumps(reordered, separators=(",", ":")).encode() + b"\n" + + +# Valid layouts of one record: only the first is what the writer emits today. +_LAYOUTS = { + "compact": _compact, + "spaced": _spaced, + "event_type_last": _event_type_last, + "event_type_past_the_lead": _event_type_past_the_lead, +} + + +@pytest.mark.unit +class TestGuessAndDecode: + @staticmethod + def _line(record: EventRecord) -> bytes: + encoder = msgspec.json.Encoder(enc_hook=EventType.encode_hook) + return encoder.encode(record) + b"\n" + + def test_issued_line_decodes_only_key_fields(self): + line = self._line( + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="u1", + conversation_id="conv1__repeat_2", + turn=3, + data=PromptData(messages=({"role": "user", "content": "hello"},)), + ) + ) + + assert guess_and_decode(line) == _IssuedKey( + sample_uuid="u1", conversation_id="conv1__repeat_2", turn=3 + ) + + def test_complete_line_decodes_full_record(self): + record = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + conversation_id="conv1__repeat_2", + turn=3, + data=TextModelOutput(output="done"), + ) + + assert guess_and_decode(self._line(record)) == record + + def test_issued_line_with_data_past_1024_bytes_decodes_in_full(self): + conversation_id = "c" * 2000 + line = self._line( + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="u1", + conversation_id=conversation_id, + turn=3, + data=PromptData(messages=({"role": "user", "content": "hello"},)), + ) + ) + assert line.find(b',"data":') > 1024 + + assert guess_and_decode(line) == _IssuedKey( + sample_uuid="u1", conversation_id=conversation_id, turn=3 + ) + + def test_other_event_returns_none(self): + line = self._line( + EventRecord( + event_type=SampleEventType.RECV_FIRST, + sample_uuid="u1", + conversation_id="conv1", + turn=3, + ) + ) + + assert guess_and_decode(line) is None + + @pytest.mark.parametrize("layout", list(_LAYOUTS)) + def test_complete_line_in_any_layout(self, layout): + record = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + conversation_id=_LONG_ID, + turn=3, + data=TextModelOutput(output="done"), + ) + + assert guess_and_decode(_LAYOUTS[layout](record)) == record + + @pytest.mark.parametrize("layout", list(_LAYOUTS)) + def test_issued_line_in_any_layout(self, layout): + record = EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="u1", + conversation_id=_LONG_ID, + turn=3, + data=PromptData(messages=({"role": "user", "content": "hello"},)), + ) + + assert guess_and_decode(_LAYOUTS[layout](record)) == _IssuedKey( + sample_uuid="u1", conversation_id=_LONG_ID, turn=3 + ) + + @pytest.mark.parametrize("layout", list(_LAYOUTS)) + def test_other_event_in_any_layout_returns_none(self, layout): + record = EventRecord( + event_type=SampleEventType.RECV_FIRST, + sample_uuid="u1", + conversation_id=_LONG_ID, + turn=3, + ) + + assert guess_and_decode(_LAYOUTS[layout](record)) is None + + def test_complete_whose_payload_mentions_the_issued_marker_stays_complete(self): + """The event type comes from the line's lead, never from payload bytes.""" + record = EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="u1", + conversation_id="c", + turn=1, + data=TextModelOutput(tool_calls=[{"event_type": "sample.issued"}]), + ) + line = self._line(record) + # Searching the first bytes for the marker would read this as ISSUED. + assert b'"event_type":"sample.issued"' in line[:256] + + assert guess_and_decode(line) == record + + def test_issued_cut_inside_a_nested_object_falls_back_to_the_whole_line(self): + """A "data" key nested before the top-level one is not a field boundary.""" + line = ( + b'{"event_type":"sample.issued","timestamp_ns":1,"sample_uuid":"u1",' + b'"conversation_id":"c1","turn":1,"meta":{"x":1,"data":2},"data":null}\n' + ) + cut = line[: line.find(b',"data":')] + b"}" + with pytest.raises(msgspec.DecodeError): + msgspec.json.decode(cut) # the optimization's cut is not valid JSON + + assert guess_and_decode(line) == _IssuedKey( + sample_uuid="u1", conversation_id="c1", turn=1 + ) + + @pytest.mark.parametrize("line", [b"\n", b" \n", b""]) + def test_blank_line_returns_none(self, line): + assert guess_and_decode(line) is None + + @pytest.mark.parametrize( + "line", + [ + b"not json\n", + b'{"event_type":"sample.complete","time\n', # truncated, compact lead + b'{"event_type": "sample.complete", "time\n', # truncated, spaced + ], + ) + def test_invalid_line_raises(self, line): + with pytest.raises(msgspec.DecodeError): + guess_and_decode(line) + + +def _write_event_lines(report_dir: Path, lines: list[bytes]) -> None: + report_dir.mkdir(parents=True, exist_ok=True) + (report_dir / "sample_idx_map.json").write_bytes( + msgspec.json.encode({"performance": {}}) + ) + (report_dir / "events.jsonl").write_bytes(b"".join(lines)) + + +@pytest.mark.unit +class TestEventLogContract: + """What the event logger writes is what the scorer reads. + + The scorer classifies log lines from their leading bytes. These tests feed it + the real ``JSONLWriter`` output, and other valid layouts of the same records, + so a change to the writer or a differently formatted log cannot silently drop + events or change a score. + """ + + @staticmethod + def _bash_tool_call(call_id: str, command: str) -> dict: + return { + "id": call_id, + "type": "function", + "function": {"name": "bash", "arguments": json.dumps({"cmd": command})}, + } + + @classmethod + def _dataset(cls) -> AgenticInferenceDataset: + return AgenticInferenceDataset( + pd.DataFrame( + [ + { + "conversation_id": "code1", + "turn": 1, + "role": "user", + "content": "run the tests", + }, + { + "conversation_id": "code1", + "turn": 2, + "role": "assistant", + "content": None, + "tool_calls": [ + cls._bash_tool_call("expected-code", "python -m pytest") + ], + }, + { + "conversation_id": "sim_1", + "turn": 1, + "role": "user", + "content": "choose the next workflow action", + }, + { + "conversation_id": "sim_1", + "turn": 2, + "role": "assistant", + "content": "expected workflow action", + "intent_codes": ["I042"], + }, + ] + ) + ) + + @classmethod + def _records(cls) -> list[EventRecord]: + """Both turns scored correctly, with an unrelated event in between.""" + return [ + EventRecord(event_type=SessionEventType.STARTED), + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="code-response", + conversation_id="code1", + turn=1, + data=PromptData(messages=({"role": "user", "content": "run it"},)), + ), + EventRecord( + event_type=SampleEventType.RECV_FIRST, + sample_uuid="code-response", + conversation_id="code1", + turn=1, + ), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="code-response", + conversation_id="code1", + turn=1, + data=TextModelOutput( + tool_calls=[cls._bash_tool_call("model-code", "python test.py")] + ), + ), + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="workflow-response", + conversation_id="sim_1", + turn=1, + ), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="workflow-response", + conversation_id="sim_1", + turn=1, + data=TextModelOutput(output="intent: I042"), + ), + EventRecord(event_type=SessionEventType.ENDED), + ] + + def _assert_scored_in_full(self, report_dir: Path) -> None: + score, repeats = AgenticInferenceInlineScorer( + "performance", self._dataset(), report_dir + ).score() + assert (score, repeats) == (1.0, 1) + scores = json.loads((report_dir / "scores.json").read_text()) + assert scores["turns"]["scored"] == 2 + + def test_every_event_type_the_writer_emits_classifies_like_a_full_decode( + self, tmp_path + ): + records = [ + *(EventRecord(event_type=t, sample_uuid="u") for t in SessionEventType), + *(EventRecord(event_type=t, sample_uuid="u") for t in SampleEventType), + *(EventRecord(event_type=t, sample_uuid="u") for t in ErrorEventType), + *self._records(), + ] + writer = JSONLWriter(tmp_path / "events", flush_interval=1) + try: + for record in records: + writer.write(record) + finally: + writer.close() + lines = (tmp_path / "events.jsonl").read_bytes().splitlines(keepends=True) + assert len(lines) == len(records) + decoder = msgspec.json.Decoder(type=EventRecord, dec_hook=EventType.decode_hook) + + for line in lines: + full = decoder.decode(line) + got = guess_and_decode(line) + if full.event_type == SampleEventType.ISSUED: + assert got == _IssuedKey( + full.sample_uuid, full.conversation_id, full.turn + ) + elif full.event_type == SampleEventType.COMPLETE: + assert got == full + else: + assert got is None + + def test_real_writer_output_is_scored(self, tmp_path): + report_dir = tmp_path / "report" + report_dir.mkdir() + (report_dir / "sample_idx_map.json").write_bytes( + msgspec.json.encode({"performance": {}}) + ) + writer = JSONLWriter(report_dir / "events", flush_interval=1) + try: + for record in self._records(): + writer.write(record) + finally: + writer.close() + + self._assert_scored_in_full(report_dir) + + @pytest.mark.parametrize("layout", list(_LAYOUTS)) + def test_any_layout_scores_the_same(self, tmp_path, layout): + report_dir = tmp_path / "report" + _write_event_lines( + report_dir, [_LAYOUTS[layout](record) for record in self._records()] + ) + + self._assert_scored_in_full(report_dir) + + def test_iter_complete_reads_every_layout_and_skips_blank_and_bad_lines( + self, tmp_path, caplog + ): + def complete(uuid: str) -> EventRecord: + return EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid=uuid, + conversation_id="c1", + turn=1, + data=TextModelOutput(output=uuid), + ) + + report_dir = tmp_path / "report" + _write_event_lines( + report_dir, + [ + _compact(complete("u1")), + _spaced(complete("u2")), + _event_type_last(complete("u3")), + b"\n", + b"not json\n", + _compact( + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="u4", + conversation_id="c1", + turn=1, + ) + ), + ], + ) + scorer = AgenticInferenceInlineScorer("performance", MagicMock(), report_dir) + + with caplog.at_level(logging.WARNING, logger=scoring_mod.logger.name): + uuids = [uuid for uuid, _ in scorer._iter_complete()] + + assert uuids == ["u1", "u2", "u3"] + assert [r.getMessage() for r in caplog.records if "line 5" in r.getMessage()] + + @pytest.mark.unit class TestAgenticInferenceInlineScorer: """AgenticInferenceInlineScorer unit tests.""" @@ -531,6 +933,139 @@ def test_scores_issued_turns_without_rounding_to_full_repeats(self, tmp_path): "scored": 3, } + def test_issued_records_with_prompts_decode_only_their_keys(self, tmp_path): + """ISSUED lines carry the whole prompt after "data"; the scorer reads only + the fields before it, whatever the prompt holds, and still joins each + repeat's completion (or its failure) to the issued turn.""" + dataset = AgenticInferenceDataset( + pd.DataFrame( + [ + { + "conversation_id": "code1", + "turn": 1, + "role": "user", + "content": "run the tests", + }, + { + "conversation_id": "code1", + "turn": 2, + "role": "assistant", + "tool_calls": [ + self._bash_tool_call("expected", "python -m pytest") + ], + }, + ] + ) + ) + # A prompt that itself mentions the event-type and "data" markers. + prompt = PromptData( + messages=( + { + "role": "user", + "content": 'x,"data":"event_type":"sample.complete" ' * 500, + }, + ) + ) + report_dir = tmp_path / "report" + self._write_report( + report_dir, + [ + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="r1", + conversation_id="code1", + turn=1, + data=prompt, + ), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="r1", + conversation_id="code1", + turn=1, + data=TextModelOutput( + tool_calls=[self._bash_tool_call("model", "python test.py")] + ), + ), + EventRecord( + event_type=SampleEventType.ISSUED, + sample_uuid="r2", + conversation_id="code1__repeat_2", + turn=1, + data=prompt, + ), + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="r2", + conversation_id="code1__repeat_2", + turn=1, + ), + ], + ) + + score, repeats = AgenticInferenceInlineScorer( + "performance", dataset, report_dir + ).score() + + assert (score, repeats) == (0.5, 2) + scores = json.loads((report_dir / "scores.json").read_text()) + assert scores["turns"]["issued"] == 2 + assert scores["turns"]["observed"] == 1 + + @pytest.mark.parametrize("late_field", ["sample_uuid", "conversation_id", "turn"]) + def test_issued_line_in_another_field_order_is_decoded_in_full( + self, tmp_path, late_field + ): + """An ISSUED line with a key field after "data" falls back to decoding + the whole line.""" + dataset = AgenticInferenceDataset( + pd.DataFrame( + [ + { + "conversation_id": "sim_1", + "turn": 1, + "role": "user", + "content": "choose the next workflow action", + }, + { + "conversation_id": "sim_1", + "turn": 2, + "role": "assistant", + "content": "expected workflow action", + "intent_codes": ["I042"], + }, + ] + ) + ) + report_dir = tmp_path / "report" + self._write_report( + report_dir, + [ + EventRecord( + event_type=SampleEventType.COMPLETE, + sample_uuid="w1", + conversation_id="sim_1", + turn=1, + data=TextModelOutput(output="intent: I042"), + ), + ], + ) + fields = {"sample_uuid": '"w1"', "conversation_id": '"sim_1"', "turn": "1"} + early = "".join( + f'"{name}":{value},' for name, value in fields.items() if name != late_field + ) + line = ( + f'{{"event_type":"sample.issued",{early}"data":null,' + f'"{late_field}":{fields[late_field]}}}\n' + ) + with (report_dir / "events.jsonl").open("a") as f: + f.write(line) + + score, repeats = AgenticInferenceInlineScorer( + "performance", dataset, report_dir + ).score() + + assert (score, repeats) == (1.0, 1) + @pytest.mark.unit class TestVBenchScorerRegistration: