1. NVLink拓扑发现与可视化

任务目标

本案例要求枚举当前主机中GPU与DPU（BlueField）的NVLink连接关系，输出连接矩阵，并生成可视化拓扑图。通过NVML（NVIDIA Management Library）编程接口实现，不依赖外部图数据库或可视化库，使用Python标准库完成数据采集与SVG拓扑生成。

环境依赖

• NVIDIA驱动（含NVML）已安装，执行 nvidia-smi 可正常输出
• DPU（BlueField）已通过NVLink连接至GPU，且DPU的NVML可见（需安装Mellanox NVML扩展）
• Python 3.6+，需安装 pynvml 库：pip install pynvml

核心实现

以下代码通过NVML枚举所有设备（GPU与DPU），查询每个设备的NVLink链路对端信息，构建连接矩阵并生成SVG拓扑图。代码可直接运行，输出JSON格式连接矩阵与拓扑SVG文件。

import pynvml
import json
import os

def get_device_type(handle):
    """判断设备类型：GPU或DPU(BlueField)"""
    try:
        name = pynvml.nvmlDeviceGetName(handle).decode()
        if 'BlueField' in name or 'DPU' in name:
            return 'DPU'
        return 'GPU'
    except:
        return 'GPU'

def enumerate_nvlink_topology():
    pynvml.nvmlInit()
    device_count = pynvml.nvmlDeviceGetCount()
    
    devices = []
    handle_map = {}  # index -> (type, name, uuid)
    
    for i in range(device_count):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        name = pynvml.nvmlDeviceGetName(handle).decode()
        uuid = pynvml.nvmlDeviceGetUUID(handle).decode()
        dtype = get_device_type(handle)
        devices.append({'index': i, 'type': dtype, 'name': name, 'uuid': uuid})
        handle_map[i] = handle
    
    # 构建连接矩阵
    n = len(devices)
    matrix = [[0]*n for _ in range(n)]
    links = []
    
    for i in range(n):
        handle = handle_map[i]
        try:
            # 获取NVLink数量
            link_count = pynvml.nvmlDeviceGetMaxLinkInfo(handle) if hasattr(pynvml, 'nvmlDeviceGetMaxLinkInfo') else 0
            # 遍历可能的NVLink链路
            for link_idx in range(link_count):
                try:
                    # 查询链路对端设备
                    remote = pynvml.nvmlDeviceGetNvLinkRemoteDeviceInfo(handle, link_idx)
                    remote_handle = remote.deviceHandle
                    # 找到对端设备索引
                    for j in range(n):
                        if handle_map[j] == remote_handle:
                            matrix[i][j] = 1
                            links.append((i, j, link_idx))
                            break
                except pynvml.NVMLError:
                    continue
        except pynvml.NVMLError:
            continue
    
    pynvml.nvmlShutdown()
    return devices, matrix, links

def generate_svg(devices, links, filename='nvlink_topology.svg'):
    """生成SVG拓扑图"""
    n = len(devices)
    # 计算布局：圆形排列
    import math
    cx, cy, radius = 200, 150, 100
    positions = {}
    
    for i, dev in enumerate(devices):
        angle = 2 * math.pi * i / n - math.pi/2
        x = cx + radius * math.cos(angle)
        y = cy + radius * math.sin(angle)
        positions[i] = (x, y)
    
    # 颜色映射
    colors = {'GPU': '#76b900', 'DPU': '#00a1e0'}
    
    svg_parts = []
    svg_parts.append(f'<svg xmlns="http://www.w3.org/2000/svg" width="400" height="300">')
    
    # 绘制连接线
    for src, dst, _ in links:
        x1, y1 = positions[src]
        x2, y2 = positions[dst]
        svg_parts.append(f'<line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" stroke="#ff6600" stroke-width="2"/>')
    
    # 绘制设备节点
    for i, dev in enumerate(devices):
        x, y = positions[i]
        color = colors.get(dev['type'], '#999999')
        svg_parts.append(f'<circle cx="{x:.1f}" cy="{y:.1f}" r="20" fill="{color}" stroke="#333" stroke-width="1.5"/>')
        # 设备标签
        label = f"{dev['type']}{dev['index']}"
        svg_parts.append(f'<text x="{x:.1f}" y="{y+5:.1f}" text-anchor="middle" font-size="10" fill="white" font-family="Arial">{label}</text>')
    
    svg_parts.append('</svg>')
    
    with open(filename, 'w') as f:
        f.write('\n'.join(svg_parts))
    
    return filename

if __name__ == '__main__':
    devices, matrix, links = enumerate_nvlink_topology()
    
    # 输出连接矩阵
    print("=== NVLink连接矩阵 ===")
    header = "     " + " ".join([f"{d['type'][0]}{d['index']:2d}" for d in devices])
    print(header)
    for i, row in enumerate(matrix):
        line = f"{devices[i]['type'][0]}{devices[i]['index']:2d}  "
        line += "  ".join(["1" if v else "." for v in row])
        print(line)
    
    # 输出JSON格式
    result = {
        'devices': devices,
        'matrix': matrix,
        'links': [{'src': s, 'dst': d, 'link': l} for s, d, l in links]
    }
    with open('nvlink_topology.json', 'w') as f:
        json.dump(result, f, indent=2)
    
    # 生成SVG
    svg_file = generate_svg(devices, links)
    print(f"\n拓扑图已保存: {svg_file}")
    print(f"连接数据已保存: nvlink_topology.json")

运行与验证

1. 保存代码为 nvlink_topology.py，执行 python nvlink_topology.py
2. 程序输出连接矩阵，每行每列对应一个设备（G=GPU，D=DPU），"1"表示存在NVLink连接，"."表示无连接
3. 验证通过条件：

• 矩阵为对称矩阵（NVLink为双向连接）
• 对角线为0（设备自身无连接）
• 矩阵中"1"的数量与 nvidia-smi nvlink -s 命令显示的链路数一致
• 生成的 nvlink_topology.svg 可在浏览器中打开，GPU节点为绿色，DPU节点为蓝色，橙色连线表示NVLink连接

常见问题

• NVML不支持DPU枚举：需安装Mellanox的NVML扩展（libnvidia-ml.so 需包含BlueField支持），或使用 nvidia-smi nvlink -s 命令交叉验证
• 无NVLink连接：确认GPU与DPU已通过NVLink线缆物理连接，且驱动版本支持NVLink互联
• 权限不足：NVML调用需要root权限，使用 sudo python nvlink_topology.py 执行