From 55c57e910e8140fee1ba1bbd2fb555521e7c7563 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 14:51:29 +0100 Subject: [PATCH 1/9] Save the GCMC topologies once per repex checkpoint --- src/somd2/runner/_base.py | 32 ++++++++++++++++++++++++++++---- src/somd2/runner/_repex.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 4 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index df46a30..fe9eefa 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2598,6 +2598,7 @@ def _checkpoint( context=None, gcmc_sampler=None, num_energy_rows=None, + save_gcmc_topology=True, ): """ Save a checkpoint file. @@ -2630,6 +2631,10 @@ def _checkpoint( The number of energy samples collected since the last checkpoint. If None, every window is assumed to have collected a full block. + save_gcmc_topology: bool + Whether to save the GCMC topologies when they are due. Replica + exchange saves them once for every replica instead. + Returns ------- @@ -2688,11 +2693,8 @@ def _checkpoint( f"{self._standard_state_correction:.6f}" ) - # Saved at the first checkpoint of each run, and the final one for the - # final coordinates in the PDB files. - if self._config.gcmc and (is_final_block or not self._gcmc_topology_saved): + if save_gcmc_topology and self._is_gcmc_topology_due(is_final_block): self._save_gcmc_topology(system) - self._gcmc_topology_saved = True if is_final_block: # Assemble and save the final trajectory. @@ -3084,6 +3086,26 @@ def _flush_energy_components(self, index): ) _pq_local.write_table(table, filepath) + def _is_gcmc_topology_due(self, is_final_block): + """ + Whether the GCMC topologies are due to be saved. They are saved at the + first checkpoint of each run, and the final one for the final + coordinates in the PDB files. + + Parameters + ---------- + + is_final_block: bool + Whether this is the final block of the simulation. + + Returns + ------- + + bool + Whether the topologies are due. + """ + return self._config.gcmc and (is_final_block or not self._gcmc_topology_saved) + def _save_gcmc_topology(self, system): """ Save the end-state topologies, including the additional waters used for @@ -3105,6 +3127,8 @@ def _save_gcmc_topology(self, system): _sr.save(mols, topology) _sr.save(mols, topology.replace(".prm7", ".pdb")) + self._gcmc_topology_saved = True + def _save_ghost_residues(self, index, gcmc_sampler): """ Record the current GCMC ghost residue indices for a window. This must diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 2c9999b..ffc1527 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1073,6 +1073,8 @@ def _prepare_replicas( lock = _FileLock(self._lock_file) with lock.acquire(timeout=self._config.timeout.to("seconds")): + self._save_replica_topology(is_final_block=False) + self._map_replicas( checkpoint_executor, "Post-equilibration checkpoint", @@ -1389,6 +1391,8 @@ def _write_cycle_checkpoint( replica_list, ) + self._save_replica_topology(is_final_block) + # Now write the new checkpoint files. self._map_replicas( checkpoint_executor, @@ -2396,6 +2400,31 @@ def _commit_replica(self, index): return index, None + def _save_replica_topology(self, is_final_block): + """ + Save the GCMC topologies from the first replica when they are due. The + replicas share the topology files, so they are saved once, before the + replicas are checkpointed, rather than by each replica. + + Parameters + ---------- + + is_final_block: bool + Whether this is the final block of the simulation. + """ + if not self._is_gcmc_topology_due(is_final_block): + return + + # Commit the replica now if it isn't already, so that _checkpoint() + # re-uses the committed system. + if 0 not in self._committed: + _, error = self._commit_replica(0) + if error is not None: + raise error + + system, _ = self._committed[0] + self._save_gcmc_topology(system) + def _checkpoint( self, index, @@ -2471,6 +2500,7 @@ def _checkpoint( speed, is_final_block=is_final_block, num_energy_rows=num_energy_rows, + save_gcmc_topology=False, ) if error is not None: From bf7ff9b8be6e38786a9fb26aae428d3300b7acd7 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 14:58:20 +0100 Subject: [PATCH 2/9] Remove the lock that serialised repex replica checkpoints --- src/somd2/runner/_repex.py | 23 +++++++++++------------ tests/runner/test_repex.py | 2 ++ 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index ffc1527..a768773 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2492,19 +2492,18 @@ def _checkpoint( system, speed = committed # Call the base class checkpoint method to save the system state. - with self._lock: - index, error = super()._checkpoint( - system, - index, - block, - speed, - is_final_block=is_final_block, - num_energy_rows=num_energy_rows, - save_gcmc_topology=False, - ) + index, error = super()._checkpoint( + system, + index, + block, + speed, + is_final_block=is_final_block, + num_energy_rows=num_energy_rows, + save_gcmc_topology=False, + ) - if error is not None: - return index, error + if error is not None: + return index, error # Delete the trajectory frames from the Sire system within the # dynamics object. A replica committed earlier, with shared diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 46bb4ec..fdebfa6 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -893,6 +893,8 @@ def test_repex_gcmc_bounded_contexts(ethane_methanol, max_contexts): "num_lambda": num_lambda, "replica_exchange": True, "max_contexts": max_contexts, + # Checkpoint replicas concurrently. + "oversubscription_factor": 2, "gcmc": True, "gcmc_selection": "resname LIG", "gcmc_frequency": "4fs", From 7a0c04b721b8a626470f4a93c86241256b9e1688 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:07:11 +0100 Subject: [PATCH 3/9] Apply the Sire thread limit in regular runner worker processes --- src/somd2/runner/_base.py | 4 ++++ src/somd2/runner/_runner.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index fe9eefa..117ccd6 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1136,6 +1136,9 @@ def _limit_sire_threads(self): # visible-devices environment variable for the chosen platform. self._gpu_devices = [] + # The maximum number of Sire CPU threads, or None if it isn't limited. + self._sire_threads = None + # Limit the number of CPU threads available to Sire when running in parallel. if self._is_gpu: # First get the total number of threads that are available to Sire. @@ -1174,6 +1177,7 @@ def _limit_sire_threads(self): _logger.info(f"Setting maximum Sire CPU threads to {sire_threads}") # Update the maximum number of threads. + self._sire_threads = sire_threads _sr.legacy.Base.set_max_num_threads(sire_threads) def _set_tuned_pme_parameters(self): diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 03ccf9c..13bcb37 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -420,6 +420,10 @@ def run_window(self, index): self._config._reset_logger(_logger) _sr.base.ProgressBar.set_silent() + # The thread limit set in the parent process doesn't carry over. + if self._sire_threads is not None: + _sr.legacy.Base.set_max_num_threads(self._sire_threads) + # Get the lambda value. lambda_value = self._lambda_values[index] From c48f3f79453914b5235ce445a9630d831552b590 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:13:40 +0100 Subject: [PATCH 4/9] Let Sire use every CPU thread while writing checkpoints --- src/somd2/runner/_base.py | 28 +++++++++++++++++++++++++++- src/somd2/runner/_repex.py | 32 ++++++++++++++++++++------------ src/somd2/runner/_runner.py | 10 ++++++++-- 3 files changed, 55 insertions(+), 15 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 117ccd6..c1ee998 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -21,6 +21,7 @@ __all__ = ["RunnerBase"] +from contextlib import contextmanager as _contextmanager from glob import glob as _glob from math import isclose as _isclose from pathlib import Path as _Path @@ -1136,8 +1137,10 @@ def _limit_sire_threads(self): # visible-devices environment variable for the chosen platform. self._gpu_devices = [] - # The maximum number of Sire CPU threads, or None if it isn't limited. + # The maximum number of Sire CPU threads during dynamics and while + # checkpointing, or None if it isn't limited. self._sire_threads = None + self._checkpoint_sire_threads = None # Limit the number of CPU threads available to Sire when running in parallel. if self._is_gpu: @@ -1180,6 +1183,29 @@ def _limit_sire_threads(self): self._sire_threads = sire_threads _sr.legacy.Base.set_max_num_threads(sire_threads) + # Checkpoints never overlap one another, so they can use every + # thread unless the user has set a limit. + if self._config.max_sire_threads is None: + self._checkpoint_sire_threads = total_threads + else: + self._checkpoint_sire_threads = sire_threads + + @_contextmanager + def _checkpoint_threads(self): + """ + Raise the maximum number of Sire CPU threads while checkpoint files + are written, restoring the limit for dynamics afterwards. + """ + if self._checkpoint_sire_threads == self._sire_threads: + yield + return + + _sr.legacy.Base.set_max_num_threads(self._checkpoint_sire_threads) + try: + yield + finally: + _sr.legacy.Base.set_max_num_threads(self._sire_threads) + def _set_tuned_pme_parameters(self): """ Tune the PME parameters for a new run, saving them so that restarts diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index a768773..446e279 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1063,16 +1063,20 @@ def _prepare_replicas( # resident so that every checkpoint file can still be written # under a single lock below. if self._is_cached: - for batch in self._replica_passes(0): - self._map_replicas( - checkpoint_executor, - "Post-equilibration commit", - self._load_and_commit, - batch, - ) + with self._checkpoint_threads(): + for batch in self._replica_passes(0): + self._map_replicas( + checkpoint_executor, + "Post-equilibration commit", + self._load_and_commit, + batch, + ) lock = _FileLock(self._lock_file) - with lock.acquire(timeout=self._config.timeout.to("seconds")): + with ( + lock.acquire(timeout=self._config.timeout.to("seconds")), + self._checkpoint_threads(), + ): self._save_replica_topology(is_final_block=False) self._map_replicas( @@ -1281,9 +1285,10 @@ def _propagate_cycle( # everything is still resident once the cycle finishes, so # committing is left to _checkpoint(). if do_checkpoint and self._is_cached: - self._map_replicas( - checkpoint_executor, "Commit", self._commit_replica, batch - ) + with self._checkpoint_threads(): + self._map_replicas( + checkpoint_executor, "Commit", self._commit_replica, batch + ) return results @@ -1382,7 +1387,10 @@ def _write_cycle_checkpoint( replica_list = list(range(self._config.num_lambda)) lock = _FileLock(self._lock_file) - with lock.acquire(timeout=self._config.timeout.to("seconds")): + with ( + lock.acquire(timeout=self._config.timeout.to("seconds")), + self._checkpoint_threads(), + ): # First backup existing checkpoint files. self._map_replicas( checkpoint_executor, diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 13bcb37..e468310 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1105,7 +1105,10 @@ def _write_post_equilibration_checkpoint(self, system, window): The window being simulated. """ lock = _FileLock(self._lock_file) - with lock.acquire(timeout=self._config.timeout.to("seconds")): + with ( + lock.acquire(timeout=self._config.timeout.to("seconds")), + self._checkpoint_threads(), + ): _, error = self._checkpoint( system, window.index, @@ -1451,7 +1454,10 @@ def _write_checkpoint(self, window, system, block, speed, is_final_block): num_energy_samples = self._num_energy_samples(window) lock = _FileLock(self._lock_file) - with lock.acquire(timeout=self._config.timeout.to("seconds")): + with ( + lock.acquire(timeout=self._config.timeout.to("seconds")), + self._checkpoint_threads(), + ): # Backup any existing checkpoint files. _, error = self._backup_checkpoint(window.index) if error is not None: From 35f06a5f0ca9630018bcbd84284f4f896b8ef7d9 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:30:42 +0100 Subject: [PATCH 5/9] Work out the Sire thread limit from every available thread --- src/somd2/runner/_base.py | 4 +++- tests/runner/test_runner.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c1ee998..bee3853 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1144,7 +1144,9 @@ def _limit_sire_threads(self): # Limit the number of CPU threads available to Sire when running in parallel. if self._is_gpu: - # First get the total number of threads that are available to Sire. + # First get the total number of threads that are available to Sire, + # resetting any limit set by an earlier runner in this process. + _sr.set_default_num_threads() total_threads = _sr.legacy.Base.get_max_num_threads() # Get the available GPU devices. Subclasses re-use this list rather diff --git a/tests/runner/test_runner.py b/tests/runner/test_runner.py index 6af940e..ca10a52 100644 --- a/tests/runner/test_runner.py +++ b/tests/runner/test_runner.py @@ -167,3 +167,32 @@ def test_energy_components_after_restart(ethane_methanol): # The times are in ns. times = list(np.round(components["time"] * 1e6, 6)) assert times == [4, 8, 12, 16, 20] + + +def test_sire_thread_limit_with_repeated_runners(monkeypatch): + """ + Validate that the Sire thread limit is worked out from every available + thread, not from the limit left by an earlier runner in the same process. + """ + import sire as sr + + from somd2.runner._base import RunnerBase + + sr.set_default_num_threads() + total_threads = sr.get_max_num_threads() + if total_threads < 2: + pytest.skip("Needs more than one CPU thread.") + + monkeypatch.setattr(RunnerBase, "_get_gpu_devices", staticmethod(lambda *args: [0])) + + runner = object.__new__(RunnerBase) + runner._is_gpu = True + runner._config = Config(max_gpus=1, oversubscription_factor=2) + + try: + for _ in range(2): + runner._limit_sire_threads() + assert runner._sire_threads == max(1, total_threads // 2) + assert runner._checkpoint_sire_threads == total_threads + finally: + sr.set_default_num_threads() From c9ab5ca96af38db268bb83e81191e319ec689f9e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:35:37 +0100 Subject: [PATCH 6/9] Reassemble repex trajectories one replica at a time --- src/somd2/runner/_repex.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 446e279..389f897 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -21,6 +21,7 @@ __all__ = ["RepexRunner"] +from contextlib import nullcontext as _nullcontext from filelock import FileLock as _FileLock from numba import njit as _njit from shutil import copyfile as _copyfile @@ -468,7 +469,8 @@ def __init__(self, system, config): from threading import Lock - # Create a lock to guard the dynamics cache. + # Create a lock to guard the committed systems, which also makes the + # replicas take turns at the final checkpoint. self._lock = Lock() # Systems committed while a replica was resident in its slot, awaiting @@ -2500,15 +2502,18 @@ def _checkpoint( system, speed = committed # Call the base class checkpoint method to save the system state. - index, error = super()._checkpoint( - system, - index, - block, - speed, - is_final_block=is_final_block, - num_energy_rows=num_energy_rows, - save_gcmc_topology=False, - ) + # The final block reassembles each trajectory, which loads the + # topology, so replicas take turns to bound the memory used. + with self._lock if is_final_block else _nullcontext(): + index, error = super()._checkpoint( + system, + index, + block, + speed, + is_final_block=is_final_block, + num_energy_rows=num_energy_rows, + save_gcmc_topology=False, + ) if error is not None: return index, error From c5c1dd97e97c538a417111717c47d6371085a3c5 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:42:01 +0100 Subject: [PATCH 7/9] Save GCMC topologies from the runners rather than the base checkpoint --- src/somd2/runner/_base.py | 8 -------- src/somd2/runner/_repex.py | 34 +++++++++++++++++++--------------- src/somd2/runner/_runner.py | 6 ++++++ 3 files changed, 25 insertions(+), 23 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index bee3853..51f2928 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2630,7 +2630,6 @@ def _checkpoint( context=None, gcmc_sampler=None, num_energy_rows=None, - save_gcmc_topology=True, ): """ Save a checkpoint file. @@ -2663,10 +2662,6 @@ def _checkpoint( The number of energy samples collected since the last checkpoint. If None, every window is assumed to have collected a full block. - save_gcmc_topology: bool - Whether to save the GCMC topologies when they are due. Replica - exchange saves them once for every replica instead. - Returns ------- @@ -2725,9 +2720,6 @@ def _checkpoint( f"{self._standard_state_correction:.6f}" ) - if save_gcmc_topology and self._is_gcmc_topology_due(is_final_block): - self._save_gcmc_topology(system) - if is_final_block: # Assemble and save the final trajectory. if self._config.save_trajectories: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 389f897..2d2ff7c 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2371,11 +2371,12 @@ def _commit_replica(self, index): Commit a replica while it is still resident in its slot, holding the result until the checkpoint files are written at the end of the cycle. - This is only used when contexts are shared between replicas. A slot - holds the positions and trajectory frames of the replica it last - hosted, so a replica has to be committed before its slot is handed on. - Deferring only the file writes keeps every checkpoint file within a - single lock, so that an external reader never sees a partial set. + When contexts are shared between replicas, a slot holds the positions + and trajectory frames of the replica it last hosted, so a replica has + to be committed before its slot is handed on. Deferring only the file + writes keeps every checkpoint file within a single lock, so that an + external reader never sees a partial set. It is also used to commit + the first replica early, when the GCMC topologies are saved from it. Parameters ---------- @@ -2425,15 +2426,19 @@ def _save_replica_topology(self, is_final_block): if not self._is_gcmc_topology_due(is_final_block): return - # Commit the replica now if it isn't already, so that _checkpoint() - # re-uses the committed system. - if 0 not in self._committed: - _, error = self._commit_replica(0) - if error is not None: - raise error - - system, _ = self._committed[0] - self._save_gcmc_topology(system) + try: + # Commit the replica now if it isn't already, so that _checkpoint() + # re-uses the committed system. + if 0 not in self._committed: + _, error = self._commit_replica(0) + if error is not None: + raise error + + system, _ = self._committed[0] + self._save_gcmc_topology(system) + except Exception as e: + _logger.error(f"Saving the GCMC topologies failed: {e}") + raise def _checkpoint( self, @@ -2512,7 +2517,6 @@ def _checkpoint( speed, is_final_block=is_final_block, num_energy_rows=num_energy_rows, - save_gcmc_topology=False, ) if error is not None: diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index e468310..594a6cc 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1109,6 +1109,9 @@ def _write_post_equilibration_checkpoint(self, system, window): lock.acquire(timeout=self._config.timeout.to("seconds")), self._checkpoint_threads(), ): + if self._is_gcmc_topology_due(is_final_block=False): + self._save_gcmc_topology(system) + _, error = self._checkpoint( system, window.index, @@ -1463,6 +1466,9 @@ def _write_checkpoint(self, window, system, block, speed, is_final_block): if error is not None: raise error + if self._is_gcmc_topology_due(is_final_block): + self._save_gcmc_topology(system) + # Write the checkpoint files. _, error = self._checkpoint( system, From ce5d6499d01a297efb12cd17279440be357552f1 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:42:40 +0100 Subject: [PATCH 8/9] Update CHANGELOG. --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a92960c..38c2f1f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ Changelog * Require `frame_frequency` to be a multiple of `gcmc_frequency` for GCMC simulations, so that every frame has a matching ghost water record. Previously some combinations silently saved frames without one [#247](https://github.com/OpenBioSim/somd2/pull/247). * Fixed the regular runner's final block when the `runtime` isn't a multiple of the `checkpoint_frequency`, which ran a single integration step rather than the remaining time, without GCMC moves. Each checkpoint now appends exactly the energies sampled since the last one, so energies are no longer duplicated, or lost on restart when `checkpoint_frequency` is zero. Blocks that aren't a whole number of GCMC or terminal flip intervals no longer overrun, and energy components are time stamped correctly after a restart [#248](https://github.com/OpenBioSim/somd2/pull/248). * Require the `runtime`, `checkpoint_frequency`, `energy_frequency`, `frame_frequency`, `gcmc_frequency` and `terminal_flip_frequency` to be multiples of the `timestep` for the regular runner, which runs in whole integration steps. GCMC and terminal flip moves now keep their frequency across checkpoints [#248](https://github.com/OpenBioSim/somd2/pull/248). +* Speed up checkpointing, which could take minutes for large systems. Sire now uses every CPU thread while checkpoint files are written, unless `max_sire_threads` is set, and replica exchange checkpoints replicas in parallel, writing the shared GCMC topologies once rather than once per replica. The regular runner's worker processes now also apply the Sire thread limit, which previously only took effect in the parent process [#251](https://github.com/OpenBioSim/somd2/pull/251). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- From 3634b27b8ab5a483e77c2351df1808a03de25240 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 19 Oct 2026 15:47:42 +0100 Subject: [PATCH 9/9] Keep the Sire thread limit set before the first runner --- src/somd2/runner/_base.py | 14 +++++++++----- tests/runner/test_runner.py | 35 ++++++++++++++++++++++++----------- 2 files changed, 33 insertions(+), 16 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 51f2928..c4937f9 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -44,6 +44,9 @@ class RunnerBase: Base class for the SOMD2 simulation runner. """ + # The number of threads available to Sire before any runner limited it. + _total_sire_threads = None + def __init__(self, system, config): """ Constructor. @@ -1142,13 +1145,14 @@ def _limit_sire_threads(self): self._sire_threads = None self._checkpoint_sire_threads = None + # Undo any limit set by an earlier runner in this process. + if RunnerBase._total_sire_threads is None: + RunnerBase._total_sire_threads = _sr.legacy.Base.get_max_num_threads() + total_threads = RunnerBase._total_sire_threads + _sr.legacy.Base.set_max_num_threads(total_threads) + # Limit the number of CPU threads available to Sire when running in parallel. if self._is_gpu: - # First get the total number of threads that are available to Sire, - # resetting any limit set by an earlier runner in this process. - _sr.set_default_num_threads() - total_threads = _sr.legacy.Base.get_max_num_threads() - # Get the available GPU devices. Subclasses re-use this list rather # than querying the environment again. self._gpu_devices = self._get_gpu_devices( diff --git a/tests/runner/test_runner.py b/tests/runner/test_runner.py index ca10a52..4ef47f8 100644 --- a/tests/runner/test_runner.py +++ b/tests/runner/test_runner.py @@ -171,28 +171,41 @@ def test_energy_components_after_restart(ethane_methanol): def test_sire_thread_limit_with_repeated_runners(monkeypatch): """ - Validate that the Sire thread limit is worked out from every available - thread, not from the limit left by an earlier runner in the same process. + Validate that the Sire thread limit is worked out from the threads that + were available before the first runner, keeping a limit set by the user, + rather than from the limit left by an earlier runner in the same process. """ import sire as sr from somd2.runner._base import RunnerBase sr.set_default_num_threads() - total_threads = sr.get_max_num_threads() - if total_threads < 2: - pytest.skip("Needs more than one CPU thread.") + if sr.get_max_num_threads() < 4: + pytest.skip("Needs at least four CPU threads.") + # A limit set by the user before the first runner is created. + user_threads = sr.get_max_num_threads() // 2 + + monkeypatch.setattr(RunnerBase, "_total_sire_threads", None) monkeypatch.setattr(RunnerBase, "_get_gpu_devices", staticmethod(lambda *args: [0])) - runner = object.__new__(RunnerBase) - runner._is_gpu = True - runner._config = Config(max_gpus=1, oversubscription_factor=2) + gpu_runner = object.__new__(RunnerBase) + gpu_runner._is_gpu = True + gpu_runner._config = Config(max_gpus=1, oversubscription_factor=2) + + cpu_runner = object.__new__(RunnerBase) + cpu_runner._is_gpu = False try: + sr.set_max_num_threads(user_threads) + for _ in range(2): - runner._limit_sire_threads() - assert runner._sire_threads == max(1, total_threads // 2) - assert runner._checkpoint_sire_threads == total_threads + gpu_runner._limit_sire_threads() + assert gpu_runner._sire_threads == user_threads // 2 + assert gpu_runner._checkpoint_sire_threads == user_threads + + # A CPU runner isn't left with the limit of the GPU runner before it. + cpu_runner._limit_sire_threads() + assert sr.get_max_num_threads() == user_threads finally: sr.set_default_num_threads()