gpus: add remote shape

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

Author: vsoch

examples/shape/google/n1-standard-8/shape_gpu_remote_numa.yaml

@@ -0,0 +1,3 @@
+default:
+  type: numa
+  locality: gpu-remote

fluxbind/bind/bind.py

@@ -195,6 +195,8 @@ 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/scripts/run_mapping.sh

@@ -24,6 +24,7 @@ fi
 # Call the fluxbind helper script to get the target location string (e.g., "core:0" or "UNBOUND")
 # It ALWAYS returns a single line in the format: BIND_LOCATION,CUDA_DEVICE_ID
 # For CPU jobs, CUDA_DEVICE_ID will be the string "NONE".
+echo fluxbind shape --file "$JOB_SHAPE_FILE" --rank "$rank" --node-id "$node_id" --local-rank "$local_rank"
 BIND_INFO=$(fluxbind shape --file "$JOB_SHAPE_FILE" --rank "$rank" --node-id "$node_id" --local-rank "$local_rank")
 
 # Exit if the helper script failed

fluxbind/shape/commands.py

@@ -29,17 +29,63 @@ def run(self, command, shell: bool = False):
 class HwlocCalcCommand(Command):
     name = "hwloc-calc"
 
-    def execute(self, args_list: list) -> str:
+    def count(self, hw_type: str, within: str = "machine:0") -> int:
         """
-        Executes hwloc-calc with a list of arguments.
-        This is safer as it avoids shell interpretation of the arguments.
+        Returns the total number of a specific hardware object.
+
+        Args:
+            hw_type: The type of object to count (e.g., "core", "numa").
+            within_object: Optional object to restrict the count to (e.g., "numa:0").
+        """
+        try:
+            args = ["--number-of", hw_type, within]
+            result_stdout = self.run([self.name] + args)
+            return int(result_stdout)
+        except (RuntimeError, ValueError) as e:
+            raise RuntimeError(f"Failed to count number of '{hw_type}': {e}")
+
+    def list_cpusets(self, hw_type: str, within: str = "machine:0") -> list[str]:
         """
-        if isinstance(args_list, str):
-            args_list = shlex.split(args_list)
+        Returns a list of cpuset strings for each object of a given type.
+
+        Args:
+            hw_type: The type of object to list (e.g., "numa").
+            within_object: Optional object to restrict the list to.
+        """
+        try:
+            # Get the indices of all objects of this type
+            args_intersect = ["--intersect", hw_type, within]
+            indices_str = self.run([self.name] + args_intersect)
+            indices = indices_str.split(",")
+
+            # Cut out early
+            if not indices or not indices[0]:
+                return []
 
-        # A more robust validation could be added here if needed,
-        command_list = [self.name] + args_list
-        return self.run(command_list, shell=False)
+            # For each index, get its specific cpuset
+            return [self.run([self.name, f"{hw_type}:{i}"]) for i in indices]
+        except (RuntimeError, ValueError) as e:
+            raise RuntimeError(f"Failed to list cpusets for '{hw_type}': {e}")
+
+    def get_cpuset(self, location: str) -> str:
+        """
+        Gets the cpuset for a single, specific location string (e.g., "pci=...", "core:0").
+        """
+        return self.run([self.name, location])
+
+    def get_object_in_set(self, cpuset: str, obj_type: str, index: int) -> str:
+        """
+        Gets the Nth object of a type within a given cpuset.
+        e.g., find the 1st 'core' within cpuset '0x00ff'.
+        """
+        # This uses the robust two-step process internally
+        all_objects_str = self.run([self.name, cpuset, "--intersect", obj_type])
+        available_indices = all_objects_str.split(",")
+        try:
+            target_index = available_indices[index]
+            return f"{obj_type}:{target_index}"
+        except IndexError:
+            raise ValueError(f"Cannot find the {index}-th '{obj_type}' in cpuset {cpuset}.")
 
 
 class NvidiaSmiCommand(Command):

fluxbind/shape/gpu.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class GPUAssignment:
+    """
+    Data structure to hold information about a rank's assigned GPU.
+    """
+
+    index: int  # logical index of the GPU
+    pci_id: str  # The PCI bus ID of the GPU
+    cuda_devices: str  # CUDA_VISIBLE_DEVICES
+
+    @classmethod
+    def for_rank(cls, local_rank, gpu_pci_ids):
+        """
+        A factory method that assigns a GPU to a given local rank
+        using a round-robin strategy.
+        """
+        if not gpu_pci_ids:
+            raise RuntimeError("Attempted to assign a GPU, but no GPUs were discovered.")
+
+        num_gpus = len(gpu_pci_ids)
+        target_gpu_index = local_rank % num_gpus
+
+        # Return the assignment
+        return cls(
+            index=target_gpu_index,
+            pci_id=gpu_pci_ids[target_gpu_index],
+            cuda_devices=str(target_gpu_index),
+        )

fluxbind/shape/shape.py

@@ -1,8 +1,8 @@
-import re
 import subprocess
 import sys
 
 import fluxbind.shape.commands as commands
+import fluxbind.shape.gpu as gpus
 import fluxbind.utils as utils
 
 
@@ -21,35 +21,24 @@ def __init__(self, filepath, machine="machine:0"):
         self.data = self.load_file(filepath)
         # This discovers and cache hardware properties on init
         # The expectation is that this is running from the node (task)
-        self.num_cores = self.discover_hardware("core")
-        self.num_pus = self.discover_hardware("pu")
+        self.num_cores = commands.hwloc_calc.count("core", within=self.machine)
+        self.num_pus = commands.hwloc_calc.count("pu", within=self.machine)
+        self.numa_node_cpusets = commands.hwloc_calc.list_cpusets("numa", within=self.machine)
         self.pus_per_core = self.num_pus // self.num_cores if self.num_cores > 0 else 0
         self.gpu_pci_ids = self.discover_gpus()
 
     def discover_gpus(self) -> list:
-        """Discovers available GPU PCI bus IDs."""
-        try:
-            return commands.nvidia_smi.get_pci_bus_ids()
-        except (RuntimeError, FileNotFoundError):
-            print("Warning: Could not discover GPUs. GPU binding unavailable.", file=sys.stderr)
-            return []
+        """
+        Discovers available GPU PCI bus IDs.
+        """
+        return commands.nvidia_smi.get_pci_bus_ids()
 
     def load_file(self, filepath):
         """
         Loads and parses the YAML shape file.
         """
         return utils.read_yaml(filepath)
 
-    def discover_hardware(self, hw_type: str) -> int:
-        """
-        Gets the number of a specific hardware object.
-        """
-        try:
-            result = commands.hwloc_calc.execute(["--number-of", hw_type, self.machine])
-            return int(result)
-        except (RuntimeError, ValueError) as e:
-            raise RuntimeError(f"Failed to determine number of '{hw_type}': {e}")
-
     @staticmethod
     def parse_range(range_str: str) -> set:
         """
@@ -101,65 +90,119 @@ def find_matching_rule(self, rank: int, node_id: int) -> dict:
                 return item["default"]
         return None
 
-    def get_gpu_local_binding(self, rule: dict, local_rank: int) -> str:
+    def get_gpu_remote_binding(self, rule: dict, local_rank: int) -> str:
         """
-        Calculates binding for a rank based on its proximity to an assigned GPU.
-        Supports an optional prefer key for user-preferred object selection.
+        Calculate a binding that is deliberately remote from an assigned GPU.
         """
-        if not self.gpu_pci_ids:
-            raise RuntimeError(
-                "Shape specifies 'locality: gpu-local', but no GPUs were discovered."
-            )
+        if not self.gpu_pci_ids or not self.numa_node_cpusets:
+            raise RuntimeError("GPU/NUMA discovery failed, cannot perform remote binding.")
+        if len(self.numa_node_cpusets) < 2:
+            raise RuntimeError("'locality: gpu-remote' requires a multi-NUMA system.")
 
         num_gpus = len(self.gpu_pci_ids)
-        hwloc_type = rule.get("type")
-        if not hwloc_type:
-            raise ValueError("Rule with 'locality: gpu-local' must also specify a 'type'.")
 
-        # 1. Assign a GPU to this rank (round-robin)
+        # This assumes round robin assignment
         target_gpu_index = local_rank % num_gpus
         cuda_devices = str(target_gpu_index)
         target_gpu_pci_id = self.gpu_pci_ids[target_gpu_index]
 
-        # 2. Get the cpuset for the GPU's locality domain
-        gpu_locality_cpuset = commands.hwloc_calc.execute([f"pci={target_gpu_pci_id}"])
-        cpu_binding_string = ""
+        # Figure out remote NUMA that aren't in local set
+        local_cpuset = commands.hwloc_calc.get_cpuset(f"pci={target_gpu_pci_id}")
+        remote_numa_cpusets = [cs for cs in self.numa_node_cpusets if cs != local_cpuset]
+        if not remote_numa_cpusets:
+            raise RuntimeError(f"Could not find a NUMA node remote from GPU {target_gpu_index}.")
 
-        # 3. Determine the final CPU binding
-        if hwloc_type in ["numa", "package", "l3cache"]:
-            cpu_binding_string = gpu_locality_cpuset
+        # if the user asks for an offset. Otherwise, just take first
+        offset = rule.get("offset", 0)
+        if offset >= len(remote_numa_cpusets):
+            raise ValueError(f"Offset {offset} is out of range.")
+        target_remote_cpuset = remote_numa_cpusets[offset]
+        return self.get_binding_within_domain(rule, local_rank, target_remote_cpuset, cuda_devices)
 
-        elif hwloc_type in ["core", "pu", "l2cache"]:
-            all_objects_in_domain_str = commands.hwloc_calc.execute(
-                [gpu_locality_cpuset, "--intersect", hwloc_type]
+    def get_gpu_local_binding(self, rule: dict, local_rank: int) -> str:
+        """
+        Calculate binding for a rank based on its proximity to an assigned GPU.
+        """
+        # Get the assignment.
+        assignment = gpus.GPUAssignment.for_rank(local_rank, self.gpu_pci_ids)
+
+        # Find the LOCAL cpuset for this assigned GPU
+        local_cpuset = commands.hwloc_calc.get_cpuset(f"pci={assignment.pci_id}")
+
+        # This is shared logic for binding within a domain
+        return self.get_binding_within_domain(
+            rule, local_rank, local_cpuset, assignment.cuda_devices
+        )
+
+    def get_gpu_remote_binding(self, rule: dict, local_rank: int) -> str:
+        """
+        Calculate a binding that is deliberately remote from an assigned GPU.
+        """
+        if len(self.numa_node_cpusets) < 2:
+            raise RuntimeError("'locality: gpu-remote' requires a multi-NUMA system.")
+
+        assignment = gpus.GPUAssignment.for_rank(local_rank, self.gpu_pci_ids)
+        local_cpuset = commands.hwloc_calc.get_cpuset(f"pci={assignment.pci_id}")
+
+        # Find all REMOTE NUMA domains
+        remote_numa_cpusets = [cs for cs in self.numa_node_cpusets if cs != local_cpuset]
+        if not remote_numa_cpusets:
+            raise RuntimeError(
+                f"Could not find a NUMA node remote from the one for GPU {assignment.index}."
             )
-            available_indices = all_objects_in_domain_str.split(",")
-            target_object_index = None
 
-            # Is the user asking to prefer a specific identifier?
-            if "prefer" in rule:
-                requested_index = str(rule["prefer"])
-                if requested_index in available_indices:
-                    target_object_index = requested_index
+        # Allow for an offset, default to 0 (the first in list, no offset)
+        offset = rule.get("offset", 0)
+        if offset >= len(remote_numa_cpusets):
+            raise ValueError(f"Offset {offset} is out of range.")
+        target_remote_cpuset = remote_numa_cpusets[offset]
+
+        # 5. Delegate to the common logic, passing the REMOTE cpuset
+        return self.get_binding_within_domain(
+            rule, local_rank, target_remote_cpuset, assignment.cuda_devices
+        )
+
+    def get_binding_within_domain(self, rule, local_rank, domain_cpuset, cuda_devices):
+        """
+        Helper to calculate a CPU binding within a given cpuset domain.
+        """
+        hwloc_type = rule.get("type")
+        num_gpus = len(self.gpu_pci_ids)
 
-            if target_object_index is None:
-                rank_index_in_gpu_group = local_rank // num_gpus
+        if hwloc_type in ["numa", "package", "l3cache"]:
+            return f"{domain_cpuset},{cuda_devices}"
+
+        elif hwloc_type in ["core", "pu", "l2cache"]:
+            target_object_index_str = None
+            if "prefer" in rule:
                 try:
-                    target_object_index = available_indices[rank_index_in_gpu_group]
-                except IndexError:
+                    requested_index = int(rule["prefer"])
+                except (ValueError, TypeError):
                     raise ValueError(
-                        f"Cannot find the {rank_index_in_gpu_group}-th '{hwloc_type}' for local_rank {local_rank} "
-                        f"within the GPU's locality. Only {len(available_indices)} are available."
+                        f"The 'prefer' key must be a simple integer, but got: {rule['prefer']}"
                     )
 
-            cpu_binding_string = f"{hwloc_type}:{target_object_index}"
-        else:
-            raise ValueError(f"Unsupported type '{hwloc_type}' for 'locality: gpu-local' binding.")
+                # Attempt to get the preferred object. If it fails, we know it's not valid.
+                try:
+                    binding_obj = commands.hwloc_calc.get_object_in_set(
+                        domain_cpuset, hwloc_type, requested_index
+                    )
+                    target_object_index_str = binding_obj
+                except Exception:
+                    print(
+                        f"Warning: Preferred index '{requested_index}' is not available in domain {domain_cpuset}. Falling back."
+                    )
 
-        if not cpu_binding_string:
-            raise RuntimeError("Failed to calculate a valid cpu_binding_string.")
+            # If no preference was given, or the preference was invalid, fall back to default.
+            if target_object_index_str is None:
+                rank_index_in_group = local_rank // num_gpus
+                target_object_index_str = commands.hwloc_calc.get_object_in_set(
+                    domain_cpuset, hwloc_type, rank_index_in_group
+                )
 
-        return f"{cpu_binding_string},{cuda_devices}"
+            return f"{target_object_index_str},{cuda_devices}"
+        else:
+            raise ValueError(f"Unsupported type '{hwloc_type}' for GPU locality binding.")
 
     def get_binding_for_rank(self, rank: int, node_id: int, local_rank: int) -> str:
         """
@@ -186,6 +229,9 @@ def get_binding_for_rank(self, rank: int, node_id: int, local_rank: int) -> str:
         # Are we doing something with GPU?
         if rule.get("bind") == "gpu-local":
             return self.get_gpu_local_binding(rule, local_rank)
+        if rule.get("bind") == "gpu-remote":
+            return self.get_gpu_remote_binding(rule, local_rank)
+
         cpu_binding_string = self.get_cpu_binding(hwloc_type, rule, local_rank)
 
         # Return the CPU binding and the "NONE" sentinel for the device.