diff --git a/CHANGELOG.md b/CHANGELOG.md index 38c2f1f..83783cf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ Changelog * 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). +* Reduce the setup time of replica exchange simulations that re-use OpenMM contexts across lambda values. Minimisation and equilibration now create their dynamics objects once per context, rather than once per replica [#255](https://github.com/OpenBioSim/somd2/pull/255). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c4937f9..ee39ba0 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -1060,11 +1060,6 @@ def _set_dynamics_kwargs(self): self._config.opencl_platform_index ) - # Set the initial constraint values. This can be modified - # by minimisation and equilibration option. - self._initial_constraint = self._config.constraint - self._initial_perturbable_constraint = self._config.perturbable_constraint - # Common kwargs shared by both dynamics and GCMC sampling. self._common_kwargs = { "cutoff": self._config.cutoff, diff --git a/src/somd2/runner/_dynamics_cache.py b/src/somd2/runner/_dynamics_cache.py index 5299c30..787e9e3 100644 --- a/src/somd2/runner/_dynamics_cache.py +++ b/src/somd2/runner/_dynamics_cache.py @@ -623,11 +623,8 @@ def _create_dynamics( f"Created dynamics object for lambda {lam:.5f} on device {device}" ) - # Leave the slot marked as holding no replica, so that the first - # call to load_replica() does the full setup (lambda value, GCMC - # parameters, ghost file and sampling statistics) rather than - # assuming the seeding replica is already fully installed. - self._slot_replica[i] = None + # The first load sets up the seeding replica in full. + self.release(i) # Give each replica its own energy trajectory. These are seeded from a # slot's own so that the "ensemble" property is carried over. A slot @@ -744,6 +741,21 @@ def set(self, slot, dynamics): """ self._dynamics[slot] = dynamics + def release(self, slot): + """ + Mark a slot as holding no replica, so that the next replica loaded into + it has its state pushed and its lambda value, GCMC water state, ghost + file and energy trajectory set, e.g. after the slot's dynamics object + is replaced or a failed step left its context in a bad state. + + Parameters + ---------- + + slot: int + The index of the slot. + """ + self._slot_replica[slot] = None + def delete(self, slot): """ Delete the dynamics object for a given slot. diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 2d2ff7c..0de187e 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -261,9 +261,22 @@ def __init__(self, system, config): # Update the timestep. timestep = self._config.equilibration_timestep - # Update the initial constraint values. - self._initial_constraint = constraint - self._initial_perturbable_constraint = perturbable_constraint + # Equilibration and production constraints, and whether the initial + # contexts must be rebuilt to use them. + production_constraints = ( + self._config.constraint, + self._config.perturbable_constraint, + ) + if self._config.equilibration_constraints: + self._equilibration_constraints = production_constraints + else: + self._equilibration_constraints = ("none", "none") + + initial_constraints = (constraint, perturbable_constraint) + self._is_equilibration_rebuild = ( + initial_constraints != self._equilibration_constraints + ) + self._is_production_rebuild = initial_constraints != production_constraints # Copy the dynamics keyword arguments. dynamics_kwargs = self._dynamics_kwargs.copy() @@ -1022,19 +1035,10 @@ def _prepare_replicas( # Minimise at each lambda value. Don't minimise on restart. if self._config.minimise and not self._is_restart: - for batch in self._dynamics_batches(replica_devices): + for batch in self._setup_batches(replica_devices): try: - for success, index, e in dynamics_executor.map( - self._minimise, - batch, - ): - if not success: - msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" - if self._config.minimisation_errors: - _logger.error(msg) - raise e - else: - _logger.warning(msg) + for results in dynamics_executor.map(self._minimise_slot, batch): + self._report_minimisation(results) except KeyboardInterrupt: dynamics_executor.shutdown(wait=False, cancel_futures=True) _logger.error("Minimisation cancelled. Exiting.") @@ -1042,17 +1046,15 @@ def _prepare_replicas( # Equilibrate the system. if self._is_equilibration and not self._is_restart: - for batch in self._dynamics_batches(replica_devices): + for batch in self._setup_batches(replica_devices): try: - for success, index, e in dynamics_executor.map( - self._equilibrate, - batch, - ): - if not success: - _logger.error( - f"Equilibration failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" - ) - raise e + for results in dynamics_executor.map(self._equilibrate_slot, batch): + for success, index, e in results: + if not success: + _logger.error( + f"Equilibration failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + ) + raise e except KeyboardInterrupt: dynamics_executor.shutdown(wait=False, cancel_futures=True) _logger.error("Equilibration cancelled. Exiting.") @@ -1671,7 +1673,7 @@ def _minimise(self, index): slot = self._dynamics_cache.slot_for(index) dynamics, gcmc_sampler = self._dynamics_cache.get(slot) - if gcmc_sampler is not None and not self._is_restart: + if gcmc_sampler is not None: gcmc_sampler.push() try: _logger.info( @@ -1685,61 +1687,126 @@ def _minimise(self, index): # Minimise. dynamics.minimise(timeout=self._config.timeout) - # If we're not equilibrating and the production constraints will change, - # then we need to rebuild the context. - if not self._is_equilibration: - constraints_changed = ( - self._initial_constraint != self._config.constraint - ) or ( - self._initial_perturbable_constraint - != self._config.perturbable_constraint + # Save the minimised state back out of the slot. + self._dynamics_cache.store_replica(index) + + except Exception as e: + # The next load restores the replica's stored state over whatever + # the failure left in the context. + self._dynamics_cache.release(self._dynamics_cache.slot_for(index)) + return False, index, e + + return True, index, None + + def _report_minimisation(self, results): + """ + Log any failed minimisations, raising the first error if they are + treated as errors. + + Parameters + ---------- + + results: list of tuple + The success, index and exception for each replica. + """ + for success, index, e in results: + if not success: + msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + if self._config.minimisation_errors: + _logger.error(msg) + raise e + else: + _logger.warning(msg) + + def _minimise_slot(self, replicas): + """ + Minimise the replicas hosted by a slot, one after another. If there is + no equilibration and the production constraints differ from those used + for minimisation, the slot is then given a production dynamics object. + + Parameters + ---------- + + replicas: list of int + The replicas hosted by the slot. + + Returns + ------- + + list of tuple + The success, index and exception for each replica. + """ + results = [self._minimise(index) for index in replicas] + + # There's no need to rebuild if a failed minimisation stops the run. + is_failed = not all(success for success, _, _ in results) + if is_failed and self._config.minimisation_errors: + return results + + if not self._is_equilibration and self._is_production_rebuild: + # Load a valid state to build from, in case the last minimisation + # failed and released the slot. + index = replicas[-1] + try: + self._dynamics_cache.load_replica(index) + self._rebuild_dynamics(index) + except Exception as e: + self._report_minimisation(results) + _logger.error( + f"Could not create the production dynamics object for " + f"{_lam_sym} = {self._lambda_values[index]:.5f}: {e}" ) + raise - if constraints_changed: - # Commit the current system. - system = dynamics.commit() + return results - # Delete the dynamics object. - self._dynamics_cache.delete(slot) + def _rebuild_dynamics(self, index, **overrides): + """ + Replace the dynamics object of a replica's slot with a new one, built + from the slot's current state, e.g. to change the constraints or + timestep. Afterwards the slot holds no replica. - # Work out the device index. - device = self._dynamics_cache.device_for(slot) + Parameters + ---------- - # Copy the dynamics keyword arguments. - dynamics_kwargs = self._dynamics_kwargs.copy() + index: int + The index of the replica. - # Overload the device and lambda value. - dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._build_lambda(index) - dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] + overrides: dict + Dynamics keyword arguments that replace those used for production. + """ + slot = self._dynamics_cache.slot_for(index) + dynamics, gcmc_sampler = self._dynamics_cache.get(slot) - # Create the production dynamics object. - dynamics = system.dynamics(**dynamics_kwargs) + # Commit the current system. + system = dynamics.commit() - # Reset the GCMC water state. The dynamics object is created from - # the original Sire system, so the water state in the context does - # not match the current GCMC water state. - if gcmc_sampler is not None: - self._reset_gcmc_sampler(gcmc_sampler, dynamics) + # Delete the current dynamics object. + self._dynamics_cache.delete(slot) - # Set the new dynamics object. The rebuilt object has a - # fresh, empty energy trajectory, but no energies are - # recorded during minimisation or equilibration, and - # load_replica() re-attaches the replica's trajectory - # before the first production block. - self._dynamics_cache.set(slot, dynamics) + # Copy the dynamics keyword arguments, overloading the device and lambda value. + dynamics_kwargs = self._dynamics_kwargs.copy() + dynamics_kwargs["device"] = self._dynamics_cache.device_for(slot) + dynamics_kwargs["lambda_value"] = self._build_lambda(index) + dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] + dynamics_kwargs.update(overrides) - _logger.info( - f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" - ) + # Create the new dynamics object. + dynamics = system.dynamics(**dynamics_kwargs) - # Save the minimised state back out of the slot. - self._dynamics_cache.store_replica(index) + # Reset the GCMC water state. The dynamics object is created from the + # Sire system, so the water state in the context does not match it. + if gcmc_sampler is not None: + self._reset_gcmc_sampler(gcmc_sampler, dynamics) - except Exception as e: - return False, index, e + # Set the new dynamics object, holding no replica, so that the next load + # sets up its state, lambda value and energy trajectory. + self._dynamics_cache.set(slot, dynamics) + self._dynamics_cache.release(slot) - return True, index, None + _logger.info( + f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" + ) def _equilibrate(self, index): """ @@ -1766,13 +1833,19 @@ def _equilibrate(self, index): _logger.info(f"Equilibrating at {_lam_sym} = {self._lambda_values[index]:.5f}") try: - # Make the replica resident in its slot. - self._dynamics_cache.load_replica(index) + # Make the replica resident in its slot, with its own clock. + clock = self._dynamics_cache.get_clock(index) + self._dynamics_cache.load_replica(index, clock=clock) # Get the dynamics object (and GCMC sampler). slot = self._dynamics_cache.slot_for(index) dynamics, gcmc_sampler = self._dynamics_cache.get(slot) + # Equilibration doesn't count towards the simulation time, so the + # clock is restored afterwards. + if clock is None: + clock = dynamics._get_clock() + if gcmc_sampler is not None: gcmc_sampler.push() try: @@ -1784,58 +1857,6 @@ def _equilibrate(self, index): finally: gcmc_sampler.pop() - # Store the current water state. - water_state = gcmc_sampler.water_state() - - # Work out whether the constraints have changed from the initial minimisation. - if self._config.minimise: - constraint = self._config.constraint - perturbable_constraint = self._config.perturbable_constraint - - if not self._config.equilibration_constraints: - constraint = "none" - perturbable_constraint = "none" - - constraints_changed = (self._initial_constraint != constraint) or ( - self._initial_perturbable_constraint != perturbable_constraint - ) - - # We need to create a new dynamics object if the constraints have changed. - if constraints_changed: - _logger.info( - f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" - ) - - # Commit the current system. - system = dynamics.commit() - - # Delete the current dynamics object. - self._dynamics_cache.delete(slot) - - # Work out the device index. - device = self._dynamics_cache.device_for(slot) - - # Copy the dynamics keyword arguments. - dynamics_kwargs = self._dynamics_kwargs.copy() - - # Overload the device and lambda value. - dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._build_lambda(index) - dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] - dynamics_kwargs["timestep"] = self._config._equilibration_timestep - dynamics_kwargs["constraint"] = constraint - dynamics_kwargs["perturbable_constraint"] = perturbable_constraint - - # Create the new dynamics object. - dynamics = system.dynamics(**dynamics_kwargs) - - # Reset the GCMC water state. - if gcmc_sampler is not None: - self._reset_gcmc_sampler(gcmc_sampler, dynamics) - - # Update the dynamics object in the cache. - self._dynamics_cache.set(slot, dynamics) - # Equilibrate. dynamics.run( self._config.equilibration_time, @@ -1846,52 +1867,73 @@ def _equilibrate(self, index): save_crash_report=self._config.save_crash_report, ) - # Commit the system. - system = dynamics.commit() - - # Reset the timer. - if self._initial_time[index].value() != 0: - system.set_time(self._initial_time[index]) - else: - system.set_time(_sr.u("0ps")) + dynamics._set_clock(clock) - # Delete the dynamics object. - self._dynamics_cache.delete(slot) - - # Work out the device index. - device = self._dynamics_cache.device_for(slot) + # Save the equilibrated state back out of the slot. + self._dynamics_cache.store_replica(index) - # Copy the dynamics keyword arguments. - dynamics_kwargs = self._dynamics_kwargs.copy() + except Exception as e: + return False, index, e - # Overload the device and lambda value. - dynamics_kwargs["device"] = device - dynamics_kwargs["lambda_value"] = self._build_lambda(index) - dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index] + return True, index, None - # Create the production dynamics object. - dynamics = system.dynamics(**dynamics_kwargs) + def _equilibrate_slot(self, replicas): + """ + Equilibrate the replicas hosted by a slot, one after another, in a + single dynamics object, then give the slot one for production. This + creates two dynamics objects per slot, rather than per replica. - # Reset the GCMC water state. The dynamics object is created from - # the original Sire system, so the water state in the context does - # not match the current GCMC water state. - if gcmc_sampler is not None: - self._reset_gcmc_sampler(gcmc_sampler, dynamics) + Parameters + ---------- - # Set the new dynamics object. - self._dynamics_cache.set(slot, dynamics) + replicas: list of int + The replicas hosted by the slot. - _logger.info( - f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}" - ) + Returns + ------- - # Save the equilibrated state back out of the slot. - self._dynamics_cache.store_replica(index) + list of tuple + The success, index and exception for each replica equilibrated. + """ + # Rebuild the slot's dynamics object for equilibration, starting from + # the first replica. + if self._is_equilibration_rebuild: + index = replicas[0] + constraint, perturbable_constraint = self._equilibration_constraints + try: + self._dynamics_cache.load_replica(index) + self._rebuild_dynamics( + index, + timestep=self._config._equilibration_timestep, + constraint=constraint, + perturbable_constraint=perturbable_constraint, + ) + except Exception as e: + _logger.error( + f"Could not create the equilibration dynamics object for " + f"{_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + ) + raise + results = [] + for index in replicas: + results.append(self._equilibrate(index)) + if not results[-1][0]: + return results + + # Rebuild the slot's dynamics object for production, from the last + # replica, which is still resident. + index = replicas[-1] + try: + self._rebuild_dynamics(index) except Exception as e: - return False, index, e + _logger.error( + f"Could not create the production dynamics object for " + f"{_lam_sym} = {self._lambda_values[index]:.5f}: {e}" + ) + raise - return True, index, None + return results def _assemble_results(self, results): """ @@ -2051,10 +2093,11 @@ def _device_batches(self, replicas, devices): batches.append(sorted(batch)) return batches - def _dynamics_batches(self, devices): + def _setup_batches(self, devices): """ - Yield batches of every replica that can be propagated concurrently: at - most one per slot, and at most 'oversubscription_factor' per GPU. + Yield batches of slots to minimise or equilibrate at once, at most + 'oversubscription_factor' per GPU. Each slot is given by the replicas + it hosts, which are set up one after another. Parameters ---------- @@ -2065,11 +2108,16 @@ def _dynamics_batches(self, devices): Yields ------ - list of int - A batch of replicas. + list of list of int + A batch of slots. """ - for batch in self._replica_passes(0): - yield from self._device_batches(batch, devices) + groups = self._dynamics_cache.groups + + # A slot's replicas share its GPU, so its first replica stands in for it. + firsts = [group[0] for group in groups] + + for batch in self._device_batches(firsts, devices): + yield [groups[self._dynamics_cache.slot_for(r)] for r in batch] def _window_samples(self): """ diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index fdebfa6..d7c0a2a 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -211,6 +211,60 @@ def _fake_slots(cache, num_slots): ] +def test_release_restores_resident_replica(): + """ + Validate that loading a replica into a released slot sets it up in full, + even if it is already resident, e.g. after the slot's dynamics object was + rebuilt at another lambda value or a failed step changed its context. + """ + from somd2.runner._dynamics_cache import DynamicsCache + + n = 2 + + cache = object.__new__(DynamicsCache) + cache._lambdas = [0.0, 1.0] + cache._rest2_scale_factors = [1.0] * n + cache._num_replicas = n + cache._update_constraints = True + cache._strided = False + cache._energy_trajectories = [None] * n + cache._ghost_files = [None] * n + cache._state_moved = [False] * n + _fake_slots(cache, n) + + cache._openmm_states = [ + {"positions": np.array([r, 0]), "velocities": np.array([r]), "box": (1, 2, 3)} + for r in range(n) + ] + cache._gcmc_states = [np.eye(n, dtype=int)[r] for r in range(n)] + + replica = 1 + slot = cache.slot_for(replica) + context = cache._dynamics[slot].context() + sampler = cache._gcmc_samplers[slot] + + cache.load_replica(replica) + + def disturb(): + context.positions = np.array([99, 0]) + context.lam = None + sampler.water = np.zeros(n, dtype=int) + + # A resident replica is left alone by a plain load. + disturb() + cache.load_replica(replica) + assert context.positions[0] == 99 + assert context.lam is None + + # After a release, the same replica is set up in full. + disturb() + cache.release(slot) + cache.load_replica(replica) + assert context.positions[0] == replica + assert context.lam == cache._lambdas[replica] + assert list(sampler.water) == list(cache._gcmc_states[replica]) + + @pytest.mark.parametrize( "slot_schedule, inverted", [ @@ -972,7 +1026,8 @@ def test_repex_gcmc_without_a_selection(ethane_methanol): @pytest.mark.skipif(not has_cuda, reason="CUDA not available.") -def test_repex_concurrent_slots(ethane_methanol): +@pytest.mark.parametrize("max_contexts", [2, None]) +def test_repex_concurrent_slots(ethane_methanol, max_contexts): """ Validate that replicas sharing a slot are never propagated concurrently. Oversubscribing exercises this on a single GPU, since the worker count is @@ -997,7 +1052,7 @@ def test_repex_concurrent_slots(ethane_methanol): "max_threads": 1, "num_lambda": num_lambda, "replica_exchange": True, - "max_contexts": 2, + "max_contexts": max_contexts, "oversubscription_factor": 2, } @@ -1012,18 +1067,34 @@ def test_repex_concurrent_slots(ethane_methanol): assert not runner._config.minimisation_constraints assert runner._config.equilibration_constraints - # Every batch of dynamics, and every pass of commits, must contain at - # most one replica per slot. - batches = list(runner._dynamics_batches(runner._replica_devices())) - batches += runner._replica_passes(0) + runner._replica_passes(1) - for batch in batches: - slots = [runner._dynamics_cache.slot_for(r) for r in batch] + # Every batch of slots set up at once must hold distinct slots, at most + # oversubscription_factor on each GPU. + cache = runner._dynamics_cache + for batch in runner._setup_batches(runner._replica_devices()): + slots = [cache.slot_for(group[0]) for group in batch] + assert len(slots) == len(set(slots)), f"batch {batch} repeats a slot" + devices = [cache.device_for(slot) for slot in slots] + for device in set(devices): + assert devices.count(device) <= config["oversubscription_factor"] + + # Every pass of commits must contain at most one replica per slot. + for batch in runner._replica_passes(0) + runner._replica_passes(1): + slots = [cache.slot_for(r) for r in batch] assert len(slots) == len(set(slots)), f"batch {batch} shares a slot" runner.run() assert (Path(tmpdir) / "repex_matrix.txt").exists() + # Equilibration doesn't advance any replica's clock, so each one + # records its first energies after one energy_frequency. + import pandas as pd + + for i in range(num_lambda): + records = pd.read_parquet(runner._filenames[i]["energy_traj"]) + times = [round(t, 6) for t in records.index.get_level_values(0)] + assert times == [0.004, 0.008, 0.012], f"replica {i} recorded {times}" + @pytest.mark.parametrize( "gpu_devices, expected",