Clean up __init__ function of the scheduler and event loop for PD (#15298)

python/sglang/srt/utils/common.py

@@ -49,7 +49,7 @@ import warnings
 from collections import OrderedDict, defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from functools import lru_cache
+from functools import lru_cache, partial
 from importlib.metadata import PackageNotFoundError, version
 from importlib.util import find_spec
 from io import BytesIO
@@ -105,22 +105,6 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-show_time_cost = False
-time_infos = {}
-
-
-def get_or_create_event_loop():
-    """Gets the running event loop or creates a new one if it doesn't exist."""
-    try:
-        return asyncio.get_running_loop()
-    except RuntimeError:
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
-        return loop
-
-
-HIP_FP8_E4M3_FNUZ_MAX = 224.0
-
 
 # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
 @lru_cache(maxsize=1)
@@ -129,6 +113,7 @@ def is_hip() -> bool:
 
 
 if is_hip():
+    HIP_FP8_E4M3_FNUZ_MAX = 224.0
     FP8_E4M3_MAX = HIP_FP8_E4M3_FNUZ_MAX
 else:
     FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -191,14 +176,6 @@ def get_cuda_version():
     return (0, 0)
 
 
-def _check(cc_major):
-    if not is_cuda():
-        return False
-    return torch.cuda.get_device_capability()[0] == cc_major and tuple(
-        map(int, torch.version.cuda.split(".")[:2])
-    ) >= (12, 3)
-
-
 @contextmanager
 def device_context(device: torch.device):
     if device.type == "cpu" and is_cpu():
@@ -213,96 +190,49 @@ def device_context(device: torch.device):
             raise ValueError(f"Unknown device module: {device}")
 
 
-is_ampere_with_cuda_12_3 = lambda: _check(8)
-is_hopper_with_cuda_12_3 = lambda: _check(9)
-
-
-@lru_cache(maxsize=1)
-def is_blackwell():
+def _check_cuda_device_version(
+    device_capability_majors: List[int], cuda_version: Tuple[int, int]
+):
     if not is_cuda():
         return False
-    return torch.cuda.get_device_capability()[0] in [10, 12]
-
-
-@lru_cache(maxsize=1)
-def is_blackwell_supported(device=None) -> bool:
-    if not is_cuda():
-        return False
-    return is_sm100_supported(device) or is_sm120_supported(device)
-
-
-@lru_cache(maxsize=1)
-def is_sm120_supported(device=None) -> bool:
-    if not is_cuda():
-        return False
-    return (torch.cuda.get_device_capability(device)[0] == 12) and (
-        torch.version.cuda >= "12.8"
+    return (
+        torch.cuda.get_device_capability()[0] in device_capability_majors
+        and tuple(map(int, torch.version.cuda.split(".")[:2])) >= cuda_version
     )
 
 
-@lru_cache(maxsize=1)
-def is_sm100_supported(device=None) -> bool:
-    if not is_cuda():
-        return False
-    return (torch.cuda.get_device_capability(device)[0] == 10) and (
-        torch.version.cuda >= "12.8"
+is_ampere_with_cuda_12_3 = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version, device_capability_majors=[8], cuda_version=(12, 3)
     )
-
-
-@lru_cache(maxsize=1)
-def is_sm90_supported(device=None) -> bool:
-    if not is_cuda():
-        return False
-    return (torch.cuda.get_device_capability(device)[0] == 9) and (
-        torch.version.cuda >= "12.3"
+)
+is_hopper_with_cuda_12_3 = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)
     )
-
-
-_warned_bool_env_var_keys = set()
-
-
-def get_bool_env_var(name: str, default: str = "false") -> bool:
-    # FIXME: move your environment variable to sglang.srt.environ
-    value = os.getenv(name, default)
-    value = value.lower()
-
-    truthy_values = ("true", "1")
-    falsy_values = ("false", "0")
-
-    if (value not in truthy_values) and (value not in falsy_values):
-        if value not in _warned_bool_env_var_keys:
-            logger.warning(
-                f"get_bool_env_var({name}) see non-understandable value={value} and treat as false"
-            )
-        _warned_bool_env_var_keys.add(value)
-
-    return value in truthy_values
-
-
-def get_int_env_var(name: str, default: int = 0) -> int:
-    # FIXME: move your environment variable to sglang.srt.environ
-    value = os.getenv(name)
-    if value is None or not value.strip():
-        return default
-    try:
-        return int(value)
-    except ValueError:
-        return default
-
-
-def get_float_env_var(name: str, default: float = 0.0) -> float:
-    # FIXME: move your environment variable to sglang.srt.environ
-    value = os.getenv(name)
-    if value is None or not value.strip():
-        return default
-    try:
-        return float(value)
-    except ValueError:
-        return default
-
-
-def support_triton(backend: str) -> bool:
-    return backend not in ["torch_native", "intel_amx"]
+)
+is_blackwell_supported = is_blackwell = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version,
+        device_capability_majors=[10, 12],
+        cuda_version=(12, 8),
+    )
+)
+is_sm120_supported = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version, device_capability_majors=[12], cuda_version=(12, 8)
+    )
+)
+is_sm100_supported = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8)
+    )
+)
+is_sm90_supported = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version, device_capability_majors=[9], cuda_version=(12, 3)
+    )
+)
 
 
 try:
@@ -364,6 +294,53 @@ def random_uuid() -> str:
     return str(uuid.uuid4().hex)
 
 
+_warned_bool_env_var_keys = set()
+
+
+def get_bool_env_var(name: str, default: str = "false") -> bool:
+    # FIXME: move your environment variable to sglang.srt.environ
+    value = os.getenv(name, default)
+    value = value.lower()
+
+    truthy_values = ("true", "1")
+    falsy_values = ("false", "0")
+
+    if (value not in truthy_values) and (value not in falsy_values):
+        if value not in _warned_bool_env_var_keys:
+            logger.warning(
+                f"get_bool_env_var({name}) see non-understandable value={value} and treat as false"
+            )
+        _warned_bool_env_var_keys.add(value)
+
+    return value in truthy_values
+
+
+def get_int_env_var(name: str, default: int = 0) -> int:
+    # FIXME: move your environment variable to sglang.srt.environ
+    value = os.getenv(name)
+    if value is None or not value.strip():
+        return default
+    try:
+        return int(value)
+    except ValueError:
+        return default
+
+
+def get_float_env_var(name: str, default: float = 0.0) -> float:
+    # FIXME: move your environment variable to sglang.srt.environ
+    value = os.getenv(name)
+    if value is None or not value.strip():
+        return default
+    try:
+        return float(value)
+    except ValueError:
+        return default
+
+
+def support_triton(backend: str) -> bool:
+    return backend not in ["torch_native", "intel_amx"]
+
+
 _ENABLE_TORCH_INFERENCE_MODE = get_bool_env_var(
     "SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false"
 )
@@ -421,6 +398,10 @@ class DynamicGradMode(_DecoratorContextManager):
             return self.__class__()
 
 
+show_time_cost = False
+time_infos = {}
+
+
 def enable_show_time_cost():
     global show_time_cost
     show_time_cost = True
@@ -3778,3 +3759,13 @@ def raise_error_or_warn(obj, strict, counter_name, message, log_interval=1000):
         if count % log_interval == 0:
             logger.warning(message)
         setattr(obj, counter_name, count + 1)
+
+
+def get_or_create_event_loop():
+    """Gets the running event loop or creates a new one if it doesn't exist."""
+    try:
+        return asyncio.get_running_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        return loop