在使用exo构建本地或分布式推理集群时，每个节点的设备能力探测是非常关键的一步。
exo 需要知道：
- 当前节点使用CPU 还是 GPU
- GPU 型号、显存大小
- 节点的大致算力能力（FLOPS）
本文结合tinygrad，介绍一套稳定、可落地、工程级安全的 Linux 设备能力探测方案，适用于：
- exo worker / node
- 本地推理
- Docker / 裸机 / 无 GPU 环境
一、为什么 exo 需要设备能力探测？
在 exo 中，每个节点都会向调度层上报自身能力，例如：
- 是否支持 CUDA
- 显存大小是否满足模型加载
- 是否能承担 FP16 / INT8 推理
如果设备探测逻辑不健壮，可能会导致：
- ❌ 节点启动即崩溃
- ❌ GPU 探测失败导致服务不可用
- ❌ Docker / 云环境下无法运行
设计原则只有一个：
GPU 是加速项，CPU 是兜底项
任何情况下，exo 节点都必须能启动
二、整体设计思路（exo 场景）
本方案基于tinygrad 的真实执行后端：
fromtinygradimportDevice Device.DEFAULT而不是单纯扫描系统硬件。
探测优先级
- NVIDIA GPU（CUDA / NV / GPU）
- AMD GPU
- CPU（最终兜底）
任何 GPU 探测失败，立即回退 CPU，不抛异常。
三、完整代码（exo 节点可直接使用）
git clone https://github.com/exo-explore/exo.git cd exo pip install -e .启动不起来 搜索linux_device_capabilities进行修改
#vim /home/michah/exo/exo/topology/device_capabilities.py async def linux_device_capabilities() -> DeviceCapabilities: import psutil from tinygrad import Device if DEBUG >= 2: print(f"tinygrad {Device.DEFAULT=}") # ----------------------------- # NVIDIA / CUDA 路径（安全版） # ----------------------------- if Device.DEFAULT in ("CUDA", "NV", "GPU"): try: import pynvml pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) gpu_raw_name = pynvml.nvmlDeviceGetName(handle).upper() gpu_name = ( gpu_raw_name.rsplit(" ", 1)[0] if gpu_raw_name.endswith("GB") else gpu_raw_name ) gpu_memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle) if DEBUG >= 2: print(f"NVIDIA device {gpu_name=} {gpu_memory_info=}") pynvml.nvmlShutdown() return DeviceCapabilities( model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info.total // 2**20, flops=CHIP_FLOPS.get( gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0), ), ) except Exception as e: # 关键：任何 NVML / 驱动 / import 错误都回退 CPU if DEBUG >= 1: print(f"[WARN] NVIDIA GPU detection failed, fallback to CPU: {e}") # ----------------------------- # AMD 路径（原样保留） # ----------------------------- if Device.DEFAULT == "AMD": try: import pyamdgpuinfo gpu_raw_info = pyamdgpuinfo.get_gpu(0) gpu_name = gpu_raw_info.name gpu_memory_info = gpu_raw_info.memory_info["vram_size"] if DEBUG >= 2: print(f"AMD device {gpu_name=} {gpu_memory_info=}") return DeviceCapabilities( model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info // 2**20, flops=CHIP_FLOPS.get( gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0), ), ) except Exception as e: if DEBUG >= 1: print(f"[WARN] AMD GPU detection failed, fallback to CPU: {e}") # ----------------------------- # CPU / 兜底路径 # ----------------------------- return DeviceCapabilities( model=f"Linux Box (Device: {Device.DEFAULT})", chip=f"Unknown Chip (Device: {Device.DEFAULT})", memory=psutil.virtual_memory().total // 2**20, flops=DeviceFlops(fp32=0, fp16=0, int8=0), )四、关键实现解析（exo 视角）
1️⃣ 为什么使用Device.DEFAULT
Device.DEFAULT代表tinygrad 实际使用的后端，而不是系统“可能存在”的设备。
这对 exo 非常重要：
exo 调度依据的是「真实可执行能力」，不是硬件清单
2️⃣ NVIDIA GPU 探测为什么必须 try / except
常见失败场景：
- Docker 容器无
/dev/nvidia* - 没装 NVIDIA Driver
- pynvml 安装了但 NVML 初始化失败
- MIG / 权限问题
因此：
exceptException:fallback to CPU这是 exo 节点稳定运行的关键。
3️⃣ GPU 名称清洗的意义
"NVIDIA RTX 3090 24GB"→"NVIDIA RTX 3090"目的：
- 避免
CHIP_FLOPS查表失败 - 统一 exo 节点能力上报口径
4️⃣ CPU 兜底是 exo 的生命线
psutil.virtual_memory()无论：
- 裸机
- 云主机
- CI
- Docker
exo 节点都必须可启动、可注册、可上报。
五、exo 中的典型使用场景
- worker 启动时自动上报能力
- 调度前判断模型是否可加载
- 异构节点（CPU / GPU 混合）统一管理
- GPU 节点异常时自动降级
