在使用 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 的真实执行后端：
python
from tinygrad import Device
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
python
Device.DEFAULT代表 tinygrad 实际使用的后端，而不是系统"可能存在"的设备。
这对 exo 非常重要：
exo 调度依据的是「真实可执行能力」，不是硬件清单
2️⃣ NVIDIA GPU 探测为什么必须 try / except
常见失败场景：
- Docker 容器无
/dev/nvidia* - 没装 NVIDIA Driver
- pynvml 安装了但 NVML 初始化失败
- MIG / 权限问题
因此：
python
except Exception:
fallback to CPU这是 exo 节点稳定运行的关键。
3️⃣ GPU 名称清洗的意义
python
"NVIDIA RTX 3090 24GB" → "NVIDIA RTX 3090"目的：
- 避免
CHIP_FLOPS查表失败 - 统一 exo 节点能力上报口径
4️⃣ CPU 兜底是 exo 的生命线
python
psutil.virtual_memory()无论：
- 裸机
- 云主机
- CI
- Docker
exo 节点都必须可启动、可注册、可上报。
五、exo 中的典型使用场景
- worker 启动时自动上报能力
- 调度前判断模型是否可加载
- 异构节点（CPU / GPU 混合）统一管理
- GPU 节点异常时自动降级