Optimize structured chat-template tokenization drain at scale

Motivation

PR #441 adds structured prompt and assistant-message token counting through tokenizer.apply_chat_template(). Unlike plain-text counting, this path currently processes structured inputs on the small in-process tokenizer thread pool, and the mixed-batch loop awaits each structured item individually. Every multi-turn prompt also renders and tokenizes the complete conversation history again.

This is intentionally simple and correct, but long-running, high-concurrency agentic workloads can accumulate structured tokenizations faster than they are processed and add time after inference finishes.

In a one-hour Qwen 3.6 agentic stress run using the complete 613-trajectory dataset and target concurrency 128:

Samples issued	Samples completed	Failed	Pending tokenizations at drain	Tokenization drain time	Pending after drain
5,531	5,403	0	408	15.299 s	0

This was not yet a correctness problem: request issuing and tokenization run in separate processes, and the unlimited default drain completed successfully. It does show where additional shutdown latency can appear at larger scale.

What “metrics drain” means

After the benchmark phase stops issuing requests and in-flight inference responses complete, the metrics aggregator must finish tokenizing any buffered ISL/OSL/TPOT inputs before publishing its final snapshot. This final metrics work is the tokenization drain; it is separate from the response drain controlled by performance_timeout_s / accuracy_timeout_s.

settings.drain.metrics_drain_timeout_s is the wall-clock budget for this tokenization drain:

• 0 (default): wait indefinitely for buffered tokenizations.
• Positive value: stop waiting after that many seconds.
• Successful drain: final snapshot has state == "complete" and n_pending_tasks == 0.
• Timed-out or failed drain: final snapshot has n_pending_tasks > 0; token metrics are incomplete and Report.complete is false.

How to measure it

Run a representative long-duration workload with structured chat inputs, realistic conversation lengths, and production-like concurrency. Keep metrics_drain_timeout_s: 0 while measuring so the drain is allowed to finish.

The metrics-aggregator log reports both endpoints:

... Draining 408 pending tokenizations...
... Tokenizations fully drained (n_pending_tasks=0)

Measure the timestamp difference between those lines and record:

1. Samples issued and completed.
2. Target concurrency and benchmark duration.
3. Pending tokenizations at drain start.
4. Tokenization drain duration.
5. Final n_pending_tasks from metrics/final_snapshot.json.
6. Any apply_chat_template failed fallbacks, which are correctness signals and should not be confused with queue backlog.

The useful operational comparison is tokenization drain time relative to benchmark duration and the configured metrics_drain_timeout_s. A growing pending count or drain duration across increasing concurrency/duration indicates that tokenization is not keeping pace.

Possible optimizations if this becomes material

1. Parallelize structured rendering

Submit structured items concurrently instead of awaiting each item serially. The smallest version can use the existing bounded thread pool. If that is insufficient, extend the process-sharded drain path so each worker loads the full tokenizer wrapper and can run apply_chat_template() for a chunk of PromptInput / MessageInput objects.

Multiprocess support must preserve:

• Input/output ordering.
• Per-item exception isolation and fallback behavior.
• Custom and slow tokenizer support.
• Dataset-provided chat templates and chat_template_kwargs.
• Bounded worker count and memory use, since each process holds a tokenizer.

2. Investigate incremental multi-turn prompt counting

Today turn N tokenizes the complete history again:

turn 1: [user_1]
turn 2: [user_1, assistant_1, tool_1, user_2]
turn 3: [user_1, assistant_1, tool_1, user_2, assistant_2, tool_2, user_3]

An incremental path could cache the previous rendered/tokenized prompt and process only the newly appended turn where the model template permits it.

This cannot assume token counts are always additive. Chat templates may rewrite generation markers, reasoning preservation, tool framing, or message-boundary whitespace when a new turn is appended. A safe design should validate that the previous serialization/tokens are an unchanged prefix of the new prompt and fall back to complete apply_chat_template() rendering whenever that invariant does not hold. Cache keys must also include tools, selected chat template, chat_template_kwargs, and other rendering inputs.

Acceptance criteria

• Add a structured-prompt concurrency test demonstrating more than one active render.
• Add a long/multi-turn benchmark or synthetic probe that reports pending count and drain duration before and after the optimization.
• Preserve token counts for existing structured-message and structured-prompt tests.
• Preserve n_pending_tasks and drain-timeout semantics.
• Document memory/throughput tradeoffs if full tokenizer wrappers are loaded in worker processes.

Related review: #441 discussion https://github.com/mlcommons/endpoints/pull/441#discussion_r3753082947