bug: fix graphic to show entire tree

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

Author: vsoch

fluxbind/graph/graph.py

@@ -278,18 +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
@@ -332,46 +332,70 @@ def translate_type(self, requested_type: str):
 
     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.
+        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}'...")
 
-        # Concrete hwloc type string for our target bind level.
-        # Hey spack, we can concretize too!
         bind_type_concrete = self.translate_type(bind_level)
 
+        # Iterate through every object found by the allocator.
         for gp, data in total_allocation:
 
-            # Case 1: The allocated object IS the type we want to bind to.
+            # 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: 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 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
 
-            # Case 3: The allocated object is a child of the type we want to bind to.
-            # Example: total_allocation=[PU:0], bind_level='core'
+            # For all other containers, 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 (e.g., if two PUs map to the same Core)
-        # and sort for deterministic assignment.
+        # 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)        
+        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.
@@ -545,7 +569,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(