feat: implement DVM allocation method for complete ranks without core…

ipsframework/resource_manager.py

@@ -371,6 +371,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 +533,97 @@ 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: Physical cores per rank; 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())
+        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 = min(ppn, available // cpp, remaining)
+            if ranks:
+                plan.append((name, ranks))
+                remaining -= ranks
+            if not remaining:
+                break
+        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.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

ipsframework/services.py

@@ -956,7 +956,7 @@ def launch_task(self, nproc: int, working_dir: str, binary: str, *args, **keywor
         manage how the binary is launched.  Keywords may be the following:
 
             * *task_ppn* : the processes per node value for this task
-            * *task_cpp* : the cores per process, only used when ``MPIRUN=srun`` commands
+            * *task_cpp* : cores per process
             * *task_gpp* : the gpus per process, only used when ``MPIRUN=srun`` commands
             * *omp* : If ``True`` the task will be launch with the correct OpenMP environment
                variables set, only used when ``MPIRUN=srun``

ipsframework/task_manager.py

@@ -460,6 +460,8 @@ def build_launch_cmd(
                             str(nproc),
                         ]
                     )
+                    if cpp:
+                        cmd += f' --map-by slot:PE={cpp}:CORECPUS --bind-to core'
                 elif smp_node:
                     cmd = ' '.join([mpicmd, nproc_flag, str(nproc)])
                 else:

tests/new/test_resource_manager.py

@@ -375,3 +375,36 @@ def test_allocations(tmpdir):
         assert lines[6] == 'core: 1  - task_id: 0  - owner: comp0'
         assert lines[7] == 'core: 2  - task_id: 0  - owner: comp0'
         assert lines[8] == 'core: 3  - task_id: 0  - owner: comp0'
+
+
+@pytest.mark.parametrize('cpp', [2, 4])
+def test_dvm_complete_ranks(tmp_path, cpp: int) -> None:
+    """Ensure ranks cannot be split across nodes with fragmented capacity.
+
+    :param tmp_path: Isolated resource report directory.
+    :param cpp: Physical cores per MPI rank.
+    :returns: None.
+    """
+    cm = mock.Mock()
+    cm.fwk_sim_name = 'test'
+    cm.sim_map = {'test': mock.Mock(sim_root=str(tmp_path))}
+    cm.get_platform_parameter.return_value = 'HOST'
+    rm = ResourceManager(mock.Mock())
+    rm.initialize(mock.Mock(), mock.Mock(), cm, cmd_nodes=2, cmd_ppn=4)
+    cm.get_platform_parameter.return_value = 'OPENMPI-DVM'
+    # One rank on each node leaves fewer than cpp free cores on each.
+    rm.get_allocation('first', 2, 1, False, False, task_ppn=1, task_cpp=3)
+    with pytest.raises(InsufficientResourcesError):
+        rm.get_allocation('second', 1, 2, False, False, task_cpp=cpp)
+    assert rm.avail_cores == 2
+    assert rm.processes == 2
+    rm.release_allocation(1, 0)
+    result = rm.get_allocation('second', 2, 3, False, False, task_cpp=cpp)
+    assert result.cores_allocated == 2 * cpp
+    assert result.ppn == min(2, 4 // cpp)
+    rm.release_allocation(3, 0)
+    assert rm.avail_cores == 8
+    for invalid in (-1, 1.5, '4'):
+        with pytest.raises(ValueError):
+            rm.get_allocation('invalid', 1, 4, False, False, task_cpp=invalid)
+    rm.reporting_file.close()

tests/new/test_task_manager.py

@@ -1056,3 +1056,65 @@ def test_build_launch_cmd_prun_uses_explicit_dvm_uri(monkeypatch):
     monkeypatch.setenv('IPS_DVM_URI_FILE', '/work/dvm.uri.nid001.42')
     cmd, _ = tm.build_launch_cmd(**kwargs)
     assert cmd.startswith('/usr/bin/prun --dvm-uri file:/work/dvm.uri.nid001.42 ')
+
+
+@pytest.mark.parametrize('cpp', [0, 1, 4])
+@pytest.mark.parametrize('whole_nodes,whole_socks', [(False, False), (False, True), (True, True)])
+def test_dvm_task_core_budget(
+    tmp_path, monkeypatch, cpp: int, whole_nodes: bool, whole_socks: bool
+) -> None:
+    """Exercise allocation, command generation, exhaustion, and release together.
+
+    :param tmp_path: Isolated resource report directory.
+    :param monkeypatch: Fixture setting the DVM rendezvous file.
+    :param cpp: Requested cores per rank, including the default sentinel.
+    :param whole_nodes: Whether to reserve whole nodes.
+    :param whole_socks: Whether to reserve whole sockets.
+    :returns: None.
+    """
+    fwk = mock.Mock()
+    cm = mock.Mock()
+    cm.fwk_sim_name = 'test'
+    cm.sim_map = {'test': mock.Mock(sim_root=str(tmp_path))}
+    cm.get_platform_parameter.return_value = 'HOST'
+    tm = TaskManager(fwk)
+    rm = ResourceManager(fwk)
+    tm.initialize(mock.Mock(), rm, cm)
+    rm.initialize(mock.Mock(), tm, cm, cmd_nodes=1, cmd_ppn=4)
+    cm.get_platform_parameter.return_value = 'OPENMPI-DVM'
+    tm.task_launch_cmd = 'mpirun'
+    tm.mpicmd = '/usr/bin/prun'
+    monkeypatch.setenv('IPS_DVM_URI_FILE', '/tmp/test-dvm-uri')
+    task_id, cmd, _, reserved = tm._init_task(
+        'component',
+        1,
+        'executable',
+        str(tmp_path),
+        0,
+        cpp,
+        False,
+        0,
+        whole_nodes,
+        whole_socks,
+        [],
+        None,
+    )
+    assert '-np 1' in cmd
+    if cpp:
+        assert f'--map-by slot:PE={cpp}:CORECPUS --bind-to core' in cmd
+    else:
+        assert '--map-by' not in cmd
+    expected = 4 if whole_nodes or whole_socks else (cpp or 1)
+    assert reserved == expected
+    assert rm.alloc_cores == expected
+    assert rm.processes == 1
+    with pytest.raises(InsufficientResourcesError):
+        rm.get_allocation('other', 1, 99, False, False, task_cpp=4)
+    assert rm.alloc_cores == expected
+    rm.release_allocation(task_id, 0)
+    assert rm.avail_cores == 4
+    assert rm.alloc_cores == rm.processes == 0
+    with pytest.raises(BadResourceRequestError):
+        rm.get_allocation('other', 1, 100, False, False, task_cpp=5)
+    assert rm.avail_cores == 4
+    rm.reporting_file.close()