From e39388f70240a8ec88de05a12e6a1fc24d86961a Mon Sep 17 00:00:00 2001
From: Ting-Hong Shieh <32212900+ting-hong-shieh@users.noreply.github.com>
Date: Thu, 13 Aug 2026 14:34:52 +0800
Subject: [PATCH] Add NIXL prefill bootstrap timeout

---
 python/sglang/srt/disaggregation/nixl/conn.py |  5 +++
 .../disaggregation/test_nixl_backend_basic.py | 32 +++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py
index 97e6251930fa..81bf3362f85d 100644
--- a/python/sglang/srt/disaggregation/nixl/conn.py
+++ b/python/sglang/srt/disaggregation/nixl/conn.py
@@ -2750,6 +2750,7 @@ def __init__(
         self._send_failed = False
         self._send_error: Optional[Exception] = None
         self._transfer_start_time: Optional[float] = None
+        self.init_time = time.time()
 
     def send(
         self,
@@ -2790,6 +2791,10 @@ def poll(self) -> KVPoll:
         if self._send_failed:
             return KVPoll.Failed  # type: ignore
         status = self.kv_mgr.check_status(self.bootstrap_room)
+        if status == KVPoll.Bootstrapping:
+            timeout_result = self._check_bootstrap_timeout()
+            if timeout_result is not None:
+                return timeout_result
         # Hold Success until all staging chunks transferred: a deferred chunk
         # can still be pending, and concluding now would drop it.
         if (
diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py
index 9136395edd10..fb52a16fd730 100644
--- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py
+++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py
@@ -340,6 +340,38 @@ def test_last_zero_page_chunk_is_sent_for_aux_only_completion(self):
         self.assertTrue(sender.should_send_kv_chunk(3, last_chunk=False))
 
 
+class TestNixlKVSenderBootstrapTimeout(CustomTestCase):
+    @patch("sglang.srt.disaggregation.nixl.conn.time.time", return_value=123.0)
+    def test_init_records_bootstrap_start_time(self, mock_time):
+        with patch(
+            "sglang.srt.disaggregation.nixl.conn.CommonKVSender.__init__",
+            return_value=None,
+        ):
+            sender = NixlKVSender(MagicMock(), "decode:8998", 11, [0], 0)
+
+        self.assertEqual(sender.init_time, 123.0)
+
+    @patch("sglang.srt.disaggregation.common.conn.time.time", return_value=20.0)
+    def test_bootstrap_timeout_records_failure_and_updates_status(self, mock_time):
+        mgr = MagicMock()
+        mgr.bootstrap_timeout = 5
+        mgr.check_status.return_value = KVPoll.Bootstrapping
+        mgr._staging_outstanding = {}
+
+        sender = object.__new__(NixlKVSender)
+        sender.kv_mgr = mgr
+        sender.bootstrap_room = 11
+        sender.init_time = 10.0
+        sender._send_failed = False
+        sender._transfer_start_time = None
+        sender._transfer_metric = MagicMock()
+
+        self.assertEqual(sender.poll(), KVPoll.Failed)
+        mgr.record_failure.assert_called_once()
+        self.assertIn("timed out", mgr.record_failure.call_args[0][1])
+        mgr.update_status.assert_called_once_with(11, KVPoll.Failed)
+
+
 class TestNixlAbortHandling(CustomTestCase):
     def _make_manager(self, request_status=None):
         mgr = object.__new__(NixlKVManager)