Skip to content

perf(agent): isolate prompt rendering from event loop - #2328

Closed
liuhao-labs wants to merge 5 commits into
THUDM:mainfrom
liuhao-labs:codex/2285-tokenizer-worker
Closed

liuhao-labs wants to merge 5 commits into
THUDM:mainfrom
liuhao-labs:codex/2285-tokenizer-worker

Conversation

@liuhao-labs

@liuhao-labs liuhao-labs commented Aug 26, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Long full-history prompt rendering blocks the agent adapter's aiohttp event loop. Move the existing _session_prompt_ids() call to one adapter-owned worker thread, retaining exact upstream rendering, saved sampled-token prefixes, and safe fallback for changed histories/tool schemas.

The worker bounds submitted running/queued work. Canceled requests retain their executor capacity until physical completion, discard queued jobs, return promptly during admission cancellation, and avoid starting SGLang or recording a trajectory after cancellation. App cleanup drains the worker. Session prefix lists/template options are snapshotted before offload; the live session is used for generation and records only flushed turns.

Bounded recent-window metrics at GET /debug/adapter_metrics include prompt length, total prompt-render latency, queue wait, worker utilization, event-loop lag, cancellations and discarded work, with P50/P95/P99 distributions.

Related to #2285.

Scope

This improves event-loop responsiveness and observability. It preserves upstream full-history rendering and canonical-boundary checks, including original sampled prompt/response token prefixes introduced by #2457. It does not remove the cumulative O(N²) tokenization cost or claim general incremental tokenization. Sampler score-centering/top-p/R3 metadata, output-token validation, and protocol behavior are unchanged.

Validation

  • PYTHONPATH=. python -m pytest -q tests/test_agent tests/test_grouped_rewards.py: 169 passed, 2 skipped against main 58b08786.
  • Direct CPU runners: worker tests 12 passed; prompt continuity 14 passed. Real in-memory BPE tests compare synchronous and worker IDs for tool continuation, changed tools and compaction, and assert both full-history and splice-boundary renders run off the event-loop thread.
  • Cancellation handoff, repeated cancellation with delayed cleanup, stable session snapshots, capacity ownership, health/metrics responsiveness and executor shutdown regressions passed.
  • The upstream sampler replay/token contracts and Responses/Anthropic/Codex harness tests passed.
  • Local synthetic continuity workload: 16,250 prompt IDs, 12 renders per mode, identical IDs. Synchronous heartbeat p95/max lag was 47.17/48.64 ms; worker p95/max lag was 0.63/1.09 ms. Total render time was 0.197/0.206 s respectively. These single-host timings show event-loop isolation, not a tokenizer throughput improvement or production guarantee.
  • Repository-pinned pre-commit, workflow regeneration, YAML and git diff --check passed. Sunabako 8-GPU/35-minute E2E registration, container init/cidfile cleanup, proxy/mount setup and artifacts remain unchanged. GPU E2E execution is not claimed locally.

Comment thread slime/agent/adapters/tokenizer_worker.py Outdated
Keep full-history chat-template rendering as the source of truth while moving it to a bounded adapter-owned worker. Add cancellation-safe cleanup and bounded queue, latency, utilization, prompt-length, and event-loop metrics.

Related to THUDM#2285

Signed-off-by: Hao Liu <liuaho6@gmail.com>
Track cleanup-owned admission cancellations explicitly so Python 3.10 keeps client cancellation distinct from worker shutdown. Shield permit acquisition and settle it before releasing a successful handoff.

Signed-off-by: Hao Liu <liuaho6@gmail.com>
Release an admission permit from the acquisition task completion callback when cancellation races with a successful handoff. Avoid waiting through repeated client cancellations and cover the occupied-permit case.

Signed-off-by: Hao Liu <liuaho6@gmail.com>
Make the repeated-cancellation regression distinguish callback cleanup from the prior waiting loop by holding the acquisition task in its cancellation handler.

Signed-off-by: Hao Liu <liuaho6@gmail.com>
Signed-off-by: Hao Liu <liuaho6@gmail.com>
@liuhao-labs
liuhao-labs force-pushed the codex/2285-tokenizer-worker branch from 256bb42 to 87c2fdb Compare October 10, 2026 15:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants