[WIP] Feature/gpu sharing factor

libensemble/executors/mpi_runner.py

@@ -124,15 +124,17 @@ def _set_gpu_cli_option(self, wresources, extra_args, gpu_setting_name, gpu_valu
     def _set_gpu_env_var(self, wresources, task, gpus_per_node, gpus_env):
         """Add GPU environment variable setting to the tasks environment"""
         jassert(wresources.matching_slots, f"Cannot assign CPUs/GPUs to non-matching slots per node {wresources.slots}")
-        slot_list = wresources.get_slots_as_string(multiplier=wresources.gpus_per_rset_per_node, limit=gpus_per_node)
-        task._add_to_env(gpus_env, slot_list)
+        gpu_list = wresources.get_gpu_ids_as_string(limit=gpus_per_node)
+        task._add_to_env(gpus_env, gpu_list)
 
     def _local_runner_set_gpus(self, task, wresources, extra_args, gpus_per_node, ppn):
         """Set default GPU setting for MPI runner"""
 
         arg_type = self.default_gpu_arg_type
         if arg_type is not None:
             gpu_value = gpus_per_node // ppn if arg_type == "option_gpus_per_task" else gpus_per_node
+            if isinstance(gpu_value, float) and gpu_value.is_integer():
+                gpu_value = int(gpu_value)
             gpu_setting_name = self.default_gpu_args[arg_type]
             jassert(gpu_setting_name is not None, f"No default gpu_setting_name for {arg_type}")
             extra_args = self._set_gpu_cli_option(wresources, extra_args, gpu_setting_name, gpu_value)
@@ -174,7 +176,7 @@ def _assign_gpus(self, task, resources, nprocs, nnodes, ppn, ngpus, extra_args,
 
         # gpus per node for this worker.
         if wresources.doihave_gpus():
-            gpus_avail_per_node = wresources.slot_count * wresources.gpus_per_rset_per_node
+            gpus_avail_per_node = len(set(wresources.gpu_ids_on_node)) * wresources.gpus_per_rset_per_node
         else:
             gpus_avail_per_node = 0
 
@@ -188,7 +190,12 @@ def _assign_gpus(self, task, resources, nprocs, nnodes, ppn, ngpus, extra_args,
                 nnodes = wresources.local_node_count
 
         if ngpus is not None:
-            gpus_req_per_node = ngpus // nnodes
+            import math
+
+            try:
+                gpus_req_per_node = math.ceil(float(ngpus) / nnodes)
+            except Exception:
+                gpus_req_per_node = ngpus // nnodes
             if gpus_req_per_node > gpus_avail_per_node:
                 logger.info(f"Asked for more GPUs per node than available - max is {gpus_avail_per_node}")
             gpus_per_node = min(gpus_req_per_node, gpus_avail_per_node)
@@ -218,6 +225,8 @@ def _assign_gpus(self, task, resources, nprocs, nnodes, ppn, ngpus, extra_args,
 
         elif gpu_setting_type in ["option_gpus_per_node", "option_gpus_per_task"]:
             gpu_value = gpus_per_node // ppn if gpu_setting_type == "option_gpus_per_task" else gpus_per_node
+            if isinstance(gpu_value, float) and gpu_value.is_integer():
+                gpu_value = int(gpu_value)
             gpu_setting_name = self.platform_info.get("gpu_setting_name", self._get_default_arg(gpu_setting_type))
             extra_args = self._set_gpu_cli_option(wresources, extra_args, gpu_setting_name, gpu_value)
 
@@ -246,8 +255,13 @@ def _get_min_nodes(self, nprocs, ppn, nnodes, ngpus, resources):
             if nprocs:
                 proc_min_nodes = (nprocs + procs_on_node - 1) // procs_on_node
             if ngpus:
-                gpus_on_node = wresources.slot_count * wresources.gpus_per_rset_per_node
-                gpu_min_nodes = (ngpus + gpus_on_node - 1) // gpus_on_node
+                gpus_on_node = len(set(wresources.gpu_ids_on_node)) * wresources.gpus_per_rset_per_node
+                import math
+
+                try:
+                    gpu_min_nodes = math.ceil(float(ngpus) / gpus_on_node) if gpus_on_node else 1
+                except Exception:
+                    gpu_min_nodes = (ngpus + gpus_on_node - 1) // gpus_on_node
 
             min_nodes = max(proc_min_nodes, gpu_min_nodes)
             nnodes = min(min_nodes, total_nodes)

libensemble/resources/resources.py

@@ -164,6 +164,7 @@ def __init__(self, libE_specs: dict, platform_info: dict = {}, top_level_dir: st
         self.num_resource_sets = libE_specs.get("num_resource_sets", None)
         self.enforce_worker_core_bounds = libE_specs.get("enforce_worker_core_bounds", False)
         self.gpus_per_group = libE_specs.get("gpus_per_group")
+        self.gpu_sharing_factor = libE_specs.get("gpu_sharing_factor", 1)
         resource_info = libE_specs.get("resource_info", {})
 
         # resource_info overrides platform

libensemble/resources/rset_resources.py

@@ -33,6 +33,8 @@ class RSetResources:
         ("group", int),  # Group ID this resource set belongs to
         ("slot", int),  # Slot ID this resource set belongs to
         ("gpus", bool),  # Does this resource set have GPUs
+        ("gpu_id", int),  # Physical GPU ID for GPU resource sets
+        ("gpu_share", int),  # Logical share within a physical GPU
         # ('pool', int),    # Pool ID (eg. separate gen/sim resources) - not yet used.
     ]
 
@@ -61,19 +63,35 @@ def __init__(self, num_workers: int, resources: GlobalResources):
         self.nodes_in_rset = len(self.split_list[0])
 
         gpus_avail_per_node = resources.gpus_avail_per_node
+        self.gpu_sharing_factor = resources.gpu_sharing_factor
         self.rsets_per_node = RSetResources.get_rsets_on_a_node(self.total_num_rsets, resources)
-        self.gpu_rsets_per_node = min(gpus_avail_per_node, self.rsets_per_node)
+        self.gpu_rsets_per_node = min(gpus_avail_per_node * self.gpu_sharing_factor, self.rsets_per_node)
         self.nongpu_rsets_per_node = self.rsets_per_node - self.gpu_rsets_per_node
 
         self.all_rsets = np.zeros(self.total_num_rsets, dtype=RSetResources.rset_dtype)
         self.all_rsets["group"], self.all_rsets["slot"], self.all_rsets["gpus"] = RSetResources.get_group_list(
-            self.split_list, gpus_avail_per_node, resources.gpus_per_group
+            self.split_list, gpus_avail_per_node, resources.gpus_per_group, self.gpu_sharing_factor
         )
+        if gpus_avail_per_node:
+            gpu_positions = RSetResources.get_gpu_positions(self.split_list)
+            self.all_rsets["gpu_id"] = np.where(
+                self.all_rsets["gpus"],
+                (gpu_positions // self.gpu_sharing_factor) % gpus_avail_per_node,
+                -1,
+            )
+            self.all_rsets["gpu_share"] = np.where(self.all_rsets["gpus"], gpu_positions % self.gpu_sharing_factor, -1)
 
         self.total_num_gpu_rsets = np.count_nonzero(self.all_rsets["gpus"])
         self.total_num_nongpu_rsets = np.count_nonzero(~self.all_rsets["gpus"])
 
-        self.gpus_per_rset_per_node = gpus_avail_per_node // self.gpu_rsets_per_node if self.gpu_rsets_per_node else 0
+        if self.gpu_rsets_per_node:
+            self.gpus_per_rset_per_node = (
+                1
+                if self.gpu_sharing_factor > 1
+                else gpus_avail_per_node // min(gpus_avail_per_node, self.rsets_per_node)
+            )
+        else:
+            self.gpus_per_rset_per_node = 0
         self.cores_per_rset_per_node = resources.physical_cores_avail_per_node // self.rsets_per_node
 
         # Oversubsribe
@@ -89,25 +107,38 @@ def __init__(self, num_workers: int, resources: GlobalResources):
         self.procs_per_rset = self.procs_per_rset_per_node * self.nodes_in_rset
 
     @staticmethod
-    def get_group_list(split_list, gpus_per_node=0, gpus_per_group=None):
-        """Return lists of group ids and slot IDs by resource set"""
+    def get_group_list(split_list, gpus_per_node=0, gpus_per_group=None, gpu_sharing_factor=1):
+        """Return group, slot, and GPU lists by resource set"""
         group = 1
         slot = 0
         group_list, slot_list, gpu_list = [], [], []
         node = split_list[0]
 
         for i in range(len(split_list)):
-            # still break on new node if gpus_per_group is set
-            if split_list[i] != node or slot == gpus_per_group:
+            # break on new node or when group size reached (account for sharing)
+            group_threshold = gpus_per_group * gpu_sharing_factor if gpus_per_group is not None else None
+            if split_list[i] != node or slot == group_threshold:
                 node = split_list[i]
                 group += 1
                 slot = 0
             group_list.append(group)
             slot_list.append(slot)
-            gpu_list.append(slot < gpus_per_node)
+            is_gpu = slot < gpus_per_node * gpu_sharing_factor
+            gpu_list.append(is_gpu)
             slot += 1
         return group_list, slot_list, gpu_list
 
+    @staticmethod
+    def get_gpu_positions(split_list):
+        """Return each resource set's node-local position, independent of scheduler groups."""
+        positions = []
+        node_counts = {}
+        for nodes in split_list:
+            node = nodes[0]
+            positions.append(node_counts.get(node, 0))
+            node_counts[node] = positions[-1] + 1
+        return np.asarray(positions)
+
     @staticmethod
     def best_split(a, n):
         """Creates the most even split of list a into n parts and return list of lists"""

libensemble/resources/scheduler.py

@@ -57,6 +57,13 @@ def __init__(self, user_resources=None, sched_opts={}):
                 Useful if setting an environment variable such as ``CUDA_VISIBLE_DEVICES``
                 to specific slots, which should match over multiple nodes.
                 Default: True
+
+            "gpu_packing" [String]:
+                GPU share packing strategy when ``gpu_sharing_factor > 1``.
+                "pack" (default) co-locates fractional tasks on the most-occupied GPU
+                to preserve fully-free GPUs for exclusive tasks. "spread" distributes
+                across least-occupied GPUs.
+                Default: "pack"
         """
 
         self.resources = user_resources or Resources.resources.resource_manager
@@ -71,6 +78,9 @@ def __init__(self, user_resources=None, sched_opts={}):
         # Process scheduler options
         self.split2fit = sched_opts.get("split2fit", True)
         self.match_slots = sched_opts.get("match_slots", True)
+        self.gpu_packing = sched_opts.get("gpu_packing", "pack")
+        if self.gpu_packing not in ("pack", "spread"):
+            raise ValueError(f"gpu_packing must be 'pack' or 'spread', got {self.gpu_packing!r}")
         self.last_use_gpus = None
 
     def assign_resources(self, rsets_req, use_gpus=None, user_params=[]):
@@ -92,6 +102,10 @@ def assign_resources(self, rsets_req, use_gpus=None, user_params=[]):
             return []
 
         self.check_total_rsets(rsets_req, use_gpus)
+
+        # GPU-share-aware path for physical-exclusive packing
+        if use_gpus and getattr(self.resources, "gpu_sharing_factor", 1) > 1:
+            return self._assign_gpu_sharing_resources(rsets_req, use_gpus, user_params)
         self.log_msg = None  # Log resource messages only when find resources
         num_groups = self.resources.num_groups
 
@@ -180,6 +194,159 @@ def assign_resources(self, rsets_req, use_gpus=None, user_params=[]):
 
         return rset_team
 
+    def _build_gpu_bins(self, valid_rsets_by_group):
+        """Build per-group GPU bins: group -> gpu_id -> list[rset_id]"""
+        gpu_bins_per_group = {}
+        for g, rlist in valid_rsets_by_group.items():
+            bins = {}
+            for rid in rlist:
+                gid = int(self.resources.rsets[rid]["gpu_id"])
+                bins.setdefault(gid, []).append(rid)
+            for gid in bins:
+                bins[gid] = sorted(bins[gid])
+            gpu_bins_per_group[g] = bins
+        return gpu_bins_per_group
+
+    def _can_group_satisfy_gpu(self, group, need, gpu_bins_per_group):
+        """Check if group can satisfy need with physical-exclusive semantics"""
+        bins = gpu_bins_per_group.get(group, {})
+        sf = getattr(self.resources, "gpu_sharing_factor", 1)
+        whole = need // sf
+        rem = need % sf
+        fully_free = sum(1 for lst in bins.values() if len(lst) == sf)
+        if fully_free < whole:
+            return False
+        if rem == 0:
+            return True
+        leftover_fully = fully_free - whole
+        if leftover_fully > 0:
+            return True
+        for lst in bins.values():
+            if len(lst) == sf:
+                continue
+            if len(lst) >= rem:
+                return True
+        total_free = sum(len(lst) for lst in bins.values())
+        return (total_free - whole * sf) >= rem
+
+    def _select_rsets_for_gpu_group(self, group, need, gpu_bins_per_group):
+        """Select rsets for one group using pack/spread"""
+        bins = gpu_bins_per_group[group]
+        sf = getattr(self.resources, "gpu_sharing_factor", 1)
+        whole = need // sf
+        rem = need % sf
+        selected = []
+        fully_free_gids = sorted([gid for gid, lst in bins.items() if len(lst) == sf])
+        if len(fully_free_gids) < whole:
+            raise InsufficientFreeResources
+        chosen_whole = fully_free_gids[:whole]
+        for gid in chosen_whole:
+            selected.extend(bins[gid])
+        remaining_bins = {gid: lst for gid, lst in bins.items() if gid not in chosen_whole}
+        if rem > 0:
+            candidates = [(len(lst), gid) for gid, lst in remaining_bins.items() if len(lst) >= rem]
+            if candidates:
+                if self.gpu_packing == "pack":
+                    candidates.sort(key=lambda x: (x[0], x[1]))
+                else:
+                    candidates.sort(key=lambda x: (-x[0], x[1]))
+                chosen_gid = candidates[0][1]
+                selected.extend(remaining_bins[chosen_gid][:rem])
+            else:
+                sorted_bins = [(len(lst), gid, lst) for gid, lst in remaining_bins.items()]
+                if self.gpu_packing == "pack":
+                    sorted_bins.sort(key=lambda x: (x[0], x[1]))
+                else:
+                    sorted_bins.sort(key=lambda x: (-x[0], x[1]))
+                need_rem = rem
+                for _, gid, lst in sorted_bins:
+                    take = min(len(lst), need_rem)
+                    selected.extend(lst[:take])
+                    need_rem -= take
+                    if need_rem == 0:
+                        break
+                if need_rem > 0:
+                    raise InsufficientFreeResources
+        return sorted(selected)
+
+    def _assign_gpu_sharing_resources(self, rsets_req, use_gpus, user_params):
+        """Physical-exclusive GPU-share-aware assignment with packing"""
+        num_groups = self.resources.num_groups
+        max_grpsize = self.resources.gpu_rsets_per_node
+        avail_rsets_by_group = self.get_avail_rsets_by_group()
+        valid_rsets_by_group = self.filter_for_rset_type(avail_rsets_by_group, use_gpus)
+        rsets_req, num_groups_req, rsets_per_group = self.calc_req_split(
+            rsets_req, max_grpsize, num_groups, extend=True
+        )
+        sorted_lengths = ResourceScheduler.get_sorted_lens(valid_rsets_by_group)
+        if num_groups_req > len(sorted_lengths) or sorted_lengths[num_groups_req - 1] < rsets_per_group:
+            if not self.split2fit or (sorted_lengths and sorted_lengths[0] == 0):
+                raise InsufficientFreeResources
+        gpu_bins_per_group = self._build_gpu_bins(valid_rsets_by_group)
+
+        def feasible_groups_for(need_per_group, ngroups):
+            feasible = [
+                g for g in valid_rsets_by_group if self._can_group_satisfy_gpu(g, need_per_group, gpu_bins_per_group)
+            ]
+
+            # sort by packing preference: pack prefers most occupied (smallest free) that still feasible
+            def sort_key(g):
+                total_free = len(valid_rsets_by_group[g])
+                return (total_free, g) if self.gpu_packing == "pack" else (-total_free, g)
+
+            feasible.sort(key=sort_key)
+            return feasible
+
+        feasible = feasible_groups_for(rsets_per_group, num_groups_req)
+        if len(feasible) < num_groups_req:
+            if not self.split2fit:
+                raise InsufficientFreeResources
+            # Try to find alternative split that is feasible
+            sorted_lengths = ResourceScheduler.get_sorted_lens(valid_rsets_by_group)
+            max_even = sorted_lengths[num_groups_req - 1] if num_groups_req <= len(sorted_lengths) else 0
+            tried = set()
+            while len(feasible) < num_groups_req:
+                key = (rsets_req, num_groups_req, rsets_per_group)
+                if key in tried:
+                    num_groups_req += 1
+                    if num_groups_req > num_groups:
+                        raise InsufficientFreeResources
+                    rsets_per_group = sorted_lengths[num_groups_req - 1] if num_groups_req <= len(sorted_lengths) else 0
+                    tried = set()
+                    continue
+                tried.add(key)
+                try:
+                    rsets_req, num_groups_req, rsets_per_group = self.calc_even_split_uneven_groups(
+                        max_even, num_groups_req, rsets_req, sorted_lengths, num_groups, user_params
+                    )
+                except InsufficientFreeResources:
+                    raise
+                feasible = feasible_groups_for(rsets_per_group, num_groups_req)
+                if len(feasible) >= num_groups_req:
+                    break
+                # if still not feasible, try larger group count
+                if rsets_per_group * num_groups_req != rsets_req or not self.check_params(user_params, num_groups_req):
+                    continue
+                # fallback: increment groups
+                if len(feasible) < num_groups_req:
+                    max_even = sorted_lengths[num_groups_req - 1] if num_groups_req <= len(sorted_lengths) else 0
+        if len(feasible) < num_groups_req:
+            raise InsufficientFreeResources
+        cand_groups = feasible[:num_groups_req]
+        rset_team = []
+        for g in cand_groups:
+            part = self._select_rsets_for_gpu_group(g, rsets_per_group, gpu_bins_per_group)
+            rset_team.extend(part)
+        rset_team = sorted(rset_team)
+        self.avail_rsets_by_group = self.filter_out_rset_team(avail_rsets_by_group, rset_team)
+        self.rsets_free -= len(rset_team)
+        self.gpu_rsets_free -= len(rset_team)
+        logger.debug(
+            f"rset_team found (gpu-sharing): Req: {rsets_req} rsets. Found: {rset_team}"
+            + f" Avail sets {self.avail_rsets_by_group}"
+        )
+        return rset_team
+
     def find_rsets_any_slots(self, valid_rsets_by_group, max_grpsize, rsets_req, ngroups, rsets_per_group):
         """Find optimal non-matching slots across groups"""
         tmp_rsets_by_group = copy.deepcopy(valid_rsets_by_group)