Skip to content

Commit 6f57a18

Browse files
lohedgeslohedges
committed
Wrap GCMC push/pop pairs in try/finally to guard against crashes.
1 parent 86e14b0 commit 6f57a18

2 files changed

Lines changed: 104 additions & 76 deletions

File tree

‎src/somd2/runner/_repex.py‎

Lines changed: 55 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -532,12 +532,14 @@ def mix_states(self):
532532

533533
# Update the water state in the GCMCSampler.
534534
self._gcmc_samplers[i].push()
535-
self._gcmc_samplers[i]._set_water_state(
536-
self._dynamics[i].context(),
537-
indices=water_idxs,
538-
states=self._gcmc_states[state][water_idxs],
539-
)
540-
self._gcmc_samplers[i].pop()
535+
try:
536+
self._gcmc_samplers[i]._set_water_state(
537+
self._dynamics[i].context(),
538+
indices=water_idxs,
539+
states=self._gcmc_states[state][water_idxs],
540+
)
541+
finally:
542+
self._gcmc_samplers[i].pop()
541543

542544
# Update the swap matrix.
543545
old_state = self._old_states[i]
@@ -820,12 +822,14 @@ def __init__(self, system, config):
820822
# Reset the GCMC water state.
821823
if gcmc_sampler is not None:
822824
gcmc_sampler.push()
823-
gcmc_sampler._set_water_state(
824-
dynamics.context(),
825-
states=self._dynamics_cache._gcmc_states[state],
826-
force=True,
827-
)
828-
gcmc_sampler.pop()
825+
try:
826+
gcmc_sampler._set_water_state(
827+
dynamics.context(),
828+
states=self._dynamics_cache._gcmc_states[state],
829+
force=True,
830+
)
831+
finally:
832+
gcmc_sampler.pop()
829833

830834
# Conversion factor for reduced potential.
831835
kT = (_sr.units.k_boltz * self._config.temperature).to(_sr.units.kcal_per_mol)
@@ -1277,13 +1281,13 @@ def _run_block(
12771281
if is_gcmc:
12781282
# Push the PyCUDA context on top of the stack.
12791283
gcmc_sampler.push()
1280-
1281-
# Perform the GCMC move.
1282-
_logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}")
1283-
gcmc_sampler.move(dynamics.context())
1284-
1285-
# Remove the PyCUDA context from the stack.
1286-
gcmc_sampler.pop()
1284+
try:
1285+
# Perform the GCMC move.
1286+
_logger.info(f"Performing GCMC move at {_lam_sym} = {lam:.5f}")
1287+
gcmc_sampler.move(dynamics.context())
1288+
finally:
1289+
# Remove the PyCUDA context from the stack.
1290+
gcmc_sampler.pop()
12871291

12881292
# Save the GCMC state.
12891293
self._dynamics_cache.save_gcmc_state(index)
@@ -1340,15 +1344,15 @@ def _minimise(self, index):
13401344
if gcmc_sampler is not None:
13411345
# Push the PyCUDA context on top of the stack.
13421346
gcmc_sampler.push()
1343-
1344-
_logger.info(
1345-
f"Pre-equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}"
1346-
)
1347-
for i in range(100):
1348-
gcmc_sampler.move(dynamics.context())
1349-
1350-
# Remove the PyCUDA context from the stack.
1351-
gcmc_sampler.pop()
1347+
try:
1348+
_logger.info(
1349+
f"Pre-equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}"
1350+
)
1351+
for i in range(100):
1352+
gcmc_sampler.move(dynamics.context())
1353+
finally:
1354+
# Remove the PyCUDA context from the stack.
1355+
gcmc_sampler.pop()
13521356

13531357
# Minimise.
13541358
dynamics.minimise(timeout=self._config.timeout)
@@ -1433,15 +1437,15 @@ def _equilibrate(self, index):
14331437
if gcmc_sampler is not None:
14341438
# Push the PyCUDA context on top of the stack.
14351439
gcmc_sampler.push()
1436-
1437-
_logger.info(
1438-
f"Equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}"
1439-
)
1440-
for i in range(100):
1441-
gcmc_sampler.move(dynamics.context())
1442-
1443-
# Remove the PyCUDA context from the stack.
1444-
gcmc_sampler.pop()
1440+
try:
1441+
_logger.info(
1442+
f"Equilibrating with GCMC moves at {_lam_sym} = {self._lambda_values[index]:.5f}"
1443+
)
1444+
for i in range(100):
1445+
gcmc_sampler.move(dynamics.context())
1446+
finally:
1447+
# Remove the PyCUDA context from the stack.
1448+
gcmc_sampler.pop()
14451449

14461450
# Store the current water state.
14471451
water_state = gcmc_sampler.water_state()
@@ -1693,14 +1697,14 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
16931697
if gcmc_sampler is not None:
16941698
# Push the PyCUDA context on top of the stack.
16951699
gcmc_sampler.push()
1696-
1697-
_logger.info(
1698-
f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} "
1699-
f"is {gcmc_sampler.num_waters()}"
1700-
)
1701-
1702-
# Remove the PyCUDA context from the stack.
1703-
gcmc_sampler.pop()
1700+
try:
1701+
_logger.info(
1702+
f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} "
1703+
f"is {gcmc_sampler.num_waters()}"
1704+
)
1705+
finally:
1706+
# Remove the PyCUDA context from the stack.
1707+
gcmc_sampler.pop()
17041708

17051709
if is_final_block:
17061710
_logger.success(f"{_lam_sym} = {lam:.5f} complete")
@@ -1823,12 +1827,12 @@ def _reset_gcmc_sampler(gcmc_sampler, dynamics):
18231827

18241828
# Push the PyCUDA context on top of the stack.
18251829
gcmc_sampler.push()
1826-
1827-
# Set the water state.
1828-
gcmc_sampler._set_water_state(dynamics.context(), force=True)
1829-
1830-
# Remove the PyCUDA context from the stack.
1831-
gcmc_sampler.pop()
1830+
try:
1831+
# Set the water state.
1832+
gcmc_sampler._set_water_state(dynamics.context(), force=True)
1833+
finally:
1834+
# Remove the PyCUDA context from the stack.
1835+
gcmc_sampler.pop()
18321836

18331837
# Re-bind the GCMC sampler to the dynamics object.
18341838
gcmc_sampler.bind_dynamics(dynamics)

‎src/somd2/runner/_runner.py‎

Lines changed: 49 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -521,8 +521,12 @@ def generate_lam_vals(lambda_base, increment=0.001):
521521
f"Equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}"
522522
)
523523

524-
for i in range(100):
525-
gcmc_sampler.move(dynamics.context())
524+
gcmc_sampler.push()
525+
try:
526+
for i in range(100):
527+
gcmc_sampler.move(dynamics.context())
528+
finally:
529+
gcmc_sampler.pop()
526530

527531
# Run without saving energies or frames.
528532
dynamics.run(
@@ -610,20 +614,24 @@ def generate_lam_vals(lambda_base, increment=0.001):
610614
if self._is_restart:
611615
from openmm.unit import angstrom
612616

613-
# First set all waters to non-ghosts.
614-
gcmc_sampler._set_water_state(
615-
dynamics.context(),
616-
states=_np.ones(len(gcmc_sampler._water_indices)),
617-
force=True,
618-
)
617+
gcmc_sampler.push()
618+
try:
619+
# First set all waters to non-ghosts.
620+
gcmc_sampler._set_water_state(
621+
dynamics.context(),
622+
states=_np.ones(len(gcmc_sampler._water_indices)),
623+
force=True,
624+
)
619625

620-
# Now set the ghost waters.
621-
gcmc_sampler._set_water_state(
622-
dynamics.context(),
623-
self._restart_ghost_waters[index],
624-
states=_np.zeros(len(gcmc_sampler._water_indices)),
625-
force=True,
626-
)
626+
# Now set the ghost waters.
627+
gcmc_sampler._set_water_state(
628+
dynamics.context(),
629+
self._restart_ghost_waters[index],
630+
states=_np.zeros(len(gcmc_sampler._water_indices)),
631+
force=True,
632+
)
633+
finally:
634+
gcmc_sampler.pop()
627635

628636
# Finally, reset the context positions to match the restart system.
629637
dynamics.context().setPositions(
@@ -634,10 +642,14 @@ def generate_lam_vals(lambda_base, increment=0.001):
634642
# the water state in the new context to match the equilibrated system.
635643
elif is_equilibrated:
636644
# Reset the water state.
637-
gcmc_sampler._set_water_state(
638-
dynamics.context(),
639-
force=True,
640-
)
645+
gcmc_sampler.push()
646+
try:
647+
gcmc_sampler._set_water_state(
648+
dynamics.context(),
649+
force=True,
650+
)
651+
finally:
652+
gcmc_sampler.pop()
641653

642654
# Set the number of neighbours used for the energy calculation.
643655
# If not None, then we add one to account for the extra windows
@@ -728,7 +740,11 @@ def generate_lam_vals(lambda_base, increment=0.001):
728740
_logger.info(
729741
f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
730742
)
731-
gcmc_sampler.move(dynamics.context())
743+
gcmc_sampler.push()
744+
try:
745+
gcmc_sampler.move(dynamics.context())
746+
finally:
747+
gcmc_sampler.pop()
732748

733749
else:
734750
dynamics.run(
@@ -816,10 +832,14 @@ def generate_lam_vals(lambda_base, increment=0.001):
816832

817833
# Log the number of waters within the GCMC sampling volume.
818834
if gcmc_sampler is not None:
819-
_logger.info(
820-
f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} "
821-
f"is {gcmc_sampler.num_waters()}"
822-
)
835+
gcmc_sampler.push()
836+
try:
837+
_logger.info(
838+
f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} "
839+
f"is {gcmc_sampler.num_waters()}"
840+
)
841+
finally:
842+
gcmc_sampler.pop()
823843

824844
if is_final_block:
825845
_logger.success(
@@ -928,7 +948,11 @@ def generate_lam_vals(lambda_base, increment=0.001):
928948
_logger.info(
929949
f"Performing GCMC move at {_lam_sym} = {lambda_value:.5f}"
930950
)
931-
gcmc_sampler.move(dynamics.context())
951+
gcmc_sampler.push()
952+
try:
953+
gcmc_sampler.move(dynamics.context())
954+
finally:
955+
gcmc_sampler.pop()
932956

933957
# Update the runtime.
934958
runtime += self._config.energy_frequency

0 commit comments

Comments
 (0)