[PD] Drain NIXL completion notifications before enforcing the WaitingForInput timeout (#32267)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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()
+        # Drain notifications before enforcing the waiting deadline. The decode
-        if timeout_result is not None:
+        # agent has no NIXL progress thread (num_threads=0), so incoming
-            return timeout_result
+        # 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:

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