From 1eee8fbdcc25b44e13bc097d5ff6ac24e8c24af4 Mon Sep 17 00:00:00 2001 From: IvanShan177 Date: Wed, 29 Jul 2026 00:20:03 +0800 Subject: [PATCH] [PD] Drain NIXL completion notifications before enforcing the WaitingForInput timeout (#32267) Co-authored-by: Claude Opus 4.8 --- python/sglang/srt/disaggregation/nixl/conn.py | 14 ++++++++++---- .../disaggregation/test_nixl_backend_basic.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 0e1548fe9..b095d7c3e 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -2744,10 +2744,11 @@ class NixlKVReceiver(CommonKVReceiver): if not self.started_transfer: return status - timeout_result = self._check_waiting_timeout() - if timeout_result is not None: - return timeout_result - + # Drain notifications before enforcing the waiting deadline. The decode + # agent has no NIXL progress thread (num_threads=0), so incoming + # completion notifications are only ingested here via + # update_transfer_status(); a completion queued by NIXL at/after the + # deadline would otherwise lose to the timeout purely by poll ordering. self.kv_mgr.update_transfer_status() if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard( @@ -2756,6 +2757,11 @@ class NixlKVReceiver(CommonKVReceiver): self.conclude_state = KVPoll.Success del self.kv_mgr.transfer_statuses[self.bootstrap_room] return self.conclude_state # type: ignore + + timeout_result = self._check_waiting_timeout() + if timeout_result is not None: + return timeout_result + return KVPoll.WaitingForInput # type: ignore def _register_kv_args(self) -> bool: diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 8f4220c47..f97ef2faa 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -576,6 +576,7 @@ class TestNixlReceiverPoll(CustomTestCase): mgr = MagicMock() mgr.waiting_timeout = 5 mgr.check_status.return_value = status + mgr.check_transfer_done.return_value = False mgr.transfer_statuses = {} mgr.addr_to_rooms_tracker = defaultdict(set) mgr.addr_to_rooms_tracker["prefill:8998"].add(11) @@ -622,6 +623,24 @@ class TestNixlReceiverPoll(CustomTestCase): self.assertIn("timed out", mgr.record_failure.call_args[0][1]) mgr.update_status.assert_called_once_with(11, KVPoll.Failed) + @patch("sglang.srt.disaggregation.nixl.conn.time.time") + def test_queued_completion_wins_over_waiting_timeout(self, mock_time): + # Past the deadline, but the completion is already queued/observed: + # draining before the timeout check must yield Success, not a false + # timeout, and must not send an abort. + mock_time.return_value = 20.0 + receiver, mgr = self._make_receiver(status=KVPoll.WaitingForInput) + receiver.started_transfer = True + receiver.init_time = 10.0 + mgr.transfer_statuses = {11: TransferStatus()} + mgr.check_transfer_done.return_value = True + + self.assertEqual(receiver.poll(), KVPoll.Success) + mgr.update_transfer_status.assert_called_once_with() + mgr.record_failure.assert_not_called() + mgr.update_status.assert_not_called() + self.assertNotIn(11, mgr.transfer_statuses) + @patch("sglang.srt.disaggregation.nixl.conn.time.time") def test_transfer_done_returns_success_and_cleans_room_state(self, mock_time): mock_time.return_value = 12.0