Cleanup: relocate temp_set_env and consolidate multi-device/CUDA helpers in common.py (#30180)

This commit is contained in:
Lianmin Zheng
2026-07-05 18:47:26 -07:00
committed by GitHub
parent 8673e85e6c
commit 9d00385b63
7 changed files with 890 additions and 872 deletions
@@ -2129,11 +2129,6 @@ def initialize_model_parallel(
logger.info(
f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}"
)
else:
if get_tensor_model_parallel_rank() == 0:
logger.info(
f"DCP disabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}"
)
attn_dp_size = attention_data_parallel_size
attn_cp_size = attention_context_model_parallel_size
-29
View File
@@ -24,35 +24,6 @@ def _default_hip() -> bool:
return False
@contextmanager
def temp_set_env(*, allow_sglang: bool = False, **env_vars: Any):
"""Temporarily set environment variables, restoring originals on exit.
By default, SGLANG_*/SGL_* keys are rejected — use ``Envs`` descriptors
for those. Pass ``allow_sglang=True`` only for special env vars that
intentionally bypass ``environ.py``.
"""
if not allow_sglang:
for key in env_vars:
if key.startswith("SGLANG_") or key.startswith("SGL_"):
raise ValueError("temp_set_env should not be used for sglang env vars")
backup = {key: os.environ.get(key) for key in env_vars}
try:
for key, value in env_vars.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = str(value)
yield
finally:
for key, value in backup.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
class EnvField:
_allow_set_name = True
+2 -1
View File
@@ -91,7 +91,7 @@ from sglang.srt.utils.common import is_cuda_alike
DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
0.8 # Reserve 20% GPU memory headroom for ModelOpt calibration
)
from sglang.srt.environ import envs, temp_set_env
from sglang.srt.environ import envs
from sglang.srt.model_loader.weight_utils import (
buffered_multi_thread_safetensors_weights_iterator,
download_safetensors_index_file_from_hf,
@@ -119,6 +119,7 @@ from sglang.srt.utils import (
rank0_log,
set_weight_attrs,
)
from sglang.srt.utils.common import temp_set_env
if TYPE_CHECKING:
from sglang.srt.configs.device_config import DeviceConfig
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -7,8 +7,8 @@ import tempfile
from pathlib import Path
from types import SimpleNamespace
from sglang.srt.environ import temp_set_env
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.common import temp_set_env
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,