[PD-Disagg] Improve KVManager init across all backends (#19240)
Signed-off-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
@@ -9,11 +9,10 @@ import struct
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import requests
|
||||
|
||||
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
|
||||
from sglang.srt.disaggregation.common.conn import (
|
||||
@@ -206,33 +205,11 @@ class MooncakeKVManager(CommonKVManager):
|
||||
threading.Thread(
|
||||
target=self.transfer_worker, args=(queue, executor), daemon=True
|
||||
).start()
|
||||
# If a timeout happens on the prefill side, it means prefill instances
|
||||
# fail to receive the KV indices from the decode instance of this request.
|
||||
# These timeout requests should be aborted to release the tree cache.
|
||||
self.bootstrap_timeout = envs.SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get()
|
||||
|
||||
self.enable_custom_mem_pool, self.custom_mem_pool_type = (
|
||||
check_mooncake_custom_mem_pool_enabled()
|
||||
)
|
||||
elif self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
self.heartbeat_failures = {}
|
||||
self.session_pool = defaultdict(requests.Session)
|
||||
self.session_pool_lock = threading.Lock()
|
||||
self.addr_to_rooms_tracker = defaultdict(set)
|
||||
self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
|
||||
# Heartbeat interval should be at least 2 seconds
|
||||
self.heartbeat_interval = max(
|
||||
envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0
|
||||
)
|
||||
# Heartbeat failure should be at least 1
|
||||
self.max_failures = max(
|
||||
envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1
|
||||
)
|
||||
self.start_decode_thread()
|
||||
# If a timeout happens on the decode side, it means decode instances
|
||||
# fail to receive the KV Cache transfer done signal after bootstrapping.
|
||||
# These timeout requests should be aborted to release the tree cache.
|
||||
self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
|
||||
|
||||
def init_engine(self):
|
||||
self.engine = get_mooncake_transfer_engine()
|
||||
|
||||
Reference in New Issue
Block a user