refactor: get topology

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

Author: thxCode

gpustack_runtime/cmds/detector.py

@@ -117,6 +117,20 @@ def run(self):
                 for topo in topologies:
                     print(format_topology_table(topo))
 
+                # Legend
+                legend_lines = [
+                    "",
+                    "Legend (from nearest to farthest):",
+                    "  X    = Self",
+                    "  LINK = Connection traversing with High-Speed Link (e.g., AMD XGMI, Ascend HCCS, NVIDIA NVLink)",
+                    "  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)",
+                ]
+                print(os.linesep.join(legend_lines))
+
 
 def format_devices_json(devs: Devices) -> str:
     return json.dumps([dev.to_dict() for dev in devs], indent=2)
@@ -259,21 +273,8 @@ def format_topology_table(topo: Topology) -> str:
         "+" + "-" * (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)
+    return os.linesep.join(header_lines + topology_lines + footer_lines)
 
 
 def format_topologies_json(topologies: list[Topology]) -> str:

gpustack_runtime/detector/__types__.py

@@ -251,21 +251,6 @@ class Topology:
     The value at index i represents the Memory set for device i.
     """
 
-    @staticmethod
-    def stringify_devices_distance(distance: int) -> str:
-        """
-        Stringify the devices distance to a human-readable format.
-
-        Args:
-            distance:
-                The distance between two devices.
-
-        Returns:
-            A string representing the distance.
-
-        """
-        return str(distance)
-
     def __init__(
         self,
         manufacturer: ManufacturerEnum,
@@ -298,7 +283,7 @@ def stringify(self) -> list[list[str]]:
         devices_info: list[list[str]] = [[]] * devices_count
         for i in range(devices_count):
             devices_info[i] = [
-                self.stringify_devices_distance(d) for d in self.devices_distances[i]
+                stringify_devices_distance(d) for d in self.devices_distances[i]
             ]
             devices_info[i] += [
                 self.devices_cpu_affinities[i]
@@ -313,6 +298,76 @@ def stringify(self) -> list[list[str]]:
         return devices_info
 
 
+class TopologyDistanceEnum(int, Enum):
+    """
+    Enum for Topology Distance Levels.
+    """
+
+    INTERNAL = 0
+    """
+    Same device.
+    """
+    LINK = 5
+    """
+    Connection traversing with High-Speed Link (e.g., AMD XGMI, Ascend HCCS, NVIDIA NVLink).
+    """
+    PIX = 10
+    """
+    Connection traversing at most a single PCIe bridge.
+    """
+    PXB = 20
+    """
+    Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge).
+    """
+    PHB = 30
+    """
+    Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU).
+    """
+    NODE = 40
+    """
+    Connection traversing PCIe and the interconnect between NUMA nodes.
+    """
+    SYS = 50
+    """
+    Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI).
+    """
+    UNKNOWN = 100
+    """
+    Unknown connection.
+    """
+
+
+def stringify_devices_distance(distance: int) -> str:
+    """
+    Stringify the devices distance to a human-readable format.
+
+    Args:
+        distance:
+            The distance between two devices.
+
+    Returns:
+        A string representing the distance.
+
+    """
+    match distance:
+        case 0:
+            return "X"
+        case 5:
+            return "LINK"
+        case 10:
+            return "PIX"
+        case 20:
+            return "PXB"
+        case 30:
+            return "PHB"
+        case 40:
+            return "NODE"
+        case 50:
+            return "SYS"
+        case _:
+            return "N/A"
+
+
 def reduce_devices_distances(
     devices_distances: list[list[int]],
 ) -> dict[int, list[int]]:

gpustack_runtime/detector/__utils__.py

@@ -673,6 +673,28 @@ def get_cpu_numa_node_mapping() -> list[int | None]:
     return mapping
 
 
+@lru_cache(maxsize=1)
+def get_numa_node_cpu_mapping() -> dict[int, list[int]]:
+    """
+    Map NUMA nodes to CPU cores.
+    The key corresponds to the NUMA node number,
+    and the value is the list of CPU cores that belong to that NUMA node.
+
+    Returns:
+        A dictionary mapping NUMA nodes to lists of CPU cores.
+
+    """
+    cpu_numa_mapping = get_cpu_numa_node_mapping()
+
+    numa_cpu_mapping: dict[int, list[int]] = {}
+    for cpu_idx, numa_node in enumerate(cpu_numa_mapping):
+        if numa_node is not None:
+            if numa_node not in numa_cpu_mapping:
+                numa_cpu_mapping[numa_node] = []
+            numa_cpu_mapping[numa_node].append(cpu_idx)
+    return numa_cpu_mapping
+
+
 @lru_cache
 def map_cpu_affinity_to_numa_node(cpu_affinity: int | str | None) -> str:
     """
@@ -722,6 +744,53 @@ def map_cpu_affinity_to_numa_node(cpu_affinity: int | str | None) -> str:
     return list_to_range_str(sorted(numa_nodes))
 
 
+@lru_cache
+def map_numa_node_to_cpu_affinity(numa_node: int | str | None) -> str:
+    """
+    Map NUMA nodes to CPU affinity.
+
+    Args:
+        numa_node:
+            The NUMA nodes as an integer bitmask or a string (e.g., "0-1,3").
+
+    Returns:
+        A comma-separated string of CPU core indices,
+        or blank string if no CPU cores are found.
+
+    """
+    if numa_node is None:
+        return ""
+
+    if isinstance(numa_node, int):
+        numa_indices = bits_to_list(numa_node)
+    else:
+        numa_indices: list[int] = []
+        for part in numa_node.split(","):
+            if "-" in part:
+                lo, hi = part.split("-")
+                lo_idx = safe_int(lo, -1)
+                hi_idx = safe_int(hi, -1)
+                if lo_idx == -1 or hi_idx == -1 or lo_idx > hi_idx:
+                    continue
+                numa_indices.extend(list(range(lo_idx, hi_idx + 1)))
+            else:
+                idx = safe_int(part, -1)
+                if idx == -1:
+                    continue
+                numa_indices.append(idx)
+
+    numa_cpu_mapping = get_numa_node_cpu_mapping()
+
+    cpu_cores: set[int] = set()
+    for numa_idx in numa_indices:
+        if numa_idx in numa_cpu_mapping:
+            cpu_cores.update(numa_cpu_mapping[numa_idx])
+    if not cpu_cores:
+        return ""
+
+    return list_to_range_str(sorted(cpu_cores))
+
+
 def bits_to_list(bits: int, offset: int = 0) -> list[int]:
     """
     Convert a bitmask to a list of set bit indices.
@@ -765,20 +834,18 @@ def bits_to_str(bits: int, offset: int = 0, prefix: str = "") -> str:
     if bits_list:
         if bits_str:
             bits_str += ","
-        bits_str += list_to_range_str(bits_list, is_sorted=True)
+        bits_str += list_to_range_str(bits_list)
 
     return bits_str
 
 
-def list_to_range_str(indices: list[int], is_sorted: bool = False) -> str:
+def list_to_range_str(indices: list[int]) -> str:
     """
     Convert a list of indices to a comma-separated string with ranges.
 
     Args:
         indices:
-            The list of indices.
-        is_sorted:
-            The indicates whether the input list is already sorted.
+            The list of indices, must be sorted in ascending order.
 
     Returns:
         A comma-separated string with ranges (e.g., "0,2-4,6").
@@ -787,17 +854,14 @@ def list_to_range_str(indices: list[int], is_sorted: bool = False) -> str:
     if not indices:
         return ""
 
-    sorted_indices = indices
-    if not is_sorted:
-        sorted_indices = sorted(set(indices))
-    if len(sorted_indices) == 1:
-        return f"{sorted_indices[0]}"
-    if len(sorted_indices) == (sorted_indices[-1] - sorted_indices[0] + 1):
-        return f"{sorted_indices[0]}-{sorted_indices[-1]}"
+    if len(indices) == 1:
+        return f"{indices[0]}"
+    if len(indices) == (indices[-1] - indices[0] + 1):
+        return f"{indices[0]}-{indices[-1]}"
 
-    start, end = sorted_indices[0], sorted_indices[0]
+    start, end = indices[0], indices[0]
     ranges: list[tuple[int, int]] = []
-    for i in sorted_indices[1:]:
+    for i in indices[1:]:
         if i == end + 1:
             end = i
         else: