From aadb9720fe6c56395bfcebb88a4dbe8148531f78 Mon Sep 17 00:00:00 2001 From: Jeremy Zhang Date: Wed, 12 Aug 2026 02:15:16 +0800 Subject: [PATCH] fix: route scheduler aborts to multi-tokenizer workers (#33940) --- .../scheduler_components/output_sender.py | 11 ++-- .../test_output_sender.py | 63 +++++++++++++++++++ 2 files changed, 69 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/managers/scheduler_components/test_output_sender.py diff --git a/python/sglang/srt/managers/scheduler_components/output_sender.py b/python/sglang/srt/managers/scheduler_components/output_sender.py index d4be68081..7ec53f696 100644 --- a/python/sglang/srt/managers/scheduler_components/output_sender.py +++ b/python/sglang/srt/managers/scheduler_components/output_sender.py @@ -12,17 +12,18 @@ class SenderWrapper: def send_output( self, output: Union[BaseReq, BaseBatchReq], - recv_obj: Optional[Union[BaseReq, BaseBatchReq]] = None, + recv_obj: Optional[object] = None, ): if self.socket is None: return + http_worker_ipc = getattr(recv_obj, "http_worker_ipc", None) if ( - isinstance(recv_obj, BaseReq) - and recv_obj.http_worker_ipc is not None + isinstance(output, BaseReq) + and http_worker_ipc is not None and output.http_worker_ipc is None ): - # handle communicator reqs for multi-http worker case - output.http_worker_ipc = recv_obj.http_worker_ipc + # Scheduler Req is not a BaseReq but carries the same return route. + output.http_worker_ipc = http_worker_ipc sock_send(self.socket, output) diff --git a/test/registered/unit/managers/scheduler_components/test_output_sender.py b/test/registered/unit/managers/scheduler_components/test_output_sender.py new file mode 100644 index 000000000..2a1cc3725 --- /dev/null +++ b/test/registered/unit/managers/scheduler_components/test_output_sender.py @@ -0,0 +1,63 @@ +import asyncio +import unittest +from unittest.mock import MagicMock, patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.managers.io_struct import AbortReq +from sglang.srt.managers.multi_tokenizer_mixin import MultiTokenizerRouter +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.scheduler_components.output_sender import SenderWrapper +from sglang.srt.sampling.sampling_params import SamplingParams + +register_cpu_ci(est_time=8, suite="base-a-test-cpu") + + +def _make_scheduler_req(http_worker_ipc: str) -> Req: + return Req( + rid="scheduler-request", + origin_input_text="prompt", + origin_input_ids=[1], + sampling_params=SamplingParams(), + http_worker_ipc=http_worker_ipc, + ) + + +class TestSenderWrapper(CustomTestCase): + @patch("sglang.srt.managers.scheduler_components.output_sender.sock_send") + def test_preserves_existing_output_route(self, mock_sock_send): + socket = MagicMock() + scheduler_req = _make_scheduler_req("ipc:///source-worker") + output = AbortReq( + rid="already-routed", + http_worker_ipc="ipc:///destination-worker", + ) + + SenderWrapper(socket).send_output(output, scheduler_req) + + self.assertEqual(output.http_worker_ipc, "ipc:///destination-worker") + mock_sock_send.assert_called_once_with(socket, output) + + @patch("sglang.srt.managers.scheduler_components.output_sender.sock_send") + def test_scheduler_abort_routes_to_origin_worker(self, mock_sock_send): + socket = MagicMock() + worker_ipc = "ipc:///tokenizer-worker-2" + output = AbortReq(rid="overflow-request") + + SenderWrapper(socket).send_output(output, _make_scheduler_req(worker_ipc)) + + routed_output = mock_sock_send.call_args.args[1] + router = MultiTokenizerRouter.__new__(MultiTokenizerRouter) + router.socket_mapping = MagicMock() + asyncio.run(router._distribute_result_to_workers(routed_output)) + + router.socket_mapping.send_output.assert_called_once_with( + worker_ipc, routed_output + ) + + +if __name__ == "__main__": + unittest.main()