From 48c2aca9ba739a0c98b059698237a133f290f07f Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Wed, 14 Jan 2026 23:06:01 +0800 Subject: [PATCH] [Env] centralize pd vars in environ.py (#16264) --- .../srt/disaggregation/mooncake/conn.py | 24 +++++++++---------- python/sglang/srt/disaggregation/nixl/conn.py | 11 ++++----- python/sglang/srt/environ.py | 9 +++++++ 3 files changed, 24 insertions(+), 20 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 32e8c0b69..6b5ae37e6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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() diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index bfe7695bd..d3e390c92 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -2,7 +2,6 @@ from __future__ import annotations import dataclasses import logging -import os import struct import threading import time @@ -23,8 +22,8 @@ from sglang.srt.disaggregation.common.conn import ( ) from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous 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 get_int_env_var logger = logging.getLogger(__name__) @@ -161,15 +160,13 @@ class NixlKVManager(CommonKVManager): # 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 - ) - self.waiting_timeout = get_int_env_var( - "SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300 + 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( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1e3cf8c61..ee156befb 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -221,6 +221,15 @@ class Envs: SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_TARGET_VERIFY = EnvInt(1) SGLANG_SCHEDULER_RECV_SKIPPER_WEIGHT_NONE = EnvInt(1) + # PD Disaggregation (runtime) + # NOTE: For SGLANG_DISAGGREGATION_THREAD_POOL_SIZE, the effective default is + # computed dynamically at runtime based on cpu_count; see disaggregation backends. + SGLANG_DISAGGREGATION_THREAD_POOL_SIZE = EnvInt(None) + SGLANG_DISAGGREGATION_QUEUE_SIZE = EnvInt(4) + SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT = EnvInt(300) + SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0) + SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) + SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) # Scheduler: others: SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period.