chore: import upstream snapshot with attribution

python/sglang/check_env.py

@@ -0,0 +1,595 @@
+"""Check environment configurations and dependency versions."""
+
+import importlib.metadata
+import os
+import resource
+import subprocess
+import sys
+from abc import abstractmethod
+from collections import OrderedDict, defaultdict
+
+import torch
+
+from sglang.srt.utils import is_hip, is_mps, is_musa, is_npu
+
+
+def is_cuda_v2():
+    return torch.version.cuda is not None
+
+
+# List of packages to check versions
+PACKAGE_LIST = [
+    "sglang",
+    "sglang-kernel",
+    "flashinfer_python",
+    "flashinfer_cubin",
+    "flashinfer_jit_cache",
+    "triton",
+    "transformers",
+    "torchao",
+    "numpy",
+    "aiohttp",
+    "fastapi",
+    "huggingface_hub",
+    "interegular",
+    "modelscope",
+    "orjson",
+    "outlines",
+    "packaging",
+    "psutil",
+    "pydantic",
+    "python-multipart",
+    "pyzmq",
+    "torchao",
+    "uvicorn",
+    "uvloop",
+    "vllm",
+    "xgrammar",
+    "openai",
+    "tiktoken",
+    "anthropic",
+    "litellm",
+    "torchcodec",
+]
+
+
+class BaseEnv:
+    """Base class for environment check"""
+
+    def __init__(self):
+        self.package_list = PACKAGE_LIST
+
+    @abstractmethod
+    def get_info(self) -> dict:
+        """
+        Get CUDA-related information if available.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_topology(self) -> dict:
+        raise NotImplementedError
+
+    def get_package_versions(self) -> dict:
+        """
+        Get versions of specified packages.
+        """
+        versions = {}
+        for package in self.package_list:
+            package_name = package.split("==")[0].split(">=")[0].split("<=")[0]
+            try:
+                version = importlib.metadata.version(package_name)
+                versions[package_name] = version
+            except ModuleNotFoundError:
+                versions[package_name] = "Module Not Found"
+        return versions
+
+    def get_device_info(self):
+        """
+        Get information about available GPU devices.
+        """
+        devices = defaultdict(list)
+        capabilities = defaultdict(list)
+        for k in range(torch.cuda.device_count()):
+            devices[torch.cuda.get_device_name(k)].append(str(k))
+            capability = torch.cuda.get_device_capability(k)
+            capabilities[f"{capability[0]}.{capability[1]}"].append(str(k))
+
+        gpu_info = {}
+        for name, device_ids in devices.items():
+            gpu_info[f"GPU {','.join(device_ids)}"] = name
+
+        if len(capabilities) == 1:
+            # All GPUs have the same compute capability
+            cap, gpu_ids = list(capabilities.items())[0]
+            gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
+        else:
+            # GPUs have different compute capabilities
+            for cap, gpu_ids in capabilities.items():
+                gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
+
+        return gpu_info
+
+    def get_hypervisor_vendor(self) -> dict:
+        try:
+            output = subprocess.check_output(["lscpu"], text=True)
+            for line in output.split("\n"):
+                if "Hypervisor vendor:" in line:
+                    return {"Hypervisor vendor:": line.split(":")[1].strip()}
+            return {}
+        except:
+            return {}
+
+    def get_ulimit_soft(self) -> dict:
+        ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
+        return {"ulimit soft": ulimit_soft}
+
+    def check_env(self):
+        """
+        Check and print environment information.
+        """
+        env_info = OrderedDict()
+        env_info["Python"] = sys.version.replace("\n", "")
+        env_info.update(self.get_info())
+        env_info["PyTorch"] = torch.__version__
+        env_info.update(self.get_package_versions())
+        env_info.update(self.get_topology())
+        env_info.update(self.get_hypervisor_vendor())
+        env_info.update(self.get_ulimit_soft())
+
+        for k, v in env_info.items():
+            print(f"{k}: {v}")
+
+
+class GPUEnv(BaseEnv):
+    """Environment checker for Nvidia GPU"""
+
+    def get_info(self):
+        cuda_info = {"CUDA available": torch.cuda.is_available()}
+
+        if cuda_info["CUDA available"]:
+            cuda_info.update(self.get_device_info())
+            cuda_info.update(self._get_cuda_version_info())
+
+        return cuda_info
+
+    def _get_cuda_version_info(self):
+        """
+        Get CUDA version information.
+        """
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        cuda_info = {"CUDA_HOME": CUDA_HOME}
+
+        if CUDA_HOME and os.path.isdir(CUDA_HOME):
+            cuda_info.update(self._get_nvcc_info())
+            cuda_info.update(self._get_cuda_driver_version())
+
+        return cuda_info
+
+    def _get_nvcc_info(self):
+        """
+        Get NVCC version information.
+        """
+        from torch.utils.cpp_extension import CUDA_HOME
+
+        try:
+            nvcc = os.path.join(CUDA_HOME, "bin/nvcc")
+            nvcc_output = (
+                subprocess.check_output(f'"{nvcc}" -V', shell=True)
+                .decode("utf-8")
+                .strip()
+            )
+            return {
+                "NVCC": nvcc_output[
+                    nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind(
+                        "Build"
+                    )
+                ].strip()
+            }
+        except subprocess.SubprocessError:
+            return {"NVCC": "Not Available"}
+
+    def _get_cuda_driver_version(self):
+        """
+        Get CUDA driver version.
+        """
+        from sglang.srt.utils.common import get_nvidia_driver_version_str
+
+        ver = get_nvidia_driver_version_str()
+        if ver is None:
+            return {"CUDA Driver Version": "Not Available"}
+        return {"CUDA Driver Version": ver}
+
+    def get_topology(self):
+        """
+        Get GPU topology information.
+        """
+        try:
+            result = subprocess.run(
+                ["nvidia-smi", "topo", "-m"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                check=True,
+            )
+            return {
+                "NVIDIA Topology": (
+                    "\n" + result.stdout if result.returncode == 0 else None
+                )
+            }
+        except subprocess.SubprocessError:
+            return {}
+
+
+class HIPEnv(BaseEnv):
+    """Environment checker for ROCm/HIP"""
+
+    def get_info(self):
+        cuda_info = {"ROCM available": torch.cuda.is_available()}
+
+        if cuda_info["ROCM available"]:
+            cuda_info.update(self.get_device_info())
+            cuda_info.update(self._get_cuda_version_info())
+
+        return cuda_info
+
+    def _get_cuda_version_info(self):
+        from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME
+
+        cuda_info = {"ROCM_HOME": ROCM_HOME}
+
+        if ROCM_HOME and os.path.isdir(ROCM_HOME):
+            cuda_info.update(self._get_hipcc_info())
+            cuda_info.update(self._get_rocm_driver_version())
+
+        return cuda_info
+
+    def _get_hipcc_info(self):
+        from torch.utils.cpp_extension import ROCM_HOME
+
+        try:
+            hipcc = os.path.join(ROCM_HOME, "bin/hipcc")
+            hipcc_output = (
+                subprocess.check_output(f'"{hipcc}" --version', shell=True)
+                .decode("utf-8")
+                .strip()
+            )
+            return {
+                "HIPCC": hipcc_output[
+                    hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang")
+                ].strip()
+            }
+        except subprocess.SubprocessError:
+            return {"HIPCC": "Not Available"}
+
+    def _get_rocm_driver_version(self):
+        try:
+            output = subprocess.check_output(
+                [
+                    "rocm-smi",
+                    "--showdriverversion",
+                    "--csv",
+                ]
+            )
+            versions = set(output.decode().strip().split("\n"))
+            versions.discard("name, value")
+            ver = versions.pop()
+            ver = ver.replace('"Driver version", ', "").replace('"', "")
+
+            return {"ROCM Driver Version": ver}
+        except subprocess.SubprocessError:
+            return {"ROCM Driver Version": "Not Available"}
+
+    def get_topology(self):
+        try:
+            result = subprocess.run(
+                ["rocm-smi", "--showtopotype"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                check=True,
+            )
+            return {
+                "AMD Topology": "\n" + result.stdout if result.returncode == 0 else None
+            }
+        except subprocess.SubprocessError:
+            return {}
+
+
+class NPUEnv(BaseEnv):
+    """Environment checker for Ascend NPU"""
+
+    EXTRA_PACKAGE_LIST = [
+        "torch_npu",
+        "sgl-kernel-npu",
+        "deep_ep",
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.package_list.extend(NPUEnv.EXTRA_PACKAGE_LIST)
+
+    def get_info(self):
+        cuda_info = {"NPU available": torch.npu.is_available()}
+        if cuda_info["NPU available"]:
+            cuda_info.update(self.get_device_info())
+            cuda_info.update(self._get_cann_version_info())
+
+        return cuda_info
+
+    def get_device_info(self):
+        """
+        Get information about available NPUs.
+        Need to override due to torch_npu interface differences.
+        """
+        devices = defaultdict(list)
+        for k in range(torch.npu.device_count()):
+            devices[torch.npu.get_device_name(k)].append(str(k))
+
+        npu_info = {}
+        for name, device_ids in devices.items():
+            npu_info[f"NPU {','.join(device_ids)}"] = name
+
+        return npu_info
+
+    def _get_cann_version_info(self):
+        cann_envs = ["ASCEND_TOOLKIT_HOME", "ASCEND_INSTALL_PATH"]
+        for var in cann_envs:
+            path = os.environ.get(var)
+            if path and os.path.exists(path):
+                CANN_HOME = path
+                break
+        else:
+            default_path = "/usr/local/Ascend/ascend-toolkit/latest"
+            CANN_HOME = default_path if os.path.exists(default_path) else None
+
+        if CANN_HOME:
+            npu_info = {"CANN_HOME": CANN_HOME}
+            npu_info.update(self._get_cann_info(CANN_HOME))
+            npu_info.update(self._get_ascend_driver_version())
+            return npu_info
+        else:
+            return {"CANN_HOME": "Not found"}
+
+    def _get_cann_info(self, CANN_HOME: str):
+        cann_info = {}
+        cann_version_file = os.path.join(CANN_HOME, "version.cfg")
+        if os.path.exists(cann_version_file):
+            with open(cann_version_file, "r", encoding="utf-8") as f:
+                f.readline()  # discard first line comment in version.cfg
+                cann_info["CANN"] = f.readline().split("[")[1].split("]")[0]
+        else:
+            cann_info["CANN"] = "Not Available"
+        try:
+            bisheng = os.path.join(CANN_HOME, "tools/bisheng_compiler/bin/bisheng")
+            if not os.path.isfile(bisheng):
+                # Check path for old CANN version
+                bisheng = os.path.join(CANN_HOME, "compiler/ccec_compiler/bin/bisheng")
+            bisheng_output = (
+                subprocess.check_output([bisheng, "--version"]).decode("utf-8").strip()
+            )
+            cann_info["BiSheng"] = bisheng_output.split("\n")[0].strip()
+        except subprocess.SubprocessError:
+            cann_info["BiSheng"] = "Not Available"
+        return cann_info
+
+    def _get_ascend_driver_version(self):
+        try:
+            output = subprocess.check_output(
+                [
+                    "npu-smi",
+                    "info",
+                    "-t",
+                    "board",
+                    "-i",
+                    "0",
+                ]
+            )
+            for line in output.decode().strip().split("\n"):
+                if "Software Version" in line:
+                    version = line.split(":")[-1].strip()
+                    break
+            else:
+                version = "Not Available"
+
+            return {"Ascend Driver Version": version}
+        except subprocess.SubprocessError:
+            return {"Ascend Driver Version": "Not Available"}
+
+    def get_topology(self):
+        try:
+            result = subprocess.run(
+                ["npu-smi", "info", "-t", "topo"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                check=True,
+            )
+            return {
+                "Ascend Topology": (
+                    "\n" + result.stdout if result.returncode == 0 else None
+                )
+            }
+        except subprocess.SubprocessError:
+            return {}
+
+
+class MUSAEnv(BaseEnv):
+    """Environment checker for MThreads GPU"""
+
+    def get_info(self):
+        musa_info = {"MUSA available": torch.musa.is_available()}
+
+        if musa_info["MUSA available"]:
+            musa_info.update(self.get_device_info())
+            musa_info.update(self._get_musa_version_info())
+
+        return musa_info
+
+    def _get_musa_version_info(self):
+        """
+        Get MUSA version information.
+        """
+        from torch_musa.utils.musa_extension import MUSA_HOME
+
+        musa_info = {"MUSA_HOME": MUSA_HOME}
+
+        if MUSA_HOME and os.path.isdir(MUSA_HOME):
+            musa_info.update(self._get_mcc_info())
+            musa_info.update(self._get_musa_driver_version())
+
+        return musa_info
+
+    def _get_mcc_info(self):
+        """
+        Get MCC version information.
+        """
+        from torch_musa.utils.musa_extension import MUSA_HOME
+
+        try:
+            mcc = os.path.join(MUSA_HOME, "bin/mcc")
+            mcc_output = (
+                subprocess.check_output(f'"{mcc}" --version', shell=True)
+                .decode("utf-8")
+                .strip()
+            )
+            return {
+                "MCC": mcc_output[
+                    mcc_output.rfind("mcc version") : mcc_output.rfind("Target")
+                ].strip()
+            }
+        except subprocess.SubprocessError:
+            return {"MCC": "Not Available"}
+
+    def _get_musa_driver_version(self):
+        """
+        Get MUSA driver version.
+        """
+        try:
+            output = subprocess.check_output(
+                [
+                    "mthreads-gmi",
+                    "-q",
+                ],
+                text=True,
+            )
+            driver_version = None
+            for line in output.splitlines():
+                if "Driver Version" in line:
+                    driver_version = line.split(":", 1)[1].strip()
+                    break
+
+            return {"MUSA Driver Version": driver_version}
+        except subprocess.SubprocessError:
+            return {"MUSA Driver Version": "Not Available"}
+
+    def get_topology(self):
+        """
+        Get GPU topology information.
+        """
+        try:
+            result = subprocess.run(
+                ["mthreads-gmi", "topo", "-m"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                check=True,
+            )
+            return {
+                "MTHREADS Topology": (
+                    "\n" + result.stdout if result.returncode == 0 else None
+                )
+            }
+        except subprocess.SubprocessError:
+            return {}
+
+
+class MPSEnv(BaseEnv):
+    """Environment checker for Apple Silicon MPS"""
+
+    EXTRA_PACKAGE_LIST = ["mlx", "mlx-lm", "mlx-metal"]
+
+    def __init__(self):
+        super().__init__()
+        self.package_list.extend(MPSEnv.EXTRA_PACKAGE_LIST)
+
+    def get_info(self):
+        import platform
+
+        info = {"MPS available": torch.backends.mps.is_available()}
+        if not info["MPS available"]:
+            return info
+
+        info["macOS Version"] = platform.mac_ver()[0]
+
+        try:
+            info["macOS Build"] = subprocess.check_output(
+                ["sw_vers", "-buildVersion"], text=True
+            ).strip()
+        except Exception:
+            info["macOS Build"] = "Not Available"
+
+        for label, key in [
+            ("Apple Silicon", "machdep.cpu.brand_string"),
+            ("Unified Memory", "hw.memsize"),
+            ("CPU Cores (Total)", "hw.ncpu"),
+        ]:
+            try:
+                info[label] = subprocess.check_output(
+                    ["sysctl", "-n", key], text=True
+                ).strip()
+            except Exception:
+                info[label] = "Not Available"
+
+        try:
+            mem_bytes = int(info["Unified Memory"])
+            info["Unified Memory"] = f"{mem_bytes / 1024**3:.1f} GB"
+        except Exception:
+            pass
+
+        for label, key in [
+            ("CPU Cores (Performance)", "hw.perflevel0.logicalcpu"),
+            ("CPU Cores (Efficiency)", "hw.perflevel1.logicalcpu"),
+        ]:
+            try:
+                info[label] = subprocess.check_output(
+                    ["sysctl", "-n", key], text=True
+                ).strip()
+            except Exception:
+                pass
+
+        # Single system_profiler call for both Metal support and GPU cores
+        info["Metal Support"] = "Not Available"
+        info["GPU Cores"] = "Not Available"
+        try:
+            sp = subprocess.check_output(
+                ["system_profiler", "SPDisplaysDataType"], text=True
+            )
+            for line in sp.splitlines():
+                line = line.strip()
+                if "Metal Support" in line or "Metal Family" in line:
+                    info["Metal Support"] = line.partition(":")[2].strip()
+                if "Total Number of Cores" in line:
+                    info["GPU Cores"] = line.partition(":")[2].strip()
+        except Exception:
+            pass
+
+        return info
+
+    def get_topology(self):
+        return {}
+
+
+if __name__ == "__main__":
+    if is_cuda_v2():
+        env = GPUEnv()
+    elif is_hip():
+        env = HIPEnv()
+    elif is_npu():
+        env = NPUEnv()
+    elif is_musa():
+        env = MUSAEnv()
+    elif is_mps():
+        env = MPSEnv()
+    env.check_env()