From 8f5ebbef8fd396da633d97065f91673faa775010 Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Thu, 29 Oct 2026 19:27:05 -0400 Subject: [PATCH 1/8] 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() From 0316f4b49d29a0a4e48b21f9e140a21779b67d19 Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Tue, 3 Nov 2026 10:44:40 -0500 Subject: [PATCH 2/8] Corrected issue whereby we were erroneously overriding the hwthreads=true --- ipsframework/task_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ipsframework/task_manager.py b/ipsframework/task_manager.py index 45479415..7b0c3d41 100644 --- a/ipsframework/task_manager.py +++ b/ipsframework/task_manager.py @@ -461,7 +461,7 @@ def build_launch_cmd( ] ) if cpp: - cmd += f' --map-by slot:PE={cpp}:CORECPUS --bind-to core' + cmd += f' --map-by slot:PE={cpp}' elif smp_node: cmd = ' '.join([mpicmd, nproc_flag, str(nproc)]) else: From cd42e177a0813a820d79ced8c31114221213fb19 Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Tue, 3 Nov 2026 10:56:41 -0500 Subject: [PATCH 3/8] DVM test harness, though it relies strongly on monkeypatching --- tests/new/test_dvm_plugin.py | 170 +++++++++++++++++++++++++++++++++++ 1 file changed, 170 insertions(+) create mode 100644 tests/new/test_dvm_plugin.py diff --git a/tests/new/test_dvm_plugin.py b/tests/new/test_dvm_plugin.py new file mode 100644 index 00000000..06dbc06b --- /dev/null +++ b/tests/new/test_dvm_plugin.py @@ -0,0 +1,170 @@ +"""DVM diagnostics and cleanup without requiring PRTE or Slurm.""" +import json +import logging +import subprocess +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from ipsframework import services + + +def records(worker): + return [json.loads(line) for line in worker.dvm_diagnostics.path.read_text().splitlines()] + + +@pytest.fixture +def running(tmp_path, monkeypatch): + monkeypatch.setenv('PMIX_SERVER_URI41', 'test-original') + monkeypatch.setenv('PRTE_MCA_ras_slurm_use_entire_allocation', '0') + monkeypatch.setenv('PRTE_MCA_rmaps_default_mapping_policy', 'core') + proc = Mock(pid=123, returncode=None) + proc.poll.side_effect = lambda: proc.returncode + + def popen(command, **kwargs): + from pathlib import Path + Path(command[-1]).write_text('test-uri\n') + return proc + + monkeypatch.setattr(services.subprocess, 'Popen', popen) + monkeypatch.setattr(services.shutil, 'which', lambda name: '/bin/' + name) + worker = SimpleNamespace(address='tcp://worker:1234') + plugin = services.DVMPlugin(logging.getLogger(__name__), tmp_path) + plugin.setup(worker) + yield plugin, worker, proc + worker.dvm_log.close() + + +def test_ready_and_graceful_shutdown(running, monkeypatch): + plugin, worker, proc = running + + def stop(*args, **kwargs): + proc.returncode = 0 + return SimpleNamespace(returncode=0) + + monkeypatch.setattr(services.subprocess, 'run', stop) + plugin.teardown(worker) + events = records(worker) + assert [r['event'] for r in events] == [ + 'launch_requested', 'launched', 'ready', 'shutdown_begin', + 'shutdown_command', 'exit_observed', 'shutdown_command_end', 'shutdown_complete', + ] + assert events[0]['environment']['PRTE_MCA_ras_slurm_use_entire_allocation'] == '1' + assert events[2]['uri'] == 'test-uri' + assert events[5]['shutdown_requested'] is True + proc.terminate.assert_not_called() + proc.kill.assert_not_called() + assert worker.dvm_log.closed + assert (worker.dvm_diagnostics.directory / 'shutdown.log').exists() + + +def test_shutdown_timeout_escalates_and_reaps(running, monkeypatch): + plugin, worker, proc = running + monkeypatch.setattr(services.subprocess, 'run', Mock(side_effect=subprocess.TimeoutExpired('pterm', 10))) + calls = [] + + def wait(timeout): + calls.append(timeout) + if len(calls) < 3: + raise subprocess.TimeoutExpired('prte', timeout) + proc.returncode = -9 + return -9 + + proc.wait.side_effect = wait + plugin.teardown(worker) + assert calls == [10, 5, 5] + proc.terminate.assert_called_once() + proc.kill.assert_called_once() + events = records(worker) + assert {'shutdown_command_failed', 'terminate_requested', 'kill_requested'} <= {e['event'] for e in events} + assert events[-1]['signal'] == 9 + assert events[-1]['dvm_return_code'] == -9 + + +def test_unexpected_exit_observed_once(running): + plugin, worker, proc = running + proc.returncode = 195 + worker.dvm_diagnostics.record('instance_end', instance='INSTANCE_1', return_code=1) + plugin.teardown(worker) + exits = [r for r in records(worker) if r['event'] == 'exit_observed'] + assert len(exits) == 1 + assert exits[0]['shutdown_requested'] is False + assert exits[0]['dvm_return_code'] == 195 + proc.terminate.assert_not_called() + + +def test_launch_failure_is_recorded_and_closed(tmp_path, monkeypatch): + monkeypatch.setenv('PRTE_MCA_ras_slurm_use_entire_allocation', '0') + monkeypatch.setenv('PRTE_MCA_rmaps_default_mapping_policy', 'core') + monkeypatch.setattr(services.subprocess, 'Popen', Mock(side_effect=OSError('missing prte'))) + worker = SimpleNamespace(address='worker') + plugin = services.DVMPlugin(logging.getLogger(__name__), tmp_path) + with pytest.raises(OSError, match='missing prte'): + plugin.setup(worker) + assert 'startup_failed' in [r['event'] for r in records(worker)] + assert worker.dvm_log.closed + + +def test_child_failure_references_dvm(tmp_path, monkeypatch): + monkeypatch.setenv('IPS_DVM_DIAGNOSTICS', '/run/dvm/id/lifecycle.jsonl') + monkeypatch.setenv('IPS_DVM_INSTANCE', 'INSTANCE_3') + proxy = Mock() + proxy.get_working_dir.return_value = str(tmp_path) + task = Mock(start_time=0, nproc=1, cores_allocated=1, binary='prun', args=(), command='prun test') + task.process.wait.return_value = 195 + proxy.task_map = {7: task} + assert services.ServicesProxy.wait_task(proxy, 7) == 195 + failure = json.loads(next(tmp_path.glob('dvm-failures.*.jsonl')).read_text()) + assert failure['instance'] == 'INSTANCE_3' + assert failure['return_code'] == 195 + assert failure['dvm_lifecycle'] == '/run/dvm/id/lifecycle.jsonl' + assert failure['dvm_state'] == 'unknown_in_child' + + +@pytest.mark.parametrize('return_code', [0, 1, 195]) +def test_instance_correlated_with_worker_state(running, monkeypatch, tmp_path, return_code): + _, worker, dvm = running + worker.name = 'worker' + worker.get_current_task = lambda: 'INSTANCE_4' + worker.log_event = Mock() + monkeypatch.setattr(services, 'get_worker', lambda: worker) + child = Mock(returncode=return_code) + + def communicate(**kwargs): + if return_code: + dvm.returncode = 195 + return '', '' + + child.communicate.side_effect = communicate + popen = Mock(return_value=child) + monkeypatch.setattr(services.subprocess, 'Popen', popen) + services.launch('ips.py', 'INSTANCE_4', tmp_path) + env = popen.call_args.kwargs['env'] + assert env['IPS_DVM_DIAGNOSTICS'] == str(worker.dvm_diagnostics.path) + assert env['IPS_DVM_INSTANCE'] == 'INSTANCE_4' + event = records(worker)[-1] + assert event['event'] == 'instance_end' + assert event['return_code'] == return_code + assert event['dvm_state'] == ('exited' if return_code else 'running') + assert worker.log_event.call_args.args[1]['state'] == ('Failed' if return_code else 'Succeeded') + + +def test_missing_shutdown_command_escalates(running, monkeypatch): + plugin, worker, proc = running + monkeypatch.setattr(services.shutil, 'which', lambda name: None) + count = 0 + + def wait(timeout): + nonlocal count + count += 1 + if count == 1: + raise subprocess.TimeoutExpired('prte', timeout) + proc.returncode = -15 + return -15 + + proc.wait.side_effect = wait + plugin.teardown(worker) + assert 'shutdown_command_unavailable' in [r['event'] for r in records(worker)] + proc.terminate.assert_called_once() + proc.kill.assert_not_called() From d795364e59954b8938f4357d5ab8d3b6d0c7197d Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Tue, 3 Nov 2026 11:17:35 -0500 Subject: [PATCH 4/8] feat: improve support for cores per task in DVM, enhancing CPU budget management --- doc/user_guides/ensembles.rst | 39 +++++++++- ipsframework/resource_manager.py | 60 ++++++++++++--- ipsframework/services.py | 44 ++++++++--- tests/new/test_dvm_plugin.py | 25 ++++-- tests/new/test_resource_manager.py | 101 ++++++++++++++++++++++++ tests/new/test_run_ensemble.py | 120 ++++++++++++++++++++++++++++- tests/new/test_task_manager.py | 4 +- 7 files changed, 361 insertions(+), 32 deletions(-) diff --git a/doc/user_guides/ensembles.rst b/doc/user_guides/ensembles.rst index 91e7779c..2f4d2d84 100644 --- a/doc/user_guides/ensembles.rst +++ b/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 diff --git a/ipsframework/resource_manager.py b/ipsframework/resource_manager.py index 843bf63d..52c9efe3 100644 --- a/ipsframework/resource_manager.py +++ b/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( @@ -552,7 +580,8 @@ def _get_dvm_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_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. @@ -564,6 +593,8 @@ def _get_dvm_allocation( 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( @@ -586,12 +617,17 @@ def _get_dvm_allocation( for sock in node.sockets if sock.avail_cores == sock.total_cores ) - ranks = min(ppn, available // cpp, remaining) + 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) @@ -612,6 +648,7 @@ def _get_dvm_allocation( 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, @@ -777,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) diff --git a/ipsframework/services.py b/ipsframework/services.py index 115af3ee..3002ce04 100644 --- a/ipsframework/services.py +++ b/ipsframework/services.py @@ -37,6 +37,7 @@ from rich.console import Console from ipsframework import ipsutil, messages +from ipsframework._internal.cpu_resources import _ensemble_worker_threads, _get_cpu_nodes from ipsframework._internal.definitions import IPS_EVENT_TYPE from ipsframework.cca_es_spec import initialize_event_service from ipsframework.component_registry import ComponentID @@ -2905,14 +2906,16 @@ def run_ensemble( `logfile` and `errfile` will write to the working directory local to the ensemble instance. - TODO be able to specify the number of cores per instance - :param template: configuration template file :param variables: a dict of variables to pass to the ensemble runs :param run_dir: in which to run the ensembles :param name: ensemble name, or string to prepend to generated instance directory and file names - :param cores_per_instance: How many cores per ensemble instances? + :param cores_per_instance: Positive CPU budget per instance, default one. + Counts physical cores, or hardware threads when ``hwthreads=True``. + Concurrent tasks inside an instance share this budget. Parent + concurrency uses allocated CPU capacity, independently of rank limits. + Manual platform CPU capacities must use the same unit. :param num_nodes: Total number of nodes to allocate for the ensemble runs. There will be one Dask worker assigned to each of these nodes. If none specified, then assign a worker to each allocated @@ -2928,6 +2931,13 @@ def run_ensemble( # This should be a unique variable across all ensembles we keep track # of in the portal This ID should only be shared by runs within an # ensemble + cores_per_instance = 1 if cores_per_instance is None else cores_per_instance + if ( + isinstance(cores_per_instance, bool) + or not isinstance(cores_per_instance, int) + or cores_per_instance < 1 + ): + raise ValueError('cores_per_instance must be a positive integer') portal_ensemble_id = str(uuid.uuid4()) use_portal = self._should_use_portal() @@ -3057,15 +3067,12 @@ def create_platform_config_file(prefix, working_dir, cores_per_instance, **kwarg # `prun` *should* use the environment variables set by the plugin # to find the DVM. platform_config['MPIRUN_VERSION'] = 'OPENMPI-DVM' + platform_config['DVM_HWTHREADS'] = hwthreads - # Set the budget of cores per instance. By default, we will give - # a single core per instance. - if cores_per_instance is not None: - platform_config['CORES_PER_NODE'] = cores_per_instance - platform_config['PROCS_PER_NODE'] = cores_per_instance - else: - platform_config['CORES_PER_NODE'] = 1 - platform_config['PROCS_PER_NODE'] = 1 + # This virtual node contains the instance's CPU budget. Rank limits + # remain separate even when their default equals that budget. + platform_config['CORES_PER_NODE'] = cores_per_instance + platform_config['PROCS_PER_NODE'] = cores_per_instance # for now each instance will always run on just one node platform_config['NODES'] = 1 @@ -3240,6 +3247,7 @@ def send_ensemble_instance_to_portal(ensemble_name: str, data_path: Path) -> Non try: # Note that we *always* use Dask to run the ensemble tasks self.logger.debug(f'Submitting {len(instances)} ensemble tasks') + self.task_pools[task_pool_name]._ensemble_budget = cores_per_instance num_submitted = self.submit_tasks( task_pool_name, block=True, @@ -3260,6 +3268,7 @@ def send_ensemble_instance_to_portal(ensemble_name: str, data_path: Path) -> Non except Exception as e: self.critical(f'Got an exception running ensemble: {e!s}') traceback.print_exc() + raise finally: # self.debug('Getting finished tasks') # exit_status = self.get_finished_tasks(task_pool_name) @@ -3523,6 +3532,7 @@ def __init__(self, name: str, services: ServicesProxy): self.dask_scheduler_file = None self.dask_client = None self.worker_event_logfile = None + self._ensemble_budget: int | None = None def _wait_any_task(self, block=True): """ @@ -3728,6 +3738,14 @@ def _make_worker_args( return base_args services: ServicesProxy = self.services + ensemble_threads = None + if self._ensemble_budget is not None: + if dask_worker_per_gpu: + raise ValueError('Ensemble CPU budgets require one Dask worker per node') + cpu_nodes = _get_cpu_nodes(services.get_config_param, hwthreads) + ensemble_threads = _ensemble_worker_threads( + [capacity for _, capacity in cpu_nodes], self._ensemble_budget + ) # Note that we use the absolute path since at some point we may # be in a different directory, which means that we otherwise would @@ -3824,6 +3842,10 @@ def _make_worker_args( nthreads = dask_ppw if dask_ppw else services.get_config_param('PROCS_PER_NODE') // gpn task_ppn = gpn task_gpp = 1 + elif ensemble_threads is not None: + nthreads = ensemble_threads + task_ppn = 1 + task_gpp = 0 else: # The number of threads per Dask worker is the number of processors # on that node divided by the cores per instance, which we're diff --git a/tests/new/test_dvm_plugin.py b/tests/new/test_dvm_plugin.py index 06dbc06b..ab312fe2 100644 --- a/tests/new/test_dvm_plugin.py +++ b/tests/new/test_dvm_plugin.py @@ -1,7 +1,9 @@ """DVM diagnostics and cleanup without requiring PRTE or Slurm.""" + import json import logging import subprocess +from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock @@ -23,7 +25,6 @@ def running(tmp_path, monkeypatch): proc.poll.side_effect = lambda: proc.returncode def popen(command, **kwargs): - from pathlib import Path Path(command[-1]).write_text('test-uri\n') return proc @@ -47,8 +48,14 @@ def stop(*args, **kwargs): plugin.teardown(worker) events = records(worker) assert [r['event'] for r in events] == [ - 'launch_requested', 'launched', 'ready', 'shutdown_begin', - 'shutdown_command', 'exit_observed', 'shutdown_command_end', 'shutdown_complete', + 'launch_requested', + 'launched', + 'ready', + 'shutdown_begin', + 'shutdown_command', + 'exit_observed', + 'shutdown_command_end', + 'shutdown_complete', ] assert events[0]['environment']['PRTE_MCA_ras_slurm_use_entire_allocation'] == '1' assert events[2]['uri'] == 'test-uri' @@ -61,7 +68,9 @@ def stop(*args, **kwargs): def test_shutdown_timeout_escalates_and_reaps(running, monkeypatch): plugin, worker, proc = running - monkeypatch.setattr(services.subprocess, 'run', Mock(side_effect=subprocess.TimeoutExpired('pterm', 10))) + monkeypatch.setattr( + services.subprocess, 'run', Mock(side_effect=subprocess.TimeoutExpired('pterm', 10)) + ) calls = [] def wait(timeout): @@ -77,7 +86,9 @@ def wait(timeout): proc.terminate.assert_called_once() proc.kill.assert_called_once() events = records(worker) - assert {'shutdown_command_failed', 'terminate_requested', 'kill_requested'} <= {e['event'] for e in events} + assert {'shutdown_command_failed', 'terminate_requested', 'kill_requested'} <= { + e['event'] for e in events + } assert events[-1]['signal'] == 9 assert events[-1]['dvm_return_code'] == -9 @@ -111,7 +122,9 @@ def test_child_failure_references_dvm(tmp_path, monkeypatch): monkeypatch.setenv('IPS_DVM_INSTANCE', 'INSTANCE_3') proxy = Mock() proxy.get_working_dir.return_value = str(tmp_path) - task = Mock(start_time=0, nproc=1, cores_allocated=1, binary='prun', args=(), command='prun test') + task = Mock( + start_time=0, nproc=1, cores_allocated=1, binary='prun', args=(), command='prun test' + ) task.process.wait.return_value = 195 proxy.task_map = {7: task} assert services.ServicesProxy.wait_task(proxy, 7) == 195 diff --git a/tests/new/test_resource_manager.py b/tests/new/test_resource_manager.py index edc68659..05387caf 100644 --- a/tests/new/test_resource_manager.py +++ b/tests/new/test_resource_manager.py @@ -408,3 +408,104 @@ def test_dvm_complete_ranks(tmp_path, cpp: int) -> None: with pytest.raises(ValueError): rm.get_allocation('invalid', 1, 4, False, False, task_cpp=invalid) rm.reporting_file.close() + + +@pytest.mark.parametrize( + 'whole_nodes,whole_socks,reserved', [(False, False, 4), (False, True, 4), (True, True, 8)] +) +def test_dvm_unequal_core_and_rank_limits( + tmp_path, whole_nodes: bool, whole_socks: bool, reserved: int +) -> None: + """Initialize from an eight-core, two-rank platform and verify allocation. + + :param tmp_path: Resource-report directory. + :param whole_nodes: Reserve complete nodes. + :param whole_socks: Reserve complete sockets. + :param reserved: Expected CPU reservation. + :returns: None. + """ + params = { + 'HOST': 'test', + 'MPIRUN_VERSION': 'OPENMPI-DVM', + 'NODE_DETECTION': 'manual', + 'CORES_PER_NODE': 8, + 'PROCS_PER_NODE': 2, + 'NODES': 1, + 'TOTAL_PROCS': 0, + 'SOCKETS_PER_NODE': 2, + 'NODE_ALLOCATION_MODE': 'SHARED', + 'GPUS_PER_NODE': 0, + } + cm = mock.Mock() + cm.fwk_sim_name = 'test' + cm.sim_map = {'test': mock.Mock(sim_root=str(tmp_path))} + cm.get_platform_parameter.side_effect = lambda key, **kwargs: params.get(key) + rm = ResourceManager(mock.Mock()) + rm.initialize(mock.Mock(), mock.Mock(), cm) + assert rm.total_cores == rm.avail_cores == 8 + assert rm.ppn == rm.max_ppn == 2 + result = rm.get_allocation('test', 1, 1, whole_nodes, whole_socks, task_cpp=4) + assert result.cores_allocated == reserved + assert result.ppn == 1 + assert rm.processes == 1 + rm.release_allocation(1, 0) + assert rm.avail_cores == 8 + result = rm.get_allocation('test', 2, 2, False, False, task_cpp=4) + assert result.cores_allocated == 8 + assert result.ppn == 2 + with pytest.raises(InsufficientResourcesError): + rm.get_allocation('blocked', 1, 3, False, False, task_cpp=1) + assert rm.avail_cores == 0 + rm.release_allocation(2, 0) + with pytest.raises(BadResourceRequestError): + rm.get_allocation('oversized', 2, 4, False, False, task_cpp=5) + assert rm.avail_cores == 8 + assert rm.processes == rm.alloc_cores == 0 + # Separate tasks may each launch one rank while sharing the CPU budget. + rm.get_allocation('first', 1, 5, False, False, task_cpp=4) + rm.get_allocation('second', 1, 6, False, False, task_cpp=4) + assert rm.alloc_cores == 8 + rm.release_allocation(5, 0) + rm.release_allocation(6, 0) + assert rm.avail_cores == 8 + rm.reporting_file.close() + + +def test_dvm_rank_limit_does_not_destroy_cpu_capacity(tmp_path) -> None: + """Track concurrent rank occupancy independently from reserved CPUs. + + :param tmp_path: Resource-report directory. + :returns: None. + """ + params = { + 'HOST': 'test', + 'MPIRUN_VERSION': 'OPENMPI-DVM', + 'NODE_DETECTION': 'manual', + 'CORES_PER_NODE': 8, + 'PROCS_PER_NODE': 2, + 'NODES': 1, + 'TOTAL_PROCS': 0, + 'SOCKETS_PER_NODE': 1, + 'NODE_ALLOCATION_MODE': 'SHARED', + 'GPUS_PER_NODE': 0, + } + cm = mock.Mock() + cm.fwk_sim_name = 'test' + cm.sim_map = {'test': mock.Mock(sim_root=str(tmp_path))} + cm.get_platform_parameter.side_effect = lambda key, **kwargs: params.get(key) + rm = ResourceManager(mock.Mock()) + rm.initialize(mock.Mock(), mock.Mock(), cm) + rm.get_allocation('first', 1, 1, False, False, task_cpp=1) + rm.get_allocation('second', 1, 2, False, False, task_cpp=1) + assert rm.avail_cores == 6 + with pytest.raises(InsufficientResourcesError): + rm.get_allocation('third', 1, 3, False, False, task_cpp=1) + rm.release_allocation(1, 0) + allocation = rm.get_allocation('third', 1, 3, False, False, task_cpp=4) + assert allocation.cores_allocated == 4 + assert rm.avail_cores == 3 + rm.release_allocation(2, 0) + rm.release_allocation(3, 0) + assert rm.avail_cores == 8 + assert not rm._dvm_rank_allocations + rm.reporting_file.close() diff --git a/tests/new/test_run_ensemble.py b/tests/new/test_run_ensemble.py index df935175..cf893ad5 100644 --- a/tests/new/test_run_ensemble.py +++ b/tests/new/test_run_ensemble.py @@ -1,5 +1,9 @@ import logging import os +from unittest import mock + +import pytest +from configobj import ConfigObj from ipsframework import ServicesProxy, TaskPool from ipsframework import services as services_module @@ -45,7 +49,11 @@ def record_submit_tasks(*args, **kwargs): submit_tasks_kwargs.append(kwargs) return 1 - monkeypatch.setattr(services, 'create_task_pool', lambda name: None) + monkeypatch.setattr( + services, + 'create_task_pool', + lambda name: services.task_pools.setdefault(name, TaskPool(name, services)), + ) monkeypatch.setattr(services, 'add_task', record_add_task) monkeypatch.setattr(services, 'submit_tasks', record_submit_tasks) monkeypatch.setattr(services, 'get_finished_tasks', lambda task_pool_name: {}) @@ -70,7 +78,7 @@ def record_submit_tasks(*args, **kwargs): 'block': True, 'use_dask': True, 'dask_nodes': 1, - 'dask_ppw': None, + 'dask_ppw': 1, 'oversubscribe': False, 'hwthreads': False, 'logfile': 'instance.out', @@ -241,3 +249,111 @@ def test_launch_timeout_kills_task(tmpdir, monkeypatch): monkeypatch.setattr(services_module, 'get_worker', DummyDaskWorker) assert services_module.launch('sleep', 'task_0', str(tmpdir), 100, timeout=1) == ('task_0', -1) + + +@pytest.mark.parametrize('budget,hwthreads', [(None, False), (8, False), (8, True)]) +def test_ensemble_generated_budget( + tmp_path, monkeypatch, budget: int | None, hwthreads: bool +) -> None: + """Keep generated instance budgets and parent scheduling in the same units. + + :param tmp_path: Directory for generated configuration files. + :param monkeypatch: Fixture replacing external task submission. + :param budget: Requested budget, or default sentinel. + :param hwthreads: Hardware-thread allocation mode. + :returns: None. + """ + template = tmp_path / 'template.config' + template.write_text('[comp]\nA = ?\n') + services = ServicesProxy(None, None, None, {'USE_PORTAL': 'False'}, None) + services.fwk = DummyFramework() + services.logger = logging.getLogger(__name__) + monkeypatch.setattr( + services, + 'create_task_pool', + lambda name: services.task_pools.setdefault(name, TaskPool(name, services)), + ) + monkeypatch.setattr(services, 'add_task', mock.Mock()) + submit = mock.Mock(return_value=1) + monkeypatch.setattr(services, 'submit_tasks', submit) + monkeypatch.setattr(services, 'remove_task_pool', mock.Mock()) + (tmp_path / 'runs').mkdir() + services.run_ensemble( + template, + {'comp': {'A': ['1']}}, + tmp_path / 'runs', + 'instance', + num_nodes=1, + cores_per_instance=budget, + hwthreads=hwthreads, + ) + config = ConfigObj(str(next(tmp_path.rglob('*_platform.config')))) + assert int(config['CORES_PER_NODE']) == (budget or 1) + assert int(config['PROCS_PER_NODE']) == (budget or 1) + assert int(config['NODES']) == 1 + assert config.as_bool('DVM_HWTHREADS') is hwthreads + assert submit.call_args.kwargs['dask_ppw'] == (budget or 1) + assert next(iter(services.task_pools.values()))._ensemble_budget == (budget or 1) + + +@pytest.mark.parametrize('capacity,budget,expected', [(128, 8, 16), (128, 1, 128), (256, 8, 32)]) +def test_ensemble_worker_command_uses_cpu_budget( + monkeypatch, capacity: int, budget: int, expected: int +) -> None: + """Size the real worker command using CPU budgets rather than rank limits. + + :param monkeypatch: Fixture mocking process launches. + :param capacity: Explicit CPU capacity in the selected unit. + :param budget: Per-instance CPU budget. + :param expected: Number of worker execution threads. + :returns: None. + """ + services = mock.Mock() + config = { + 'NODE_DETECTION': 'manual', + 'CORES_PER_NODE': capacity, + 'PROCS_PER_NODE': 2, + 'NODES': 1, + } + services.get_config_param.side_effect = lambda key, **kwargs: config.get(key) + services.launch_task.side_effect = RuntimeError('stop at worker launch') + pool = TaskPool('ensemble', services) + pool._ensemble_budget = budget + pool.dask_scheduler = ['dask', 'scheduler'] + pool.DaskWorker = ['dask', 'worker'] + monkeypatch.setattr(services_module.subprocess, 'Popen', mock.Mock()) + with pytest.raises(RuntimeError, match='stop at worker launch'): + pool.submit_dask_tasks(dask_ppw=budget) + args = services.launch_task.call_args.args + assert args[args.index('--nthreads') + 1] == str(expected) + + +def test_oversized_ensemble_budget_fails_before_launch(monkeypatch) -> None: + """Reject an oversized budget before starting a scheduler or worker. + + :param monkeypatch: Fixture mocking subprocess creation. + :returns: None. + """ + services = mock.Mock() + config = {'NODE_DETECTION': 'manual', 'CORES_PER_NODE': 128, 'NODES': 1} + services.get_config_param.side_effect = lambda key, **kwargs: config.get(key) + pool = TaskPool('ensemble', services) + pool._ensemble_budget = 129 + popen = mock.Mock() + monkeypatch.setattr(services_module.subprocess, 'Popen', popen) + with pytest.raises(ValueError, match='exceeds'): + pool.submit_dask_tasks() + popen.assert_not_called() + services.launch_task.assert_not_called() + + +@pytest.mark.parametrize('budget', [0, -1, True, 1.5]) +def test_run_ensemble_rejects_invalid_budget(budget) -> None: + """Reject invalid instance budgets before generating any files. + + :param budget: Invalid CPU budget. + :returns: None. + """ + services = ServicesProxy(None, None, None, {'USE_PORTAL': 'False'}, None) + with pytest.raises(ValueError, match='positive integer'): + services.run_ensemble('unused', {}, 'unused', 'instance', cores_per_instance=budget) diff --git a/tests/new/test_task_manager.py b/tests/new/test_task_manager.py index fdfe8162..7a1be344 100644 --- a/tests/new/test_task_manager.py +++ b/tests/new/test_task_manager.py @@ -1101,7 +1101,9 @@ def test_dvm_task_core_budget( ) assert '-np 1' in cmd if cpp: - assert f'--map-by slot:PE={cpp}:CORECPUS --bind-to core' in cmd + assert f'--map-by slot:PE={cpp} ' in cmd + assert 'CORECPUS' not in cmd + assert '--bind-to core' not in cmd else: assert '--map-by' not in cmd expected = 4 if whole_nodes or whole_socks else (cpp or 1) From ad3f5e93cb65c8b176bd4f32880bfa9a3b7a6e74 Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Tue, 3 Nov 2026 11:41:32 -0500 Subject: [PATCH 5/8] feat: add CPU budget management independent of MPI rank limits and local topology --- ipsframework/_internal/cpu_resources.py | 107 +++++++++++++++++ tests/new/test_cpu_resources.py | 148 ++++++++++++++++++++++++ 2 files changed, 255 insertions(+) create mode 100644 ipsframework/_internal/cpu_resources.py create mode 100644 tests/new/test_cpu_resources.py diff --git a/ipsframework/_internal/cpu_resources.py b/ipsframework/_internal/cpu_resources.py new file mode 100644 index 00000000..a59056f6 --- /dev/null +++ b/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 diff --git a/tests/new/test_cpu_resources.py b/tests/new/test_cpu_resources.py new file mode 100644 index 00000000..f75b4af1 --- /dev/null +++ b/tests/new/test_cpu_resources.py @@ -0,0 +1,148 @@ +"""CPU-budget tests use explicit scheduler data, never the test host's hwloc.""" + +from unittest import mock + +import pytest + +from ipsframework._internal.cpu_resources import ( + _ensemble_worker_threads, + _expand_slurm_counts, + _get_cpu_nodes, +) + + +def _parameters(**overrides): + """Create a configuration lookup. + + :param overrides: Platform settings to override. + :returns: Lookup accepting optional service keyword arguments. + """ + values = {'NODE_DETECTION': 'manual', 'CORES_PER_NODE': 8, 'NODES': 2, 'PROCS_PER_NODE': 2} + values.update(overrides) + return lambda name, **kwargs: values.get(name) + + +@pytest.mark.parametrize('hwthreads', [False, True]) +def test_manual_budget_is_independent_of_slots(hwthreads: bool) -> None: + """Manual budgets already use the selected CPU unit. + + :param hwthreads: Selected allocation mode. + :returns: None. + """ + assert _get_cpu_nodes(_parameters(), hwthreads) == [('dummynode0', 8), ('dummynode1', 8)] + + +@pytest.mark.parametrize('hwthreads,expected', [(False, [128, 32]), (True, [256, 64])]) +def test_slurm_restricted_allocation(monkeypatch, hwthreads: bool, expected: list[int]) -> None: + """Restrict capacity using remote Slurm topology and granted CPU counts. + + :param monkeypatch: Environment and subprocess fixture. + :param hwthreads: Selected allocation mode. + :param expected: Expected per-node capacities. + :returns: None. + """ + monkeypatch.setenv('SLURM_JOB_NODELIST', 'nid[1-2]') + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '256,64') + topology = '\n'.join( + f'NodeName=nid{i} CPUTot=256 Sockets=2 CoresPerSocket=64 ThreadsPerCore=2' for i in (1, 2) + ) + query = mock.Mock(side_effect=['nid1\nnid2\n', topology]) + monkeypatch.setattr('ipsframework._internal.cpu_resources.subprocess.check_output', query) + assert _get_cpu_nodes( + _parameters(NODE_DETECTION='slurm_env', CORES_PER_NODE=128), hwthreads + ) == list(zip(['nid1', 'nid2'], expected, strict=True)) + assert query.call_count == 2 + + +@pytest.mark.parametrize('hwthreads,expected', [(False, 4), (True, 8)]) +def test_slurm_core_counted_cpus(monkeypatch, hwthreads: bool, expected: int) -> None: + """Support Slurm configured to count cores rather than SMT threads. + + :param monkeypatch: Environment and subprocess fixture. + :param hwthreads: Selected allocation mode. + :param expected: Expected capacity. + :returns: None. + """ + monkeypatch.setenv('SLURM_JOB_NODELIST', 'node1') + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '4') + monkeypatch.setattr( + 'ipsframework._internal.cpu_resources.subprocess.check_output', + mock.Mock( + side_effect=[ + 'node1', + 'NodeName=node1 CPUTot=8 Sockets=2 CoresPerSocket=4 ThreadsPerCore=2', + ] + ), + ) + assert _get_cpu_nodes(_parameters(NODE_DETECTION='slurm_env'), hwthreads) == [ + ('node1', expected) + ] + + +def test_slurm_counts() -> None: + """Validate compressed CPU counts. + + :returns: None. + """ + assert _expand_slurm_counts('256(x2),64') == [256, 256, 64] + for invalid in ('0', '2(x0)', 'abc', '2,'): + with pytest.raises(ValueError): + _expand_slurm_counts(invalid) + + +@pytest.mark.parametrize( + 'capacities,budget,expected', + [([128, 128], 8, 16), ([128], 1, 128), ([256], 8, 32), ([128, 32], 8, 4)], +) +def test_ensemble_concurrency(capacities: list[int], budget: int, expected: int) -> None: + """Compute concurrency in the chosen CPU unit. + + :param capacities: Available CPUs per worker. + :param budget: Per-instance CPU budget. + :param expected: Safe uniform concurrency. + :returns: None. + """ + assert _ensemble_worker_threads(capacities, budget) == expected + + +@pytest.mark.parametrize('budget', [0, -1, True, 1.5, 129]) +def test_invalid_ensemble_budget(budget) -> None: + """Reject invalid or oversized instance budgets. + + :param budget: Invalid per-instance budget. + :returns: None. + """ + with pytest.raises(ValueError): + _ensemble_worker_threads([128], budget) + + +def test_slurm_missing_topology_is_not_local_fallback(monkeypatch) -> None: + """Missing remote topology must not substitute the Mac's CPU count. + + :param monkeypatch: Fixture supplying deliberately incomplete Slurm data. + :returns: None. + """ + monkeypatch.setenv('SLURM_JOB_NODELIST', 'node1') + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '8') + monkeypatch.setattr( + 'ipsframework._internal.cpu_resources.subprocess.check_output', + mock.Mock(side_effect=['node1', 'NodeName=node1 CPUTot=256']), + ) + with pytest.raises(ValueError, match='Missing Slurm CPU topology'): + _get_cpu_nodes(_parameters(NODE_DETECTION='slurm_env')) + + +def test_slurm_mismatched_counts(monkeypatch) -> None: + """Reject mismatched scheduler node/count records. + + :param monkeypatch: Fixture supplying malformed Slurm data. + :returns: None. + """ + monkeypatch.setenv('SLURM_JOB_NODELIST', 'node[1-2]') + monkeypatch.setenv('SLURM_JOB_CPUS_PER_NODE', '8') + monkeypatch.setattr( + 'ipsframework._internal.cpu_resources.subprocess.check_output', + mock.Mock(return_value='node1 node2'), + ) + with pytest.raises(ValueError, match='different lengths'): + _get_cpu_nodes(_parameters(NODE_DETECTION='slurm_env')) From 013808f1649fe09eaa85272b0b63c14b54e00c3c Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Tue, 3 Nov 2026 15:01:39 -0500 Subject: [PATCH 6/8] feat: enhance task initialization with error handling and resource rollback --- ipsframework/task_manager.py | 197 ++++++++++++++++++--------------- tests/new/test_task_manager.py | 111 +++++++++++++++++++ 2 files changed, 218 insertions(+), 90 deletions(-) diff --git a/ipsframework/task_manager.py b/ipsframework/task_manager.py index 7b0c3d41..697fbe34 100644 --- a/ipsframework/task_manager.py +++ b/ipsframework/task_manager.py @@ -330,42 +330,63 @@ def _init_task( task_cpp=tcpp, task_gpp=tgpp, ) - self.fwk.debug('RM: get_allocation() returned %s', str(allocation)) + try: + self.fwk.debug('RM: get_allocation() returned %s', str(allocation)) - if allocation.partial_node or allocation.accurate_nodes: - nodes = ','.join(allocation.nodelist) - else: - nodes = '' + if allocation.partial_node or allocation.accurate_nodes: + nodes = ','.join(allocation.nodelist) + else: + nodes = '' + + (cmd, env_update) = self.build_launch_cmd( + nproc, + binary, + cmd_args, + working_dir, + allocation.ppn, + allocation.max_ppn, + nodes, + allocation.accurate_nodes, + allocation.partial_node, + task_id, + allocation.cpp, + omp, + tgpp, + allocation.corelist, + launch_cmd_extra_args, + ) - (cmd, env_update) = self.build_launch_cmd( - nproc, - binary, - cmd_args, - working_dir, - allocation.ppn, - allocation.max_ppn, - nodes, - allocation.accurate_nodes, - allocation.partial_node, - task_id, - allocation.cpp, - omp, - tgpp, - allocation.corelist, - launch_cmd_extra_args, - ) + self.curr_task_table[task_id] = { + 'component': caller_id, + 'status': 'init_task', + 'binary': binary, + 'nproc': nproc, + 'args': cmd_args, + 'launch_cmd': cmd, + 'env_update': env_update, + } + + return (task_id, cmd, env_update, allocation.cores_allocated) + except Exception: + self._rollback_tasks([task_id]) + raise + + def _rollback_tasks(self, task_ids: list[int]) -> None: + """Release reservations after task initialization fails. - self.curr_task_table[task_id] = { - 'component': caller_id, - 'status': 'init_task', - 'binary': binary, - 'nproc': nproc, - 'args': cmd_args, - 'launch_cmd': cmd, - 'env_update': env_update, - } + Attempt every release, logging cleanup failures without masking the + original initialization exception. Keep task records if release fails. - return (task_id, cmd, env_update, allocation.cores_allocated) + :param task_ids: Task identifiers reserved by the failed operation. + :returns: None. + """ + for task_id in task_ids: + try: + self.resource_mgr.release_allocation(task_id, -1) + except Exception: + self.fwk.exception('TM: Failed to roll back task %d', task_id) + else: + self.curr_task_table.pop(task_id, None) def build_launch_cmd( self, @@ -749,65 +770,61 @@ def init_task_pool(self, init_task_msg: messages.ServiceRequestMessage): caller_id = init_task_msg.sender_id task_dict = init_task_msg.args[0] ret_dict = {} - for task_name in task_dict: - # handle for task related things - task_init = task_dict[task_name] + try: + for task_name in task_dict: + # handle for task related things + task_init = task_dict[task_name] + + try: + ret_dict[task_name] = self._init_task( + caller_id, + task_init.nproc, + task_init.binary, + task_init.working_dir, + task_init.tppn, + task_init.tcpp, + task_init.omp, + task_init.tgpp, + task_init.wnodes, + task_init.wsocks, + task_init.cmd_args, + task_init.launch_cmd_extra_args, + ) + except InsufficientResourcesError: + continue + except BadResourceRequestError as e: + self.fwk.error( + 'There has been a fatal error, %s requested %d too many processors in task %d', + caller_id, + e.deficit, + e.task_id, + ) + raise + except ResourceRequestMismatchError as e: + self.fwk.error( + 'There has been a fatal error, %s requested too few processors per node to launch task %d (request: procs = %d, ppn = %d)', + caller_id, + e.task_id, + e.nproc, + e.ppn, + ) + raise + except GpuResourceRequestMismatchError as e: + self.fwk.error( + 'There has been a fatal error, %s requested too many GPUs per node to launch task %d (requested: ppn = %d, gpp = %d)', + caller_id, + e.task_id, + e.ppn, + e.gpp, + ) + raise + except Exception: + self.fwk.exception('TM:init_task_pool(): Allocation exception') + raise - try: - ret_dict[task_name] = self._init_task( - caller_id, - task_init.nproc, - task_init.binary, - task_init.working_dir, - task_init.tppn, - task_init.tcpp, - task_init.omp, - task_init.tgpp, - task_init.wnodes, - task_init.wsocks, - task_init.cmd_args, - task_init.launch_cmd_extra_args, - ) - except InsufficientResourcesError: - continue - except BadResourceRequestError as e: - self.fwk.error( - 'There has been a fatal error, %s requested %d too many processors in task %d', - caller_id, - e.deficit, - e.task_id, - ) - for task_id, _, _, _ in ret_dict.values(): - self.resource_mgr.release_allocation(task_id, -1) - del self.curr_task_table[task_id] - raise - except ResourceRequestMismatchError as e: - self.fwk.error( - 'There has been a fatal error, %s requested too few processors per node to launch task %d (request: procs = %d, ppn = %d)', - caller_id, - e.task_id, - e.nproc, - e.ppn, - ) - for task_id, _, _, _ in ret_dict.values(): - self.resource_mgr.release_allocation(task_id, -1) - del self.curr_task_table[task_id] - raise - except GpuResourceRequestMismatchError as e: - self.fwk.error( - 'There has been a fatal error, %s requested too many GPUs per node to launch task %d (requested: ppn = %d, gpp = %d)', - caller_id, - e.task_id, - e.ppn, - e.gpp, - ) - for task_id, _, _, _ in ret_dict.values(): - self.resource_mgr.release_allocation(task_id, -1) - del self.curr_task_table[task_id] - raise - except Exception: - self.fwk.exception('TM:init_task_pool(): Allocation exception') - raise + except Exception: + self._rollback_tasks([task_id for task_id, _, _, _ in ret_dict.values()]) + raise return ret_dict diff --git a/tests/new/test_task_manager.py b/tests/new/test_task_manager.py index 7a1be344..4bf773ae 100644 --- a/tests/new/test_task_manager.py +++ b/tests/new/test_task_manager.py @@ -1120,3 +1120,114 @@ def test_dvm_task_core_budget( rm.get_allocation('other', 1, 100, False, False, task_cpp=5) assert rm.avail_cores == 4 rm.reporting_file.close() + + +@pytest.fixture +def dvm_pool_manager(tmp_path, monkeypatch): + """Yield a task manager with four synthetic CPU resources. + + :param tmp_path: Directory for the resource report. + :param monkeypatch: Fixture supplying a synthetic DVM URI. + :yields: Task manager using real resource accounting without hardware discovery. + """ + fwk = mock.Mock() + cm = mock.Mock() + cm.fwk_sim_name = 'test' + cm.sim_map = {'test': mock.Mock(sim_root=str(tmp_path))} + settings = {'MPIRUN_VERSION': 'OPENMPI-DVM', 'TOTAL_PROCS': 4} + cm.get_platform_parameter.side_effect = lambda name, **kwargs: settings.get(name) + tm = TaskManager(fwk) + rm = ResourceManager(fwk) + tm.initialize(mock.Mock(), rm, cm) + rm.initialize(mock.Mock(), tm, cm, cmd_nodes=1, cmd_ppn=4) + tm.task_launch_cmd = 'mpirun' + tm.mpicmd = '/usr/bin/prun' + monkeypatch.setenv('IPS_DVM_URI_FILE', '/tmp/test-dvm-uri') + yield tm + rm.reporting_file.close() + + +def _pool_request(**tasks: TaskInit) -> ServiceRequestMessage: + """Build a pool initialization request. + + :param tasks: Named task initialization records. + :returns: Service request containing the supplied tasks. + """ + return ServiceRequestMessage('id', 'id', 'c', 'init_task_pool', tasks) + + +def _pool_task(cpp: int = 1) -> TaskInit: + """Construct a one-rank task using partial-node allocation. + + :param cpp: CPU resources to reserve for the rank. + :returns: Task initialization record. + """ + return TaskInit(1, 'exe', '/dir', 0, cpp, 0, False, False, False, False, (), None) + + +@pytest.mark.parametrize('failure', ['validation', 'command']) +def test_dvm_pool_fatal_error_restores_resources(dvm_pool_manager, failure: str) -> None: + """Roll back a failed submission while preserving a prior task. + + :param dvm_pool_manager: Manager with synthetic DVM resources. + :param failure: Whether validation or command construction fails. + :returns: None. + """ + tm = dvm_pool_manager + rm = tm.resource_mgr + prior = tm.init_task_pool(_pool_request(prior=_pool_task()))['prior'][0] + prior_table = tm.curr_task_table.copy() + prior_ranks = rm._dvm_rank_allocations.copy() + if failure == 'command': + error = RuntimeError('command construction failed') + with mock.patch.object(tm, 'build_launch_cmd', side_effect=[('exe', None), error]): + with pytest.raises(RuntimeError) as caught: + tm.init_task_pool(_pool_request(first=_pool_task(), second=_pool_task())) + assert caught.value is error + else: + with pytest.raises(ValueError, match='nonnegative integer task_cpp'): + tm.init_task_pool(_pool_request(first=_pool_task(), second=_pool_task(-1))) + assert tm.curr_task_table == prior_table + assert rm._dvm_rank_allocations == prior_ranks + assert (rm.avail_cores, rm.alloc_cores, rm.processes) == (3, 1, 1) + assert sum(node.avail_cores for node in rm.nodes.values()) == 3 + # Reuse every recovered core while the unrelated task is still reserved. + recovered = tm.init_task_pool(_pool_request(recovered=_pool_task(3)))['recovered'][0] + assert (rm.avail_cores, rm.alloc_cores, rm.processes) == (0, 4, 2) + for task_id in (prior, recovered): + tm.finish_task(ServiceRequestMessage('id', 'id', 'c', 'finish_task', task_id, 0)) + assert not tm.curr_task_table + assert not rm._dvm_rank_allocations + assert (rm.avail_cores, rm.alloc_cores, rm.processes) == (4, 0, 0) + + +def test_dvm_pool_exhaustion_preserves_successes(dvm_pool_manager) -> None: + """Keep successful tasks when another task must wait for capacity. + + :param dvm_pool_manager: Manager with synthetic DVM resources. + :returns: None. + """ + tm = dvm_pool_manager + result = tm.init_task_pool(_pool_request(first=_pool_task(4), waiting=_pool_task())) + assert set(result) == {'first'} + task_id = result['first'][0] + assert set(tm.curr_task_table) == {task_id} + assert set(tm.resource_mgr._dvm_rank_allocations) == {task_id} + assert tm.resource_mgr.alloc_cores == 4 + tm.finish_task(ServiceRequestMessage('id', 'id', 'c', 'finish_task', task_id, 0)) + assert tm.resource_mgr.avail_cores == 4 + + +def test_task_rollback_attempts_all_releases() -> None: + """Continue cleanup after a release failure, retaining its task record. + + :returns: None. + """ + tm = TaskManager(mock.Mock()) + tm.resource_mgr = mock.Mock() + tm.resource_mgr.release_allocation.side_effect = [RuntimeError('release failed'), True] + tm.curr_task_table = {1: {}, 2: {}} + tm._rollback_tasks([1, 2]) + assert tm.resource_mgr.release_allocation.call_args_list == [mock.call(1, -1), mock.call(2, -1)] + assert set(tm.curr_task_table) == {1} + tm.fwk.exception.assert_called_once() From cb99a025dac44cac480f12405780f886d9ba0e3c Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Thu, 5 Nov 2026 12:36:43 -0500 Subject: [PATCH 7/8] Fixed heading for ensembles.rst --- doc/user_guides/ensembles.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/user_guides/ensembles.rst b/doc/user_guides/ensembles.rst index 2f4d2d84..72917c51 100644 --- a/doc/user_guides/ensembles.rst +++ b/doc/user_guides/ensembles.rst @@ -439,7 +439,7 @@ 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 From 94c6cdbaa9ca2df7ba6d47ec598b07b3bb4007ab Mon Sep 17 00:00:00 2001 From: Mark Coletti Date: Thu, 5 Nov 2026 12:37:39 -0500 Subject: [PATCH 8/8] Fixed cores_per_instance check where variable would never be a bool. Repaired portal ID comment so that it's now adjacent once again to the code to which it refers. --- ipsframework/services.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/ipsframework/services.py b/ipsframework/services.py index 3002ce04..353781a6 100644 --- a/ipsframework/services.py +++ b/ipsframework/services.py @@ -2928,16 +2928,15 @@ def run_ensemble( :returns: a list of dicts mapping created subdirs to simulation names and their parameters """ + cores_per_instance = 1 if cores_per_instance is None \ + else cores_per_instance + + if cores_per_instance < 1: + raise ValueError('cores_per_instance must be a positive integer') + # This should be a unique variable across all ensembles we keep track # of in the portal This ID should only be shared by runs within an # ensemble - cores_per_instance = 1 if cores_per_instance is None else cores_per_instance - if ( - isinstance(cores_per_instance, bool) - or not isinstance(cores_per_instance, int) - or cores_per_instance < 1 - ): - raise ValueError('cores_per_instance must be a positive integer') portal_ensemble_id = str(uuid.uuid4()) use_portal = self._should_use_portal()