Fix issue #250

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
 --------------------------------------------------------------------------------------

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
@@ -43,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.
@@ -1136,11 +1140,19 @@ def _limit_sire_threads(self):
         # visible-devices environment variable for the chosen platform.
         self._gpu_devices = []
 
+        # 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
+
+        # 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.
-            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(
@@ -1174,8 +1186,32 @@ 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)
 
+            # 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
@@ -2688,12 +2724,6 @@ 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):
-                self._save_gcmc_topology(system)
-                self._gcmc_topology_saved = True
-
             if is_final_block:
                 # Assemble and save the final trajectory.
                 if self._config.save_trajectories:
@@ -3084,6 +3114,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 +3155,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

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
@@ -1063,16 +1065,22 @@ 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(
                     checkpoint_executor,
                     "Post-equilibration checkpoint",
@@ -1279,9 +1287,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
 
@@ -1380,7 +1389,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,
@@ -1389,6 +1401,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,
@@ -2357,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
         ----------
@@ -2396,6 +2411,35 @@ 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
+
+        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,
         index,
@@ -2463,7 +2507,9 @@ def _checkpoint(
                 system, speed = committed
 
             # Call the base class checkpoint method to save the system state.
-            with self._lock:
+            # 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,
@@ -2473,8 +2519,8 @@ def _checkpoint(
                     num_energy_rows=num_energy_rows,
                 )
 
-                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

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]
 
@@ -1101,7 +1105,13 @@ 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(),
+        ):
+            if self._is_gcmc_topology_due(is_final_block=False):
+                self._save_gcmc_topology(system)
+
             _, error = self._checkpoint(
                 system,
                 window.index,
@@ -1447,12 +1457,18 @@ 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:
                 raise error
 
+            if self._is_gcmc_topology_due(is_final_block):
+                self._save_gcmc_topology(system)
+
             # Write the checkpoint files.
             _, error = self._checkpoint(
                 system,

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",

tests/runner/test_runner.py

@@ -167,3 +167,45 @@ 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 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()
+    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]))
+
+    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):
+            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()