[PD] Consolidate shared logic into common backend (#25979)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
Shangming Cai
@@ -25,7 +25,10 @@ from sglang.srt.disaggregation.base.conn import (
|
||||
KVPoll,
|
||||
KVTransferMetric,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
filter_kv_indices_for_cp_rank,
|
||||
)
|
||||
from sglang.srt.distributed import get_pp_group, get_world_group
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
@@ -594,6 +597,92 @@ class CommonKVManager(BaseKVManager):
|
||||
|
||||
return src_kv_ptrs, sliced_dst
|
||||
|
||||
def _start_heartbeat_checker_thread(self):
|
||||
"""Start the heartbeat checker thread for Decode worker."""
|
||||
|
||||
def heartbeat_checker():
|
||||
while True:
|
||||
time.sleep(self.heartbeat_interval)
|
||||
with self.connection_lock:
|
||||
addresses = list(self.prefill_info_table.keys())
|
||||
|
||||
for bootstrap_addr in addresses:
|
||||
session = None
|
||||
try:
|
||||
with self.session_pool_lock:
|
||||
session = self.session_pool[bootstrap_addr]
|
||||
response = session.get(
|
||||
f"http://{bootstrap_addr}/health",
|
||||
timeout=(2, 3),
|
||||
headers={"Connection": "keep-alive"},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
self.heartbeat_failures[bootstrap_addr] = 0
|
||||
self._on_heartbeat_success(bootstrap_addr)
|
||||
else:
|
||||
logger.info(
|
||||
f"Attempting to reconnect to {bootstrap_addr}..."
|
||||
)
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
except Exception:
|
||||
logger.info(f"Attempting to reconnect to {bootstrap_addr}...")
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
|
||||
if (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0)
|
||||
>= self.max_failures
|
||||
):
|
||||
self._handle_node_failure(bootstrap_addr)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
|
||||
threading.Thread(target=heartbeat_checker, daemon=True).start()
|
||||
|
||||
def _on_heartbeat_success(self, bootstrap_addr: str):
|
||||
"""Hook called on successful heartbeat. Override for backend-specific cleanup."""
|
||||
pass
|
||||
|
||||
def _handle_node_failure(self, failed_bootstrap_addr: str):
|
||||
"""Handle failure of a prefill node."""
|
||||
with self.connection_lock:
|
||||
keys_to_remove = [
|
||||
k for k in self.connection_pool if k.startswith(failed_bootstrap_addr)
|
||||
]
|
||||
for k in keys_to_remove:
|
||||
del self.connection_pool[k]
|
||||
self.prefill_info_table.pop(failed_bootstrap_addr, None)
|
||||
|
||||
possible_affected_rooms = self.addr_to_rooms_tracker.get(
|
||||
failed_bootstrap_addr, []
|
||||
)
|
||||
self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None)
|
||||
|
||||
affected_rooms = []
|
||||
for room in possible_affected_rooms:
|
||||
if (
|
||||
room in self.request_status
|
||||
and self.check_status(room) != KVPoll.Success
|
||||
):
|
||||
self.record_failure(
|
||||
room,
|
||||
f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})",
|
||||
)
|
||||
self.update_status(room, KVPoll.Failed)
|
||||
affected_rooms.append(room)
|
||||
|
||||
logger.error(
|
||||
f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), "
|
||||
f"{len(affected_rooms)} requests affected"
|
||||
)
|
||||
|
||||
|
||||
class CommonKVSender(BaseKVSender):
|
||||
def __init__(
|
||||
@@ -614,6 +703,7 @@ class CommonKVSender(BaseKVSender):
|
||||
self._transfer_num_state_indices = 0
|
||||
# inner state
|
||||
self.curr_idx = 0
|
||||
self.init_time: Optional[float] = None
|
||||
if self.kv_mgr.is_dummy_cp_rank:
|
||||
# Non-authoritative CP ranks are dummy participants.
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput)
|
||||
@@ -667,10 +757,10 @@ class CommonKVSender(BaseKVSender):
|
||||
)
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return 0
|
||||
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0
|
||||
return num_pages > 0 or last_chunk
|
||||
|
||||
def get_transfer_metric(self) -> KVTransferMetric:
|
||||
total_bytes = self._transfer_num_kv_indices * self.kv_mgr.kv_item_lens_sum
|
||||
@@ -691,6 +781,36 @@ class CommonKVSender(BaseKVSender):
|
||||
if component_indices is not None:
|
||||
self._transfer_num_state_indices += len(component_indices)
|
||||
|
||||
def _prepare_send_indices(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
state_indices: Optional[List] = None,
|
||||
) -> Tuple[npt.NDArray[np.int32], slice, bool, bool]:
|
||||
"""Common pre-processing for send(): index tracking and CP-rank handling.
|
||||
|
||||
Returns:
|
||||
(kv_indices, index_slice, is_last_chunk, should_skip)
|
||||
If should_skip is True, the caller should return immediately.
|
||||
"""
|
||||
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
||||
self.curr_idx += len(kv_indices)
|
||||
is_last_chunk = self.curr_idx == self.num_kv_indices
|
||||
|
||||
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
||||
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
||||
self.kv_mgr,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
)
|
||||
elif self.kv_mgr.is_dummy_cp_rank:
|
||||
if not is_last_chunk:
|
||||
return kv_indices, index_slice, is_last_chunk, True
|
||||
else:
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
||||
return kv_indices, index_slice, is_last_chunk, True
|
||||
|
||||
return kv_indices, index_slice, is_last_chunk, False
|
||||
|
||||
def send(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
@@ -698,6 +818,25 @@ class CommonKVSender(BaseKVSender):
|
||||
):
|
||||
pass
|
||||
|
||||
def _check_bootstrap_timeout(self) -> Optional[KVPoll]:
|
||||
if self.init_time is None:
|
||||
return None
|
||||
elapsed = time.time() - self.init_time
|
||||
if elapsed < self.kv_mgr.bootstrap_timeout:
|
||||
return None
|
||||
logger.warning_once(
|
||||
"Some requests timed out when bootstrapping, "
|
||||
"which means prefill instances fail to receive the KV indices from the decode instance of this request. "
|
||||
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
||||
)
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s "
|
||||
f"in KVPoll.Bootstrapping",
|
||||
)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
return KVPoll.Failed
|
||||
|
||||
def poll(self) -> KVPoll:
|
||||
pass
|
||||
|
||||
@@ -737,6 +876,7 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
self.kv_mgr = mgr
|
||||
self.conclude_state: Optional[KVPoll] = None
|
||||
self.require_staging: bool = False
|
||||
self.init_time: Optional[float] = None
|
||||
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
|
||||
|
||||
@@ -906,6 +1046,24 @@ class CommonKVReceiver(BaseKVReceiver):
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
def _check_waiting_timeout(self) -> Optional[KVPoll]:
|
||||
if self.init_time is None:
|
||||
return None
|
||||
elapsed = time.time() - self.init_time
|
||||
if elapsed < self.kv_mgr.waiting_timeout:
|
||||
return None
|
||||
logger.warning_once(
|
||||
"Some requests fail to receive KV Cache transfer done signal after bootstrapping. "
|
||||
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
||||
)
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s "
|
||||
f"in KVPoll.WaitingForInput",
|
||||
)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
return KVPoll.Failed
|
||||
|
||||
def failure_exception(self):
|
||||
raise Exception("Fake KVReceiver Exception")
|
||||
|
||||
|
||||
@@ -1,12 +1,27 @@
|
||||
import ctypes
|
||||
import dataclasses
|
||||
import struct
|
||||
import threading
|
||||
from collections import deque
|
||||
from typing import List, Tuple
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TransferKVChunk:
|
||||
"""Work unit for KV cache transfer from prefill to decode."""
|
||||
|
||||
room: int
|
||||
prefill_kv_indices: npt.NDArray[np.int32]
|
||||
index_slice: slice
|
||||
is_last_chunk: bool
|
||||
prefill_aux_index: Optional[int]
|
||||
state_indices: Optional[List]
|
||||
chunk_id: Optional[int] = None
|
||||
|
||||
|
||||
def pack_list_of_buffers(buffers: List[bytes]) -> bytes:
|
||||
if not buffers:
|
||||
return b""
|
||||
@@ -59,6 +74,26 @@ class FastQueue:
|
||||
return self._buf.popleft()
|
||||
|
||||
|
||||
class AuxDataCodec:
|
||||
"""Handles serialization and deserialization of auxiliary data buffers."""
|
||||
|
||||
@staticmethod
|
||||
def serialize_data_from_buffer(src_addr, data_length):
|
||||
"""Serialize data from memory buffer to bytes."""
|
||||
buffer = (ctypes.c_byte * data_length).from_address(src_addr)
|
||||
return bytes(buffer)
|
||||
|
||||
@staticmethod
|
||||
def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
|
||||
"""Deserialize bytes into target memory buffer."""
|
||||
dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
|
||||
item_len = kv_args.aux_item_lens[buffer_index]
|
||||
dst_addr = dst_aux_ptr + item_len * aux_index
|
||||
buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
|
||||
buffer[:] = data
|
||||
return
|
||||
|
||||
|
||||
def group_concurrent_contiguous(
|
||||
src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32]
|
||||
) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]:
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import ctypes
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
@@ -29,7 +28,9 @@ from sglang.srt.disaggregation.common.staging_handler import (
|
||||
StagingTransferInfo,
|
||||
)
|
||||
from sglang.srt.disaggregation.common.utils import (
|
||||
AuxDataCodec,
|
||||
FastQueue,
|
||||
TransferKVChunk,
|
||||
group_concurrent_contiguous,
|
||||
pack_int_lists,
|
||||
unpack_int_lists,
|
||||
@@ -37,10 +38,7 @@ from sglang.srt.disaggregation.common.utils import (
|
||||
from sglang.srt.disaggregation.mooncake.utils import (
|
||||
check_mooncake_custom_mem_pool_enabled,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
filter_kv_indices_for_cp_rank,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
@@ -64,17 +62,6 @@ class KVTransferError(Exception):
|
||||
return f"KVTransferError(bootstrap_room={self.bootstrap_room}): {self.failure_reason}"
|
||||
|
||||
|
||||
# prefill
|
||||
@dataclasses.dataclass
|
||||
class TransferKVChunk:
|
||||
room: int
|
||||
prefill_kv_indices: npt.NDArray[np.int32]
|
||||
index_slice: slice
|
||||
is_last_chunk: bool
|
||||
prefill_aux_index: Optional[int]
|
||||
state_indices: Optional[List]
|
||||
|
||||
|
||||
# decode
|
||||
@dataclasses.dataclass
|
||||
class TransferInfo:
|
||||
@@ -162,26 +149,6 @@ class KVArgsRegisterInfo:
|
||||
)
|
||||
|
||||
|
||||
class AuxDataCodec:
|
||||
"""Handles serialization and deserialization of auxiliary data buffers"""
|
||||
|
||||
@staticmethod
|
||||
def serialize_data_from_buffer(src_addr, data_length):
|
||||
"""Serialize data from memory buffer to bytes"""
|
||||
buffer = (ctypes.c_byte * data_length).from_address(src_addr)
|
||||
return bytes(buffer)
|
||||
|
||||
@staticmethod
|
||||
def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
|
||||
"""Deserialize bytes into target memory buffer"""
|
||||
dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
|
||||
item_len = kv_args.aux_item_lens[buffer_index]
|
||||
dst_addr = dst_aux_ptr + item_len * aux_index
|
||||
buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
|
||||
buffer[:] = data
|
||||
return
|
||||
|
||||
|
||||
class MooncakeKVManager(CommonKVManager):
|
||||
AUX_DATA_HEADER = b"AUX_DATA"
|
||||
|
||||
@@ -1478,62 +1445,8 @@ class MooncakeKVManager(CommonKVManager):
|
||||
)
|
||||
self.update_status(bootstrap_room, status)
|
||||
|
||||
def heartbeat_checker():
|
||||
while True:
|
||||
time.sleep(self.heartbeat_interval)
|
||||
with self.connection_lock:
|
||||
addresses = list(self.prefill_info_table.keys())
|
||||
|
||||
for bootstrap_addr in addresses:
|
||||
session = None
|
||||
try:
|
||||
with self.session_pool_lock:
|
||||
session = self.session_pool[bootstrap_addr]
|
||||
response = session.get(
|
||||
f"http://{bootstrap_addr}/health",
|
||||
timeout=(2, 3),
|
||||
headers={"Connection": "keep-alive"},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
self.heartbeat_failures[bootstrap_addr] = 0
|
||||
|
||||
current_rooms = self.addr_to_rooms_tracker[
|
||||
bootstrap_addr
|
||||
].copy()
|
||||
|
||||
for bootstrap_room in current_rooms:
|
||||
# Remove KVPoll.Success requests from the tracker
|
||||
if bootstrap_room not in self.request_status:
|
||||
self.addr_to_rooms_tracker[bootstrap_addr].discard(
|
||||
bootstrap_room
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Attempting to reconnect to {bootstrap_addr}..."
|
||||
)
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
except Exception:
|
||||
logger.info(f"Attempting to reconnect to {bootstrap_addr}...")
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
|
||||
if (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0)
|
||||
>= self.max_failures
|
||||
):
|
||||
self._handle_node_failure(bootstrap_addr)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
|
||||
threading.Thread(target=decode_thread).start()
|
||||
threading.Thread(target=heartbeat_checker).start()
|
||||
self._start_heartbeat_checker_thread()
|
||||
|
||||
def add_transfer_request(
|
||||
self,
|
||||
@@ -1583,6 +1496,13 @@ class MooncakeKVManager(CommonKVManager):
|
||||
def get_session_id(self):
|
||||
return self.engine.get_session_id()
|
||||
|
||||
def _on_heartbeat_success(self, bootstrap_addr: str):
|
||||
current_rooms = self.addr_to_rooms_tracker[bootstrap_addr].copy()
|
||||
for bootstrap_room in current_rooms:
|
||||
# Remove KVPoll.Success requests from the tracker
|
||||
if bootstrap_room not in self.request_status:
|
||||
self.addr_to_rooms_tracker[bootstrap_addr].discard(bootstrap_room)
|
||||
|
||||
def _run_one_probe_pass(self) -> None:
|
||||
with self.session_lock:
|
||||
snapshot = list(self.failed_sessions)
|
||||
@@ -1666,33 +1586,15 @@ class MooncakeKVSender(CommonKVSender):
|
||||
self.conclude_state = None
|
||||
self.init_time = time.time()
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0 or last_chunk
|
||||
|
||||
def send(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
state_indices: Optional[List] = None,
|
||||
):
|
||||
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
||||
self.curr_idx += len(kv_indices)
|
||||
is_last_chunk = self.curr_idx == self.num_kv_indices
|
||||
|
||||
# Special handling for cp
|
||||
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
||||
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
||||
self.kv_mgr,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
kv_indices, index_slice, is_last_chunk, should_skip = (
|
||||
self._prepare_send_indices(kv_indices, state_indices)
|
||||
)
|
||||
elif self.kv_mgr.is_dummy_cp_rank:
|
||||
if not is_last_chunk:
|
||||
return
|
||||
else:
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
||||
if should_skip:
|
||||
return
|
||||
|
||||
if not is_last_chunk:
|
||||
@@ -1719,21 +1621,9 @@ class MooncakeKVSender(CommonKVSender):
|
||||
if status in (KVPoll.Success, KVPoll.Failed):
|
||||
self.conclude_state = status
|
||||
elif status == KVPoll.Bootstrapping:
|
||||
if self.init_time is not None:
|
||||
now = time.time()
|
||||
elapsed = now - self.init_time
|
||||
if elapsed >= self.kv_mgr.bootstrap_timeout:
|
||||
logger.warning_once(
|
||||
"Some requests timed out when bootstrapping, "
|
||||
"which means prefill instances fail to receive the KV indices from the decode instance of this request. "
|
||||
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
||||
)
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.Bootstrapping",
|
||||
)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
return KVPoll.Failed
|
||||
timeout_result = self._check_bootstrap_timeout()
|
||||
if timeout_result is not None:
|
||||
return timeout_result
|
||||
|
||||
return status
|
||||
else:
|
||||
@@ -1819,12 +1709,6 @@ class MooncakeKVReceiver(CommonKVReceiver):
|
||||
]
|
||||
)
|
||||
|
||||
def init(
|
||||
self,
|
||||
prefill_dp_rank: int,
|
||||
):
|
||||
super().init(prefill_dp_rank)
|
||||
|
||||
def send_metadata(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
@@ -1874,33 +1758,20 @@ class MooncakeKVReceiver(CommonKVReceiver):
|
||||
self.init_time = time.time()
|
||||
|
||||
def poll(self) -> KVPoll:
|
||||
if self.conclude_state is None:
|
||||
if self.conclude_state is not None:
|
||||
return self.conclude_state
|
||||
|
||||
status = self.kv_mgr.check_status(self.bootstrap_room)
|
||||
if status in (KVPoll.Success, KVPoll.Failed):
|
||||
self.conclude_state = status
|
||||
elif status == KVPoll.WaitingForInput:
|
||||
if self.init_time is not None:
|
||||
now = time.time()
|
||||
elapsed = now - self.init_time
|
||||
if elapsed >= self.kv_mgr.waiting_timeout:
|
||||
logger.warning_once(
|
||||
"Some requests fail to receive KV Cache transfer done signal after bootstrapping. "
|
||||
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
||||
)
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput",
|
||||
)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
return KVPoll.Failed
|
||||
timeout_result = self._check_waiting_timeout()
|
||||
if timeout_result is not None:
|
||||
return timeout_result
|
||||
|
||||
return status
|
||||
|
||||
else:
|
||||
return self.conclude_state
|
||||
|
||||
def failure_exception(self):
|
||||
# Explicitly set the status to failure since this request has failed in another rank
|
||||
if self.conclude_state is None:
|
||||
self.conclude_state = KVPoll.Failed
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import dataclasses
|
||||
import logging
|
||||
import os
|
||||
@@ -33,11 +32,11 @@ from sglang.srt.disaggregation.common.conn import (
|
||||
CommonKVReceiver,
|
||||
CommonKVSender,
|
||||
)
|
||||
from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
filter_kv_indices_for_cp_rank,
|
||||
from sglang.srt.disaggregation.common.utils import (
|
||||
AuxDataCodec,
|
||||
group_concurrent_contiguous,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils.common import get_int_env_var
|
||||
from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto
|
||||
@@ -176,22 +175,6 @@ class KVArgsRegisterInfo:
|
||||
)
|
||||
|
||||
|
||||
class AuxDataCodec:
|
||||
@staticmethod
|
||||
def serialize_data_from_buffer(src_addr, data_length):
|
||||
buffer = (ctypes.c_byte * data_length).from_address(src_addr)
|
||||
return bytes(buffer)
|
||||
|
||||
@staticmethod
|
||||
def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
|
||||
dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
|
||||
item_len = kv_args.aux_item_lens[buffer_index]
|
||||
dst_addr = dst_aux_ptr + item_len * aux_index
|
||||
buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
|
||||
buffer[:] = data
|
||||
return
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TPSliceConfig:
|
||||
page_size: int
|
||||
@@ -1132,7 +1115,7 @@ class MoriKVManager(CommonKVManager):
|
||||
bootstrap_room: int,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
index_slice: slice,
|
||||
is_last: bool,
|
||||
is_last_chunk: bool,
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[npt.NDArray[np.int32]] = None,
|
||||
) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]:
|
||||
@@ -1163,7 +1146,7 @@ class MoriKVManager(CommonKVManager):
|
||||
self.update_status(bootstrap_room, KVPoll.Failed)
|
||||
return [], list(transfer_infos.values())
|
||||
targets.append(TransferTarget(info=info, peer_info=peer_info))
|
||||
if is_last:
|
||||
if is_last_chunk:
|
||||
target_infos_snapshot = list(transfer_infos.values())
|
||||
|
||||
result_statuses: List[TransferStatus] = []
|
||||
@@ -1179,7 +1162,7 @@ class MoriKVManager(CommonKVManager):
|
||||
)
|
||||
|
||||
if (
|
||||
is_last
|
||||
is_last_chunk
|
||||
and state_indices is not None
|
||||
and not info.is_dummy
|
||||
and self.state_mem_descs
|
||||
@@ -1191,7 +1174,7 @@ class MoriKVManager(CommonKVManager):
|
||||
)
|
||||
|
||||
if (
|
||||
is_last
|
||||
is_last_chunk
|
||||
and aux_index is not None
|
||||
and info.dst_aux_index >= 0
|
||||
and self.pp_group.is_last_rank
|
||||
@@ -1212,7 +1195,7 @@ class MoriKVManager(CommonKVManager):
|
||||
)
|
||||
return result_statuses, target_infos_snapshot
|
||||
|
||||
if is_last:
|
||||
if is_last_chunk:
|
||||
with self.transfer_lock:
|
||||
# Keep transfer_infos alive until sender.clear() so abort/failure
|
||||
# paths can still recover notification targets after posting.
|
||||
@@ -1243,38 +1226,28 @@ class MoriKVSender(CommonKVSender):
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
state_indices: Optional[List] = None,
|
||||
):
|
||||
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
||||
self.curr_idx += len(kv_indices)
|
||||
is_last = self.curr_idx == self.num_kv_indices
|
||||
|
||||
# Special handling for cp
|
||||
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
||||
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
||||
self.kv_mgr,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
kv_indices, index_slice, is_last_chunk, should_skip = (
|
||||
self._prepare_send_indices(kv_indices, state_indices)
|
||||
)
|
||||
elif self.kv_mgr.is_dummy_cp_rank:
|
||||
if not is_last:
|
||||
return
|
||||
else:
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
||||
if should_skip:
|
||||
return
|
||||
|
||||
normalized_state = _normalize_state_indices(state_indices) if is_last else None
|
||||
normalized_state = (
|
||||
_normalize_state_indices(state_indices) if is_last_chunk else None
|
||||
)
|
||||
statuses, infos = self.kv_mgr.add_transfer_request(
|
||||
self.bootstrap_room,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
is_last,
|
||||
aux_index=self.aux_index if is_last else None,
|
||||
is_last_chunk,
|
||||
aux_index=self.aux_index if is_last_chunk else None,
|
||||
state_indices=normalized_state,
|
||||
)
|
||||
self.transfer_statuses.extend(statuses)
|
||||
self._record_transfer_indices(kv_indices, None)
|
||||
if infos is not None:
|
||||
self.pending_infos = infos
|
||||
if is_last:
|
||||
if is_last_chunk:
|
||||
self.sent_last_chunk = True
|
||||
self._maybe_finalize_if_room_failed()
|
||||
|
||||
@@ -1295,15 +1268,9 @@ class MoriKVSender(CommonKVSender):
|
||||
status = self.kv_mgr.check_status(self.bootstrap_room)
|
||||
|
||||
if status == KVPoll.Bootstrapping:
|
||||
elapsed = time.time() - self.init_time
|
||||
if elapsed >= self.kv_mgr.bootstrap_timeout:
|
||||
reason = (
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s "
|
||||
"waiting for decode handshake"
|
||||
)
|
||||
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
self._finalize_failure(reason)
|
||||
timeout_result = self._check_bootstrap_timeout()
|
||||
if timeout_result is not None:
|
||||
self._finalize_failure()
|
||||
return KVPoll.Failed
|
||||
return status
|
||||
|
||||
@@ -1499,14 +1466,10 @@ class MoriKVReceiver(CommonKVReceiver):
|
||||
self.conclude_state = status
|
||||
return status
|
||||
|
||||
if status == KVPoll.WaitingForInput and self.init_time is not None:
|
||||
elapsed = time.time() - self.init_time
|
||||
if elapsed >= self.kv_mgr.waiting_timeout:
|
||||
reason = f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s waiting for KV transfer"
|
||||
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
return KVPoll.Failed
|
||||
if status == KVPoll.WaitingForInput:
|
||||
timeout_result = self._check_waiting_timeout()
|
||||
if timeout_result is not None:
|
||||
return timeout_result
|
||||
|
||||
return status
|
||||
|
||||
|
||||
@@ -26,14 +26,12 @@ from sglang.srt.disaggregation.common.conn import (
|
||||
from sglang.srt.disaggregation.common.staging_handler import StagingRegisterInfo
|
||||
from sglang.srt.disaggregation.common.utils import (
|
||||
FastQueue,
|
||||
TransferKVChunk,
|
||||
group_concurrent_contiguous,
|
||||
pack_int_lists,
|
||||
unpack_int_lists,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import (
|
||||
DisaggregationMode,
|
||||
filter_kv_indices_for_cp_rank,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
@@ -104,17 +102,6 @@ class TransferInfo:
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TransferKVChunk:
|
||||
room: int
|
||||
prefill_kv_indices: npt.NDArray[np.int32]
|
||||
index_slice: slice
|
||||
is_last: bool
|
||||
chunk_id: int
|
||||
prefill_aux_index: Optional[int]
|
||||
state_indices: Optional[List]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class KVArgsRegisterInfo:
|
||||
"""Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread."""
|
||||
@@ -176,7 +163,7 @@ class TransferStatus:
|
||||
received_kvs_per_pp: Dict[int, Set[int]] = dataclasses.field(
|
||||
default_factory=lambda: defaultdict(set)
|
||||
)
|
||||
# Expected chunk count per pp_rank (set when is_last=True): {pp_rank: expected_count}
|
||||
# Expected chunk count per pp_rank (set when is_last_chunk=True): {pp_rank: expected_count}
|
||||
expected_kvs_per_pp: Dict[int, int] = dataclasses.field(default_factory=dict)
|
||||
# Number of PP ranks expected to send data.
|
||||
num_pp_ranks_expected: Optional[int] = None
|
||||
@@ -186,12 +173,8 @@ class TransferStatus:
|
||||
received_state_per_pp: Set[int] = dataclasses.field(default_factory=set)
|
||||
# Whether state data is expected (set based on state_type).
|
||||
expects_state: bool = False
|
||||
# Mark as failed
|
||||
is_failure: bool = False
|
||||
|
||||
def is_done(self):
|
||||
if self.is_failure:
|
||||
return True
|
||||
if self.num_pp_ranks_expected is None or not self.received_aux:
|
||||
return False
|
||||
# If state data is expected, check all PP ranks have sent it
|
||||
@@ -209,9 +192,6 @@ class TransferStatus:
|
||||
return False
|
||||
return True
|
||||
|
||||
def is_failed(self):
|
||||
return self.is_failure
|
||||
|
||||
|
||||
class NixlKVManager(CommonKVManager):
|
||||
def __init__(
|
||||
@@ -471,92 +451,6 @@ class NixlKVManager(CommonKVManager):
|
||||
)
|
||||
self._staging_ctx.prefetched_rooms.add(room)
|
||||
|
||||
def _start_heartbeat_checker_thread(self):
|
||||
"""
|
||||
Start the heartbeat checker thread for Decode worker.
|
||||
TODO (smor): unite nixl heartbeat checker with mooncake's.
|
||||
"""
|
||||
|
||||
def heartbeat_checker():
|
||||
while True:
|
||||
time.sleep(self.heartbeat_interval)
|
||||
with self.connection_lock:
|
||||
addresses = list(self.prefill_info_table.keys())
|
||||
|
||||
for bootstrap_addr in addresses:
|
||||
session = None
|
||||
try:
|
||||
with self.session_pool_lock:
|
||||
session = self.session_pool[bootstrap_addr]
|
||||
response = session.get(
|
||||
f"http://{bootstrap_addr}/health",
|
||||
timeout=(2, 3),
|
||||
headers={"Connection": "keep-alive"},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
self.heartbeat_failures[bootstrap_addr] = 0
|
||||
|
||||
else:
|
||||
logger.info(
|
||||
f"Attempting to reconnect to {bootstrap_addr}..."
|
||||
)
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
except Exception:
|
||||
logger.info(f"Attempting to reconnect to {bootstrap_addr}...")
|
||||
self.heartbeat_failures[bootstrap_addr] = (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
||||
)
|
||||
|
||||
if (
|
||||
self.heartbeat_failures.get(bootstrap_addr, 0)
|
||||
>= self.max_failures
|
||||
):
|
||||
self._handle_node_failure(bootstrap_addr)
|
||||
with self.session_pool_lock:
|
||||
if bootstrap_addr in self.session_pool:
|
||||
del self.session_pool[bootstrap_addr]
|
||||
|
||||
threading.Thread(target=heartbeat_checker, daemon=True).start()
|
||||
|
||||
def _handle_node_failure(self, failed_bootstrap_addr):
|
||||
"""Handle failure of a prefill node."""
|
||||
with self.connection_lock:
|
||||
keys_to_remove = [
|
||||
k for k in self.connection_pool if k.startswith(failed_bootstrap_addr)
|
||||
]
|
||||
for k in keys_to_remove:
|
||||
del self.connection_pool[k]
|
||||
self.prefill_info_table.pop(failed_bootstrap_addr, None)
|
||||
|
||||
possible_affected_rooms = self.addr_to_rooms_tracker.get(
|
||||
failed_bootstrap_addr, []
|
||||
)
|
||||
self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None)
|
||||
|
||||
# Mark all pending transfers associated with the failed node as failed
|
||||
affected_rooms = []
|
||||
for room in possible_affected_rooms:
|
||||
if (
|
||||
room in self.transfer_statuses
|
||||
and not self.transfer_statuses[room].is_done()
|
||||
):
|
||||
# Mark the transfer as failed
|
||||
self.transfer_statuses[room].is_failure = True
|
||||
affected_rooms.append(room)
|
||||
|
||||
logger.error(
|
||||
f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), "
|
||||
f"{len(affected_rooms)} transfers affected"
|
||||
)
|
||||
for room in possible_affected_rooms:
|
||||
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.get(bootstrap_room, KVPoll.WaitingForInput)
|
||||
|
||||
@@ -606,7 +500,7 @@ class NixlKVManager(CommonKVManager):
|
||||
|
||||
# Skip KV RDMA transfer when there are no pages to send
|
||||
# (e.g., decode-side radix cache matched the entire prefix).
|
||||
# Aux data is still sent below when is_last=True.
|
||||
# Aux data is still sent below when is_last_chunk=True.
|
||||
if len(kv_chunk.prefill_kv_indices) > 0:
|
||||
chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
|
||||
|
||||
@@ -659,7 +553,7 @@ class NixlKVManager(CommonKVManager):
|
||||
if kv_xfer_handle is None:
|
||||
notif = (
|
||||
f"{req.room}_kv_{kv_chunk.chunk_id}"
|
||||
f"_{int(kv_chunk.is_last)}_{self.kv_args.engine_rank}"
|
||||
f"_{int(kv_chunk.is_last_chunk)}_{self.kv_args.engine_rank}"
|
||||
)
|
||||
if self.is_mla_backend or (
|
||||
decode_tp_size == self.attn_tp_size
|
||||
@@ -688,7 +582,7 @@ class NixlKVManager(CommonKVManager):
|
||||
|
||||
handles.append(kv_xfer_handle)
|
||||
|
||||
if kv_chunk.is_last:
|
||||
if kv_chunk.is_last_chunk:
|
||||
dst_info = self.decode_kv_args_table[req.agent_name]
|
||||
if kv_chunk.state_indices:
|
||||
state_xfer_handles = self.maybe_send_extra(
|
||||
@@ -739,7 +633,7 @@ class NixlKVManager(CommonKVManager):
|
||||
break
|
||||
time.sleep(0)
|
||||
|
||||
if kv_chunk.is_last:
|
||||
if kv_chunk.is_last_chunk:
|
||||
self.update_status(room, KVPoll.Success)
|
||||
# Drop per-room state on Success (parity with mooncake
|
||||
# transfer_worker; staging prefetch sets are NIXL-only).
|
||||
@@ -1265,7 +1159,7 @@ class NixlKVManager(CommonKVManager):
|
||||
return (None, True)
|
||||
|
||||
notif_tag = (
|
||||
f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last)}"
|
||||
f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last_chunk)}"
|
||||
f"_{self.kv_args.engine_rank}_{chunk_idx}"
|
||||
f"_{page_start}_{num_pages}_{req.agent_name}"
|
||||
)
|
||||
@@ -1574,13 +1468,13 @@ class NixlKVManager(CommonKVManager):
|
||||
bootstrap_room: int,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
index_slice: slice,
|
||||
is_last: bool,
|
||||
is_last_chunk: bool,
|
||||
chunk_id: int,
|
||||
aux_index: Optional[int] = None,
|
||||
state_indices: Optional[List] = None,
|
||||
):
|
||||
assert self.disaggregation_mode == DisaggregationMode.PREFILL
|
||||
assert not is_last or (is_last and aux_index is not None)
|
||||
assert not is_last_chunk or (is_last_chunk and aux_index is not None)
|
||||
|
||||
# Prefetch STAGING_REQ to decode before enqueueing so decode has
|
||||
# already allocated staging by the time the worker picks up the
|
||||
@@ -1601,7 +1495,7 @@ class NixlKVManager(CommonKVManager):
|
||||
room=bootstrap_room,
|
||||
prefill_kv_indices=kv_indices,
|
||||
index_slice=index_slice,
|
||||
is_last=is_last,
|
||||
is_last_chunk=is_last_chunk,
|
||||
chunk_id=chunk_id,
|
||||
prefill_aux_index=aux_index,
|
||||
state_indices=state_indices,
|
||||
@@ -1628,9 +1522,9 @@ class NixlKVManager(CommonKVManager):
|
||||
tag = components[1]
|
||||
if tag == "kv":
|
||||
chunk_id = int(components[2])
|
||||
is_last = bool(int(components[3]))
|
||||
is_last_chunk = bool(int(components[3]))
|
||||
pp_rank = int(components[4]) if len(components) > 4 else 0
|
||||
self._track_kv_arrival(room, chunk_id, is_last, pp_rank)
|
||||
self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank)
|
||||
elif tag == "stg":
|
||||
self._handle_stg_notification(components, room)
|
||||
elif tag == "aux":
|
||||
@@ -1647,13 +1541,13 @@ class NixlKVManager(CommonKVManager):
|
||||
Format: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name}
|
||||
"""
|
||||
chunk_id = int(components[2])
|
||||
is_last = bool(int(components[3]))
|
||||
is_last_chunk = bool(int(components[3]))
|
||||
pp_rank = int(components[4])
|
||||
chunk_idx = int(components[5])
|
||||
page_start = int(components[6])
|
||||
num_pages = int(components[7])
|
||||
agent_name = components[8] if len(components) > 8 else ""
|
||||
self._track_kv_arrival(room, chunk_id, is_last, pp_rank)
|
||||
self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank)
|
||||
self._handle_staging_chunk_arrived(
|
||||
room, chunk_idx, page_start, num_pages, agent_name
|
||||
)
|
||||
@@ -1683,10 +1577,12 @@ class NixlKVManager(CommonKVManager):
|
||||
):
|
||||
self._maybe_submit_last_scatter(room)
|
||||
|
||||
def _track_kv_arrival(self, room: int, chunk_id: int, is_last: bool, pp_rank: int):
|
||||
def _track_kv_arrival(
|
||||
self, room: int, chunk_id: int, is_last_chunk: bool, pp_rank: int
|
||||
):
|
||||
"""Update transfer status tracking for a kv chunk arrival."""
|
||||
self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add(chunk_id)
|
||||
if is_last:
|
||||
if is_last_chunk:
|
||||
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = chunk_id + 1
|
||||
if self.transfer_statuses[room].num_pp_ranks_expected is None:
|
||||
self.transfer_statuses[room].num_pp_ranks_expected = (
|
||||
@@ -1827,12 +1723,6 @@ class NixlKVSender(CommonKVSender):
|
||||
self._send_error: Optional[Exception] = None
|
||||
self._transfer_start_time: Optional[float] = None
|
||||
|
||||
def pop_decode_prefix_len(self) -> int:
|
||||
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
|
||||
|
||||
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
||||
return num_pages > 0 or last_chunk
|
||||
|
||||
def send(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
@@ -1841,22 +1731,10 @@ class NixlKVSender(CommonKVSender):
|
||||
if self._send_failed:
|
||||
return
|
||||
|
||||
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
||||
self.curr_idx += len(kv_indices)
|
||||
is_last = self.curr_idx == self.num_kv_indices
|
||||
|
||||
# Special handling for cp
|
||||
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
||||
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
||||
self.kv_mgr,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
kv_indices, index_slice, is_last_chunk, should_skip = (
|
||||
self._prepare_send_indices(kv_indices, state_indices)
|
||||
)
|
||||
elif self.kv_mgr.is_dummy_cp_rank:
|
||||
if not is_last:
|
||||
return
|
||||
else:
|
||||
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
||||
if should_skip:
|
||||
return
|
||||
|
||||
if self._transfer_start_time is None and (
|
||||
@@ -1868,14 +1746,14 @@ class NixlKVSender(CommonKVSender):
|
||||
self.bootstrap_room,
|
||||
kv_indices,
|
||||
index_slice,
|
||||
is_last,
|
||||
is_last_chunk,
|
||||
self.chunk_id,
|
||||
self.aux_index,
|
||||
state_indices,
|
||||
)
|
||||
self._record_transfer_indices(kv_indices, state_indices)
|
||||
self.chunk_id += 1
|
||||
if is_last:
|
||||
if is_last_chunk:
|
||||
self.has_sent = True
|
||||
|
||||
def poll(self) -> KVPoll:
|
||||
@@ -1892,9 +1770,6 @@ class NixlKVSender(CommonKVSender):
|
||||
)
|
||||
return status
|
||||
|
||||
def clear(self):
|
||||
super().clear()
|
||||
|
||||
def failure_exception(self):
|
||||
if self._send_error is not None:
|
||||
raise self._send_error
|
||||
@@ -1915,12 +1790,6 @@ class NixlKVReceiver(CommonKVReceiver):
|
||||
super().__init__(mgr, bootstrap_addr, bootstrap_room)
|
||||
self.init_time = None
|
||||
|
||||
def init(
|
||||
self,
|
||||
prefill_dp_rank: int,
|
||||
):
|
||||
super().init(prefill_dp_rank)
|
||||
|
||||
def send_metadata(
|
||||
self,
|
||||
kv_indices: npt.NDArray[np.int32],
|
||||
@@ -1997,30 +1866,15 @@ class NixlKVReceiver(CommonKVReceiver):
|
||||
if not self.started_transfer:
|
||||
return status
|
||||
|
||||
now = time.time()
|
||||
elapsed = now - self.init_time
|
||||
|
||||
if elapsed >= self.kv_mgr.waiting_timeout:
|
||||
logger.error(f"Request {self.bootstrap_room} waiting_timeout")
|
||||
self.kv_mgr.record_failure(
|
||||
self.bootstrap_room,
|
||||
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput",
|
||||
)
|
||||
self.conclude_state = KVPoll.Failed
|
||||
return KVPoll.Failed
|
||||
timeout_result = self._check_waiting_timeout()
|
||||
if timeout_result is not None:
|
||||
return timeout_result
|
||||
|
||||
self.kv_mgr.update_transfer_status()
|
||||
if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore
|
||||
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard(
|
||||
self.bootstrap_room
|
||||
)
|
||||
# Check if the transfer failed
|
||||
if self.kv_mgr.transfer_statuses[self.bootstrap_room].is_failed():
|
||||
self.conclude_state = KVPoll.Failed
|
||||
logger.error(
|
||||
f"Transfer for room {self.bootstrap_room} failed due to node failure"
|
||||
)
|
||||
else:
|
||||
self.conclude_state = KVPoll.Success
|
||||
del self.kv_mgr.transfer_statuses[self.bootstrap_room]
|
||||
return self.conclude_state # type: ignore
|
||||
|
||||
Reference in New Issue
Block a user