bug: fix graphic to show entire tree

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

Author: vsoch

fluxbind/graph/graph.py

@@ -108,6 +108,10 @@ def build_graph(self, element, parent_gp=None):
     def discover_devices(self):
         """
         Discover different kinds of GPU and NIC.
+
+        TODO: a better design would be to discover devices, then annotate the graph
+        nodes with device type. Then when we search for gpu or nic, we can just
+        search by that attribute. Right now we search for PCIDev, and then filter.
         """
         # Don't assume only one vendor of GPU.
         for vendor, command in {"NVIDIA": commands.nvidia_smi, "AMD": commands.rocm_smi}.items():
@@ -180,7 +184,7 @@ def add_latency_edges(matrix, indexes):
                             weight=matrix[i][j],  # The weight is the latency value from the matrix.
                         )
 
-        # --- First, try to parse the modern hwloc v2.x format ---
+        # First, try to parse newer hwloc v2.x format
         for dist_el in root.iter("distances2"):
             try:
                 # nbobjs is how many objects are in the matrix (e.g., 2 for 2 NUMA nodes).
@@ -274,15 +278,18 @@ def match_resources(self, jobspec, allocated_gps=None):
         log.debug(f"Successfully found a slot with {len(final_allocation)} objects.")
         return final_allocation
 
-    def sort_by_affinity(self, candidates, affinity, allocated):
+    def sort_by_affinity(self, candidates, affinity, allocated, domain_gp):
+        """
+        Sort list of candidates by affinity so we get closest one.
+        """
         target_type = self.translate_type(affinity.get("type"))
         if not target_type:
             log.warning("Affinity spec missing 'type'.")
             return candidates
-        machine_gp = self.find_objects(type="Machine")[0][0]
-
-        # Affinity targets can be anywhere, so search from Machine
-        targets = self.get_available_children(machine_gp, target_type, allocated)
+        
+        # Search within the domain we were provided, not across the machine
+        log.debug(f"    -> Searching for affinity target '{target_type}' within the current domain.")
+        targets = self.get_available_children(domain_gp, target_type, allocated)        
         if not targets:
             log.warning(f"Affinity target '{target_type}' not found.")
             return candidates
@@ -323,19 +330,82 @@ def translate_type(self, requested_type: str):
         # capitalizing the word (e.g., 'l3cache' -> 'L3cache').
         return mapping.get(requested_type.lower(), requested_type.capitalize())
 
+    def find_bindable_leaves(self, total_allocation, bind_level):
+        """
+        Transforms a list of allocated resources into a final list of bindable
+        nodes by correctly handling all relationships for EACH allocated object.
+        """
+        leaf_nodes = []
+        log.debug(f"Transforming {len(total_allocation)} allocated objects to bind_level '{bind_level}'...")
+        
+        bind_type_concrete = self.translate_type(bind_level)
+
+        # Iterate through every object found by the allocator.
+        for gp, data in total_allocation:
+            
+            # Case 1 (Identity): The object is already the type we want to bind to.
+            if data.get('type') == bind_type_concrete:
+                leaf_nodes.append((gp, data))
+                continue
+
+            # Case 2 (Container): The object is a container. Find its children of the target type.
+            found_descendants = False
+            container_type = data.get("type")
+
+            # If the container is a NUMANode, the relationship is via the parent Package.
+            if container_type == "NUMANode":
+                package = self.get_ancestor_of_type(gp, "Package")
+                if package:
+                    descendants = self.get_descendants(package[0], type=bind_type_concrete)
+                    if descendants:
+                        leaf_nodes.extend(descendants)
+                        found_descendants = True
+            
+            # If the container is a PCIDev, the relationship is via NUMA locality.
+            elif container_type == "PCIDev":
+                numa_idx = data.get("numa_os_index")
+                if numa_idx is not None:
+                    all_bindable = self.find_objects(type=bind_type_concrete)
+                    local_nodes = [node for node in all_bindable if node[1].get("numa_os_index") == numa_idx]
+                    if local_nodes:
+                        leaf_nodes.extend(local_nodes)
+                        found_descendants = True
+            
+            # For all other containers (like Package), the relationship is simple hierarchy.
+            else:
+                descendants = self.get_descendants(gp, type=bind_type_concrete)
+                if descendants:
+                    leaf_nodes.extend(descendants)
+                    found_descendants = True
+            
+            # If we successfully found descendants, onto the next!
+            if found_descendants:
+                continue
+
+            # Last case (Child): If it's not the right type and not a container, it might be a child.
+            # Find its ancestor of the target type (e.g., allocated PU, bind to Core).
+            ancestor = self.get_ancestor_of_type(gp, bind_type_concrete)
+            if ancestor:
+                leaf_nodes.append(ancestor)
+
+        # De-duplicate the final list and sort for deterministic assignment.
+        leaf_nodes = list({gp: (gp, data) for gp, data in leaf_nodes}.values())
+        leaf_nodes.sort(key=self.get_sort_key_for_node)
+        
+        log.debug(f"Transformation resulted in {len(leaf_nodes)} unique bindable leaf nodes.")
+        return leaf_nodes
+
     def summarize(self, nodes):
         """
         Given a set of nodes in the graph (a set of resources) print a textual visual.
         """
         for gp, data in nodes:
             p_info = ""
             if data["type"] in ["Core", "PU"]:
-                # BUGFIX: Changed self.topology to self
                 package = self.get_ancestor_of_type(gp, "Package")
                 if package:
                     p_info = f"Package:{package[1].get('os_index')} -> "
                 if data["type"] == "PU":
-                    # BUGFIX: Changed self.topology to self
                     core = self.get_ancestor_of_type(gp, "Core")
                     if core:
                         p_info += f"Core:{core[1].get('os_index')} -> "
@@ -393,22 +463,7 @@ def get_available_children(self, parent_gp, child_type, allocated):
 
         parent_type = parent_node.get("type")
         child_type_lower = child_type.lower()
-
-        candidate_gp_set = set()
-        if child_type_lower == "gpu":
-            for gpu_info in self.gpus:
-                if (pci_id := gpu_info.get("pci_bus_id")) and (
-                    matches := self.find_objects(type="PCIDev", pci_busid=pci_id)
-                ):
-                    candidate_gp_set.add(matches[0][0])
-        elif child_type_lower == "nic":
-            for nic_gp, _ in self.nics:
-                candidate_gp_set.add(nic_gp)
-        else:
-            for gp, _ in self.find_objects(type=child_type):
-                candidate_gp_set.add(gp)
-
-        all_candidates = [(gp, self.graph.nodes[gp]) for gp in candidate_gp_set]
+        all_candidates = self.find_objects(type=child_type)
         log.debug(
             f"    -   -> Found {len(all_candidates)} total unique system-wide candidates for '{child_type}'."
         )
@@ -422,14 +477,14 @@ def get_available_children(self, parent_gp, child_type, allocated):
 
             # Rule 1: Relationship to NUMA node is through shared PACKAGE parent (for Cores) or LOCALITY (for Devices)
             if parent_type == "NUMANode":
-                if child_type_lower in ["gpu", "nic"]:
-                    if data.get("numa_os_index") == parent_node.get("os_index"):
-                        is_valid_child = True
-                elif child_type_lower in ["core", "pu"]:
+                if child_type_lower in ["core", "pu"]:
                     package_of_numa = self.get_ancestor_of_type(parent_gp, "Package")
                     if package_of_numa and nx.has_path(self.hierarchy_view, package_of_numa[0], gp):
                         is_valid_child = True
-
+                elif child_type_lower == "pcidev":
+                    if data.get("numa_os_index") == parent_node.get("os_index"):
+                        is_valid_child = True
+            
             # Rule 2 (NEW): Relationship of a Core/PU to a Device is through shared NUMA LOCALITY
             elif parent_type == "PCIDev" and child_type_lower in ["core", "pu"]:
                 parent_numa_idx = parent_node.get("numa_os_index")
@@ -449,8 +504,9 @@ def get_available_children(self, parent_gp, child_type, allocated):
         return available
 
     def find_objects(self, **attributes):
-        # STOPED HERE, trace this and see if gp is globally identifier, add comments.
-        # then go back to making distances/distances2 function and add comments.
+        """
+        Search nodes in the graph for a specific attribute (or more than one)
+        """
         return [
             (gp, data)
             for gp, data in self.graph.nodes(data=True)
@@ -468,9 +524,23 @@ def find_assignment_recursive(self, request, domain_gp, allocated, depth=0):
         domain_info = f"{domain_node.get('type')}:{domain_node.get('os_index', domain_node.get('pci_busid', domain_gp))}"
         log.debug(f"{indent}[ENTER] find_assignment(req={request}, domain={domain_info})")
 
-        req_type, count = self.translate_type(request["type"]), request["count"]
+        # This can also be gpu/nic
+        raw_request_type = request['type']
+        req_type, count = self.translate_type(raw_request_type), request["count"]
 
+        # If the req_type is gpu or nic, this isn't an actual type in the graph - it is PCIDev.
         candidates = self.get_available_children(domain_gp, req_type, allocated)
+
+        # Now we handle the type of the pcidev request and filter candidates to those devices.
+        if raw_request_type.lower() == 'gpu':
+            gpu_bus_ids = {g['pci_bus_id'] for g in self.gpus}
+            candidates = [node for node in candidates if node[1].get('pci_busid') in gpu_bus_ids]
+            log.debug(f"{indent}  -> Filtered for 'gpu', {len(candidates)} candidates remain.")
+        elif raw_request_type.lower() == 'nic':
+            nic_bus_ids = {n[1]['pci_busid'] for n in self.nics}
+            candidates = [node for node in candidates if node[1].get('pci_busid') in nic_bus_ids]
+            log.debug(f"{indent}  -> Filtered for 'nic', {len(candidates)} candidates remain.")
+
         log.debug(f"{indent}  -> Found {len(candidates)} initial candidates for '{req_type}'.")
 
         affinity_spec = request.get("affinity")
@@ -496,7 +566,7 @@ def find_assignment_recursive(self, request, domain_gp, allocated, depth=0):
                 self.last_affinity_target = (target_gp, domain_node)
             else:
                 log.debug(f"{indent}  -> Sorting candidates by GLOBAL affinity to {affinity_spec}")
-                candidates = self.sort_by_affinity(candidates, affinity_spec, allocated)
+                candidates = self.sort_by_affinity(candidates, affinity_spec, allocated, domain_gp)
 
         if len(candidates) < count:
             log.debug(
@@ -550,6 +620,9 @@ def find_assignment_recursive(self, request, domain_gp, allocated, depth=0):
         return None
 
     def get_descendants(self, gp_index, **filters):
+        """
+        Given a global position index, return descendents that match a filter.
+        """
         if gp_index not in self.graph:
             return []
         desc = list(nx.descendants(self.hierarchy_view, gp_index))
@@ -648,14 +721,17 @@ def create_ordered_gpus(self):
         precompute_numa_affinities are done.
         """
         ordered_gpus = []
+
+        # The gpus we found were discovered with nvidia/rocm-smi and we need
+        # to map to things in the graph.
         for gpu_info in self.gpus:
             pci_id = gpu_info.get("pci_bus_id")
             if not pci_id:
                 continue
 
             # Find the corresponding PCIDev object in our graph
             # Note: We now store and search for types in lowercase
-            matches = self.find_objects(type="pcidev", pci_busid=pci_id)
+            matches = self.find_objects(type="PCIDev", pci_busid=pci_id)
             if not matches:
                 log.warning(
                     f"Could not find a graph object for discovered GPU with PCI ID: {pci_id}"

fluxbind/graph/graphic.py

@@ -37,15 +37,13 @@ def _build_contextual_subgraph(self):
         if not self.assigned_nodes:
             return nx.DiGraph()
 
-        # Step 1: Identify the type of resource we are drawing (e.g., 'core', 'pu').
         leaf_type = self.assigned_nodes[0][1].get("type")
         if not leaf_type:
             return nx.DiGraph()
 
-        # Step 2: Find a common parent container for the allocated resources.
         first_node_gp = self.assigned_nodes[0][0]
 
-        # Use the existing, correct helper function.
+        # We respect your capitalization fix here.
         parent = self.topology.get_ancestor_of_type(
             first_node_gp, "Package"
         ) or self.topology.get_ancestor_of_type(first_node_gp, "NUMANode")
@@ -56,29 +54,19 @@ def _build_contextual_subgraph(self):
         elif leaf_type in ["Package", "NUMANode", "Machine"]:
             search_domain_gp = first_node_gp
 
-        # Step 3: Get all sibling nodes of the same type within that context.
         if search_domain_gp:
             all_siblings = self.topology.get_descendants(search_domain_gp, type=leaf_type)
             if not all_siblings and leaf_type in ["Package", "NUMANode"]:
                 all_siblings = self.assigned_nodes
         else:
             all_siblings = self.assigned_nodes
 
-        # Step 4: Build the final set of nodes to draw.
         nodes_to_draw_gps = set()
         for gp, _ in all_siblings:
             nodes_to_draw_gps.add(gp)
             nodes_to_draw_gps.update(nx.ancestors(self.topology.hierarchy_view, gp))
 
         final_subgraph = self.topology.graph.subgraph(nodes_to_draw_gps).copy()
-
-        # Filter out types we don't want to see, for clarity.
-        nodes_to_remove = [
-            gp
-            for gp, data in final_subgraph.nodes(data=True)
-            if data.get("type") not in ["Core", "PU"]
-        ]
-        final_subgraph.remove_nodes_from(nodes_to_remove)
         return final_subgraph
 
     def draw(self, filename: str):