chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[PD-Disagg] Deduplicate common KVManager methods into CommonKVManager (#19205)

Browse Source 
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Liangsheng Yin
2026-02-23 18:34:21 -08:00
committed by GitHub
parent 8aeb16f3fc
commit ea7ef63e6d
4 changed files with 20 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
python/sglang/srt/disaggregation/common/conn.py
View File

@@ -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).

28
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -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,

20
python/sglang/srt/disaggregation/mori/conn.py
View File

@@ -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)

18
python/sglang/srt/disaggregation/nixl/conn.py
View File

@@ -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(