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).
@@ -275,6 +279,9 @@ def match_resources(self, jobspec, allocated_gps=None):
         return final_allocation
 
     def sort_by_affinity(self, candidates, affinity, allocated):
+        """
+        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'.")
@@ -323,6 +330,48 @@ 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):
+        """
+        Given an allocation, find bindable leaf nodes.
+
+        This function handles three cases for each allocated resource:
+        1. Identity: The resource is already the desired bindable type.
+        2. Descendants: The resource is a container (e.g., Package) of the desired type.
+        3. Ancestors: The resource is a child (e.g., PU) of the desired type.
+        """
+        leaf_nodes = []
+        log.debug(f"Transforming {len(total_allocation)} allocated objects to bind_level '{bind_level}'...")
+        
+        # Concrete hwloc type string for our target bind level.
+        # Hey spack, we can concretize too!
+        bind_type_concrete = self.translate_type(bind_level)
+
+        for gp, data in total_allocation:
+            
+            # Case 1: The allocated object IS the type we want to bind to.
+            if data.get('type') == bind_type_concrete:
+                leaf_nodes.append((gp, data))
+                continue
+
+            # Case 2: The allocated object is a container of the type we want to bind to.
+            descendants = self.get_descendants(gp, type=bind_type_concrete)
+            if descendants:
+                 leaf_nodes.extend(descendants)
+                 continue
+            
+            # Case 3: The allocated object is a child of the type we want to bind to.
+            # Example: total_allocation=[PU:0], bind_level='core'
+            ancestor = self.get_ancestor_of_type(gp, bind_type_concrete)
+            if ancestor:
+                leaf_nodes.append(ancestor)
+
+        # De-duplicate the final list (e.g., if two PUs map to the same Core)
+        # 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.
@@ -393,22 +442,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 +456,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 +483,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 +503,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")
@@ -550,6 +599,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 +700,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):