[v0.5.15 19/21] Retain KV ownership until Mooncake transfer quiescence

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

@@ -12,6 +12,7 @@
     MooncakeKVManager,
     MooncakeKVReceiver,
     MooncakeKVSender,
+    _submit_transfer_futures,
 )
 from sglang.srt.utils.network import get_local_ip_auto
 
@@ -156,21 +157,11 @@ def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
             return self._transfer_data(mooncake_session_id, transfer_blocks)
 
         if self.enable_custom_mem_pool:
-            futures = [
-                executor.submit(
-                    process_layer,
-                    src_ptr,
-                    dst_ptr,
-                    item_len,
-                )
+            calls = [
+                (process_layer, (src_ptr, dst_ptr, item_len))
                 for (src_ptr, dst_ptr, item_len) in layers_params
             ]
-            for future in concurrent.futures.as_completed(futures):
-                status = future.result()
-                if status != 0:
-                    for f in futures:
-                        f.cancel()
-                    return status
+            return _submit_transfer_futures(executor, calls)
         else:
             # Combining all layers' params in one batch transfer is more efficient
             # compared to using multiple threads

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

@@ -156,6 +156,15 @@ def failure_exception(self):
         """
         ...
 
+    def begin_failure_quiescence(self):
+        pass
+
+    def is_transfer_quiesced(self) -> bool:
+        return True
+
+    def is_failure_quiescing(self) -> bool:
+        return False
+
 
 class BaseKVReceiver(ABC):
     @abstractmethod
@@ -215,6 +224,15 @@ def abort(self):
         """
         pass
 
+    def begin_failure_quiescence(self):
+        pass
+
+    def is_transfer_quiesced(self) -> bool:
+        return True
+
+    def is_failure_quiescing(self) -> bool:
+        return False
+
 
 class BaseKVBootstrapServer(ABC):
     @abstractmethod

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

@@ -75,6 +75,7 @@ class PrefillServerInfo:
     page_size: Optional[int]
     kv_cache_dtype: Optional[str]
     follow_bootstrap_room: bool
+    capabilities: Optional[List[str]] = None
 
     # Pre-computed rank mapping (set by try_ensure_parallel_info on decode side)
     target_tp_rank: Optional[int] = None
@@ -94,16 +95,19 @@ def __post_init__(self):
             str(self.kv_cache_dtype) if self.kv_cache_dtype is not None else None
         )
         self.follow_bootstrap_room = bool(self.follow_bootstrap_room)
+        self.capabilities = list(self.capabilities or [])
 
 
 @dataclasses.dataclass
 class PrefillRankInfo:
     rank_ip: str
     rank_port: int
+    capabilities: Optional[List[str]] = None
 
     def __post_init__(self):
         self.rank_ip = str(self.rank_ip)
         self.rank_port = int(self.rank_port)
+        self.capabilities = list(self.capabilities or [])
 
 
 class CommonKVManager(BaseKVManager):
@@ -421,6 +425,9 @@ def register_to_bootstrap(self):
             "kv_cache_dtype": self.server_args.kv_cache_dtype,
             "load_balance_method": self.server_args.load_balance_method,
         }
+        capabilities = self.get_bootstrap_capabilities()
+        if capabilities:
+            payload["capabilities"] = capabilities
 
         max_retries, initial_delay, max_delay = 5, 1.0, 30.0
         for attempt in range(max_retries):
@@ -450,6 +457,9 @@ def register_to_bootstrap(self):
             f"Prefill instance failed to register to bootstrap server after {max_retries} retries"
         )
 
+    def get_bootstrap_capabilities(self) -> List[str]:
+        return []
+
     def _connect(self, endpoint: str, is_ipv6: bool = False):
         with self._socket_lock:
             sock = self._socket_cache.get(endpoint)
@@ -1230,6 +1240,7 @@ def __init__(self, host: str, port: int):
         self.page_size = None
         self.kv_cache_dtype: Optional[str] = None
         self.follow_bootstrap_room: Optional[bool] = None
+        self.capabilities: List[str] = []
         self.prefill_port_table: Dict[
             int, Dict[int, Dict[int, Dict[int, PrefillRankInfo]]]
         ] = {}
@@ -1296,6 +1307,7 @@ async def _handle_route_put(self, request: web.Request):
         rank_port = int(data["rank_port"])
         page_size = int(data["page_size"])
         kv_cache_dtype = data["kv_cache_dtype"]
+        capabilities = list(data.get("capabilities") or [])
 
         if self.attn_tp_size is None:
             self.attn_tp_size = attn_tp_size
@@ -1315,6 +1327,9 @@ async def _handle_route_put(self, request: web.Request):
         if self.kv_cache_dtype is None and kv_cache_dtype is not None:
             self.kv_cache_dtype = kv_cache_dtype
 
+        if capabilities:
+            self.capabilities = capabilities
+
         if self.follow_bootstrap_room is None:
             load_balance_method = data.get(
                 "load_balance_method", "follow_bootstrap_room"
@@ -1335,6 +1350,7 @@ async def _handle_route_put(self, request: web.Request):
             tp_group_table[pp_rank] = PrefillRankInfo(
                 rank_ip=rank_ip,
                 rank_port=rank_port,
+                capabilities=capabilities,
             )
 
             self._registered_count += 1
@@ -1384,8 +1400,12 @@ async def _handle_route_get(self, request: web.Request):
                     if self.follow_bootstrap_room is not None
                     else True
                 ),
+                capabilities=self.capabilities,
             )
-            return web.json_response(dataclasses.asdict(info), status=200)
+            response = dataclasses.asdict(info)
+            if not response["capabilities"]:
+                response.pop("capabilities")
+            return web.json_response(response, status=200)
 
         if not self._is_ready():
             return web.Response(
@@ -1407,7 +1427,10 @@ async def _handle_route_get(self, request: web.Request):
                 status=404,
             )
 
-        return web.json_response(dataclasses.asdict(bootstrap_info), status=200)
+        response = dataclasses.asdict(bootstrap_info)
+        if not response["capabilities"]:
+            response.pop("capabilities")
+        return web.json_response(response, status=200)
 
     async def _handle_register_dp_rank(self, request: web.Request):
         data = await request.json()