refactor: get topology

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

Author: thxCode

gpustack_runtime/detector/__utils__.py

@@ -24,6 +24,14 @@ class PCIDevice:
     """
     Path to the PCI device in sysfs.
     """
+    root: str
+    """
+    Root of the PCI device.
+    """
+    switches: list[str]
+    """
+    Switches of the PCI device.
+    """
     address: str
     """
     Address of the PCI device.
@@ -64,7 +72,7 @@ def get_pci_devices(
         vendor = [vendor]
 
     for dev_path in sysfs_pci_path.iterdir():
-        dev_address = dev_path.name
+        dev_address = dev_path.name.lower()
         if address and dev_address not in address:
             continue
 
@@ -92,10 +100,19 @@ def get_pci_devices(
         if dev_class is None or dev_config is None:
             continue
 
+        dev_path_resolved = dev_path.resolve()
+        dev_switches = []
+        dev_root = dev_path_resolved.parent
+        while (dev_root != sysfs_pci_path) and (dev_root.name.count(":") == 2):
+            dev_switches.append(dev_root.name)
+            dev_root = dev_root.parent
+
         pci_devices.append(
             PCIDevice(
                 vendor=dev_vendor,
-                path=str(dev_path),
+                path=str(dev_path_resolved),
+                root=dev_root.name,
+                switches=dev_switches,
                 address=dev_address,
                 class_=dev_class,
                 config=dev_config,
@@ -105,6 +122,40 @@ def get_pci_devices(
     return pci_devices
 
 
+def compare_pci_devices(
+    dev_a: PCIDevice | None,
+    dev_b: PCIDevice | None,
+) -> int:
+    """
+    Compare two PCI devices.
+
+    Args:
+        dev_a:
+            The first PCI device.
+        dev_b:
+            The second PCI device.
+
+    Returns:
+        -1 if devices have different roots,
+        0 if devices have the same root but different switches,
+        1 if devices have the same root and same switches.
+
+    """
+    if dev_a and dev_b:
+        is_same_root = dev_a.root == dev_b.root
+        is_same_switch = is_same_root and len(dev_a.switches) == len(dev_b.switches)
+        if is_same_switch:
+            for sw_a, sw_b in zip(dev_a.switches, dev_b.switches, strict=False):
+                if sw_a != sw_b:
+                    is_same_switch = False
+                    break
+        if is_same_switch:
+            return 1
+        if is_same_root:
+            return 0
+    return -1
+
+
 @dataclass
 class DeviceFile:
     path: str

gpustack_runtime/detector/amd.py

@@ -11,6 +11,7 @@
 from .__utils__ import (
     PCIDevice,
     byte_to_mebibyte,
+    compare_pci_devices,
     get_brief_version,
     get_numa_node_by_bdf,
     get_pci_devices,
@@ -217,7 +218,7 @@ def detect(self) -> Devices | None:
 
                 with contextlib.suppress(pyamdsmi.AmdSmiException):
                     dev_bdf = pyamdsmi.amdsmi_get_gpu_device_bdf(dev)
-                    dev_appendix["bdf"] = str(dev_bdf).lower()
+                    dev_appendix["bdf"] = dev_bdf
 
                 ret.append(
                     Device(
@@ -287,20 +288,48 @@ def get_device_handle(dev: Device):
             return devs_mapping.get(dev.index)
 
         try:
+            pci_devices = self.detect_pci_devices()
+
+            def distance_pci_devices(bdf_a: str, bdf_b: str) -> TopologyDistanceEnum:
+                """
+                Compute distance between two PCI devices by their BDFs.
+
+                Args:
+                    bdf_a:
+                        The BDF of the first PCI device.
+                    bdf_b:
+                        The BDF of the second PCI device.
+
+                Returns:
+                    The TopologyDistanceEnum representing the distance.
+
+                """
+                pcid_a = pci_devices.get(bdf_a, None)
+                pcid_b = pci_devices.get(bdf_b, None)
+
+                score = compare_pci_devices(pcid_a, pcid_b)
+                if score > 0:
+                    return TopologyDistanceEnum.PIX
+                if score == 0:
+                    return TopologyDistanceEnum.PXB
+                return TopologyDistanceEnum.PHB
+
             pyamdsmi.amdsmi_init()
 
+            # Get NUMA and CPU affinities.
             for i, dev_i in enumerate(devices):
-                dev_i_handle = get_device_handle(dev_i)
-
                 # Get affinity with PCIe BDF if possible.
                 if dev_i_bdf := dev_i.appendix.get("bdf", ""):
-                    numa_node = get_numa_node_by_bdf(dev_i_bdf)
-                    topology.devices_numa_affinities[i] = numa_node
+                    topology.devices_numa_affinities[i] = get_numa_node_by_bdf(
+                        dev_i_bdf,
+                    )
                     topology.devices_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
-                        numa_node,
+                        topology.devices_numa_affinities[i],
                     )
                 # Otherwise, get affinity via AMD SMI.
                 if not topology.devices_cpu_affinities[i]:
+                    dev_i_handle = get_device_handle(dev_i)
+
                     # Get NUMA affinity.
                     try:
                         dev_i_numa_node = pyamdsmi.amdsmi_topo_get_numa_node_number(
@@ -315,41 +344,50 @@ def get_device_handle(dev: Device):
                         )
                     # Get CPU affinity.
                     topology.devices_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
-                        numa_node=topology.devices_numa_affinities[i],
+                        topology.devices_numa_affinities[i],
                     )
 
-                # Get distances to other devices.
+            # Get distances to other devices.
+            for i, dev_i in enumerate(devices):
+                dev_i_handle = get_device_handle(dev_i)
+
                 for j, dev_j in enumerate(devices):
-                    if i == j:
-                        continue
-                    if topology.devices_distances[i][j] != 0:
+                    if (
+                        dev_i.index == dev_j.index
+                        or topology.devices_distances[i][j] != 0
+                    ):
                         continue
 
                     dev_j_handle = get_device_handle(dev_j)
 
                     distance = TopologyDistanceEnum.UNK
                     try:
-                        link_type = pyamdsmi.amdsmi_topo_get_link_type(
+                        link = pyamdsmi.amdsmi_topo_get_link_type(
                             dev_i_handle,
                             dev_j_handle,
                         )
-                        match int(link_type.type):
+                        link_type = link.get("type", -1)
+                        link_hops = link.get("hops", -1)
+                        match link_type:
                             case pyamdsmi.AMDSMI_LINK_TYPE_INTERNAL:
                                 distance = TopologyDistanceEnum.SELF
-                            case pyamdsmi.AMDSMI_LINK_TYPE_XGMI:
-                                distance = TopologyDistanceEnum.LINK
-                            # For PCIe links,
-                            # further distinguish between PHB and SYS based on NUMA affinity.
                             case pyamdsmi.AMDSMI_LINK_TYPE_PCIE:
-                                if link_type == pyamdsmi.AMDSMI_LINK_TYPE_PCIE:
-                                    dev_i_numa, dev_j_numa = (
-                                        topology.devices_numa_affinities[i],
-                                        topology.devices_numa_affinities[j],
+                                dev_i_numa, dev_j_numa = (
+                                    topology.devices_numa_affinities[i],
+                                    topology.devices_numa_affinities[j],
+                                )
+                                if dev_i_numa and dev_i_numa == dev_j_numa:
+                                    distance = distance_pci_devices(
+                                        dev_i.appendix.get("bdf", ""),
+                                        dev_j.appendix.get("bdf", ""),
                                     )
-                                    if dev_i_numa != "" and dev_i_numa == dev_j_numa:
-                                        distance = TopologyDistanceEnum.PHB
-                                    else:
-                                        distance = TopologyDistanceEnum.SYS
+                                else:
+                                    distance = TopologyDistanceEnum.SYS
+                            case pyamdsmi.AMDSMI_LINK_TYPE_XGMI:
+                                distance = TopologyDistanceEnum.LINK
+                            case _:
+                                if link_hops == 0:
+                                    distance = TopologyDistanceEnum.SELF
                     except pyamdsmi.AmdSmiException:
                         debug_log_exception(
                             logger,
@@ -361,7 +399,7 @@ def get_device_handle(dev: Device):
                     topology.devices_distances[i][j] = distance
                     topology.devices_distances[j][i] = distance
         except pyamdsmi.AmdSmiException:
-            debug_log_exception(logger, "Failed to get topology")
+            debug_log_exception(logger, "Failed to fetch topology")
             raise
         except Exception:
             debug_log_exception(logger, "Failed to process topology fetching")

gpustack_runtime/detector/ascend.py

@@ -213,7 +213,7 @@ def detect(self) -> Devices | None:
                             dev_card_id,
                             dev_device_id,
                         )
-                        dev_appendix["bdf"] = str(dev_bdf).lower()
+                        dev_appendix["bdf"] = dev_bdf
 
                     ret.append(
                         Device(
@@ -276,10 +276,11 @@ def get_topology(self, devices: Devices | None = None) -> Topology | None:
 
                 # Get affinity with PCIe BDF if possible.
                 if dev_i_bdf := dev_i.appendix.get("bdf", ""):
-                    numa_node = get_numa_node_by_bdf(dev_i_bdf)
-                    topology.devices_numa_affinities[i] = numa_node
+                    topology.devices_numa_affinities[i] = get_numa_node_by_bdf(
+                        dev_i_bdf,
+                    )
                     topology.devices_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
-                        numa_node,
+                        topology.devices_numa_affinities[i],
                     )
                 # Otherwise, get affinity via DCMI.
                 if not topology.devices_cpu_affinities[i]:
@@ -298,14 +299,15 @@ def get_topology(self, devices: Devices | None = None) -> Topology | None:
                         )
                     # Get NUMA affinity.
                     topology.devices_numa_affinities[i] = map_cpu_affinity_to_numa_node(
-                        cpu_affinity=topology.devices_cpu_affinities[i],
+                        topology.devices_cpu_affinities[i],
                     )
 
                 # Get distances to other devices.
                 for j, dev_j in enumerate(devices):
-                    if i == j:
-                        continue
-                    if topology.devices_distances[i][j] != 0:
+                    if (
+                        dev_i.index == dev_j.index
+                        or topology.devices_distances[i][j] != 0
+                    ):
                         continue
 
                     dev_j_card_id = dev_j.appendix["card_id"]