refactor: get topology

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

Author: thxCode

gpustack_runtime/detector/__utils__.py

@@ -695,6 +695,30 @@ def get_numa_node_cpu_mapping() -> dict[int, list[int]]:
     return numa_cpu_mapping
 
 
+@lru_cache(maxsize=128)
+def get_numa_node_by_bdf(bdf: str) -> str:
+    """
+    Get the NUMA node for a given PCI device BDF (Bus:Device.Function) address.
+
+    Args:
+        bdf:
+            The PCI device BDF address (e.g., "0000:00:1f.0").
+
+    Returns:
+        The NUMA node index if found, otherwise blank string.
+
+    """
+    if bdf:
+        with contextlib.suppress(Exception):
+            with Path(f"/sys/bus/pci/devices/{bdf}/numa_node").open("r") as f:
+                s = f.read().strip()
+            numa_node = int(s)
+            if numa_node >= 0:
+                return str(numa_node)
+
+    return ""
+
+
 @lru_cache
 def map_cpu_affinity_to_numa_node(cpu_affinity: int | str | None) -> str:
     """
@@ -715,6 +739,8 @@ def map_cpu_affinity_to_numa_node(cpu_affinity: int | str | None) -> str:
     if isinstance(cpu_affinity, int):
         cpu_indices = bitmask_to_list(cpu_affinity)
     else:
+        if not cpu_affinity:
+            return ""
         cpu_indices: list[int] = []
         for part in cpu_affinity.split(","):
             if "-" in part:
@@ -764,6 +790,8 @@ def map_numa_node_to_cpu_affinity(numa_node: int | str | None) -> str:
     if isinstance(numa_node, int):
         numa_indices = bitmask_to_list(numa_node)
     else:
+        if not numa_node:
+            return ""
         numa_indices: list[int] = []
         for part in numa_node.split(","):
             if "-" in part:
@@ -810,29 +838,6 @@ def bitmask_to_list(bitmask: int, offset: int = 0) -> list[int]:
     return indices
 
 
-def bitmask_to_str(bitmask_list: list) -> str:
-    """
-    Convert a bitmask to a comma-separated string of set bit indices.
-
-    Args:
-        bitmask_list:
-            An integer list stores each item in bitmask.
-
-    Returns:
-        If a bitmask are contiguous, returns a range string (e.g., "2-5"),
-        Otherwise, returns a comma-separated string of indices (e.g., "0,2-4").
-
-    """
-    bits_lists = []
-    offset = 0
-    for bitmask in bitmask_list:
-        if bitmask != 0:
-            bits_lists.extend(bitmask_to_list(bitmask, offset))
-        offset += get_bits_size()
-
-    return list_to_range_str(sorted(bits_lists))
-
-
 def list_to_range_str(indices: list[int]) -> str:
     """
     Convert a list of indices to a comma-separated string with ranges.
@@ -872,3 +877,26 @@ def list_to_range_str(indices: list[int]) -> str:
     range_str = ",".join(range_str_parts)
 
     return range_str
+
+
+def bitmask_to_str(bitmask_list: list) -> str:
+    """
+    Convert a bitmask to a comma-separated string of set bit indices.
+
+    Args:
+        bitmask_list:
+            An integer list stores each item in bitmask.
+
+    Returns:
+        If a bitmask are contiguous, returns a range string (e.g., "2-5"),
+        Otherwise, returns a comma-separated string of indices (e.g., "0,2-4").
+
+    """
+    bits_lists = []
+    offset = 0
+    for bitmask in bitmask_list:
+        if bitmask != 0:
+            bits_lists.extend(bitmask_to_list(bitmask, offset))
+        offset += get_bits_size()
+
+    return list_to_range_str(sorted(bits_lists))

gpustack_runtime/detector/amd.py

@@ -12,6 +12,7 @@
     PCIDevice,
     byte_to_mebibyte,
     get_brief_version,
+    get_numa_node_by_bdf,
     get_pci_devices,
     get_utilization,
     map_numa_node_to_cpu_affinity,
@@ -216,7 +217,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)
+                    dev_appendix["bdf"] = str(dev_bdf).lower()
 
                 ret.append(
                     Device(
@@ -277,10 +278,8 @@ def get_topology(self, devices: Devices | None = None) -> Topology | None:
         devs_mapping = None
 
         def get_device_handle(dev: Device):
-            if "bdf" in dev.appendix:
-                return pyamdsmi.amdsmi_get_processor_handle_from_bdf(
-                    dev.appendix["bdf"],
-                )
+            if bdf := dev.appendix.get("bdf", None):
+                return pyamdsmi.amdsmi_get_processor_handle_from_bdf(bdf)
             nonlocal devs_mapping
             if devs_mapping is None:
                 devs = pyamdsmi.amdsmi_get_processor_handles()
@@ -293,24 +292,32 @@ def get_device_handle(dev: Device):
             for i, dev_i in enumerate(devices):
                 dev_i_handle = get_device_handle(dev_i)
 
-                # Get NUMA affinity.
-                try:
-                    dev_i_numa_node = pyamdsmi.amdsmi_topo_get_numa_node_number(
-                        dev_i_handle,
+                # 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_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
+                        numa_node,
                     )
-                    topology.devices_numa_affinities[i] = str(dev_i_numa_node)
-                except pyamdsmi.AmdSmiException:
-                    debug_log_exception(
-                        logger,
-                        "Failed to get NUMA affinity for device %d",
-                        dev_i.index,
+                # Otherwise, get affinity via AMD SMI.
+                if not topology.devices_cpu_affinities[i]:
+                    # Get NUMA affinity.
+                    try:
+                        dev_i_numa_node = pyamdsmi.amdsmi_topo_get_numa_node_number(
+                            dev_i_handle,
+                        )
+                        topology.devices_numa_affinities[i] = str(dev_i_numa_node)
+                    except pyamdsmi.AmdSmiException:
+                        debug_log_exception(
+                            logger,
+                            "Failed to get NUMA affinity for device %d",
+                            dev_i.index,
+                        )
+                    # Get CPU affinity.
+                    topology.devices_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
+                        numa_node=topology.devices_numa_affinities[i],
                     )
 
-                # Get CPU affinity.
-                topology.devices_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
-                    numa_node=topology.devices_numa_affinities[i],
-                )
-
                 # Get distances to other devices.
                 for j, dev_j in enumerate(devices):
                     if i == j:

gpustack_runtime/detector/ascend.py

@@ -18,9 +18,11 @@
 from .__utils__ import (
     PCIDevice,
     get_brief_version,
+    get_numa_node_by_bdf,
     get_pci_devices,
     get_utilization,
     map_cpu_affinity_to_numa_node,
+    map_numa_node_to_cpu_affinity,
 )
 
 logger = logging.getLogger(__name__)
@@ -111,9 +113,6 @@ def detect(self) -> Devices | None:
 
             _, card_list = pydcmi.dcmi_get_card_list()
             for dev_card_id in card_list:
-                device_id_max_in_card, _, _ = pydcmi.dcmi_get_device_id_in_card(
-                    dev_card_id,
-                )
                 device_num_in_card = pydcmi.dcmi_get_device_num_in_card(dev_card_id)
                 for dev_device_id in range(device_num_in_card):
                     dev_is_vgpu = False
@@ -193,7 +192,7 @@ def detect(self) -> Devices | None:
                         "vgpu": dev_is_vgpu,
                         "card_id": dev_card_id,
                         "device_id": dev_device_id,
-                        "device_id_max": device_id_max_in_card - 1,
+                        "device_id_max": device_num_in_card - 1,
                     }
 
                     dev_roce_ip, dev_roce_mask, dev_roce_gateway = (
@@ -209,6 +208,13 @@ def detect(self) -> Devices | None:
                     if dev_roce_gateway:
                         dev_appendix["roce_gateway"] = str(dev_roce_gateway)
 
+                    with contextlib.suppress(pydcmi.DCMIError):
+                        dev_bdf = pydcmi.dcmi_get_device_bdf(
+                            dev_card_id,
+                            dev_device_id,
+                        )
+                        dev_appendix["bdf"] = str(dev_bdf).lower()
+
                     ret.append(
                         Device(
                             manufacturer=self.manufacturer,
@@ -265,39 +271,58 @@ def get_topology(self, devices: Devices | None = None) -> Topology | None:
             pydcmi.dcmi_init()
 
             for i, dev_i in enumerate(devices):
-                # Get CPU affinity.
-                try:
-                    cpu_affinity = pydcmi.dcmi_get_affinity_cpu_info_by_device_id(
-                        dev_i.appendix["card_id"],
-                        dev_i.appendix["device_id"],
+                dev_i_card_id = dev_i.appendix["card_id"]
+                dev_i_device_id = dev_i.appendix["device_id"]
+
+                # 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_cpu_affinities[i] = map_numa_node_to_cpu_affinity(
+                        numa_node,
                     )
-                    topology.devices_cpu_affinities[i] = cpu_affinity
-                except pydcmi.DCMIError:
-                    debug_log_exception(
-                        slogger,
-                        "Failed to get CPU affinity for device %d",
-                        dev_i.index,
+                # Otherwise, get affinity via DCMI.
+                if not topology.devices_cpu_affinities[i]:
+                    # Get CPU affinity.
+                    try:
+                        cpu_affinity = pydcmi.dcmi_get_affinity_cpu_info_by_device_id(
+                            dev_i.appendix["card_id"],
+                            dev_i.appendix["device_id"],
+                        )
+                        topology.devices_cpu_affinities[i] = cpu_affinity
+                    except pydcmi.DCMIError:
+                        debug_log_exception(
+                            slogger,
+                            "Failed to get CPU affinity for device %d",
+                            dev_i.index,
+                        )
+                    # Get NUMA affinity.
+                    topology.devices_numa_affinities[i] = map_cpu_affinity_to_numa_node(
+                        cpu_affinity=topology.devices_cpu_affinities[i],
                     )
 
-                # Get NUMA affinity.
-                topology.devices_numa_affinities[i] = map_cpu_affinity_to_numa_node(
-                    cpu_affinity=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:
                         continue
 
+                    dev_j_card_id = dev_j.appendix["card_id"]
+                    dev_j_device_id = dev_j.appendix["device_id"]
+
+                    # If two devices are the same card,
+                    # skip distance calculation.
+                    if dev_i_card_id == dev_j_card_id:
+                        continue
+
                     distance = TopologyDistanceEnum.UNK
                     try:
                         topo_type = pydcmi.dcmi_get_topo_info_by_device_id(
-                            dev_i.appendix["card_id"],
-                            dev_i.appendix["device_id"],
-                            dev_j.appendix["card_id"],
-                            dev_j.appendix["device_id"],
+                            dev_i_card_id,
+                            dev_i_device_id,
+                            dev_j_card_id,
+                            dev_j_device_id,
                         )
                         distance = _TOPO_TYPE_DISTANCE_MAPPING.get(topo_type, distance)
                     except pydcmi.DCMIError: