[Env] centralize pd vars in environ.py (#16264)
This commit is contained in:
@@ -32,8 +32,9 @@ from sglang.srt.disaggregation.mooncake.utils import (
|
||||
check_mooncake_custom_mem_pool_enabled,
|
||||
)
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import format_tcp_address, get_int_env_var, is_valid_ipv6_address
|
||||
from sglang.srt.utils import format_tcp_address, is_valid_ipv6_address
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -170,11 +171,12 @@ class MooncakeKVManager(CommonKVManager):
|
||||
self.session_lock = threading.Lock()
|
||||
# Determine the number of threads to use for kv sender
|
||||
cpu_count = os.cpu_count()
|
||||
transfer_thread_pool_size = get_int_env_var(
|
||||
"SGLANG_DISAGGREGATION_THREAD_POOL_SIZE",
|
||||
min(max(4, int(0.5 * cpu_count) // 8), 12),
|
||||
transfer_thread_pool_size = (
|
||||
envs.SGLANG_DISAGGREGATION_THREAD_POOL_SIZE.get()
|
||||
)
|
||||
transfer_queue_size = get_int_env_var("SGLANG_DISAGGREGATION_QUEUE_SIZE", 4)
|
||||
if transfer_thread_pool_size is None:
|
||||
transfer_thread_pool_size = min(max(4, int(0.5 * cpu_count) // 8), 12)
|
||||
transfer_queue_size = envs.SGLANG_DISAGGREGATION_QUEUE_SIZE.get()
|
||||
self.transfer_queues: List[FastQueue] = [
|
||||
FastQueue() for _ in range(transfer_queue_size)
|
||||
]
|
||||
@@ -195,9 +197,7 @@ class MooncakeKVManager(CommonKVManager):
|
||||
# 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 = get_int_env_var(
|
||||
"SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300
|
||||
)
|
||||
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()
|
||||
@@ -210,19 +210,17 @@ class MooncakeKVManager(CommonKVManager):
|
||||
self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set)
|
||||
# Heartbeat interval should be at least 2 seconds
|
||||
self.heartbeat_interval = max(
|
||||
float(os.getenv("SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL", 5.0)), 2.0
|
||||
envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0
|
||||
)
|
||||
# Heartbeat failure should be at least 1
|
||||
self.max_failures = max(
|
||||
get_int_env_var("SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE", 2), 1
|
||||
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 = get_int_env_var(
|
||||
"SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300
|
||||
)
|
||||
self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get()
|
||||
|
||||
self.failure_records: Dict[int, str] = {}
|
||||
self.failure_lock = threading.Lock()
|
||||
|
||||
Reference in New Issue
Block a user