265 add support for cores per task when using prun

doc/user_guides/ensembles.rst

@@ -438,10 +438,47 @@ not necessarily the time it occurred. A running process does not establish
 DVM responsiveness. Abrupt worker or node loss may leave no final record.
 Components must still check ``wait_task()`` results and propagate failures.
 
+CPU budgets and rank limits
+---------------------------
+
+``cores_per_instance`` is the CPU budget available to each subordinate IPS
+instance, not a request for that many MPI ranks. It defaults to one and must
+be a positive integer. A task launched with ``nproc=2, task_cpp=4`` requires
+eight processing elements from this budget.
+
+The generated single-node platform sets both ``CORES_PER_NODE`` and
+``PROCS_PER_NODE`` to the instance budget by default. These settings remain
+independent: an explicitly configured DVM instance with eight cores and a
+two-process limit can run two four-core ranks. Lowering the process limit
+does not discard CPU capacity. CPU reservations and concurrent rank occupancy
+are tracked and released separately. ``TOTAL_PROCS``, when set, also limits
+concurrent ranks across the instance's nodes.
+
+Parent scheduling allows ``allocated CPU capacity // cores_per_instance``
+instances concurrently per worker. A 128-core node can therefore run sixteen
+eight-core instances at once. All workers use the concurrency supported by
+the smallest allocated node budget. An instance budget exceeding that capacity
+is rejected before workers are launched. Generic non-ensemble Dask task pools
+retain their existing ``dask_ppw`` behavior.
+
+In core mode, budgets count physical cores. With ``hwthreads=True``, budgets
+count hardware threads. Manual platforms must declare ``CORES_PER_NODE`` in
+the selected unit; generated instance platforms do this automatically.
+For Slurm platforms, ``CORES_PER_NODE`` describes physical capacity and IPS
+uses ``SLURM_JOB_CPUS_PER_NODE`` plus ``scontrol show node`` topology to bound
+and convert the allocation. IPS does not inspect the login host's or macOS
+host's topology to infer compute-node capacity. Partial allocations are
+conservatively rounded down in physical-core mode. Missing scheduler topology
+is an error, rather than permission to use an entire machine.
+
+This CPU-budget discovery path currently supports ``NODE_DETECTION=manual``
+and ``slurm_env``. Other launchers retain their existing resource-discovery
+behavior. Instances remain on a single node, and the DVM controls actual
+binding; the instance configuration describes a budget rather than CPU IDs.
+
 Limitations
 -----------
 
-- The method signature indicates `cores_per_instance` is not yet fully implemented
 - Parameter substitution is limited to simple string replacement
 
 See Also

ipsframework/_internal/cpu_resources.py

@@ -0,0 +1,107 @@
+"""CPU budgets independent of MPI rank limits and local machine topology."""
+
+import os
+import re
+import subprocess
+from collections.abc import Callable
+from typing import Any
+
+
+def _expand_slurm_counts(value: str) -> list[int]:
+    """Expand Slurm's per-node CPU counts.
+
+    :param value: A count expression such as ``256(x2),128``.
+    :returns: One positive CPU count per allocated node.
+    :raises ValueError: If the expression is malformed.
+    """
+    counts = []
+    for group in value.split(','):
+        match = re.fullmatch(r'([1-9][0-9]*)(?:\(x([1-9][0-9]*)\))?', group.strip())
+        if match is None:
+            raise ValueError(f'Invalid SLURM_JOB_CPUS_PER_NODE: {value!r}')
+        counts.extend([int(match[1])] * int(match[2] or 1))
+    return counts
+
+
+def _get_cpu_nodes(
+    get_parameter: Callable[..., Any], hwthreads: bool = False
+) -> list[tuple[str, int]]:
+    """Read allocated CPU budgets without using process-slot counts as capacity.
+
+    Manual configurations declare budgets in the selected processing-element
+    unit. Slurm CPU counts are converted using the compute nodes' scheduler
+    topology, never the topology of the host running IPS. For partial Slurm
+    allocations, core-mode capacity is conservatively rounded down.
+
+    :param get_parameter: Configuration lookup accepting ``silent=True``.
+    :param hwthreads: Whether the allocation unit is a hardware thread.
+    :returns: Node names and allocatable processing-element counts.
+    :raises ValueError: If capacity or scheduler topology cannot be established.
+    """
+    detection = get_parameter('NODE_DETECTION', silent=True)
+    configured = int(get_parameter('CORES_PER_NODE', silent=True) or 0)
+    if configured < 1:
+        raise ValueError('CPU budgeting requires a positive CORES_PER_NODE')
+    if detection == 'manual':
+        nodes = int(get_parameter('NODES', silent=True) or 1)
+        if nodes < 1:
+            raise ValueError('NODES must be positive')
+        return [(f'dummynode{i}', configured) for i in range(nodes)]
+    if detection != 'slurm_env':
+        raise ValueError('CPU budgeting currently requires manual or slurm_env node detection')
+
+    nodelist = os.environ.get('SLURM_JOB_NODELIST') or os.environ.get('SLURM_NODELIST')
+    counts = os.environ.get('SLURM_JOB_CPUS_PER_NODE')
+    if not nodelist or not counts:
+        raise ValueError('Slurm CPU budgeting requires the job node list and CPUs per node')
+    names = subprocess.check_output(
+        ['scontrol', 'show', 'hostnames', nodelist], encoding='UTF-8'
+    ).split()
+    allocated = _expand_slurm_counts(counts)
+    if len(names) != len(allocated):
+        raise ValueError('Slurm node list and CPU counts have different lengths')
+    topology = subprocess.check_output(
+        ['scontrol', 'show', 'node', nodelist, '--oneliner'], encoding='UTF-8'
+    )
+    topologies = {}
+    for line in topology.splitlines():
+        fields = dict(part.split('=', 1) for part in line.split() if '=' in part)
+        if 'NodeName' in fields:
+            topologies[fields['NodeName']] = fields
+    result = []
+    for name, granted in zip(names, allocated, strict=True):
+        fields = topologies.get(name, {})
+        try:
+            tpc = int(fields['ThreadsPerCore'])
+            physical = int(fields['Sockets']) * int(fields['CoresPerSocket'])
+            scheduled = int(fields['CPUTot'])
+        except (KeyError, ValueError) as exc:
+            raise ValueError(f'Missing Slurm CPU topology for {name}') from exc
+        if min(tpc, physical, scheduled) < 1 or scheduled not in (physical, physical * tpc):
+            raise ValueError(f'Inconsistent Slurm CPU topology for {name}')
+        # Slurm may count schedulable CPUs as cores or hardware threads.
+        scheduled_per_core = scheduled // physical
+        if hwthreads:
+            capacity = min(configured, physical) * tpc
+            capacity = min(capacity, granted * tpc // scheduled_per_core)
+        else:
+            capacity = min(configured, physical, granted // scheduled_per_core)
+        if capacity < 1:
+            raise ValueError(f'No complete processing element available on {name}')
+        result.append((name, capacity))
+    return result
+
+
+def _ensemble_worker_threads(capacities: list[int], budget: int) -> int:
+    """Compute a safe uniform concurrency for one Dask worker per node.
+
+    :param capacities: Allocated CPU budgets of the worker nodes.
+    :param budget: Processing elements reserved for each ensemble instance.
+    :returns: Instances allowed concurrently on each worker.
+    :raises ValueError: If the budget is invalid or cannot fit every worker.
+    """
+    if isinstance(budget, bool) or not isinstance(budget, int) or budget < 1:
+        raise ValueError('cores_per_instance must be a positive integer')
+    if not capacities or min(capacities) < budget:
+        raise ValueError('cores_per_instance exceeds the allocated CPU capacity of a worker node')
+    return min(capacities) // budget

ipsframework/resource_manager.py

@@ -8,6 +8,7 @@
 from collections import namedtuple
 from math import ceil
 
+from ._internal.cpu_resources import _get_cpu_nodes
 from .ips_es_spec import EventManager
 from .ips_exceptions import (
     BadResourceRequestError,
@@ -108,6 +109,15 @@ def initialize(self, data_mngr, task_mngr, config_mngr, cmd_nodes=0, cmd_ppn=0):
         self.TM = task_mngr
         self.CM = config_mngr
         self.node_alloc_mode = self.CM.get_platform_parameter('NODE_ALLOCATION_MODE')
+        self._dvm = (
+            str(self.CM.get_platform_parameter('MPIRUN_VERSION', silent=True)).upper()
+            == 'OPENMPI-DVM'
+        )
+
+        self._dvm_rank_allocations: dict[int, dict[str, int]] = {}
+        self._dvm_total_ranks = (
+            int(self.CM.get_platform_parameter('TOTAL_PROCS', silent=True) or 0) if self._dvm else 0
+        )
 
         rfile_name = os.path.join(self.CM.sim_map[self.CM.fwk_sim_name].sim_root, 'resource_usage')
         # SIMYAN: try to safely make the directory...
@@ -137,13 +147,31 @@ def initialize(self, data_mngr, task_mngr, config_mngr, cmd_nodes=0, cmd_ppn=0):
             self.host = self.CM.get_platform_parameter('HOST')
             self.fwk.debug('RM: Host = %s', self.host)
             try:
-                (
-                    list_of_nodes,
-                    self.cores_per_node,
-                    self.sockets_per_node,
-                    self.max_ppn,
-                    self.accurate_nodes,
-                ) = get_resource_list(self.CM, self.host)
+                if self._dvm:
+                    list_of_nodes = _get_cpu_nodes(
+                        self.CM.get_platform_parameter,
+                        str(self.CM.get_platform_parameter('DVM_HWTHREADS', silent=True)).lower()
+                        == 'true',
+                    )
+                    # Process limits restrict ranks, never the CPU inventory.
+                    self.max_ppn = int(self.CM.get_platform_parameter('PROCS_PER_NODE'))
+                    if self.max_ppn <= 0:
+                        self.max_ppn = max(count for _, count in list_of_nodes)
+                    self.cores_per_node = max(count for _, count in list_of_nodes)
+                    self.sockets_per_node = int(
+                        self.CM.get_platform_parameter('SOCKETS_PER_NODE') or 1
+                    )
+                    self.accurate_nodes = (
+                        self.CM.get_platform_parameter('NODE_DETECTION') != 'manual'
+                    )
+                else:
+                    (
+                        list_of_nodes,
+                        self.cores_per_node,
+                        self.sockets_per_node,
+                        self.max_ppn,
+                        self.accurate_nodes,
+                    ) = get_resource_list(self.CM, self.host)
                 self.fwk.warning('RM: list_of_nodes = %s', str(list_of_nodes))
                 self.fwk.warning('RM: max_ppn = %d ', int(self.max_ppn))
                 if self.accurate_nodes is True and not self.CM.get_platform_parameter(
@@ -167,7 +195,7 @@ def initialize(self, data_mngr, task_mngr, config_mngr, cmd_nodes=0, cmd_ppn=0):
             if user_ppn <= self.max_ppn:
                 self.ppn = user_ppn
                 for i, (node, count) in enumerate(list_of_nodes):
-                    if count > self.ppn:
+                    if not self._dvm and count > self.ppn:
                         list_of_nodes[i] = (node, self.ppn)
                 self.fwk.warning('Using user set procs per node: %d', user_ppn)
             else:
@@ -182,7 +210,7 @@ def initialize(self, data_mngr, task_mngr, config_mngr, cmd_nodes=0, cmd_ppn=0):
                 if user_ppn <= self.max_ppn:
                     self.ppn = user_ppn
                     for i, (node, count) in enumerate(list_of_nodes):
-                        if count > self.ppn:
+                        if not self._dvm and count > self.ppn:
                             list_of_nodes[i] = (node, self.ppn)
                 else:
                     self.fwk.warning(
@@ -371,6 +399,14 @@ def get_allocation(
             whole_nodes = True
             whole_socks = True
 
+        if (
+            str(self.CM.get_platform_parameter('MPIRUN_VERSION', silent=True)).upper()
+            == 'OPENMPI-DVM'
+        ):
+            return self._get_dvm_allocation(
+                comp_id, nproc, task_id, whole_nodes, whole_socks, ppn, task_cpp, task_gpp
+            )
+
         # Are there enough cores to satisfy the request?
         # Returns the list of nodes that fit the bill
         allocation_possible = False
@@ -525,6 +561,106 @@ def get_allocation(
                     cores_allocated=cores_allocated,
                 )
 
+    def _get_dvm_allocation(
+        self,
+        comp_id: str,
+        nproc: int,
+        task_id: int,
+        whole_nodes: bool,
+        whole_socks: bool,
+        ppn: int,
+        task_cpp: int,
+        task_gpp: int,
+    ) -> Allocation:
+        """Reserve complete ranks for a DVM launch without reducing their core count.
+
+        :param comp_id: Identifier of the requesting component.
+        :param nproc: Number of MPI ranks requested.
+        :param task_id: Identifier used to track and release the allocation.
+        :param whole_nodes: Whether to reserve entire nodes.
+        :param whole_socks: Whether to reserve entire sockets.
+        :param ppn: Maximum ranks per node requested by the caller.
+        :param task_cpp: Processing elements per rank in the configured CPU mode;
+            zero retains launcher defaults.
+        :param task_gpp: GPUs per rank.
+        :returns: Allocation containing rank counts and reserved physical cores.
+        :raises ValueError: If the rank or core count is invalid.
+        :raises BadResourceRequestError: If the request exceeds total capacity.
+        :raises InsufficientResourcesError: If capacity is temporarily occupied.
+        """
+        if nproc < 1 or not isinstance(task_cpp, int) or task_cpp < 0:
+            raise ValueError('DVM tasks require positive nproc and nonnegative integer task_cpp')
+        cpp = task_cpp or 1
+        ppn = min(ppn, self.cores_per_node // cpp)
+        total_capacity = sum(min(ppn, node.total_cores // cpp) for node in self.nodes.values())
+        if self._dvm_total_ranks > 0:
+            total_capacity = min(total_capacity, self._dvm_total_ranks)
+        requested_nodes = ceil(nproc / max(ppn, 1))
+        if nproc > total_capacity:
+            raise BadResourceRequestError(
+                comp_id, task_id, requested_nodes, max(1, requested_nodes - len(self.nodes))
+            )
+        if not self.check_gpus(ppn, task_gpp):
+            raise GpuResourceRequestMismatchError(comp_id, task_id, ppn, task_gpp, self.gpn)
+
+        # Plan first so a failed request never consumes a partial allocation.
+        plan = []
+        remaining = nproc
+        for name in self.avail_nodes:
+            node = self.nodes[name]
+            available = node.avail_cores
+            if whole_nodes:
+                available = available if available == node.total_cores else 0
+            elif whole_socks:
+                available = sum(
+                    sock.avail_cores
+                    for sock in node.sockets
+                    if sock.avail_cores == sock.total_cores
+                )
+            ranks_in_use = sum(
+                allocation.get(name, 0) for allocation in self._dvm_rank_allocations.values()
+            )
+            ranks = min(ppn, available // cpp, remaining, max(0, self.max_ppn - ranks_in_use))
+            if ranks:
+                plan.append((name, ranks))
+                remaining -= ranks
+            if not remaining:
+                break
+        if self._dvm_total_ranks > 0:
+            remaining = max(remaining, nproc - max(0, self._dvm_total_ranks - self.processes))
+        if remaining:
+            raise InsufficientResourcesError(
+                comp_id, task_id, requested_nodes, ceil(remaining / ppn)
+            )
+
+        entries = []
+        cores_allocated = 0
+        for name, ranks in plan:
+            node = self.nodes[name]
+            count, cores = node.allocate(whole_nodes, whole_socks, task_id, comp_id, ranks * cpp)
+            entries.append((name, cores))
+            cores_allocated += count
+            if name not in self.alloc_nodes:
+                self.alloc_nodes.append(name)
+            if not node.avail_cores:
+                self.avail_nodes.remove(name)
+        self.processes += nproc
+        self.alloc_cores += cores_allocated
+        self.avail_cores -= cores_allocated
+        self.active_tasks[task_id] = (comp_id, nproc, cores_allocated)
+        self._dvm_rank_allocations[task_id] = dict(plan)
+        self.report_rm_status('allocation for DVM task %d' % task_id)
+        return Allocation(
+            partial_node=not whole_nodes,
+            nodelist=[name for name, _ in plan],
+            corelist=None if whole_nodes else entries,
+            ppn=ppn,
+            max_ppn=self.max_ppn,
+            cpp=task_cpp,
+            accurate_nodes=self.accurate_nodes,
+            cores_allocated=cores_allocated,
+        )
+
     def check_whole_node_cap(self, nproc, ppn):
         """
         Determine if it is currently possible to allocate *nproc* processes
@@ -678,6 +814,7 @@ def release_allocation(self, task_id, status):
         self.avail_cores += num_cores
         self.alloc_cores -= num_cores
         self.processes -= nproc
+        self._dvm_rank_allocations.pop(task_id, None)
 
         self.report_rm_status('released nodes for task %d' % task_id)