📚 课程目录第1章 NVLink拓扑感知的通信模式识别

1. NVLink拓扑感知的通信模式识别

任务说明

本案例目标：给定一个多GPU节点的NVLink拓扑描述文件，识别其中存在的通信模式（如环形、全连接、混合拓扑），并输出每个GPU可用的NVLink对端列表及拓扑类型判定。该识别结果是后续通信调度与优化（如选择最优通信路径、避免共享链路拥塞）的基础。

拓扑输入格式

使用简化的邻接矩阵描述NVLink连接关系。矩阵为N×N（N为GPU数量），元素为1表示对应两个GPU之间存在直接NVLink连接，0表示无直接连接。矩阵对角线固定为0。输入文件为纯文本，每行N个以空格分隔的整数。

实现代码

import sys
from collections import deque

def read_topology(filepath):
    """读取NVLink拓扑邻接矩阵"""
    with open(filepath, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    matrix = [list(map(int, line.split())) for line in lines]
    n = len(matrix)
    # 校验矩阵合法性
    for i in range(n):
        if len(matrix[i]) != n:
            raise ValueError(f"第{i+1}行长度不为{n}")
        if matrix[i][i] != 0:
            raise ValueError(f"对角线元素必须为0，第{i+1}行错误")
        for j in range(n):
            if matrix[i][j] != matrix[j][i]:
                raise ValueError(f"矩阵不对称: ({i},{j}) != ({j},{i})")
    return matrix

def get_neighbors(matrix, gpu_id):
    """返回指定GPU的直接NVLink邻居列表（升序）"""
    n = len(matrix)
    return [j for j in range(n) if matrix[gpu_id][j] == 1]

def classify_topology(matrix):
    """
    识别拓扑类型:
    - 'ring': 每个GPU度数为2，且整体构成单一环
    - 'fully_connected': 每个GPU与其他所有GPU直连
    - 'hybrid': 不属于以上两种，但存在至少一条NVLink
    - 'disconnected': 无任何NVLink连接
    """
    n = len(matrix)
    degrees = [sum(row) for row in matrix]
    
    # 无连接
    if all(d == 0 for d in degrees):
        return 'disconnected'
    
    # 全连接: 每个GPU度数为n-1
    if all(d == n-1 for d in degrees):
        return 'fully_connected'
    
    # 环形: 每个度数=2 且 图连通且每个节点度数为2 → 单一环
    if all(d == 2 for d in degrees):
        # 检查是否连通（从0出发BFS）
        visited = [False] * n
        q = deque([0])
        visited[0] = True
        while q:
            u = q.popleft()
            for v in get_neighbors(matrix, u):
                if not visited[v]:
                    visited[v] = True
                    q.append(v)
        if all(visited):
            return 'ring'
    
    return 'hybrid'

def main():
    if len(sys.argv) != 2:
        print("用法: python nvlink_topology_analyzer.py <topology_file>")
        sys.exit(1)
    
    matrix = read_topology(sys.argv[1])
    n = len(matrix)
    topo_type = classify_topology(matrix)
    
    print(f"GPU数量: {n}")
    print(f"拓扑类型: {topo_type}")
    print("各GPU的NVLink对端:")
    for i in range(n):
        neighbors = get_neighbors(matrix, i)
        print(f"  GPU {i}: {neighbors if neighbors else '无'}")

if __name__ == "__main__":
    main()

输入样例

创建文件 topo_4gpu_ring.txt，内容如下（4个GPU构成环形拓扑）：

运行与验证

在终端执行：

python nvlink_topology_analyzer.py topo_4gpu_ring.txt

预期输出（由代码逻辑直接推出）：

GPU数量: 4
拓扑类型: ring
各GPU的NVLink对端:
  GPU 0: [1, 3]
  GPU 1: [0, 2]
  GPU 2: [1, 3]
  GPU 3: [0, 2]

判定标准：程序输出与上述预期一致，且拓扑类型正确识别为ring。可另建全连接（所有非对角线元素为1）和混合拓扑文件测试，观察分类结果是否符合定义。

拓扑识别逻辑说明

分类依据：

ring：每个GPU恰好有2个NVLink邻居，且从任一GPU出发沿连接可遍历全部GPU（单一连通环）；
fully_connected：每个GPU与其他所有GPU直连（度数为N-1）；
hybrid：存在连接但不满足上述两种结构（如部分GPU有3个邻居、部分有1个，或存在多个独立环）；
disconnected：所有GPU之间无任何NVLink。

验证方法

准备三个额外测试文件验证分类正确性：

全连接（topo_4gpu_full.txt）：4×4矩阵，除对角线外全为1，预期输出fully_connected；
混合拓扑（topo_4gpu_hybrid.txt）：GPU0连接GPU1、GPU2、GPU3，GPU1仅连接GPU0，GPU2仅连接GPU0，GPU3仅连接GPU0，预期输出hybrid（GPU1/2/3度数为1，GPU0度数为3）；
无连接（topo_4gpu_none.txt）：全0矩阵，预期输出disconnected。

运行命令替换文件名即可。观察输出中的拓扑类型与各GPU邻居列表是否符合上述预期，全部一致则判定通过。