feat: get topology

Signed-off-by: thxCode <[email protected]>

Author: thxCode

gpustack_runtime/__main__.py

@@ -18,6 +18,7 @@
     DeleteWorkloadSubCommand,
     DetectDevicesSubCommand,
     ExecWorkloadSubCommand,
+    GetDevicesTopologySubCommand,
     GetWorkloadSubCommand,
     InspectWorkloadSubCommand,
     ListImagesSubCommand,
@@ -68,6 +69,7 @@ def main():
     ExecWorkloadSubCommand.register(subcommand_parser)
     InspectWorkloadSubCommand.register(subcommand_parser)
     DetectDevicesSubCommand.register(subcommand_parser)
+    GetDevicesTopologySubCommand.register(subcommand_parser)
     ListImagesSubCommand.register(subcommand_parser)
     SaveImagesSubCommand.register(subcommand_parser)
     CopyImagesSubCommand.register(subcommand_parser)

gpustack_runtime/cmds/__init__.py

@@ -11,7 +11,7 @@
     ListWorkloadsSubCommand,
     LogsWorkloadSubCommand,
 )
-from .detector import DetectDevicesSubCommand
+from .detector import DetectDevicesSubCommand, GetDevicesTopologySubCommand
 from .images import (
     CopyImagesSubCommand,
     ListImagesSubCommand,
@@ -29,6 +29,7 @@
     "DeleteWorkloadsSubCommand",
     "DetectDevicesSubCommand",
     "ExecWorkloadSubCommand",
+    "GetDevicesTopologySubCommand",
     "GetWorkloadSubCommand",
     "InspectWorkloadSubCommand",
     "ListImagesSubCommand",

gpustack_runtime/cmds/detector.py

@@ -5,7 +5,13 @@
 import time
 from typing import TYPE_CHECKING
 
-from ..detector import Devices, detect_devices
+from ..detector import (
+    Devices,
+    Topology,
+    detect_devices,
+    get_devices_topologies,
+    group_devices_by_manufacturer,
+)
 from .__types__ import SubCommand
 
 if TYPE_CHECKING:
@@ -49,27 +55,67 @@ def __init__(self, args: Namespace):
         self.watch = args.watch
 
     def run(self):
-        try:
-            while True:
-                devs: Devices = detect_devices(fast=False)
-                print("\033[2J\033[H", end="")
-                match self.format.lower():
-                    case "json":
-                        print(format_devices_json(devs))
-                    case _:
-                        # Group devs by manufacturer
-                        grouped_devs: dict[str, Devices] = {}
-                        for dev in devs:
-                            if dev.manufacturer not in grouped_devs:
-                                grouped_devs[dev.manufacturer] = []
-                            grouped_devs[dev.manufacturer].append(dev)
-                        for manu in sorted(grouped_devs.keys()):
-                            print(format_devices_table(grouped_devs[manu]))
-                if not self.watch:
-                    break
-                time.sleep(self.watch)
-        except KeyboardInterrupt:
+        while True:
+            devs: Devices = detect_devices(fast=False)
             print("\033[2J\033[H", end="")
+            match self.format.lower():
+                case "json":
+                    print(format_devices_json(devs))
+                case _:
+                    # Group devices by manufacturer.
+                    group_devs = group_devices_by_manufacturer(devs)
+                    if not group_devs:
+                        print("No GPUs detected.")
+                    else:
+                        # Print each group separately.
+                        for devs in group_devs.values():
+                            print(format_devices_table(devs))
+            if not self.watch:
+                break
+            time.sleep(self.watch)
+
+
+class GetDevicesTopologySubCommand(SubCommand):
+    """
+    Command to detect GPUs topology.
+    """
+
+    format: str = "table"
+
+    @staticmethod
+    def register(parser: _SubParsersAction):
+        topo_parser = parser.add_parser(
+            "topology",
+            help="Detect GPUs topology",
+            aliases=["topo"],
+        )
+
+        topo_parser.add_argument(
+            "--format",
+            type=str,
+            choices=["table", "json"],
+            default="table",
+            help="Output format",
+        )
+
+        topo_parser.set_defaults(func=GetDevicesTopologySubCommand)
+
+    def __init__(self, args: Namespace):
+        self.format = args.format
+
+    def run(self):
+        topologies = get_devices_topologies(fast=False)
+        print("\033[2J\033[H", end="")
+        if not topologies:
+            print("No GPU topology information available.")
+            return
+
+        match self.format.lower():
+            case "json":
+                print(format_topologies_json(topologies))
+            case _:
+                for topo in topologies:
+                    print(format_topology_table(topo))
 
 
 def format_devices_json(devs: Devices) -> str:
@@ -151,3 +197,84 @@ def format_devices_table(devs: Devices) -> str:
 
     # Combine all parts
     return os.linesep.join(header_lines + device_lines + footer_lines)
+
+
+def format_topology_table(topo: Topology) -> str:
+    content = topo.stringify()
+
+    # Column headers
+    col_headers = [str(topo.manufacturer).upper()]
+    col_headers += ["Device " + str(idx) for idx in range(len(content))]
+    col_headers += [
+        "CPU Affinity",
+        "NUMA Affinity",
+    ]
+    # Gather all rows to determine max width for each column
+    rows = []
+    for row_idx, row_devs in enumerate(content):
+        row = ["Device " + str(row_idx), *row_devs]
+        rows.append(row)
+
+    # Calculate max width for each column
+    col_widths = [len(header) for header in col_headers]
+    for row in rows:
+        for i, cell in enumerate(row):
+            col_widths[i] = max(col_widths[i], len(str(cell)))
+
+    # Add padding
+    col_widths = [w + 2 for w in col_widths]
+
+    # Calculate table width
+    width = sum(col_widths) + len(col_widths) + 1
+
+    # Header section
+    header_lines = [
+        "+" + "-" * (width - 2) + "+",
+    ]
+
+    # Column header line
+    col_header_line = "|"
+    for i, header in enumerate(col_headers):
+        col_header_line += f" {header.center(col_widths[i] - 2)} |"
+    header_lines.append(col_header_line)
+
+    # Separator line
+    separator = "|" + "|".join(["-" * w for w in col_widths]) + "|"
+    header_lines.append(separator)
+
+    # Topology rows
+    topology_lines = []
+    for row in rows:
+        row_line = "|"
+        for j, data in enumerate(row):
+            cell = str(data)
+            # Truncate if too long
+            if len(cell) > col_widths[j] - 2:
+                cell = cell[: col_widths[j] - 5] + "..."
+            row_line += f" {cell.ljust(col_widths[j] - 2)} |"
+        topology_lines.append(row_line)
+
+    # Footer section
+    footer_lines = [
+        "+" + "-" * (width - 2) + "+",
+    ]
+
+    # Legend
+    legend_lines = [
+        "",
+        "Legend (from nearest to farthest):",
+        "  X    = Self",
+        "  LINK = Connection traversing with High-Speed Link (e.g., NVIDIA NVLink, Ascend HCCS)",
+        "  PIX  = Connection traversing at most a single PCIe bridge",
+        "  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)",
+        "  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)",
+        "  NODE = Connection traversing PCIe and the interconnect between NUMA nodes",
+        "  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)",
+    ]
+
+    # Combine all parts
+    return os.linesep.join(header_lines + topology_lines + footer_lines + legend_lines)
+
+
+def format_topologies_json(topologies: list[Topology]) -> str:
+    return json.dumps([topo.to_dict() for topo in topologies], indent=2)

gpustack_runtime/deployer/__init__.py

@@ -51,6 +51,14 @@
     DockerDeployer(),
     KubernetesDeployer(),
 ]
+"""
+List of all deployers.
+"""
+
+_DEPLOYERS_MAP: dict[str, Deployer] = {dep.name: dep for dep in _DEPLOYERS}
+"""
+Mapping from deployer name to deployer.
+"""
 
 
 def supported_list() -> list[Deployer]:

gpustack_runtime/detector/__init__.py

@@ -9,8 +9,10 @@
     Device,
     Devices,
     ManufacturerEnum,
+    Topology,
     backend_to_manufacturer,
     manufacturer_to_backend,
+    reduce_devices_distances,
     supported_backends,
     supported_manufacturers,
 )
@@ -35,6 +37,16 @@
     MThreadsDetector(),
     NVIDIADetector(),
 ]
+"""
+List of all detectors.
+"""
+
+_DETECTORS_MAP: dict[ManufacturerEnum, Detector] = {
+    det.manufacturer: det for det in _DETECTORS
+}
+"""
+Mapping from manufacturer to detector.
+"""
 
 
 def supported_list() -> list[Detector]:
@@ -113,14 +125,82 @@ def detect_devices(fast: bool = True) -> Devices:
     return devices
 
 
+def get_devices_topologies(
+    devices: Devices | None = None,
+    fast: bool = True,
+) -> list[Topology]:
+    """
+    Get the topology information of the given devices.
+
+    Args:
+        devices:
+            A list of devices to get the topology information from.
+            If None, detects devices automatically.
+        fast:
+            If True, return topologies from the first supported detector.
+            Otherwise, return topologies from all supported detectors.
+
+    Returns:
+        A list of Topology objects for each manufacturer group.
+
+    """
+    if devices is None:
+        devices = detect_devices(fast=fast)
+
+    topologies: list[Topology] = []
+
+    # Group devices by manufacturer.
+    group_devices = group_devices_by_manufacturer(devices)
+    if not group_devices:
+        return topologies
+
+    # Get topology for each group.
+    for manu, devs in group_devices.items():
+        det = _DETECTORS_MAP.get(manu)
+        if det is not None:
+            topo = det.get_topology(devs)
+            if topo:
+                topologies.append(topo)
+            if fast and topologies:
+                return topologies
+
+    return topologies
+
+
+def group_devices_by_manufacturer(
+    devices: Devices | None,
+) -> dict[ManufacturerEnum, Devices]:
+    """
+    Group devices by their manufacturer.
+
+    Args:
+        devices:
+            A list of devices to be grouped.
+
+    Returns:
+        A dictionary mapping each manufacturer to its corresponding list of devices.
+
+    """
+    group_devices: dict[ManufacturerEnum, Devices] = {}
+    for dev in devices or []:
+        if dev.manufacturer not in group_devices:
+            group_devices[dev.manufacturer] = []
+        group_devices[dev.manufacturer].append(dev)
+    return group_devices
+
+
 __all__ = [
     "Device",
     "Devices",
     "ManufacturerEnum",
+    "Topology",
     "backend_to_manufacturer",
     "detect_backend",
     "detect_devices",
+    "get_devices_topologies",
+    "group_devices_by_manufacturer",
     "manufacturer_to_backend",
+    "reduce_devices_distances",
     "supported_backends",
     "supported_list",
     "supported_manufacturers",