From 8f5ebbef8fd396da633d97065f91673faa775010 Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Thu, 29 Oct 2026 19:27:05 -0400 Subject: [PATCH] feat: implement DVM allocation method for complete ranks without core reduction --- ipsframework/resource_manager.py | 99 ++++++++++++++++++++++++++++++ ipsframework/services.py | 2 +- ipsframework/task_manager.py | 2 + tests/new/test_resource_manager.py | 33 ++++++++++ tests/new/test_task_manager.py | 62 +++++++++++++++++++ 5 files changed, 197 insertions(+), 1 deletion(-) diff --git a/ipsframework/resource_manager.py b/ipsframework/resource_manager.py index e7f72c4e..843bf63d 100644 --- a/ipsframework/resource_manager.py +++ b/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 diff --git a/ipsframework/services.py b/ipsframework/services.py index e46e87e1..115af3ee 100644 --- a/ipsframework/services.py +++ b/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`` diff --git a/ipsframework/task_manager.py b/ipsframework/task_manager.py index e3c19736..45479415 100644 --- a/ipsframework/task_manager.py +++ b/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: diff --git a/tests/new/test_resource_manager.py b/tests/new/test_resource_manager.py index 4a0e4cea..edc68659 100644 --- a/tests/new/test_resource_manager.py +++ b/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() diff --git a/tests/new/test_task_manager.py b/tests/new/test_task_manager.py index de563c96..fdfe8162 100644 --- a/tests/new/test_task_manager.py +++ b/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()