Cleanup: relocate temp_set_env and consolidate multi-device/CUDA helpers in common.py (#30180)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+885
-834
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user