disaggregation: Fix FakeKVSender queue accumulation (#28652)

Signed-off-by: KP, Jeeja <jeeja.kp@intel.com>

python/sglang/srt/disaggregation/fake/conn.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from typing import List, Optional
 
 import numpy as np
@@ -13,6 +14,7 @@ from sglang.srt.disaggregation.base.conn import (
     KVTransferMetric,
 )
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.environ import envs
 from sglang.srt.server_args import ServerArgs
 
 logger = logging.getLogger(__name__)
@@ -46,13 +48,26 @@ class FakeKVSender(BaseKVSender):
         req_has_disagg_prefill_dp_rank: bool = False,
     ):
         self.kv_mgr = mgr
+        self.bootstrap_room = bootstrap_room
+        # Set by any chunk, not only the last one: nothing is transferred, so a
+        # chunk that never comes cannot change the outcome.
         self.has_sent = False
         self.conclude_state: Optional[KVPoll] = None
+        # Read here rather than off kv_mgr: a FAKE_BOOTSTRAP_HOST req on a real
+        # backend pairs this sender with that backend's KVManager, which carries
+        # no waiting_timeout in prefill mode.
+        self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
+        self.inited = False
+        self.waiting_since: Optional[float] = None
 
     def poll(self) -> KVPoll:
         if self.conclude_state is not None:
             return self.conclude_state
+
         if not self.has_sent:
+            timeout_result = self._check_waiting_timeout()
+            if timeout_result is not None:
+                return timeout_result
             # Assume handshake completed instantly
             return KVPoll.WaitingForInput
 
@@ -61,18 +76,47 @@ class FakeKVSender(BaseKVSender):
         self.conclude_state = KVPoll.Success
         return KVPoll.Success
 
+    def _check_waiting_timeout(self) -> Optional[KVPoll]:
+        # A send() that never comes must not pin the prefill inflight queue forever.
+        # No deadline before init(): the request is still in the bootstrap queue.
+        if not self.inited:
+            return None
+        if self.waiting_since is None:
+            # Clock starts at the first poll after init(), not at init() itself:
+            # the scheduler stops polling while a request queues and computes
+            # prefill. Monotonic, so an NTP step cannot fail a healthy request.
+            self.waiting_since = time.monotonic()
+            return None
+        elapsed = time.monotonic() - self.waiting_since
+        if elapsed < self.waiting_timeout:
+            return None
+        logger.warning_once(
+            "Some FakeKVSender requests fail to receive a KV chunk after bootstrapping. "
+            "If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
+        )
+        logger.debug(
+            f"FakeKVSender for room {self.bootstrap_room} timed out after {elapsed:.1f}s "
+            "in KVPoll.WaitingForInput; no KV chunk was ever sent."
+        )
+        self.conclude_state = KVPoll.Failed
+        return KVPoll.Failed
+
     def get_transfer_metric(self) -> KVTransferMetric:
         return KVTransferMetric()
 
     def init(
         self,
-        kv_indices: list[int],
+        num_kv_indices: int,
         aux_index: Optional[int] = None,
     ):
+        self.inited = True
         logger.debug(
-            f"FakeKVSender init with kv_indices: {kv_indices}, aux_index: {aux_index}"
+            f"FakeKVSender init with num_kv_indices: {num_kv_indices}, aux_index: {aux_index}"
         )
-        pass
+
+    def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
+        # A zero-page last chunk must still send: poll() only concludes after send().
+        return num_pages > 0 or last_chunk
 
     def send(
         self,
@@ -112,6 +156,9 @@ class FakeKVReceiver(BaseKVReceiver):
         if not self.bootstrap_done:
             return KVPoll.Bootstrapping
         if not self.has_sent_metadata:
+            # No deadline needed here, unlike FakeKVSender: send_metadata() is
+            # unconditional once the decode side preallocates, and waiting for
+            # KV space is not a stalled transfer.
             return KVPoll.WaitingForInput
         logger.debug("FakeKVReceiver poll success")
         self.conclude_state = KVPoll.Success