diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index 5b5bc4ed6..30560c0d0 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -165,7 +165,7 @@ class DataParallelController: self.load_snapshot_reader = create_load_snapshot_reader( server_args, port_args, - caller="dp_controller", + caller="DataParallelController", ) self._last_refresh_time = 0.0 diff --git a/python/sglang/srt/managers/load_snapshot.py b/python/sglang/srt/managers/load_snapshot.py index 917e1d744..ade1ed737 100644 --- a/python/sglang/srt/managers/load_snapshot.py +++ b/python/sglang/srt/managers/load_snapshot.py @@ -28,7 +28,8 @@ transport backends are supported: Shared memory does not work across nodes, so multi-node DP attention requires the ZMQ transport. The ``ZmqShmLoadSnapshotReader`` on node 0 receives snapshots from all schedulers via zmq PUSH/PULL and writes them -into the local SHM file. All readers (tokenizer, dp_controller) on +into the local SHM file. All readers (TokenizerManager, +DataParallelController) on node 0 then read from SHM. ``zmq_reader_owner()`` decides which process on node 0 binds the zmq @@ -89,30 +90,49 @@ def should_use_zmq(server_args) -> bool: _LOAD_AWARE_METHODS = frozenset({"total_requests", "total_tokens"}) +def _tokenizer_load_snapshot_owner_caller(server_args) -> str: + """The caller that plays the tokenizer-side zmq owner role. + + In multi-tokenizer mode (``tokenizer_worker_num > 1``) there are N + independent ``TokenizerWorker`` processes that would all try to bind the + same zmq PULL endpoint. Instead, the single ``MultiTokenizerRouter`` + process owns the socket (polls zmq -> SHM) and every worker reads SHM. + """ + if server_args.tokenizer_worker_num > 1: + return "MultiTokenizerRouter" + return "TokenizerManager" + + def zmq_reader_owner(server_args, caller: str) -> bool: """Decide which process owns the zmq PULL socket. - Exactly one of ``"dp_controller"`` or ``"tokenizer"`` must return True - when zmq mode is active. The owner polls zmq -> SHM; the other reads SHM. + Exactly one of ``"DataParallelController"``, ``"TokenizerManager"``, or + ``"MultiTokenizerRouter"`` must return True when zmq mode is active. The + owner polls zmq -> SHM; the others read SHM. Rules: - - Non-zero node_rank: no tokenizer, dp_controller only launches - schedulers and waits -> nobody owns it. - - dp_size == 1: no dp_controller exists -> tokenizer owns it. - - dp_size > 1, load-aware method: dp_controller polls on every - dispatch via refresh_load_budget() -> dp_controller owns it. - - dp_size > 1, round-robin / other: dp_controller never reads - load data -> tokenizer owns it (polls on /v1/loads calls). + - Non-zero node_rank: no TokenizerManager, DataParallelController only + launches schedulers and waits -> nobody owns it. + - dp_size == 1: no DataParallelController exists -> tokenizer-side owner + owns it. + - dp_size > 1, load-aware method: DataParallelController polls on every + dispatch via refresh_load_budget() -> DataParallelController owns it. + - dp_size > 1, round-robin / other: DataParallelController never reads + load data -> tokenizer-side owner owns it (polls on /v1/loads calls). + + The tokenizer-side owner is the ``"MultiTokenizerRouter"`` caller in + multi-tokenizer mode, otherwise the ``"TokenizerManager"`` caller. """ if not should_use_zmq(server_args): return False if server_args.node_rank != 0: return False + tokenizer_owner = _tokenizer_load_snapshot_owner_caller(server_args) if server_args.dp_size == 1: - return caller == "tokenizer" + return caller == tokenizer_owner if server_args.load_balance_method.lower() in _LOAD_AWARE_METHODS: - return caller == "dp_controller" - return caller == "tokenizer" + return caller == "DataParallelController" + return caller == tokenizer_owner # --------------------------------------------------------------------------- @@ -619,6 +639,22 @@ class ZmqShmLoadSnapshotReader: "load snapshot shm write failed for rank %d: %s", dp_rank, e ) + def fileno(self) -> int: + """Edge-triggered fd that becomes readable when zmq messages arrive. + + Lets an owner process register the reader with an event loop and drain + it via ``poll()`` instead of polling on a timer. + """ + return self._socket.getsockopt(self._zmq.FD) + + def poll(self) -> None: + """Drain the zmq PULL socket into SHM. + + Public entry point so an owner process (e.g. MultiTokenizerRouter) can + keep SHM fresh without touching internals. + """ + self._poll() + def read(self, dp_rank: int) -> Optional[LoadSnapshot]: self._poll() return self._shm_reader.read(dp_rank) @@ -683,8 +719,9 @@ def create_load_snapshot_reader(server_args, port_args, caller: str): """Create a load snapshot reader. Args: - caller: ``"dp_controller"`` or ``"tokenizer"`` -- determines who - binds the zmq PULL socket when zmq mode is active. + caller: ``"DataParallelController"``, ``"TokenizerManager"``, or + ``"MultiTokenizerRouter"`` -- determines who binds the zmq PULL + socket when zmq mode is active. """ dp_size = server_args.dp_size if zmq_reader_owner(server_args, caller): diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 5f6cc6309..d675c1245 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -50,6 +50,10 @@ from sglang.srt.managers.io_struct import ( PauseGenerationReqInput, TokenizerWorkerRegistration, ) +from sglang.srt.managers.load_snapshot import ( + create_load_snapshot_reader, + zmq_reader_owner, +) from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.server_args import PortArgs, ServerArgs from sglang.srt.utils import ( @@ -381,6 +385,19 @@ class MultiTokenizerRouter: self._handle_task = asyncio.run_coroutine_threadsafe( print_exception_wrapper(self.handle_loop), self._loop ) + + # In multi-tokenizer mode the N TokenizerWorker processes cannot each + # bind the zmq PULL socket used for load snapshots, so the single + # MultiTokenizerRouter process owns it (zmq -> SHM) and the workers + # read SHM only. Drain it event-driven via the socket's fd instead of + # polling on a timer. + self.load_snapshot_reader = None + if zmq_reader_owner(server_args, "MultiTokenizerRouter"): + self.load_snapshot_reader = create_load_snapshot_reader( + server_args, port_args, caller="MultiTokenizerRouter" + ) + self._loop.call_soon_threadsafe(self._register_load_snapshot_reader) + self.disaggregation_bootstrap_server = start_disagg_service(self.server_args) # Worker IPC names for pause/continue broadcasting @@ -391,6 +408,20 @@ class MultiTokenizerRouter: def _run_loop(self): self._loop.run_forever() + def _register_load_snapshot_reader(self): + """Drain zmq load snapshots into SHM whenever the PULL socket is readable. + + zmq exposes an edge-triggered fd; ``poll()`` drains it until empty, which + also re-arms the fd, so TokenizerWorkers reading SHM stay up to date + without any timer. + """ + assert self.load_snapshot_reader is not None + self._loop.add_reader( + self.load_snapshot_reader.fileno(), self.load_snapshot_reader.poll + ) + # Drain anything already queued before the fd was registered. + self.load_snapshot_reader.poll() + async def router_worker_obj(self): """Forward path: workers → scheduler, with pause/continue broadcast.""" while True: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 580f0725f..cb53d1844 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -385,7 +385,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self.load_snapshot_reader = create_load_snapshot_reader( self.server_args, port_args, - caller="tokenizer", + caller="TokenizerManager", ) def init_running_status(self): diff --git a/test/registered/unit/managers/test_load_snapshot_backends.py b/test/registered/unit/managers/test_load_snapshot_backends.py index 7b6a49c86..71babf607 100644 --- a/test/registered/unit/managers/test_load_snapshot_backends.py +++ b/test/registered/unit/managers/test_load_snapshot_backends.py @@ -16,6 +16,7 @@ from sglang.srt.managers.load_snapshot import ( create_load_snapshot_reader, create_load_snapshot_writer, should_use_zmq, + zmq_reader_owner, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel @@ -217,13 +218,16 @@ class TestFactoryFunctions(CustomTestCase): dp_size=1, load_balance_method="round_robin", node_rank=0, + tokenizer_worker_num=1, ) port_args = SimpleNamespace(instance_id="test_shm_factory") writer = create_load_snapshot_writer( server_args, port_args, dp_size=1, dp_rank=0 ) self.assertIsInstance(writer, ShmLoadSnapshotWriter) - reader = create_load_snapshot_reader(server_args, port_args, caller="tokenizer") + reader = create_load_snapshot_reader( + server_args, port_args, caller="TokenizerManager" + ) self.assertIsInstance(reader, ShmLoadSnapshotReader) reader.close() writer.close() @@ -240,6 +244,7 @@ class TestFactoryFunctions(CustomTestCase): dp_size=1, load_balance_method="round_robin", node_rank=0, + tokenizer_worker_num=1, ) port_args = SimpleNamespace(instance_id="test_zmq_factory") os.environ["SGLANG_LOAD_SNAPSHOT_USE_ZMQ"] = "1" @@ -249,7 +254,7 @@ class TestFactoryFunctions(CustomTestCase): ) self.assertIsInstance(writer, ZmqLoadSnapshotWriter) reader = create_load_snapshot_reader( - server_args, port_args, caller="tokenizer" + server_args, port_args, caller="TokenizerManager" ) self.assertIsInstance(reader, ZmqShmLoadSnapshotReader) reader.close() @@ -270,6 +275,71 @@ class TestFactoryFunctions(CustomTestCase): self.assertFalse(should_use_zmq(args)) +class TestZmqReaderOwner(CustomTestCase): + """At most one process binds the zmq PULL socket across all callers.""" + + CALLERS = ("TokenizerManager", "MultiTokenizerRouter", "DataParallelController") + + @staticmethod + def _args(**overrides): + base = dict( + enable_dp_attention=True, + nnodes=2, + node_rank=0, + dp_size=1, + load_balance_method="round_robin", + tokenizer_worker_num=1, + ) + base.update(overrides) + return SimpleNamespace(**base) + + def _owners(self, args): + return {c for c in self.CALLERS if zmq_reader_owner(args, c)} + + def test_zmq_disabled_no_owner(self): + args = self._args(enable_dp_attention=False, nnodes=1) + self.assertEqual(self._owners(args), set()) + + def test_non_zero_node_rank_no_owner(self): + args = self._args(node_rank=1, dp_size=4, tokenizer_worker_num=8) + self.assertEqual(self._owners(args), set()) + + def test_tokenizer_manager_owns_when_dp1(self): + self.assertEqual(self._owners(self._args(dp_size=1)), {"TokenizerManager"}) + + def test_multi_tokenizer_router_owns_in_multi_tokenizer_dp1(self): + args = self._args(dp_size=1, tokenizer_worker_num=8) + self.assertEqual(self._owners(args), {"MultiTokenizerRouter"}) + + def test_multi_tokenizer_router_owns_in_multi_tokenizer_round_robin(self): + args = self._args(dp_size=4, tokenizer_worker_num=8) + self.assertEqual(self._owners(args), {"MultiTokenizerRouter"}) + + def test_data_parallel_controller_owns_load_aware(self): + for method in ("total_tokens", "total_requests"): + args = self._args( + dp_size=4, tokenizer_worker_num=8, load_balance_method=method + ) + self.assertEqual(self._owners(args), {"DataParallelController"}) + + def test_tokenizer_manager_owns_dp4_round_robin(self): + args = self._args(dp_size=4, tokenizer_worker_num=1) + self.assertEqual(self._owners(args), {"TokenizerManager"}) + + def test_at_most_one_owner_across_configs(self): + for dp_size in (1, 4): + for tw in (1, 8): + for method in ("round_robin", "total_tokens", "total_requests"): + for node_rank in (0, 1): + args = self._args( + dp_size=dp_size, + tokenizer_worker_num=tw, + load_balance_method=method, + node_rank=node_rank, + ) + self.assertLessEqual(len(self._owners(args)), 1, args) + + class TestZmqAddr(CustomTestCase): def test_ipc_for_single_node(self): port_args = SimpleNamespace(instance_id="myinstance")