diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 99a1286..47d5ff3 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 pathlib import Path as _Path from shutil import copyfile as _copyfile @@ -42,6 +43,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. @@ -1008,11 +1012,19 @@ def __init__(self, system, config): # 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( @@ -1046,8 +1058,32 @@ def __init__(self, system, config): _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) + @property def _is_abfe_bound(self): """ @@ -2468,12 +2504,6 @@ def new_rows(df): 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: @@ -2825,6 +2855,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 @@ -2846,6 +2896,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 828739b..339d524 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 @@ -1524,7 +1525,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 @@ -1632,6 +1634,61 @@ def _safe_batches(self, num_workers, cycle=0): for i in range(ceil(len(batch) / num_workers)): yield batch[i * num_workers : (i + 1) * num_workers] + def _map_replicas(self, executor, step, func, replicas, *args): + """ + Run a step for each replica in a thread pool, letting every call + finish before raising the first error, so that files written under a + lock are all written before it is released. An interrupt exits + straight away, without waiting. + + Two replicas that share a slot must never be processed at the same + time, so the replicas must come from a single pass, unless the step + doesn't use the slot's context. + + Parameters + ---------- + + executor: concurrent.futures.ThreadPoolExecutor + The thread pool to run the step in. + + step: str + The name of the step, for logging. + + func: callable + The step, which returns the index of the replica and the exception + raised, if any. An exception it raises is treated the same way. + + replicas: list of int + The replicas to run the step for. + + args: iterable + The remaining arguments, an iterable of one per replica for each. + """ + from concurrent.futures import wait + + try: + futures = [ + executor.submit(func, *arguments) for arguments in zip(replicas, *args) + ] + wait(futures) + except KeyboardInterrupt: + executor.shutdown(wait=False, cancel_futures=True) + _logger.error(f"{step} cancelled. Exiting.") + _sys.exit(1) + + outcomes = [] + for replica, future in zip(replicas, futures): + error = future.exception() + outcomes.append((replica, error) if error else future.result()) + + for index, error in outcomes: + if error is not None: + _logger.error( + f"{step} failed for {_lam_sym} = " + f"{self._lambda_values[index]:.5f}: {error}" + ) + raise error + def _set_num_slots(self): """ Work out the number of OpenMM contexts (slots) to create, validating @@ -1860,51 +1917,35 @@ def run(self): # resident so that every checkpoint file can still be written # under a single lock below. if self._is_cached: - for batch in self._safe_batches(num_checkpoint_workers): - try: - for index, error in checkpoint_executor.map( + with self._checkpoint_threads(): + for batch in self._safe_batches(num_checkpoint_workers): + self._map_replicas( + checkpoint_executor, + "Post-equilibration commit", self._load_and_commit, batch, - ): - if error is not None: - msg = ( - f"Post-equilibration commit failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown(wait=False, cancel_futures=True) - _logger.error("Post-equilibration commit cancelled. Exiting.") - _sys.exit(1) + ) 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) + for j in range(num_checkpoint_batches): replicas = replica_list[ j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - try: - for index, error in checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(-1), - repeat(cycles), - ): - if error is not None: - msg = ( - f"Post-equilibration checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}:\n{error}" - ) - _logger.error(msg) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown(wait=False, cancel_futures=True) - _logger.error( - "Post-equilibration checkpoint cancelled. Exiting." - ) - _sys.exit(1) + self._map_replicas( + checkpoint_executor, + "Post-equilibration checkpoint", + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(-1), + repeat(cycles), + ) # Current block number. block = self._start_block @@ -2037,28 +2078,18 @@ def run(self): # context per replica everything is still resident once the # cycle finishes, so committing is left to _checkpoint(). if do_checkpoint and self._is_cached: - for j in range(ceil(len(batch) / num_checkpoint_workers)): - replicas = batch[ - j * num_checkpoint_workers : (j + 1) - * num_checkpoint_workers - ] - try: - for index, error in checkpoint_executor.map( + with self._checkpoint_threads(): + for j in range(ceil(len(batch) / num_checkpoint_workers)): + replicas = batch[ + j * num_checkpoint_workers : (j + 1) + * num_checkpoint_workers + ] + self._map_replicas( + checkpoint_executor, + "Commit", self._commit_replica, replicas, - ): - if error: - _logger.error( - f"Commit failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True ) - _logger.error("Commit cancelled. Exiting.") - _sys.exit(1) # Assemble an energy matrix from the results. _logger.info("Assembling energy matrix") @@ -2090,7 +2121,10 @@ def run(self): # Acquire the file lock to ensure that the checkpoint files are # in a consistent state if read by another process. - 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. for j in range(num_checkpoint_batches): # Get the indices of the replicas in this batch. @@ -2098,23 +2132,14 @@ def run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - try: - for index, error in checkpoint_executor.map( - self._backup_checkpoint, - replicas, - ): - if error: - _logger.error( - f"Backup failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Backup cancelled. Exiting.") - _sys.exit(1) + self._map_replicas( + checkpoint_executor, + "Backup", + self._backup_checkpoint, + replicas, + ) + + self._save_replica_topology(i == cycles - 1) # Now write the new checkpoint files. for j in range(num_checkpoint_batches): @@ -2123,27 +2148,16 @@ def run(self): j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers ] - try: - for index, error in checkpoint_executor.map( - self._checkpoint, - replicas, - repeat(self._lambda_values), - repeat(block), - repeat(num_blocks + int(rem > 0)), - repeat(i == cycles - 1), - ): - if error: - _logger.error( - f"Checkpoint failed for {_lam_sym} = " - f"{self._lambda_values[index]:.5f}: {error}" - ) - raise error - except KeyboardInterrupt: - checkpoint_executor.shutdown( - wait=False, cancel_futures=True - ) - _logger.error("Checkpoint cancelled. Exiting.") - _sys.exit(1) + self._map_replicas( + checkpoint_executor, + "Checkpoint", + self._checkpoint, + replicas, + repeat(self._lambda_values), + repeat(block), + repeat(num_blocks + int(rem > 0)), + repeat(i == cycles - 1), + ) self._save_repex_state(final=not is_checkpoint) @@ -2740,11 +2754,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 ---------- @@ -2779,6 +2794,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, lambdas, block, num_blocks, is_final_block=False): """ Checkpoint the simulation. @@ -2834,7 +2878,9 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False): 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, block, speed, is_final_block=is_final_block ) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index d094086..87f634b 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -419,6 +419,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] @@ -935,7 +939,13 @@ def generate_lam_vals(lambda_base, increment=0.001): # after an early production crash doesn't need to re-equilibrate. if is_equilibrated: 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, index, @@ -1180,13 +1190,19 @@ def generate_lam_vals(lambda_base, increment=0.001): # Acquire the file lock to ensure that the checkpoint files are # in a consistent state if read by another process. - 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. index, error = self._backup_checkpoint(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. index, error = self._checkpoint( system, @@ -1319,13 +1335,19 @@ def generate_lam_vals(lambda_base, increment=0.001): # Acquire the file lock to ensure that the checkpoint files are # in a consistent state if read by another process. - 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. index, error = self._backup_checkpoint(index) if error is not None: raise error + if self._is_gcmc_topology_due(is_final_block=True): + self._save_gcmc_topology(system) + # Write the checkpoint files. index, error = self._checkpoint( system, @@ -1535,7 +1557,10 @@ def generate_lam_vals(lambda_base, increment=0.001): # Acquire the file lock to ensure that the checkpoint files are # in a consistent state if read by another process. - 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. index, error = self._backup_checkpoint(index) @@ -1544,6 +1569,9 @@ def generate_lam_vals(lambda_base, increment=0.001): _logger.error(msg) raise RuntimeError(msg) + if self._is_gcmc_topology_due(is_final_block=True): + self._save_gcmc_topology(system) + # Write the checkpoint files. index, error = self._checkpoint( system, diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index c9232d6..ef8e578 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -499,6 +499,58 @@ def acquire(self, *args, **kwargs): assert len(acquisitions) == 2 +@pytest.mark.skipif(not has_cuda, reason="CUDA not available.") +def test_repex_checkpoint_failure_keeps_lock(ethane_methanol, monkeypatch): + """ + Validate that when one replica's checkpoint fails, the others finish + writing while the file lock is still held, before the error is raised. + """ + import time + + from filelock import FileLock, Timeout + + with tempfile.TemporaryDirectory() as tmpdir: + config = { + "runtime": "4fs", + "restart": False, + "output_directory": tmpdir, + "energy_frequency": "4fs", + "checkpoint_frequency": "4fs", + "frame_frequency": "4fs", + "platform": "cuda", + "max_threads": 1, + "num_lambda": 4, + "replica_exchange": True, + # Enough workers to checkpoint every replica at once. + "oversubscription_factor": 4, + "num_checkpoint_workers": 4, + } + + runner = RepexRunner(ethane_methanol, Config(**config)) + + held = [] + + def checkpoint(index, *args, **kwargs): + if index == 0: + return index, RuntimeError("checkpoint failed") + time.sleep(0.5) + lock = FileLock(runner._lock_file) + try: + lock.acquire(timeout=0) + lock.release() + held.append(False) + except Timeout: + held.append(True) + return index, None + + monkeypatch.setattr(runner, "_checkpoint", checkpoint) + + with pytest.raises(RuntimeError, match="checkpoint failed"): + runner.run() + + assert held == [True, True, True] + + @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") @pytest.mark.parametrize( "runtime, checkpoint_frequency, expected", @@ -566,6 +618,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",