1. NVLink拓扑发现与可视化

NVLink拓扑发现与可视化 学习卡片

任务说明

本案例目标为：在装有NVIDIA GPU（Volta及以上架构，支持NVLink）的服务器上，通过系统工具与NVML库识别GPU之间的NVLink连接关系，并生成拓扑可视化图。最终交付物为一张反映真实物理连接的拓扑图（SVG格式）及对应的连接关系清单。

实现方案

采用两层实现：命令行快速探测（使用nvidia-smi）与NVML编程接口（获取精确的P2P链路与带宽信息）。两者互为验证，确保拓扑数据准确。核心逻辑为：枚举所有GPU设备，对每一对GPU查询NVML的P2P链路状态与带宽，筛选出有效NVLink连接，再输出为Graphviz DOT格式并渲染为SVG。

代码实现


#!/usr/bin/env python3
"""
NVLink拓扑发现与可视化工具
依赖: pynvml (pip install pynvml), graphviz (apt install graphviz 或 pip install graphviz)
"""
import os
import subprocess
import tempfile
from graphviz import Digraph
from pynvml import (
    nvmlInit, nvmlShutdown, nvmlDeviceGetHandleByIndex,
    nvmlDeviceGetCount, nvmlDeviceGetName,
    nvmlDeviceGetP2PStatus, nvmlDeviceGetP2PLinkInfo,
    NVML_P2P_CAPS_INDEX_LINK, NVML_P2P_LINK_STATUS_ACTIVE,
    NVML_P2P_CAPS_INDEX_PCI, NVML_P2P_CAPS_INDEX_PHYSICAL
)

def get_nvlink_topology():
    """通过NVML获取NVLink连接矩阵"""
    nvmlInit()
    try:
        device_count = nvmlDeviceGetCount()
        if device_count < 2:
            print("检测到少于2块GPU，无法构成NVLink拓扑")
            return None

        # 收集设备信息
        devices = []
        for i in range(device_count):
            handle = nvmlDeviceGetHandleByIndex(i)
            name = nvmlDeviceGetName(handle)
            devices.append({'index': i, 'name': name, 'handle': handle})

        # 构建连接矩阵
        links = []  # (gpu_i, gpu_j, link_count, bandwidth_MBps)
        for i in range(device_count):
            for j in range(i+1, device_count):
                handle_i = devices[i]['handle']
                handle_j = devices[j]['handle']
                try:
                    # 查询P2P链路状态
                    status = nvmlDeviceGetP2PStatus(
                        handle_i, handle_j, NVML_P2P_CAPS_INDEX_LINK
                    )
                    if status == NVML_P2P_LINK_STATUS_ACTIVE:
                        # 获取链路信息（NVLink链路数与总带宽）
                        link_info = nvmlDeviceGetP2PLinkInfo(
                            handle_i, handle_j,
                            NVML_P2P_CAPS_INDEX_LINK
                        )
                        # link_info包含num_links与pcie_path等信息
                        # 此处简化：每条NVLink链路带宽为25GB/s（H100）或根据设备调整
                        # 实际带宽可从link_info中解析，但不同驱动版本字段不同
                        # 为保持兼容，使用nvidia-smi辅助获取带宽
                        links.append((i, j, link_info.num_links))
                except Exception as e:
                    # 部分驱动不支持P2P状态查询，回退到nvidia-smi
                    print(f"NVML查询GPU {i}-{j}失败: {e}，尝试nvidia-smi")
                    pass

        # 使用nvidia-smi补充带宽信息（更稳定）
        try:
            smi_output = subprocess.check_output(
                ['nvidia-smi', 'nvlink', '-s'], text=True
            )
            # 解析nvidia-smi输出补充带宽
            # 格式示例: GPU0-GPU1: 4 NVLink Links, 100.0 GB/s
            for line in smi_output.splitlines():
                if 'NVLink Links' in line:
                    parts = line.split(':')[0].strip()  # "GPU0-GPU1"
                    gpu_pair = parts.split('-')
                    if len(gpu_pair) == 2:
                        g0 = int(gpu_pair[0].replace('GPU', ''))
                        g1 = int(gpu_pair[1].replace('GPU', ''))
                        bw_str = line.split(',')[1].replace('GB/s', '').strip()
                        bw = float(bw_str)
                        # 更新links中的带宽
                        for idx, (a, b, _) in enumerate(links):
                            if (a == g0 and b == g1) or (a == g1 and b == g0):
                                links[idx] = (a, b, bw)
        except Exception as e:
            print(f"nvidia-smi带宽获取失败: {e}，使用默认带宽")

        return devices, links
    finally:
        nvmlShutdown()

def render_topology(devices, links, output_file='nvlink_topology'):
    """生成SVG拓扑图"""
    if not devices or not links:
        print("无有效拓扑数据")
        return

    dot = Digraph(comment='NVLink Topology', format='svg')
    dot.attr(rankdir='LR', bgcolor='#1a1a2e')
    dot.attr('node', shape='box', style='filled',
             fontname='Helvetica', fontsize='12')

    # 添加GPU节点
    for dev in devices:
        label = f"GPU{dev['index']}\n{dev['name']}"
        dot.node(f"gpu{dev['index']}", label,
                 fillcolor='#16213e', fontcolor='#e94560')

    # 添加NVLink连接
    for (i, j, bw) in links:
        # 带宽单位：GB/s（来自nvidia-smi）或链路数（来自NVML）
        if isinstance(bw, float):
            label = f"NVLink\n{bw:.0f} GB/s"
        else:
            label = f"NVLink\n{bw} links"
        dot.edge(f"gpu{i}", f"gpu{j}", label=label,
                 color='#00d2ff', penwidth='2')

    # 渲染SVG
    dot.render(output_file, cleanup=True)
    print(f"拓扑图已生成: {output_file}.svg")

def verify_with_nvidia_smi():
    """使用nvidia-smi验证拓扑（独立验证路径）"""
    print("=== nvidia-smi NVLink 验证 ===")
    try:
        result = subprocess.run(
            ['nvidia-smi', 'nvlink', '-s'],
            capture_output=True, text=True, check=True
        )
        print(result.stdout)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"nvidia-smi执行失败: {e.stderr}")
        return None

if __name__ == '__main__':
    # 第一步：nvidia-smi快速验证
    smi_result = verify_with_nvidia_smi()

    # 第二步：NVML精确拓扑
    result = get_nvlink_topology()
    if result:
        devices, links = result
        print("\n=== NVML检测到的NVLink连接 ===")
        for (i, j, bw) in links:
            if isinstance(bw, float):
                print(f"GPU{i} <-> GPU{j}: {bw:.1f} GB/s")
            else:
                print(f"GPU{i} <-> GPU{j}: {bw}条链路")
        render_topology(devices, links)
    else:
        print("NVML未检测到有效拓扑，请检查GPU是否支持NVLink")

运行与验证

运行方式：


# 安装依赖
pip install pynvml graphviz
# Ubuntu/Debian需安装graphviz系统包: sudo apt install graphviz

# 执行脚本
python3 nvlink_topology.py

验证方法与判定标准：

  1. 命令行验证：执行nvidia-smi nvlink -s，观察输出中是否列出GPU间的NVLink连接及带宽。若输出为空或提示“不支持”，则当前环境无NVLink或驱动不支持。
  2. 脚本输出验证：运行脚本后，终端应打印出与nvidia-smi一致的GPU连接对（编号与带宽数值需匹配）。若NVML与nvidia-smi结果不一致，以nvidia-smi为准排查驱动问题。
  3. SVG产物验证：打开生成的nvlink_topology.svg，图中节点数应等于nvidia-smi -L列出的GPU数量；每一条边对应一条实际NVLink连接，且边上的带宽标注与nvidia-smi nvlink -s输出一致。
  4. 判定通过标准：SVG中展示的连接关系与nvidia-smi nvlink -s输出的连接对完全一致（GPU编号、连接存在性），且无遗漏或多余连接。

拓扑可视化示例

下图展示4块GPU通过NVLink全互联的典型拓扑（实际生成图取决于当前服务器硬件）：

GPU0 A100 GPU1 A100 GPU2 A100 GPU3 A100 25 GB/s 25 GB/s 25 GB/s 25 GB/s 25 GB/s PCIe NVLink连接 PCIe通道

说明：上图为4块A100 GPU全互联拓扑示意。实际执行脚本时，SVG将根据当前服务器真实硬件生成——若GPU间无NVLink物理连接，则对应边不会出现；若为部分互联（如仅相邻GPU连接），则边数相应减少。判定拓扑正确性的唯一依据是nvidia-smi nvlink -s的实测输出与SVG边集合完全一致。

← 上一章 📖 返回目录 下一章 →