From ea7ef63e6d359fe74cbad5099a76a4f14f3a2806 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 23 Feb 2026 18:34:21 -0800 Subject: [PATCH] [PD-Disagg] Deduplicate common KVManager methods into CommonKVManager (#19205) Co-authored-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 20 +++++++++++++ .../srt/disaggregation/mooncake/conn.py | 28 ------------------- python/sglang/srt/disaggregation/mori/conn.py | 20 ------------- python/sglang/srt/disaggregation/nixl/conn.py | 18 ------------ 4 files changed, 20 insertions(+), 66 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 286132d4f..a71c722dc 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -99,6 +99,8 @@ class CommonKVManager(BaseKVManager): logger.debug(f"kv manager bind to {zmq_bind_host}:{self.rank_port}") self.request_status: Dict[int, KVPoll] = {} + self.failure_records: Dict[int, str] = {} + self.failure_lock = threading.Lock() if self.disaggregation_mode == DisaggregationMode.PREFILL: self.register_to_bootstrap() @@ -115,6 +117,24 @@ class CommonKVManager(BaseKVManager): f"Unsupported DisaggregationMode: {self.disaggregation_mode}" ) + def check_status(self, bootstrap_room: int) -> KVPoll: + return self.request_status[bootstrap_room] + + def update_status(self, bootstrap_room: int, status: KVPoll): + if bootstrap_room not in self.request_status: + self.request_status[bootstrap_room] = status + else: + if status == KVPoll.Failed: + self.request_status[bootstrap_room] = KVPoll.Failed + else: + self.request_status[bootstrap_room] = max( + self.request_status[bootstrap_room], status + ) + + def record_failure(self, bootstrap_room: int, failure_reason: str): + with self.failure_lock: + self.failure_records[bootstrap_room] = failure_reason + def ensure_parallel_info(self, bootstrap_addr: str) -> bool: """Fetch and cache prefill parallel info if not yet available. Returns True if info is available (cached or freshly fetched). diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 09b265490..d05fae9a6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -14,7 +14,6 @@ from typing import Dict, List, Optional, Set, Tuple import numpy as np import numpy.typing as npt import requests -import zmq from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll from sglang.srt.disaggregation.common.conn import ( @@ -235,9 +234,6 @@ class MooncakeKVManager(CommonKVManager): # These timeout requests should be aborted to release the tree cache. self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get() - self.failure_records: Dict[int, str] = {} - self.failure_lock = threading.Lock() - def init_engine(self): self.engine = get_mooncake_transfer_engine() @@ -1095,25 +1091,6 @@ class MooncakeKVManager(CommonKVManager): ) ) - def check_status(self, bootstrap_room: int): - return self.request_status[bootstrap_room] - - def update_status(self, bootstrap_room: int, status: KVPoll): - if bootstrap_room not in self.request_status: - self.request_status[bootstrap_room] = status - else: - # NOTE: status is only allowed to be incremented unless it is KVPoll.Failed - if status == KVPoll.Failed: - self.request_status[bootstrap_room] = KVPoll.Failed - else: - self.request_status[bootstrap_room] = max( - self.request_status[bootstrap_room], status - ) - - def record_failure(self, bootstrap_room: int, failure_reason: str): - with self.failure_lock: - self.failure_records[bootstrap_room] = failure_reason - def get_session_id(self): return self.engine.get_session_id() @@ -1242,11 +1219,6 @@ class MooncakeKVSender(CommonKVSender): class MooncakeKVReceiver(CommonKVReceiver): - _ctx = zmq.Context() - _socket_cache = {} - _socket_locks = {} - _global_lock = threading.Lock() - def __init__( self, mgr: MooncakeKVManager, diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 55c2d481a..830ca9e79 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -191,8 +191,6 @@ class MoriKVManager(CommonKVManager): self.kv_mem_descs: List[MemoryDesc] = [] self.aux_mem_descs: List[MemoryDesc] = [] self.state_mem_descs: List[MemoryDesc] = [] - self.failure_records: Dict[int, str] = {} - self.failure_lock = threading.Lock() self.transfer_lock = threading.Lock() self._register_local_buffers() if self.disaggregation_mode == DisaggregationMode.PREFILL: @@ -293,24 +291,6 @@ class MoriKVManager(CommonKVManager): ) self.state_mem_descs.append(desc) - def check_status(self, bootstrap_room: int): - return self.request_status[bootstrap_room] - - def update_status(self, bootstrap_room: int, status: KVPoll): - if bootstrap_room not in self.request_status: - self.request_status[bootstrap_room] = status - else: - if status == KVPoll.Failed: - self.request_status[bootstrap_room] = KVPoll.Failed - else: - self.request_status[bootstrap_room] = max( - self.request_status[bootstrap_room], status - ) - - def record_failure(self, bootstrap_room: int, failure_reason: str) -> None: - with self.failure_lock: - self.failure_records[bootstrap_room] = failure_reason - def _handle_register_message(self, payload: List[bytes]) -> None: try: register_info = KVArgsRegisterInfo.from_zmq(payload) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 457022a70..8a8807cef 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -300,24 +300,6 @@ class NixlKVManager(CommonKVManager): logger.error(f"Let room {room} be failed due to prefill down") self.update_status(room, KVPoll.Failed) - def check_status(self, bootstrap_room: int): - return self.request_status[bootstrap_room] - - def update_status(self, bootstrap_room: int, status: KVPoll): - if bootstrap_room not in self.request_status: - self.request_status[bootstrap_room] = status - else: - # NOTE: status is only allowed to be incremented unless it is KVPoll.Failed - if status == KVPoll.Failed: - self.request_status[bootstrap_room] = KVPoll.Failed - else: - self.request_status[bootstrap_room] = max( - self.request_status[bootstrap_room], status - ) - - def record_failure(self, bootstrap_room: int, failure_reason: str): - pass - def register_buffer_to_engine(self): kv_addrs = [] for kv_data_ptr, kv_data_len in zip(