python将xml格式文件转成png或者pdf格式

本文主要介绍运行NCCL代码时输出的xml文件该如何转成更加容易观看的图格式

如下是举例，服务器上的PCIE相关的topo xml 文件

<system version="1">
  <cpu numaid="1" affinity="ffffff00,0000ffff,ff000000" arch="x86_64" vendor="GenuineIntel" familyid="6" modelid="106">
    <pci busid="0000:98:00.0" class="0x060400" vendor="0x1bd4" device="0x3004" subsystem_vendor="0x0000" subsystem_device="0x0000" link_speed="16.0 GT/s PCIe" link_width="16">
      <pci busid="0000:a6:00.0" class="0x030200" vendor="0x10de" device="0x2235" subsystem_vendor="0x10de" subsystem_device="0x145a" link_speed="16.0 GT/s PCIe" link_width="16">
        <gpu dev="0" sm="86" rank="0" gdr="1"/>
      </pci>
      <pci busid="0000:a5:00.0" class="0x020000" vendor="0x15b3" device="0x1021" subsystem_vendor="0x15b3" subsystem_device="0x0023" link_speed="16.0 GT/s PCIe" link_width="16">
        <nic>
          <net name="mlx5_0" dev="0" speed="200000" port="1" latency="0.000000" guid="0xbe8cb30003ae6d94" maxconn="131072" gdr="1"/>
        </nic>
      </pci>
    </pci>
  </cpu>
</system>

将该文件转换成图格式

import os
import sys
import logging
import argparse
import xml.etree.ElementTree as ET
from graphviz import Graph


def build_topology(graph, parent_element, parent_node=None):
    dev = parent_element.tag
    res = None
    if dev == 'cpu':
        res = parent_element.get('numaid').replace(':', '_')
    elif dev == 'pci':
        res = parent_element.get('busid').replace(':', '_')
    elif dev == 'gpu':
        res = parent_element.get('dev').replace(':', '_')
    elif dev == 'nvlink':
        res = parent_element.get('target').replace(':', '_')
    elif dev == 'net':
        res = parent_element.get('name').replace(':', '_')
    
    node = dev
    if res is not None:
        node = dev + '_' + res
    
    all_attributes = {attr: parent_element.get(attr) for attr in parent_element.attrib}
    str_attr = node + '\n\n' + str(all_attributes).replace(',', '\n')
    graph.node(node, shape='box', label=str_attr)
    
    if parent_node is not None:
        graph.edge(parent_node, node)
    
    for child_element in parent_element:
        build_topology(graph, child_element, node)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Parse xml file to pdf or png.')
    parser.add_argument('-i', '--input', required=True, help='Input file path')
    parser.add_argument('-o', '--output', default='xml', help='Output file name')
    args = parser.parse_args()

    input_file_path = args.input
    output_file_name = args.output


	# 解析XML数据
	tree = ET.parse(input_file_path)
	# 获取 XML 文档对象的根结点 Element
	root_element = tree.getroot()
	
	graph = Graph()
	build_topology(graph, root_element)
	
	# 绘制拓扑图
	graph.render(output_file_name, format='png')
	# graph.view()
	print('Finished')