[fix] /pause_generation and /continue_generation wrong for --tokenizer-worker-num > 1

[maocheng23]

Summary

Resubmit of #24445 (which was reverted in #24461 due to a _check_all_req_types import-time failure). This version ships the rename inline so io_struct loads cleanly.

Fix /pause_generation and /continue_generation not propagating to all tokenizer workers when --tokenizer-worker-num > 1. Previously each HTTP request only flipped is_pause on the single worker that handled the request, leaving other workers in an inconsistent state and causing /generate requests routed to the stuck worker(s) to hang forever.

Fixes #21235.

Difference from #24445

• TokenizerWorkerRegisterReq → TokenizerWorkerRegistration.
  This is a worker→router control message, not a scheduler IO type. io_struct._check_all_req_types enforces (at import time) that any class ending in Req/Input/Output must subclass BaseReq/BaseBatchReq. Renaming away from that suffix is more semantically appropriate than retrofitting BaseReq inheritance. This was the cause of the post-merge CI failure that prompted the revert.

Root cause

With --tokenizer-worker-num=N, there are N independent TokenizerWorker processes (subclasses of TokenizerManager), each holding its own is_pause flag (tokenizer_manager.py:416). The HTTP load balancer dispatches each request to one worker, so:

• /pause_generation only sets is_pause = True on 1 of N workers (tokenizer_manager.py:1486-1490).
• /continue_generation only sets is_pause = False on 1 of N workers — likely a different one (tokenizer_manager.py:1501-1505).

After a pause+continue cycle, with probability (N-1)/N one worker remains stuck is_pause = True. The per-worker flag gates incoming requests at tokenizer_manager.py:542-543 (await self.is_pause_cond.wait_for(lambda: not self.is_pause)), so any subsequent request the load balancer routes to the stuck worker hangs forever. With PD disaggregation and 8 workers, 15/64 of requests need to traverse the stuck worker on either prefill or decode and hang waiting for KV transfer.

The shared scheduler's _engine_paused flag was toggled correctly (it's a single process), so generation resumed globally — but per-worker request gating was broken.

Fix

Route pause/continue through MultiTokenizerRouter, which fans out to every registered tokenizer manager IPC.

• TokenizerWorkerRegistration — each TokenizerWorker registers its IPC name with the router on startup.
• Router intercepts PauseGenerationReqInput / ContinueGenerationReqInput in router_worker_obj, emits PauseContinueBroadcast to every registered worker IPC, then forwards to the scheduler (skipping for mode="abort", which drains via the originator's existing abort polling). We only need to send to one scheduler and it will broadcast to others — the router forwards once to scheduler rank 0, which fans out to all TP/PP/DP ranks internally via the existing scheduler-side broadcast path.
• Workers apply the broadcast under the existing is_pause_cond and call notify_all() on continue.
• Originator worker awaits a Future resolved when its own broadcast lands, so the HTTP response only returns after the local flag is flipped (other workers are flipped in parallel).

Total added latency: one IPC round-trip (worker → router → worker), sub-ms on a single host.

Test plan

• Repro from [Bug] /pause_generation and /continue_generation wrong for --tokenizer-worker-num > 1 #21235 (16 concurrent /generate while paused, assert none complete; resume; assert all complete).
• test/registered/tokenizer/test_multi_tokenizer.py passes with --tokenizer-worker-num > 1.
• test/registered/unit/managers/test_scheduler_pause_generation.py passes.
• test/registered/language/test_srt_backend.py passes (this was the regression on [fix] /pause_generation and /continue_generation wrong for --tokenizer-worker-num > 1 #24445 caused by the naming check; verifies the rename fixes it).
• test/registered/rl/test_pause_generation_tensor_consistency.py passes.

Notes / follow-ups

• Concurrent pause/continue on the same worker: the originator stores a single _pause_continue_future; concurrent ops landing on the same worker process can overwrite it. Sequential usage (the bug repro) is fine; an op_id-keyed future map would harden this.
• Abort mode: only the originator polls its own model_update_lock. Scheduler-side abort_request clears global scheduler state, but per-worker tokenizer-side in-flight state on non-originator workers isn't drained by the originator's poll. Consider broadcasting the abort poll if that becomes load-bearing.
• Worker registration timing: relies on all TokenizerWorkers sending TokenizerWorkerRegistration before HTTP traffic begins. Standard startup ordering handles this.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Qiaolin-Yu]

/tag-and-rerun-ci

[Qiaolin-Yu]

/tag-and-rerun-ci

[lawrence-harmonic]

Does this prevent /pause_generation from being reordered with /update_weights_from_distributed when they are sent to different workers? This was an issue I had to account for in #21237 by introducing _ensure_paused_or_model_locked

Right now the updates weights only checks is_paused once, if it is False at the beginning, then it will wait on model_update_lock forever.

The race:

1. Client sends /pause_generation to worker A and /update_weights_from_distributed to worker B roughly concurrently.
2. Worker A enters pause_generation → sends PauseGenerationReqInput to router → router broadcasts PauseContinueBroadcast(is_pause=True) to both A and B. The broadcast reaches B via router → IPC → recv_loop → _result_dispatcher → schedule task → _apply_pause_continue_broadcast →
   acquire is_pause_cond → set is_pause=True. Many hops.
3. Worker B enters update_weights_from_distributed (tokenizer_control_mixin.py:426-433):
   async with self.is_pause_cond:
   is_paused = self.is_pause # one-shot read
   if is_paused: ...
   if not is_paused:
   async with self.model_update_lock.writer_lock:
   ...
4. If B grabs is_pause_cond and reads is_pause before the broadcast task lands, it sees False and commits to the writer-lock path.

Consequence: B then waits on the writer lock until its existing readers (in-flight generates holding reader_lock at tokenizer_manager.py:546) drain. Not literally "forever," but in non-abort pause mode the scheduler lets in-flight batches run to completion, so the wait is bounded
only by the longest live generation — defeating the whole point of the "if paused, skip the lock and update immediately" fast path. In abort mode it's worse, because the broadcast intentionally doesn't go to the scheduler (router_worker_obj skips the forward for abort), and only
the originator (A) polls model_update_lock.is_locked() on its own process — A's poll doesn't observe readers held on B.