Repository navigation
Reduce the runtime of benchmarking draining phase - #556
Conversation
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
github-actions
Bot
added
the
size/normal
zihaok
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #556 +/- ##
=======================================
Coverage ? 82.62%
=======================================
Files ? 170
Lines ? 24317
Branches ? 0
=======================================
Hits ? 20092
Misses ? 4225
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
7bd9f13 to
28f6d9e
Compare
viraatc
left a comment
viraatc
left a comment
There was a problem hiding this comment.
Thanks, the batching, index/offset alignment and the scoring.py fast path all check out. Three inline notes: the RAYON_NUM_THREADS leak into the audit aggregator (requesting changes on this one), a question on peak RSS for the unbounded drain, and a count-parity note for render-then-encode.
| return None | ||
| try: | ||
| with BatchTokenizer(tokenizer_name, live_workers=1, n_workers=0) as tokenizer: | ||
| # A batch is encoded across the tokenizer's rayon threads, which |
There was a problem hiding this comment.
This sets RAYON_NUM_THREADS to the full CPU count in the CLI process environment, via BatchTokenizer.__init__'s setdefault. With audit configured, run_audit runs next in the same process (commands/benchmark/cli.py:89). Its metrics aggregator is started with Popen (launcher.py:122) and inherits that value, so the aggregator's own setdefault from --metrics-tokenizer-workers (metrics_aggregator/__main__.py:262) is a no-op. Its live-lane encodes then use every core, including the ones the client workers are pinned to, during the TEST04 load phases. Base leaked "1", so this was already wrong, but now in the direction that contends with the client. Could the env write move out of BatchTokenizer.__init__ into the aggregator's __main__, and finalize leave it unset? Rayon's default already sizes itself to the process affinity.
There was a problem hiding this comment.
moved setdefault to main
| # Render serially on the tokenizer thread (the event loop stays free), | ||
| # then encode the texts in one batch: the bounded live lane mid-run, | ||
| # the shard pool at the end-of-run drain. | ||
| rendered = await loop.run_in_executor( |
There was a problem hiding this comment.
At the end-of-run drain, drain_all hands the whole buffer to count_batch_async. This now renders every PromptInput / MessageInput, pickles them to the shards and holds every encoding at once, where base held one item. Agentic prompts are full histories, so peak RSS scales with backlog × ISL. What was the aggregator's peak RSS in the c18k run? Slicing structured by _LIVE_FLUSH_MAX_ITEMS would match the 1024 batching this PR already uses in full_run_osl and the existing guard in accuracy.py:117-119. Each slice still fans out to all shards, and it puts cancellation points back, so a drain timeout stays bounded.
nv-alicheng
left a comment
nv-alicheng
left a comment
There was a problem hiding this comment.
Review Council
5 findings for PR #556: Reduce the runtime of benchmarking draining phase
Severity counts: critical 0, high 0, medium 4, low 1.
| id | file | line | severity | dims | models | summary |
|---|---|---|---|---|---|---|
| 1 | src/inference_endpoint/evaluation/scoring.py | 96 | medium | code-quality, coverage, design, functionality | claude,codex,private-lm | Byte-marker classification can decode a COMPLETE as an issued key and silently drop its model output |
| 2 | src/inference_endpoint/commands/benchmark/full_run_osl.py | 293 | medium | coverage | claude,codex,private-lm | The new finalize-side CPU probe and tokenizer wiring are untested |
| 3 | src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py | 673 | medium | functionality | codex | Drain renders the whole structured backlog in one unstoppable thread task |
| 4 | src/inference_endpoint/commands/benchmark/full_run_osl.py | 130 | medium | code-quality, design | private-lm | Batch outcomes are folded positionally without a length check |
| 5 | src/inference_endpoint/evaluation/scoring.py | 89 | low | code-quality, coverage, design, docs-readability | claude,codex,private-lm | Docstring says the event type is guessed from 1024 bytes; the classifier inspects 256 |
Coverage gaps
- none
Not anchored to the diff
- none
Already on the PR
src/inference_endpoint/async_utils/services/metrics_aggregator/token_metrics.py:673Drain renders the whole structured backlog in one unstoppable thread task
| Raises: | ||
| msgspec.DecodeError: the line is not a valid event. | ||
| """ | ||
| head = line[:_EVENT_TYPE_PREFIX_BYTES] |
There was a problem hiding this comment.
[claude,codex,private-lm] medium (code-quality, coverage, design, functionality): Byte-marker classification can decode a COMPLETE as an issued key and silently drop its model output
guess_and_decode classifies a line by searching for the exact ISSUED marker bytes anywhere in the first 256 bytes, and it tests that branch before the COMPLETE branch — unlike the previous full decode, payload bytes can therefore flip the classification, so a COMPLETE record whose payload contains the unescaped bytes "event_type":"sample.issued" (e.g. a dict-valued tool-call argument) is decoded as _IssuedKey, discarding its data and scoring the turn missing/0 with no warning. The new prompt-marker test cannot catch this: the marker there sits inside a JSON string (msgspec escapes the quotes) and _IssuedKey ignores unknown fields, so its equality assertion would pass even with a broken guess. Classification is exact-byte coupled to the current wire format, so any drift silently returns None — dropping issued turns from the accuracy denominator, with the same silent failure in Scorer._iter_complete — and the partial cut is chosen by presence checks rather than a JSON boundary, so an unlucky ,"data": match raises DecodeError that the caller skips as a malformed line. Anchor classification on the line's leading event_type (as the module comment claims) and make the partial decode a pure optimization with a total full-line fallback; derive the marker/prefix constants once from the event enum instead of duplicating them in two modules.
Root cause: Case A — the misclassification sits on the same byte-guessing helper as the design/quality findings; classify from the line's leading event_type, centralize the marker and prefix constants next to EventRecord, and fall back to a full-line decode whenever the fast path cannot be applied.
There was a problem hiding this comment.
One concrete format-compatibility case here: {"event_type": "sample.complete", ...} is valid JSON but does not match _COMPLETE_MARKER because of the space after the colon. _iter_complete() silently skips it, and guess_and_decode() similarly returns None for spaced ISSUED/COMPLETE markers. Moving event_type beyond the first 256 bytes has the same effect. The previous full JSON decode accepted these layouts.
The current JSONLWriter emits the expected compact format, but the reader now depends on that formatting without a decoding fallback or warning. The existing ISSUED fallback only runs after the exact event marker has matched, so it does not cover these cases.
Could we keep the byte check as a fast path and fall back to JSON decoding when the prefix cannot be classified, then filter by the decoded event type? Please add whitespace/reordered-field cases and a contract test that feeds actual JSONLWriter output into the scorer, so future writer changes cannot silently drop events or alter scores.
| # live_workers caps (an exported RAYON_NUM_THREADS still wins): one per | ||
| # CPU this process may run on, so a taskset-pinned client stays in its set. | ||
| try: | ||
| threads = len(get_all_online_cpus()) or 1 |
There was a problem hiding this comment.
[claude,codex,private-lm] medium (coverage): The new finalize-side CPU probe and tokenizer wiring are untested
Every full_run_osl_for_report test returns None before tokenizer construction (missing performance map, tokenizer name, or events file), so none of the new code at lines 292–303 runs under test: neither get_all_online_cpus(), its UnsupportedPlatformError fallback, the or 1 guards for an empty CPU list or cpu_count() is None, nor live_workers=threads and count_batch_sync being handed to compute_full_run_osl. A regression here — a renamed exception, a probe returning [], or a broken callback — is swallowed by the outer except Exception, silently dropping the whole full-run OSL block. Add a wrapper test with a valid events file and performance map, patching get_all_online_cpus (raise/empty/normal) and BatchTokenizer to capture live_workers and the injected count_batch callback.
| def count_pending() -> None: | ||
| nonlocal n_errors, first_error | ||
| try: | ||
| outcomes = count_batch(pending) |
There was a problem hiding this comment.
[private-lm] medium (code-quality, design): Batch outcomes are folded positionally without a length check
count_pending trusts the injected count_batch callback to return exactly one outcome per pending turn but never checks the length before folding the list into lengths/n_errors and clearing the buffer. A counter that returns fewer outcomes silently drops the unaccounted turns while n_errors stays 0, and partial — derived only from n_errors/n_undecodable/n_missing — stays False, so the block can be published as complete with a mean over an unknown subset, exactly what the partial gate exists to prevent. The sibling consumer of the same contract, TokenBatchQueue._flush (token_metrics.py:869-875), raises RuntimeError on a length mismatch. Capture len(pending) before clearing and charge the batch (or mark it partial) when the result length disagrees.
Root cause: Case A — the unvalidated fold is both the design and code-quality finding; validate len(outcomes) against the drained count and charge the difference to n_errors, mirroring TokenBatchQueue._flush.
| def guess_and_decode(line: bytes) -> EventRecord | _IssuedKey | None: | ||
| """Decode the ISSUED and COMPLETE events of one ``events.jsonl`` line. | ||
|
|
||
| The event type is guessed from the leading 1024 bytes, and other events return |
There was a problem hiding this comment.
[claude,codex,private-lm] low (code-quality, coverage, design, docs-readability): Docstring says the event type is guessed from 1024 bytes; the classifier inspects 256
_EVENT_TYPE_PREFIX_BYTES is 256 (line 67), and head = line[:_EVENT_TYPE_PREFIX_BYTES] (line 96) is what the ISSUED/COMPLETE checks search; the docstring's 1024 is a different quantity — the bound for the ,"data": search, 4 * _EVENT_TYPE_PREFIX_BYTES (line 99). A reader trusting the docstring will mis-reason about which lines the fast path recognizes: a marker past byte 256 falls through to None with no error, which the docstring implies cannot happen. Say 256 or reference _EVENT_TYPE_PREFIX_BYTES, and describe the data-search window separately if it is worth mentioning.
Review Council5 findings for PR #556: Reduce the runtime of benchmarking draining phase Severity counts: critical 0, high 0, medium 4, low 1.
Coverage gaps
Not anchored to the diff
Already on the PR
|
|
Please note related PR #561 |
github-actions
Bot
added
size/large
arekay-nv
left a comment
arekay-nv
left a comment
There was a problem hiding this comment.
Minor issues, but can be addressed in followup.
| result = [ | ||
| (index, c) for (index, _), c in zip(chunk, counts, strict=True) | ||
| ] | ||
| except Exception as exc: # noqa: BLE001 - isolate this chunk. |
There was a problem hiding this comment.
[Claude] medium (error-handling): A single crashed shard now fails the whole agentic drain
Chunks are bound to a shard at submit time (ex = procs[k % len(procs)], L654), and a failure is recorded and dropped here. If one shard dies (OOM on a large agentic PromptInput, a segfault in a trust_remote_code tokenizer), its executor goes broken: the in-flight chunk fails, and every later chunk dealt to it fails at submit with BrokenProcessPool. Nothing recounts those items, so _flush(live=False) raises, the drain ends with n_pending_tasks > 0, and run_benchmark fails the run. On main, agentic drains counted structured items in-process, so a shard crash could not lose them.
Root cause: chunks are assigned statically. Pulling chunks from a shared cursor as shards finish (~2 in flight per shard), and re-queuing a failed chunk on a surviving shard (or _count_structured_inproc as a last resort), fixes this and also removes stragglers. With round-robin, the shard dealt the long late-turn prompts finishes last while the others idle.
| 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 |
There was a problem hiding this comment.
[Codex] medium (data-integrity): Batched text counts include padding when the tokenizer presets it
encode_lengths returns len(item.ids). With BatchLongest padding configured in tokenizer.json, a batch call pads every item to the longest one. Codex reproduced [8, 8] for "a" and "abcdefgh" with the repo's char tokenizer, vs [1, 8] from the previous one-text-at-a-time call. Full-run OSL then depends on which turns share a batch and on batch_size. The aggregator's text lanes already batch through encode_lengths on main, so they share the issue, but this PR extends it to finalize. Fix: call no_padding() (and no_truncation()) on the backend once when it is loaded (_load_tokenizer / _init_worker), or count sum(item.attention_mask).
| 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 |
There was a problem hiding this comment.
[Claude] low (error-handling): A text-encode failure voids the whole batch and orphans the structured futures
If _encode_lengths_inproc(texts) raises (or the strict=True zip does), the exception leaves count_batch_sync after the structured futures were submitted on L816. They keep running unobserved, and compute_full_run_osl.count_pending charges every turn in the batch, including structured turns that counted fine. Everywhere else, this method isolates failures per slot (count_one, case _:). Fix: wrap the text encode, assign the exception to text_slots only, then collect the futures.
| if tok is None: | ||
| raise RuntimeError("tokenizer worker tokenizer unavailable") | ||
| out: list[int | Exception] = [] | ||
| for item in items: |
There was a problem hiding this comment.
[Claude, Lead] medium (performance): The structured drain uses about 1 of each shard's 8 cores
Each shard is ProcessPoolExecutor(max_workers=1) pinned to a CORES_PER_WORKER = 8 block, and this loop counts its chunk sequentially on one Python thread. Jinja rendering holds the GIL, and encoding one rendered sequence uses a single rayon thread, because rayon parallelizes across batch items. So the structured drain runs on about ⌈(CPUs − 2) / 8⌉ cores (~17–18 on a 144-CPU host) while the rest of each block idles. The finalize path in this PR shows the same workload keeps scaling with threads (_MAX_COUNT_THREADS).
Root cause: one shard size, tuned for batched text encoding, is reused for a one-item-at-a-time workload. A separate worker count for the structured drain fixes this: one process per core, or a thread pool of len(core_set) inside the shard, like count_batch_sync. See the co-location note on L730 for why that count should be configurable.
| ``count_message_tokens`` / ``count_prompt_tokens`` logic, so a sample | ||
| gets the same number whichever lane handles it. | ||
| """ | ||
| if self._procs and not live: |
There was a problem hiding this comment.
[Lead] medium (design): The drain can't be bounded when the client shares a node with the server
This sends the structured drain to the shard pool. That pool has no knob (DESIGN.md: "The shard pool has no knob"). It sizes itself from cgroup_clamped_cpus(), which deliberately widens past the process's inherited affinity mask, so taskset/numactl on the client doesn't confine it; only a cgroup cpuset does. Each shard keeps one busy thread in every 8-core block across all sockets, including blocks the server's CPU-side processes may use. Before this PR the shards ran only briefly (text items). Now they run for the whole structured drain (minutes at c6144/c18k). That is harmless when the server is idle, but not in multi-client runs, sweeps whose runs overlap, or while the server is still finishing requests that were cut off at the max duration. Suggest a cap such as settings.metrics_drain_workers, or honoring the inherited mask. The same knob would let dedicated client nodes use more than one core per block (see L380).
|
|
||
| # 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] |
There was a problem hiding this comment.
[Claude] low (concurrency): First path that shares one PreTrainedTokenizerFast across many threads
Up to 32 pool threads run apply_chat_template(tokenize=True) while the calling thread runs backend.encode_batch_fast. On main, finalize counted on one thread, and the aggregator awaits structured items one at a time. The __call__ path runs set_truncation_and_padding. When tokenizer.json ships a truncation or padding config, that takes a mutable borrow (no_truncation / no_padding) on the Rust tokenizer. If another thread is inside encode_batch with the GIL released at that moment, tokenizers raises RuntimeError: Already borrowed. The affected turns are counted as n_errors, and the OSL block is marked partial. The window is only the first concurrent calls, and a 20-trial local probe with truncation enabled did not reproduce it. Normalizing the backend state once before fanning out would close it; that is the same no_padding() / no_truncation() call suggested on L819.
| 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 |
There was a problem hiding this comment.
[Cursor] low (performance): Without a fast backend, the structured drain stays on this in-process path
_setup_shards returns without creating shards when self._text_backend is None, so _procs is empty and the drain takes this branch: one GIL-bound item at a time, the pre-PR speed. Structured counting needs only apply_chat_template, so tokenizers without a Rust backend (custom trust_remote_code slow tokenizers) are exactly the ones that miss the speedup. Gate shard creation on loading the tokenizer, not on the text backend. The comment here is also inaccurate: this branch also serves the drain whenever there are no shards.
| _WORKER_TEXT_BACKEND = getattr(_WORKER_TOKENIZER, "backend_tokenizer", None) | ||
| try: | ||
| _WORKER_PREFIX_LEN, _WORKER_BASELINE = chat_template_baseline(_WORKER_TOKENIZER) | ||
| except Exception: |
There was a problem hiding this comment.
[Cursor, Quality] low (data-integrity): The shard's baseline fallback is silent and differs from the parent's
The comment says "Same fallback as the parent", but the parent logs with logger.exception and this swallows the error. Live, drain and finalize share count_message_tokens so that a turn gets the same count whichever lane counts it. If a shard's baseline fails where the parent's succeeded (e.g. a transient trust_remote_code load problem in the child), that shard keeps the assistant frame in its message counts and nothing reports it. Log it, or fail shard warmup (e.g. have _worker_ready return the baseline so the parent can compare).
| return encode_lengths(backend, texts) | ||
|
|
||
|
|
||
| def _worker_warn_fallback(exc: Exception, impact: str) -> None: |
There was a problem hiding this comment.
[Claude, Quality] low (error-handling): Shard-side diagnostics are much weaker than the parent's
A spawn child has no logging config, so this warning reaches stderr through logging.lastResort with no logger name, timestamp or traceback, once per shard. _worker_count_structured also flattens each failure to RuntimeError(f"{type(exc).__name__}: {exc}"), so the parent's logger.exception shows the parent's traceback, not the shard's. A systematic template failure during a 20-minute drain will be hard to diagnose. Consider returning the formatted traceback string with the error. Also consider sharing one warn-once-per-error-type helper with BatchTokenizer._warn_template_fallback, which this function copies.
| 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 |
There was a problem hiding this comment.
[Claude] low (testing): The existing _init_worker test now leaks these globals
test_init_worker_overrides_inherited_cap_with_block_size (test_token_metrics.py:748) calls _init_worker in the pytest process. _init_worker now also writes _WORKER_TOKENIZER, _WORKER_PREFIX_LEN and _WORKER_BASELINE, which that test doesn't monkeypatch. A fake tokenizer and a (0, 0) baseline therefore leak into later tests, so results depend on test order. Monkeypatch the _WORKER_* globals there, or keep a single _WORKER_STATE object so there is one name to patch.
What does this PR do?
This code addresses 3 main slowness coming from the drain phase:
Benchmarked with one of the submission model on high concurrency cases, c6144 and c18k, the total finish time went from 60minutes, 90 minutes to 13 minutes and 20 minutes.
Type of change
Related issues
Testing
Checklist