refactor: get topology

with test modification

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

Author: thxCode

gpustack_runtime/cmds/detector.py

@@ -128,6 +128,7 @@ def run(self):
                     "  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)",
+                    "  N/A  = Unknown or unavailable information",
                 ]
                 print(os.linesep.join(legend_lines))
 

gpustack_runtime/detector/__types__.py

@@ -303,9 +303,9 @@ class TopologyDistanceEnum(int, Enum):
     Enum for Topology Distance Levels.
     """
 
-    INTERNAL = 0
+    SELF = 0
     """
-    Same device.
+    Self connection.
     """
     LINK = 5
     """
@@ -331,7 +331,7 @@ class TopologyDistanceEnum(int, Enum):
     """
     Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI).
     """
-    UNKNOWN = 100
+    UNK = 100
     """
     Unknown connection.
     """
@@ -350,19 +350,19 @@ def stringify_devices_distance(distance: int) -> str:
 
     """
     match distance:
-        case 0:
+        case TopologyDistanceEnum.SELF:
             return "X"
-        case 5:
+        case TopologyDistanceEnum.LINK:
             return "LINK"
-        case 10:
+        case TopologyDistanceEnum.PIX:
             return "PIX"
-        case 20:
+        case TopologyDistanceEnum.PXB:
             return "PXB"
-        case 30:
+        case TopologyDistanceEnum.PHB:
             return "PHB"
-        case 40:
+        case TopologyDistanceEnum.NODE:
             return "NODE"
-        case 50:
+        case TopologyDistanceEnum.SYS:
             return "SYS"
         case _:
             return "N/A"
@@ -450,7 +450,7 @@ def detect(self) -> Devices | None:
         """
         raise NotImplementedError
 
-    def get_topology(self, devices: Devices | None) -> Topology | None:  # noqa: ARG002
+    def get_topology(self, devices: Devices | None = None) -> Topology | None:  # noqa: ARG002
         """
         Get the Topology object between the given devices.
 

gpustack_runtime/detector/amd.py

@@ -250,7 +250,7 @@ def detect(self) -> Devices | None:
 
         return ret
 
-    def get_topology(self, devices: Devices | None) -> Topology | None:
+    def get_topology(self, devices: Devices | None = None) -> Topology | None:
         """
         Get the Topology object between AMD GPUs.
 
@@ -320,15 +320,15 @@ def get_device_handle(dev: Device):
 
                     dev_j_handle = get_device_handle(dev_j)
 
-                    distance = TopologyDistanceEnum.UNKNOWN
+                    distance = TopologyDistanceEnum.UNK
                     try:
                         link_type = pyamdsmi.amdsmi_topo_get_link_type(
                             dev_i_handle,
                             dev_j_handle,
                         )
                         match int(link_type.type):
                             case pyamdsmi.AMDSMI_LINK_TYPE_INTERNAL:
-                                distance = TopologyDistanceEnum.INTERNAL
+                                distance = TopologyDistanceEnum.SELF
                             case pyamdsmi.AMDSMI_LINK_TYPE_XGMI:
                                 distance = TopologyDistanceEnum.LINK
                             # For PCIe links,

gpustack_runtime/detector/ascend.py

@@ -27,7 +27,7 @@
 slogger = logger.getChild("internal")
 
 _TOPO_TYPE_DISTANCE_MAPPING: dict[int, int] = {
-    pydcmi.DCMI_TOPO_TYPE_SELF: TopologyDistanceEnum.INTERNAL,
+    pydcmi.DCMI_TOPO_TYPE_SELF: TopologyDistanceEnum.SELF,
     pydcmi.DCMI_TOPO_TYPE_HCCS: TopologyDistanceEnum.LINK,  # Traversing via high-speed interconnect, RoCE, etc.
     pydcmi.DCMI_TOPO_TYPE_PIX: TopologyDistanceEnum.PIX,  # Traversing via a single PCIe bridge.
     pydcmi.DCMI_TOPO_TYPE_PXB: TopologyDistanceEnum.PXB,  # Traversing via multiple PCIe bridges without PCIe Host Bridge.
@@ -111,6 +111,9 @@ 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
@@ -190,6 +193,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,
                     }
 
                     dev_roce_ip, dev_roce_mask, dev_roce_gateway = (
@@ -233,7 +237,7 @@ def detect(self) -> Devices | None:
 
         return ret
 
-    def get_topology(self, devices: Devices | None) -> Topology | None:
+    def get_topology(self, devices: Devices | None = None) -> Topology | None:
         """
         Get the Topology object between Ascend NPUs.
 
@@ -258,6 +262,8 @@ def get_topology(self, devices: Devices | None) -> Topology | None:
         )
 
         try:
+            pydcmi.dcmi_init()
+
             for i, dev_i in enumerate(devices):
                 # Get CPU affinity.
                 try:
@@ -285,7 +291,7 @@ def get_topology(self, devices: Devices | None) -> Topology | None:
                     if topology.devices_distances[i][j] != 0:
                         continue
 
-                    distance = TopologyDistanceEnum.UNKNOWN
+                    distance = TopologyDistanceEnum.UNK
                     try:
                         topo_type = pydcmi.dcmi_get_topo_info_by_device_id(
                             dev_i.appendix["card_id"],

gpustack_runtime/detector/nvidia.py

@@ -393,7 +393,7 @@ def detect(self) -> Devices | None:  # noqa: PLR0915
 
         return ret
 
-    def get_topology(self, devices: Devices | None) -> Topology | None:
+    def get_topology(self, devices: Devices | None = None) -> Topology | None:
         """
         Get the Topology object between NVIDIA GPUs.
 
@@ -485,7 +485,7 @@ def get_topology(self, devices: Devices | None) -> Topology | None:
 
                     dev_j_handle = pynvml.nvmlDeviceGetHandleByUUID(dev_j.uuid)
 
-                    distance = TopologyDistanceEnum.UNKNOWN
+                    distance = TopologyDistanceEnum.UNK
                     try:
                         distance = pynvml.nvmlDeviceGetTopologyCommonAncestor(
                             dev_i_handle,