[karverma-amd]

What does this PR do?

Makes the metrics aggregator's end-of-run drain count structured inputs
(MessageInput / PromptInput) on the tokenizer shard processes instead of
one at a time on the in-process thread.

The problem

When a run ends, the aggregator drains every pending tokenization before it
writes the final snapshot and result_summary.json. Plain-text inputs already
fan out across the ProcessPoolExecutor shards, but structured inputs went
through _token_count_message / _token_count_prompt one at a time on the
in-process thread. Those call apply_chat_template, which is Jinja and Python
and holds the GIL, so the whole structured drain ran on a single core.

Agentic runs are almost entirely structured inputs, with long multi-turn
prompts. On a large agentic replay the drain ran for hours with one core
pegged and the rest idle. Meanwhile the session showed every sample complete,
the endpoint was idle, and no summary had been written, so the run looked hung
at the tail.

The fix

• Move the structured counting into module-level count_message_tokens /
  count_prompt_tokens (plus chat_template_baseline). The in-process lane
  and the shards share one implementation, so both produce the same counts.
• The shard initializer now loads the full tokenizer and the chat-template
  baseline alongside the fast text backend. It falls back to (0, 0) the same
  way the parent does.
• During the drain (live=False), BatchTokenizer._fan_out_structured deals
  structured items to the shards in chunks of 64, round-robin, so the long
  prompts spread across shards instead of landing on one.
• A failed chunk, such as a dead worker, marks only its own items as errors.
  The worker also returns per-item exceptions instead of raising, so one bad
  input does not fail its chunk.
• Drain progress is logged every 30 s, since this phase can otherwise run with
  no output.
• The live (mid-run) lane is unchanged and never uses the shards.

With this change the same drain finished in minutes and the run wrote its
summary normally.

Comparison with #556

We also tried #556 on the same workload in place of this change. #556 speeds
up the encode step: it renders every pending item to text on one thread and
then encodes the batch in parallel. On our workload the time goes into
rendering the chat template, which stays serial in #556, and holding all of
the rendered prompts at once made the aggregator's memory grow without bound.
The aggregator was killed by the OOM killer partway through the drain, so the
run never finished.

This PR keeps memory flat because each shard renders and counts one small
chunk at a time and returns only integers. It does not touch the
full_run_osl.py or scoring.py changes in #556, which are independent of
this one.

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Related: #556

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

New tests in test_token_metrics.py (TestStructuredDrainFanOut):

• shard counts match in-process counts for structured inputs spanning several
  chunks, with chunks dealt round-robin;
• a mixed batch (token IDs, text, message, prompt) keeps its input order;
• a failed chunk marks only its own items as errors;
• the live lane never uses the shards;
• the worker returns per-item errors instead of raising.

pytest -m unit: 2724 passed, 6 skipped. Manually verified on a large agentic
run: the drain completed and result_summary.json was written.

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

Pre-commit: ruff, ruff-format, mypy, license-header and the file hygiene hooks
pass. The prettier hook could not install its Node environment locally, and
this PR changes no files that hook formats.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[mrmhodak]

@zihaok @viraatc: Please take a look, this is related to #556

[zihaok]

Hi Karan, thanks for noticing, we also did realize the potential OOM issue. #556 aims to addressing 3 level of drain slowness. Will compare the solution between this and 556

[zihaok]

Hi, could you try the updated 556 which incorporate 561 with some syntax refactoring to merge with more related changes