test: add count

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

Author: vsoch

fluxbind/bind/bind.py

@@ -187,6 +187,8 @@ def set_flags(self, cmd):
         """
         Set command flags.
         """
+        # CPU and GPU affinity MUST stay off
+        cmd += ['-o', 'cpu-affinity=off', '-o', 'gpu-affinity=off']
         if self.tasks is not None:
             cmd += ["-n", str(self.tasks)]
         if self.cpu_affinity is not None:
@@ -195,8 +197,6 @@ def set_flags(self, cmd):
             cmd += ["-o", f"gpu-affinity={self.gpu_affinity}"]
         if self.cores_per_task is not None:
             cmd += ["--cores-per-task", str(self.cores_per_task)]
-        if self.gpus_per_task is not None:
-            cmd += ["--gpus-per-task", str(self.gpus_per_task)]
         if self.tasks_per_core is not None:
             cmd += ["--tasks-per-core", str(self.tasks_per_core)]
         if self.taskmap is not None:

fluxbind/shape/commands.py

@@ -95,6 +95,11 @@ def get_object_in_set(self, cpuset: str, obj_type: str, index: int) -> str:
         """
         list_cmd = f"{self.name} '{cpuset}' --intersect {obj_type}"
         all_indices_str = self.run(list_cmd, shell=True)
+        
+        # Special case asking for all.
+        if index == "all":
+            return all_indices_str
+
         available_indices = all_indices_str.split(",")
         try:
             target_index = available_indices[index]

fluxbind/shape/gpu.py

@@ -1,30 +1,29 @@
 from dataclasses import dataclass
 
+
 @dataclass
 class GPUAssignment:
     """
     A data structure to hold information about a rank's assigned GPU(s).
     Instances are created via the for_rank() classmethod.
     """
-    indices: list[int]          # The logical indices in the ordered list (e.g., [4, 5])
-    pci_ids: list[str]          # The corresponding PCI bus IDs of the GPUs
-    numa_indices: set[int]      # The set of unique NUMA nodes these GPUs are on
-    cuda_devices: str           # The final string for CUDA_VISIBLE_DEVICES (e.g., "4,5")
+
+    indices: list[int]  # The logical indices in the ordered list (e.g., [4, 5])
+    pci_ids: list[str]  # The corresponding PCI bus IDs of the GPUs
+    numa_indices: set[int]  # The set of unique NUMA nodes these GPUs are on
+    cuda_devices: str  # The final string for CUDA_VISIBLE_DEVICES (e.g., "4,5")
 
     @classmethod
     def for_rank(
-        cls, 
-        local_rank: int, 
-        gpus_per_task: int, 
-        ordered_gpus: list[dict]
+        cls, local_rank: int, gpus_per_task: int, ordered_gpus: list[dict]
     ) -> "GPUAssignment":
         """
         A factory method that assigns a slice of GPUs to a given local rank
         from a pre-ordered, topology-aware list of all GPUs.
         """
         if not ordered_gpus:
             raise RuntimeError("Attempted to assign a GPU, but no GPUs were discovered.")
-        
+
         start_idx = local_rank * gpus_per_task
         end_idx = start_idx + gpus_per_task
 
@@ -35,13 +34,13 @@ def for_rank(
             )
 
         assigned_gpu_slice = ordered_gpus[start_idx:end_idx]
-        
+
         # The global indices for CUDA_VISIBLE_DEVICES are their positions in the ordered list
         assigned_indices = list(range(start_idx, end_idx))
-        
+
         return cls(
             indices=assigned_indices,
-            pci_ids=[gpu['pci_id'] for gpu in assigned_gpu_slice],
-            numa_indices={gpu['numa_index'] for gpu in assigned_gpu_slice},
-            cuda_devices=",".join(map(str, assigned_indices))
-        )
+            pci_ids=[gpu["pci_id"] for gpu in assigned_gpu_slice],
+            numa_indices={gpu["numa_index"] for gpu in assigned_gpu_slice},
+            cuda_devices=",".join(map(str, assigned_indices)),
+        )

fluxbind/shape/shape.py

@@ -137,13 +137,13 @@ def get_gpu_local_binding(self, rule: dict, local_rank: int, gpus_per_task: int)
         Calculate a 'gpu-local' binding using the topology-aware ordered GPU list.
         """
         assignment = gpus.GPUAssignment.for_rank(local_rank, gpus_per_task, self.ordered_gpus)
-        
+
         # The CPU domain is the union of NUMA nodes for the assigned GPUs.
         domain_locations = [f"numa:{i}" for i in assignment.numa_indices]
         domain = " ".join(domain_locations)
         cpu_binding_string = self.get_binding_in_gpu_domain(rule, local_rank, gpus_per_task, domain)
         return f"{cpu_binding_string};{assignment.cuda_devices}"
-    
+
     def get_gpu_remote_binding(self, rule: dict, local_rank: int, gpus_per_task: int) -> str:
         """
         Calculates a 'gpu-remote' binding using the topology-aware ordered GPU list.
@@ -155,14 +155,16 @@ def get_gpu_remote_binding(self, rule: dict, local_rank: int, gpus_per_task: int
         # Find all remote NUMA domains relative to the set of local domains.
         all_numa_indices = set(range(len(self.numa_node_cpusets)))
         remote_numa_indices = sorted(list(all_numa_indices - assignment.numa_indices))
-        
+
         if not remote_numa_indices:
-            raise RuntimeError(f"Cannot find a remote NUMA node for rank {local_rank}; its GPUs span all NUMA domains.")
-        
-        offset = rule.get('offset', 0)
+            raise RuntimeError(
+                f"Cannot find a remote NUMA node for rank {local_rank}; its GPUs span all NUMA domains."
+            )
+
+        offset = rule.get("offset", 0)
         if offset >= len(remote_numa_indices):
             raise ValueError(f"Offset {offset} is out of range for remote NUMA domains.")
-            
+
         target_remote_numa_idx = remote_numa_indices[offset]
         domain = f"numa:{target_remote_numa_idx}"
 
@@ -171,7 +173,7 @@ def get_gpu_remote_binding(self, rule: dict, local_rank: int, gpus_per_task: int
 
     def get_binding_in_gpu_domain(
         self, rule: dict, local_rank: int, gpus_per_task: int, domain: str
-    ) -> str:
+    ):
         """
         A dedicated binding engine for GPU jobs. It applies user preferences within a calculated domain
         (e.g., "numa:0" or "numa:0 numa:1").
@@ -184,25 +186,58 @@ def get_binding_in_gpu_domain(
             # If a broad type is requested, the binding is the domain itself.
             return domain
 
-        if "prefer" in rule:
-            try:
-                requested_index = int(rule["prefer"])
-                # Validate by attempting to get the object.
-                return commands.hwloc_calc.get_object_in_set(domain, hwloc_type, requested_index)
-            except (ValueError, RuntimeError, TypeError):
-                print(
-                    f"Warning: Preferred index '{rule['prefer']}' invalid/not in domain '{domain}'. Falling back.",
-                    file=sys.stderr,
+        elif hwloc_type in ["core", "pu", "l2cache"]:
+
+            # Get the number of objects to select, defaulting to 1.
+            count = rule.get("count", 1)
+
+            all_indices_in_domain = commands.hwloc_calc.get_object_in_set(
+                domain, hwloc_type, "all"
+            ).split(",")
+            if not all_indices_in_domain or not all_indices_in_domain[0]:
+                raise RuntimeError(f"No objects of type '{hwloc_type}' found in domain '{domain}'.")
+
+            if "prefer" in rule:
+                if count > 1:
+                    raise ValueError("'prefer' and 'count > 1' cannot be used together.")
+                try:
+                    requested_index = str(int(rule["prefer"]))
+                    if requested_index in all_indices_in_domain:
+                        return f"{hwloc_type}:{requested_index}"
+                    else:
+                        print(
+                            f"Warning: Preferred index '{requested_index}' not available in domain '{domain}'. Falling back.",
+                            file=sys.stderr,
+                        )
+                except (ValueError, TypeError):
+                    raise ValueError(
+                        f"The 'prefer' key must be a simple integer, but got: {rule['prefer']}"
+                    )
+
+            # Default assignment: Calculate the slice of objects for this rank.
+            # We need to know this rank's turn on the current domain.
+            num_domains = len(domain.split())
+            rank_turn_in_domain = local_rank // num_domains
+
+            start_index = rank_turn_in_domain * count
+            end_index = start_index + count
+
+            if end_index > len(all_indices_in_domain):
+                raise ValueError(
+                    f"Not enough '{hwloc_type}' objects in domain '{domain}' to satisfy request "
+                    f"for {count} objects for rank {local_rank} (needs up to index {end_index-1}, "
+                    f"only {len(all_indices_in_domain)} available)."
                 )
 
-        # Default assignment: Rank's Nth turn for a resource of this type within its GPU group.
-        # This is the correct index for packing sub-objects within a domain.
-        index = local_rank // gpus_per_task if gpus_per_task > 0 else local_rank
+            # Get the slice of object indices.
+            target_indices_slice = all_indices_in_domain[start_index:end_index]
 
-        # For certain patterns like interleave or spread, the index calculation
-        # would need to be more complex, but for a simple packed pattern this is the logic.
-        # Let's assume a simple packed logic for now as pattern is not yet implemented here.
-        return commands.hwloc_calc.get_object_in_set(domain, hwloc_type, index)
+            # Construct a space-separated list of location objects.
+            # e.g., "core:0 core:1 core:2 core:3 core:4 core:5"
+            binding_locations = [f"{hwloc_type}:{i}" for i in target_indices_slice]
+            return " ".join(binding_locations)
+        else:
+            raise ValueError(f"Unsupported type '{hwloc_type}' for GPU binding.")
 
     def get_binding_for_rank(self, rank, node_id, local_rank, gpus_per_task=None) -> str:
         """