From a6fd37cbc8eb4eeea9d00bf74ec8cc496c2a1ac9 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 20 Oct 2026 14:30:45 +0100
Subject: [PATCH 1/4] Build setup contexts once per slot rather than once per
 replica.

---
 CHANGELOG.md               |   1 +
 src/somd2/runner/_repex.py | 321 +++++++++++++++++++------------------
 tests/runner/test_repex.py |  23 ++-
 3 files changed, 184 insertions(+), 161 deletions(-)

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/_repex.py b/src/somd2/runner/_repex.py
index 2d2ff7c..f9cab4e 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1022,19 +1022,17 @@ 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):
+                        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)
                 except KeyboardInterrupt:
                     dynamics_executor.shutdown(wait=False, cancel_futures=True)
                     _logger.error("Minimisation cancelled. Exiting.")
@@ -1042,17 +1040,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.")
@@ -1685,61 +1681,92 @@ 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)
 
-                if constraints_changed:
-                    # Commit the current system.
-                    system = dynamics.commit()
+        except Exception as e:
+            return False, index, e
 
-                    # Delete the dynamics object.
-                    self._dynamics_cache.delete(slot)
+        return True, index, None
 
-                    # Work out the device index.
-                    device = self._dynamics_cache.device_for(slot)
+    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.
 
-                    # Copy the dynamics keyword arguments.
-                    dynamics_kwargs = self._dynamics_kwargs.copy()
+        Parameters
+        ----------
 
-                    # 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]
+        replicas: list of int
+            The replicas hosted by the slot.
 
-                    # Create the production dynamics object.
-                    dynamics = system.dynamics(**dynamics_kwargs)
+        Returns
+        -------
 
-                    # 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)
+        list of tuple
+            The success, index and exception for each replica.
+        """
+        results = [self._minimise(index) for index in replicas]
 
-                    # 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)
+        constraints_changed = (self._initial_constraint != self._config.constraint) or (
+            self._initial_perturbable_constraint != self._config.perturbable_constraint
+        )
 
-                    _logger.info(
-                        f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
-                    )
+        if not self._is_equilibration and constraints_changed:
+            try:
+                self._rebuild_dynamics(replicas[-1])
+            except Exception as e:
+                results.append((False, replicas[-1], e))
 
-            # Save the minimised state back out of the slot.
-            self._dynamics_cache.store_replica(index)
+        return results
 
-        except Exception as e:
-            return False, index, e
+    def _rebuild_dynamics(self, index, **overrides):
+        """
+        Replace the dynamics object of a replica's slot with a new one, built
+        from the replica's state, e.g. to change the constraints or timestep.
+        The replica must be resident in its slot.
 
-        return True, index, None
+        Parameters
+        ----------
+
+        index: int
+            The index of the replica.
+
+        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)
+
+        # Commit the current system.
+        system = dynamics.commit()
+
+        # Delete the current dynamics object.
+        self._dynamics_cache.delete(slot)
+
+        # 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)
+
+        # Create the new dynamics object.
+        dynamics = system.dynamics(**dynamics_kwargs)
+
+        # 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)
+
+        # Set the new dynamics object. Its energy trajectory is empty, but none
+        # are recorded during setup, and load_replica() re-attaches each one.
+        self._dynamics_cache.set(slot, dynamics)
+
+        _logger.info(
+            f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
+        )
 
     def _equilibrate(self, index):
         """
@@ -1773,6 +1800,10 @@ def _equilibrate(self, index):
             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.
+            clock = dynamics._get_clock()
+
             if gcmc_sampler is not None:
                 gcmc_sampler.push()
                 try:
@@ -1784,58 +1815,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 +1825,76 @@ def _equilibrate(self, index):
                 save_crash_report=self._config.save_crash_report,
             )
 
-            # Commit the system.
-            system = dynamics.commit()
+            dynamics._set_clock(clock)
 
-            # Reset the timer.
-            if self._initial_time[index].value() != 0:
-                system.set_time(self._initial_time[index])
-            else:
-                system.set_time(_sr.u("0ps"))
+            # Save the equilibrated state back out of the slot.
+            self._dynamics_cache.store_replica(index)
 
-            # Delete the dynamics object.
-            self._dynamics_cache.delete(slot)
+        except Exception as e:
+            return False, index, e
 
-            # Work out the device index.
-            device = self._dynamics_cache.device_for(slot)
+        return True, index, None
 
-            # Copy the dynamics keyword arguments.
-            dynamics_kwargs = self._dynamics_kwargs.copy()
+    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.
 
-            # 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]
+        Parameters
+        ----------
 
-            # Create the production dynamics object.
-            dynamics = system.dynamics(**dynamics_kwargs)
+        replicas: list of int
+            The replicas hosted by the slot.
 
-            # 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)
+        Returns
+        -------
 
-            # Set the new dynamics object.
-            self._dynamics_cache.set(slot, dynamics)
+        list of tuple
+            The success, index and exception for each replica equilibrated.
+        """
+        # Work out whether the constraints have changed from the initial minimisation.
+        constraints_changed = False
+        if self._config.minimise:
+            constraint = self._config.constraint
+            perturbable_constraint = self._config.perturbable_constraint
 
-            _logger.info(
-                f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
+            if not self._config.equilibration_constraints:
+                constraint = "none"
+                perturbable_constraint = "none"
+
+            constraints_changed = (self._initial_constraint != constraint) or (
+                self._initial_perturbable_constraint != perturbable_constraint
             )
 
-            # Save the equilibrated state back out of the slot.
-            self._dynamics_cache.store_replica(index)
+        # Rebuild the slot's dynamics object for equilibration, starting from
+        # the first replica.
+        if constraints_changed:
+            try:
+                self._dynamics_cache.load_replica(replicas[0])
+                self._rebuild_dynamics(
+                    replicas[0],
+                    timestep=self._config._equilibration_timestep,
+                    constraint=constraint,
+                    perturbable_constraint=perturbable_constraint,
+                )
+            except Exception as e:
+                return [(False, replicas[0], e)]
+
+        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.
+        try:
+            self._rebuild_dynamics(replicas[-1])
         except Exception as e:
-            return False, index, e
+            results.append((False, replicas[-1], e))
 
-        return True, index, None
+        return results
 
     def _assemble_results(self, results):
         """
@@ -2051,10 +2054,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 +2069,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..8b6bf31 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -972,7 +972,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 +998,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,9 +1013,12 @@ 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()))
+        # Every batch of slots set up at once, and every pass of commits, must
+        # contain at most one replica per slot.
+        batches = [
+            [group[0] for group in batch]
+            for batch in runner._setup_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]
@@ -1024,6 +1028,15 @@ def test_repex_concurrent_slots(ethane_methanol):
 
         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",

From 9378efa05df2bf63bda47a45ef54b830dfdecb9d Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 20 Oct 2026 14:39:44 +0100
Subject: [PATCH 2/4] Fix minor issues in per-slot setup.

---
 src/somd2/runner/_base.py           |  5 --
 src/somd2/runner/_dynamics_cache.py |  8 ++-
 src/somd2/runner/_repex.py          | 93 ++++++++++++++++++-----------
 tests/runner/test_repex.py          | 22 ++++---
 4 files changed, 78 insertions(+), 50 deletions(-)

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..c57626b 100644
--- a/src/somd2/runner/_dynamics_cache.py
+++ b/src/somd2/runner/_dynamics_cache.py
@@ -908,7 +908,7 @@ def set_states(self, states):
         """
         self._states = states
 
-    def load_replica(self, replica, clock=None):
+    def load_replica(self, replica, clock=None, force=False):
         """
         Make a replica resident in its slot, ready for a dynamics block.
 
@@ -932,6 +932,10 @@ def load_replica(self, replica, clock=None):
         clock: dict
             The simulation clock to restore, as returned by
             Dynamics._get_clock(). If None, the slot's clock is left alone.
+
+        force: bool
+            Whether to push the replica's stored state even if it is already
+            resident, e.g. when a failed step left the context in a bad state.
         """
         slot = self._replica_slot[replica]
         dynamics = self._dynamics[slot]
@@ -942,7 +946,7 @@ def load_replica(self, replica, clock=None):
 
         # The slot needs new positions and velocities if it is being handed a
         # different replica, or if the last mix moved this replica's state.
-        if resident != replica or self._state_moved[replica]:
+        if resident != replica or self._state_moved[replica] or force:
             self._apply_openmm_state(dynamics.context(), self._openmm_states[replica])
 
             # Positions have changed underneath the context, so any cached
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index f9cab4e..4da908d 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()
@@ -1667,7 +1680,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(
@@ -1709,15 +1722,19 @@ def _minimise_slot(self, replicas):
         """
         results = [self._minimise(index) for index in replicas]
 
-        constraints_changed = (self._initial_constraint != self._config.constraint) or (
-            self._initial_perturbable_constraint != self._config.perturbable_constraint
-        )
-
-        if not self._is_equilibration and constraints_changed:
+        if not self._is_equilibration and self._is_production_rebuild:
+            # Rebuild from the last replica's stored state, which is its
+            # starting state if its minimisation failed.
+            index = replicas[-1]
             try:
-                self._rebuild_dynamics(replicas[-1])
+                self._dynamics_cache.load_replica(index, force=True)
+                self._rebuild_dynamics(index)
             except Exception as e:
-                results.append((False, replicas[-1], e))
+                _logger.error(
+                    f"Could not create the production dynamics object for "
+                    f"{_lam_sym} = {self._lambda_values[index]:.5f}: {e}"
+                )
+                raise
 
         return results
 
@@ -1760,6 +1777,15 @@ def _rebuild_dynamics(self, index, **overrides):
         if gcmc_sampler is not None:
             self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
+        # A context built at the shared constraint lambda value is moved to the
+        # replica's own, since load_replica() skips this for a resident replica.
+        if dynamics_kwargs["lambda_value"] != self._lambda_values[index]:
+            dynamics.set_lambda(
+                self._lambda_values[index],
+                rest2_scale=self._rest2_scale_factors[index],
+                update_constraints=self._config.update_constraints,
+            )
+
         # Set the new dynamics object. Its energy trajectory is empty, but none
         # are recorded during setup, and load_replica() re-attaches each one.
         self._dynamics_cache.set(slot, dynamics)
@@ -1793,8 +1819,9 @@ 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)
@@ -1802,7 +1829,8 @@ def _equilibrate(self, index):
 
             # Equilibration doesn't count towards the simulation time, so the
             # clock is restored afterwards.
-            clock = dynamics._get_clock()
+            if clock is None:
+                clock = dynamics._get_clock()
 
             if gcmc_sampler is not None:
                 gcmc_sampler.push()
@@ -1853,33 +1881,25 @@ def _equilibrate_slot(self, replicas):
         list of tuple
             The success, index and exception for each replica equilibrated.
         """
-        # Work out whether the constraints have changed from the initial minimisation.
-        constraints_changed = False
-        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
-            )
-
         # Rebuild the slot's dynamics object for equilibration, starting from
         # the first replica.
-        if constraints_changed:
+        if self._is_equilibration_rebuild:
+            index = replicas[0]
+            constraint, perturbable_constraint = self._equilibration_constraints
             try:
-                self._dynamics_cache.load_replica(replicas[0])
+                self._dynamics_cache.load_replica(index)
                 self._rebuild_dynamics(
-                    replicas[0],
+                    index,
                     timestep=self._config._equilibration_timestep,
                     constraint=constraint,
                     perturbable_constraint=perturbable_constraint,
                 )
             except Exception as e:
-                return [(False, replicas[0], 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:
@@ -1889,10 +1909,15 @@ def _equilibrate_slot(self, replicas):
 
         # Rebuild the slot's dynamics object for production, from the last
         # replica, which is still resident.
+        index = replicas[-1]
         try:
-            self._rebuild_dynamics(replicas[-1])
+            self._rebuild_dynamics(index)
         except Exception as e:
-            results.append((False, replicas[-1], e))
+            _logger.error(
+                f"Could not create the production dynamics object for "
+                f"{_lam_sym} = {self._lambda_values[index]:.5f}: {e}"
+            )
+            raise
 
         return results
 
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 8b6bf31..2adfe48 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -1013,15 +1013,19 @@ def test_repex_concurrent_slots(ethane_methanol, max_contexts):
         assert not runner._config.minimisation_constraints
         assert runner._config.equilibration_constraints
 
-        # Every batch of slots set up at once, and every pass of commits, must
-        # contain at most one replica per slot.
-        batches = [
-            [group[0] for group in batch]
-            for batch in runner._setup_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()

From eeb94000c7459bd7ba5be2cbb4df692265071919 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 20 Oct 2026 15:44:52 +0100
Subject: [PATCH 3/4] Release rebuilt or failed slots so the next load sets
 them up in full.

---
 src/somd2/runner/_dynamics_cache.py | 26 ++++++++----
 src/somd2/runner/_repex.py          | 61 +++++++++++++++++------------
 2 files changed, 56 insertions(+), 31 deletions(-)

diff --git a/src/somd2/runner/_dynamics_cache.py b/src/somd2/runner/_dynamics_cache.py
index c57626b..181a930 100644
--- a/src/somd2/runner/_dynamics_cache.py
+++ b/src/somd2/runner/_dynamics_cache.py
@@ -731,7 +731,8 @@ def get(self, slot):
 
     def set(self, slot, dynamics):
         """
-        Set the dynamics object for a given slot.
+        Set the dynamics object for a given slot. The slot then holds no
+        replica, so the next one loaded into it is set up in full.
 
         Parameters
         ----------
@@ -743,6 +744,21 @@ def set(self, slot, dynamics):
             The dynamics object.
         """
         self._dynamics[slot] = dynamics
+        self.release(slot)
+
+    def release(self, slot):
+        """
+        Mark a slot as holding no replica, so that the next replica loaded
+        into it has its state, lambda value and GCMC water state pushed in
+        full, e.g. after a failed step left the context in a bad state.
+
+        Parameters
+        ----------
+
+        slot: int
+            The index of the slot.
+        """
+        self._slot_replica[slot] = None
 
     def delete(self, slot):
         """
@@ -908,7 +924,7 @@ def set_states(self, states):
         """
         self._states = states
 
-    def load_replica(self, replica, clock=None, force=False):
+    def load_replica(self, replica, clock=None):
         """
         Make a replica resident in its slot, ready for a dynamics block.
 
@@ -932,10 +948,6 @@ def load_replica(self, replica, clock=None, force=False):
         clock: dict
             The simulation clock to restore, as returned by
             Dynamics._get_clock(). If None, the slot's clock is left alone.
-
-        force: bool
-            Whether to push the replica's stored state even if it is already
-            resident, e.g. when a failed step left the context in a bad state.
         """
         slot = self._replica_slot[replica]
         dynamics = self._dynamics[slot]
@@ -946,7 +958,7 @@ def load_replica(self, replica, clock=None, force=False):
 
         # The slot needs new positions and velocities if it is being handed a
         # different replica, or if the last mix moved this replica's state.
-        if resident != replica or self._state_moved[replica] or force:
+        if resident != replica or self._state_moved[replica]:
             self._apply_openmm_state(dynamics.context(), self._openmm_states[replica])
 
             # Positions have changed underneath the context, so any cached
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 4da908d..75e54ca 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1038,14 +1038,7 @@ def _prepare_replicas(
             for batch in self._setup_batches(replica_devices):
                 try:
                     for results in dynamics_executor.map(self._minimise_slot, batch):
-                        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)
+                        self._report_minimisation(results)
                 except KeyboardInterrupt:
                     dynamics_executor.shutdown(wait=False, cancel_futures=True)
                     _logger.error("Minimisation cancelled. Exiting.")
@@ -1698,10 +1691,33 @@ def _minimise(self, index):
             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
@@ -1722,14 +1738,20 @@ def _minimise_slot(self, replicas):
         """
         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:
-            # Rebuild from the last replica's stored state, which is its
-            # starting state if its minimisation failed.
+            # 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, force=True)
+                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}"
@@ -1741,8 +1763,8 @@ def _minimise_slot(self, replicas):
     def _rebuild_dynamics(self, index, **overrides):
         """
         Replace the dynamics object of a replica's slot with a new one, built
-        from the replica's state, e.g. to change the constraints or timestep.
-        The replica must be resident in its slot.
+        from the slot's current state, e.g. to change the constraints or
+        timestep. Afterwards the slot holds no replica.
 
         Parameters
         ----------
@@ -1777,17 +1799,8 @@ def _rebuild_dynamics(self, index, **overrides):
         if gcmc_sampler is not None:
             self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
-        # A context built at the shared constraint lambda value is moved to the
-        # replica's own, since load_replica() skips this for a resident replica.
-        if dynamics_kwargs["lambda_value"] != self._lambda_values[index]:
-            dynamics.set_lambda(
-                self._lambda_values[index],
-                rest2_scale=self._rest2_scale_factors[index],
-                update_constraints=self._config.update_constraints,
-            )
-
-        # Set the new dynamics object. Its energy trajectory is empty, but none
-        # are recorded during setup, and load_replica() re-attaches each one.
+        # Set the new dynamics object. The slot then holds no replica, so the
+        # next load sets up its state, lambda value and energy trajectory.
         self._dynamics_cache.set(slot, dynamics)
 
         _logger.info(

From 85e86fa5f3ad27a44871a79bde6ecc78c972541a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 20 Oct 2026 15:50:19 +0100
Subject: [PATCH 4/4] Release slots explicitly and test that release restores a
 replica.

---
 src/somd2/runner/_dynamics_cache.py | 18 ++++------
 src/somd2/runner/_repex.py          |  5 +--
 tests/runner/test_repex.py          | 54 +++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/src/somd2/runner/_dynamics_cache.py b/src/somd2/runner/_dynamics_cache.py
index 181a930..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
@@ -731,8 +728,7 @@ def get(self, slot):
 
     def set(self, slot, dynamics):
         """
-        Set the dynamics object for a given slot. The slot then holds no
-        replica, so the next one loaded into it is set up in full.
+        Set the dynamics object for a given slot.
 
         Parameters
         ----------
@@ -744,13 +740,13 @@ def set(self, slot, dynamics):
             The dynamics object.
         """
         self._dynamics[slot] = dynamics
-        self.release(slot)
 
     def release(self, slot):
         """
-        Mark a slot as holding no replica, so that the next replica loaded
-        into it has its state, lambda value and GCMC water state pushed in
-        full, e.g. after a failed step left the context in a bad state.
+        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
         ----------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 75e54ca..0de187e 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1799,9 +1799,10 @@ def _rebuild_dynamics(self, index, **overrides):
         if gcmc_sampler is not None:
             self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
-        # Set the new dynamics object. The slot then holds no replica, so the
-        # next load sets up its state, lambda value and energy trajectory.
+        # 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)
 
         _logger.info(
             f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 2adfe48..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",
     [