Mooncake/mooncake-transfer-engine/benchmarks/hetero/topology.py

#!/usr/bin/env python3
# Copyright 2024 KVCache.AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Render the host's memory-fabric topology: GPUs and their NVLink mesh, NICs,
NUMA nodes, and the probed transport capabilities, as a text summary and a
Graphviz DOT file.

Sources, each optional and degraded gracefully:
  - nvidia-smi topo -m        (GPU<->GPU links, GPU<->NIC, NUMA affinity)
  - numactl --hardware        (NUMA nodes)
  - the conformance.json probe report (per-backend capability), if present in
    MC_FABRIC_RESULTS.

Writes <results>/topology.dot (+ topology.png if Graphviz `dot` is installed)."""

import json
import os
import re
import shutil
import subprocess
import sys

RESULTS = os.environ.get("MC_FABRIC_RESULTS", os.path.abspath("fabric-results"))


def run(cmd):
    try:
        return subprocess.run(cmd, capture_output=True, text=True,
                              timeout=20).stdout
    except Exception:
        return ""


def parse_nvidia_topo():
    """Return (gpus, links, gpu_numa) from `nvidia-smi topo -m`."""
    out = run(["nvidia-smi", "topo", "-m"])
    if not out:
        return [], {}, {}
    lines = [l for l in out.splitlines() if l.strip()]
    header = None
    rows = {}
    gpu_numa = {}
    for line in lines:
        cells = re.split(r"\s{1,}", line.strip())
        if cells and cells[0].startswith("GPU") and header is None and \
                "GPU0" in line and cells[0] == "GPU0":
            pass
        if cells[0].startswith("GPU") and re.match(r"GPU\d+$", cells[0]):
            name = cells[0]
            rows[name] = cells[1:]
        if header is None and line.lstrip().startswith("GPU0"):
            header = re.split(r"\s{1,}", line.strip())
    gpus = sorted(rows.keys(), key=lambda g: int(g[3:]))
    links = {}
    for g, cells in rows.items():
        for j, gj in enumerate(gpus):
            if j < len(cells) and cells[j].startswith("NV"):
                a, b = sorted([g, gj])
                if a != b:
                    links[(a, b)] = cells[j]  # e.g. NV18
    # NUMA affinity column: look for a numeric NUMA id per GPU row.
    for g, cells in rows.items():
        for c in cells:
            if re.match(r"^\d+$", c):
                gpu_numa[g] = int(c)
                break
    return gpus, links, gpu_numa


def parse_numa():
    out = run(["numactl", "--hardware"])
    m = re.search(r"available:\s*(\d+)\s*nodes", out)
    return int(m.group(1)) if m else 0


def probed_caps():
    path = os.path.join(RESULTS, "conformance.json")
    if not os.path.exists(path):
        return []
    try:
        data = json.load(open(path))
    except Exception:
        return []
    return [e.get("capability", {}) for e in data]


def main():
    os.makedirs(RESULTS, exist_ok=True)
    gpus, links, gpu_numa = parse_nvidia_topo()
    numa_nodes = parse_numa()
    caps = probed_caps()

    print("=== Memory fabric topology ===")
    if gpus:
        nvlink = sorted({v for v in links.values()})
        print(f"GPUs: {len(gpus)}  NVLink mesh edges: {len(links)} "
              f"({','.join(nvlink)})")
    else:
        print("GPUs: none detected (nvidia-smi unavailable)")
    if numa_nodes:
        print(f"NUMA nodes: {numa_nodes}")
    if caps:
        print("Probed transports:")
        for c in caps:
            tag = "EMULATED" if c.get("emulated") else "real"
            print(f"  {c.get('name', '?'):12s} {c.get('kind', '?'):7s} {tag}")

    # Graphviz DOT.
    dot = ["graph fabric {", '  rankdir=LR;', '  node [fontsize=10];']
    by_numa = {}
    for g in gpus:
        by_numa.setdefault(gpu_numa.get(g, 0), []).append(g)
    for node, members in sorted(by_numa.items()):
        dot.append(f'  subgraph cluster_numa{node} {{')
        dot.append(f'    label="NUMA node {node}"; color="#888888";')
        for g in members:
            dot.append(f'    {g} [shape=box, style=filled, '
                       f'fillcolor="#cfe8cf"];')
        dot.append("  }")
    for (a, b), kind in sorted(links.items()):
        dot.append(f'  {a} -- {b} [color="#2ca02c", penwidth=2, '
                   f'label="{kind}"];')
    # Fabric capability legend nodes.
    for c in caps:
        name = c.get("name", "?").replace(":", "_").replace(".", "_")
        color = "#f4cccc" if c.get("emulated") else "#d9ead3"
        label = f"{c.get('name','?')}\\n{c.get('kind','?')}"
        dot.append(f'  fab_{name} [shape=ellipse, style=filled, '
                   f'fillcolor="{color}", label="{label}"];')
    dot.append("}")

    dot_path = os.path.join(RESULTS, "topology.dot")
    with open(dot_path, "w") as f:
        f.write("\n".join(dot) + "\n")
    print(f"\nDOT -> {dot_path}")

    if shutil.which("dot"):
        png = os.path.join(RESULTS, "topology.png")
        try:
            subprocess.run(["dot", "-Tpng", dot_path, "-o", png], timeout=30,
                           check=True)
            print(f"PNG -> {png}")
        except Exception as exc:
            print(f"(graphviz render skipped: {exc})")
    else:
        print("(install graphviz `dot` to render topology.png)")
    return 0


if __name__ == "__main__":
    sys.exit(main())