diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 844c847af..60f4cc279 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -6,8 +6,9 @@ import logging import socket import threading import time +from collections import defaultdict from functools import cache -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Set, Tuple, Union import numpy as np import numpy.typing as npt @@ -107,11 +108,32 @@ class CommonKVManager(BaseKVManager): self.transfer_infos = {} self.decode_kv_args_table = {} self.pp_group = get_pp_group() + # 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() elif self.disaggregation_mode == DisaggregationMode.DECODE: self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {} self.connection_lock = threading.Lock() self.required_prefill_response_num_table: Dict[int, int] = {} self.prefill_info_table: Dict[str, PrefillServerInfo] = {} + self.heartbeat_failures: Dict[str, int] = {} + self.session_pool: Dict = defaultdict(requests.Session) + self.session_pool_lock = threading.Lock() + self.addr_to_rooms_tracker: Dict[str, Set[int]] = 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 + ) + # 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() else: raise ValueError( f"Unsupported DisaggregationMode: {self.disaggregation_mode}" diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d05fae9a6..706d32fcc 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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() diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 830ca9e79..9e0d2d5b9 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -7,8 +7,7 @@ import os import struct import threading import time -from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple +from typing import Dict, List, Optional, Tuple import msgspec import numpy as np @@ -194,16 +193,8 @@ class MoriKVManager(CommonKVManager): self.transfer_lock = threading.Lock() self._register_local_buffers() if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.bootstrap_timeout = get_int_env_var( - "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300 - ) self._start_bootstrap_thread() elif self.disaggregation_mode == DisaggregationMode.DECODE: - self.waiting_timeout = get_int_env_var( - "SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300 - ) - self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set) - self.addr_to_rooms_tracker = defaultdict(set) self.room_to_bootstrap_addr: Dict[int, str] = {} self._start_decode_thread() diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 8a8807cef..279dfab90 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -11,7 +11,6 @@ from typing import Dict, List, Optional, Set 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 ( @@ -193,21 +192,6 @@ class NixlKVManager(CommonKVManager): self.transfer_statuses: Dict[int, TransferStatus] = defaultdict( TransferStatus ) - self.heartbeat_failures = {} - self.session_pool = defaultdict(requests.Session) - self.session_pool_lock = threading.Lock() - self.addr_to_rooms_tracker = defaultdict(set) - self.connection_lock = threading.Lock() - - # 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.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get() self._start_heartbeat_checker_thread() else: raise ValueError(