|
| 1 | +""" |
| 2 | +@author: cunyue |
| 3 | +@file: hardware.py |
| 4 | +@time: 2024/11/18 15:12 |
| 5 | +@description: 硬件信息采集 |
| 6 | +""" |
| 7 | + |
| 8 | +import json |
| 9 | +import multiprocessing |
| 10 | +import platform |
| 11 | +import subprocess |
| 12 | + |
| 13 | +import psutil |
| 14 | +import pynvml |
| 15 | + |
| 16 | + |
| 17 | +def get_hardware_info(): |
| 18 | + """ |
| 19 | + 采集硬件信息，包括CPU、GPU、内存、硬盘等 |
| 20 | + """ |
| 21 | + info = { |
| 22 | + "memory": get_memory_size(), |
| 23 | + "cpu": get_cpu_info(), |
| 24 | + "gpu": { |
| 25 | + "nvidia": get_nvidia_gpu_info(), |
| 26 | + }, |
| 27 | + "soc": { |
| 28 | + "apple": get_apple_chip_info(), |
| 29 | + }, |
| 30 | + } |
| 31 | + return info |
| 32 | + |
| 33 | + |
| 34 | +# ---------------------------------- cpu信息 ---------------------------------- |
| 35 | + |
| 36 | + |
| 37 | +def get_cpu_info(): |
| 38 | + """获取 CPU 信息""" |
| 39 | + info = {"brand": None, "cores": None} |
| 40 | + |
| 41 | + # 获取 CPU 品牌, 根据不同操作系统调用不同的函数 |
| 42 | + if platform.system() == "Windows": |
| 43 | + info["brand"] = get_cpu_brand_windows() |
| 44 | + elif platform.system() == "Linux": |
| 45 | + info["brand"] = get_cpu_brand_linux() |
| 46 | + else: |
| 47 | + # 其他情况，暂时不支持 |
| 48 | + # 苹果芯片单独处理 |
| 49 | + return None |
| 50 | + try: |
| 51 | + # 获取 CPU 核心数 |
| 52 | + info["cores"] = multiprocessing.cpu_count() |
| 53 | + except Exception: # noqa |
| 54 | + pass |
| 55 | + |
| 56 | + return info |
| 57 | + |
| 58 | + |
| 59 | +def get_cpu_brand_windows(): |
| 60 | + try: |
| 61 | + # 使用 WMIC 命令获取 CPU 品牌 |
| 62 | + result = subprocess.run(["wmic", "cpu", "get", "name"], capture_output=True, text=True) |
| 63 | + cpu_brand = result.stdout.strip().split("\n")[-1].strip() |
| 64 | + return cpu_brand |
| 65 | + except Exception: # noqa |
| 66 | + return None |
| 67 | + |
| 68 | + |
| 69 | +def get_cpu_brand_linux(): |
| 70 | + try: |
| 71 | + # 使用 lscpu 命令获取 CPU 品牌 |
| 72 | + result = subprocess.run(["lscpu"], capture_output=True, text=True) |
| 73 | + for line in result.stdout.split("\n"): |
| 74 | + if "Model name:" in line: |
| 75 | + cpu_brand = line.split(":")[1].strip() |
| 76 | + return cpu_brand |
| 77 | + return None |
| 78 | + except Exception: # noqa |
| 79 | + return None |
| 80 | + |
| 81 | + |
| 82 | +# ---------------------------------- 内存信息 ---------------------------------- |
| 83 | + |
| 84 | + |
| 85 | +def get_memory_size(): |
| 86 | + """获取内存大小""" |
| 87 | + try: |
| 88 | + # 获取系统总内存大小 |
| 89 | + mem = psutil.virtual_memory() |
| 90 | + total_memory = round(mem.total / (1024**3)) # 单位为GB |
| 91 | + return total_memory |
| 92 | + except Exception: # noqa |
| 93 | + return |
| 94 | + |
| 95 | + |
| 96 | +# ---------------------------------- gpu信息 ---------------------------------- |
| 97 | + |
| 98 | + |
| 99 | +def get_nvidia_gpu_info(): |
| 100 | + """获取 GPU 信息""" |
| 101 | + |
| 102 | + def get_cuda_version(): |
| 103 | + """获取 CUDA 版本""" |
| 104 | + try: |
| 105 | + output = subprocess.check_output(["nvcc", "--version"]).decode("utf-8") |
| 106 | + for line in output.split('\n'): |
| 107 | + if "release" in line: |
| 108 | + version = line.split("release")[-1].strip().split(" ")[0][:-1] |
| 109 | + return version |
| 110 | + except Exception: # noqa |
| 111 | + return None |
| 112 | + |
| 113 | + info = {"driver": None, "cores": None, "type": [], "memory": [], "cuda": None} |
| 114 | + try: |
| 115 | + pynvml.nvmlInit() |
| 116 | + except Exception: # noqa |
| 117 | + return None |
| 118 | + |
| 119 | + try: |
| 120 | + # 获取 NVIDIA 驱动版本信息 |
| 121 | + nv_driver = pynvml.nvmlSystemGetDriverVersion() |
| 122 | + if isinstance(nv_driver, bytes): |
| 123 | + nv_driver = nv_driver.decode("utf-8") |
| 124 | + info["driver"] = nv_driver |
| 125 | + |
| 126 | + # 获取 CUDA 版本 |
| 127 | + info["cuda"] = get_cuda_version() |
| 128 | + |
| 129 | + # 获取 NVIDIA GPU 数量 |
| 130 | + info["cores"] = pynvml.nvmlDeviceGetCount() |
| 131 | + # 遍历每个 GPU，获取 GPU 信息 |
| 132 | + for i in range(info["cores"]): |
| 133 | + handle = pynvml.nvmlDeviceGetHandleByIndex(i) |
| 134 | + # 获取 GPU 型号 |
| 135 | + gpu_name = pynvml.nvmlDeviceGetName(handle) # types: bytes | str |
| 136 | + if isinstance(gpu_name, bytes): # Fix for pynvml 早期版本，关联 issue: #605 |
| 137 | + gpu_name = gpu_name.decode("utf-8") |
| 138 | + info["type"].append(gpu_name) |
| 139 | + # 获取 GPU 的总显存, 单位为GB |
| 140 | + info["memory"].append(round(pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024**3))) |
| 141 | + |
| 142 | + except pynvml.NVMLError: |
| 143 | + pass |
| 144 | + finally: |
| 145 | + # 结束 NVML |
| 146 | + pynvml.nvmlShutdown() |
| 147 | + return info |
| 148 | + |
| 149 | + |
| 150 | +# ---------------------------------- apple信息 ---------------------------------- |
| 151 | + |
| 152 | + |
| 153 | +def get_apple_chip_info(): |
| 154 | + if "mac" not in platform.platform().lower(): |
| 155 | + return None |
| 156 | + info = {"cpu": None, "gpu": None, "memory": None, "type": None} |
| 157 | + |
| 158 | + # 使用system_profiler命令以JSON格式获取GPU信息 |
| 159 | + try: |
| 160 | + result = subprocess.run(["system_profiler", "SPHardwareDataType", "-json"], capture_output=True, text=True) |
| 161 | + gpu_name = json.loads(result.stdout)["SPHardwareDataType"][0]["chip_type"] |
| 162 | + memory = json.loads(result.stdout)["SPHardwareDataType"][0]["physical_memory"] |
| 163 | + memory = str(memory).lower().replace("gb", "") |
| 164 | + # TODO: 获取GPU信息 |
| 165 | + info["type"] = gpu_name |
| 166 | + info["memory"] = memory |
| 167 | + except Exception: # noqa |
| 168 | + return None |
| 169 | + try: |
| 170 | + info["cpu"] = multiprocessing.cpu_count() |
| 171 | + except Exception: # noqa |
| 172 | + pass |
| 173 | + return info |
0 commit comments