pub_soft/openfe
3
0
Fork 0
mirror of https://github.com/OpenFreeEnergy/openfe.git synced 2026-06-04 14:14:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

Safely close reporter and clear GPU contexts to avoid UnboundLocalErrors (#1846)

Browse Source 
* Safely close reporter and clear GPU contexts to avoid UnboundLocalErrors

Added try/except block to safely close things in case there's an unbound variable.

* small fix

small fix

* Also improve deletion for the hybridtop protocol

Added error handling for reporter closure and context clearing.

* Add news item

* Update news/issue-1845.rst

Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com>

---------

Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com>
This commit is contained in:
Irfan Alibay
2026-02-16 16:14:16 +00:00
committed by GitHub
parent 89d1948894
commit 53d3dad90f
3 changed files with 68 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
news/issue-1845.rst Normal file
View File

@@ -0,0 +1,24 @@
**Added:**
* <news item>
**Changed:**
* <news item>
**Deprecated:**
* <news item>
**Removed:**
* <news item>
**Fixed:**
* Fixed a bug in Protocol termination for the HybridTop and AFE Protocols
which would unnecessarily declare an ``UnboundLocalError``.
**Security:**
* <news item>

37
src/openfe/protocols/openmm_afe/base_afe_units.py
View File

@@ -1319,24 +1319,31 @@ class BaseAbsoluteMultiStateSimulationUnit(gufe.ProtocolUnit, AbsoluteUnitMixin)
)
finally:
# close reporter when you're done to prevent file handle clashes
reporter.close()
# Have to wrap this in a try/except, because we might
# be in a situation where the reporter or sampler weren't created
try:
# Order is reporter, contexts, sampler, integrator
reporter.close() # close to prevent file handle clashes
# clear GPU context
# Note: use cache.empty() when openmmtools #690 is resolved
for context in list(sampler.energy_context_cache._lru._data.keys()):
del sampler.energy_context_cache._lru._data[context]
for context in list(sampler.sampler_context_cache._lru._data.keys()):
del sampler.sampler_context_cache._lru._data[context]
# cautiously clear out the global context cache too
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
del openmmtools.cache.global_context_cache._lru._data[context]
# clear GPU context
# Note: use cache.empty() when openmmtools #690 is resolved
for context in list(sampler.energy_context_cache._lru._data.keys()):
del sampler.energy_context_cache._lru._data[context]
for context in list(sampler.sampler_context_cache._lru._data.keys()):
del sampler.sampler_context_cache._lru._data[context]
# cautiously clear out the global context cache too
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
del openmmtools.cache.global_context_cache._lru._data[context]
del sampler.sampler_context_cache, sampler.energy_context_cache
del sampler.sampler_context_cache, sampler.energy_context_cache
# Keep these around in a dry run so we can inspect things
if not dry:
del integrator, sampler
# Keep these around in a dry run so we can inspect things
if not dry:
# At this point we know the sampler exists, so we del the integrator
# first since it's associated with the sampler
del integrator, sampler
except UnboundLocalError:
pass
if not dry:
nc = self.shared_basepath / settings["output_settings"].output_filename

38
src/openfe/protocols/openmm_rfe/hybridtop_units.py
View File

@@ -1227,25 +1227,31 @@ class HybridTopologyMultiStateSimulationUnit(gufe.ProtocolUnit, HybridTopologyUn
dry=dry,
)
finally:
# close reporter when you're done, prevent
# file handle clashes
reporter.close()
# Have to wrap this in a try/except, because we might
# be in a situation where the reporter or sampler wasn't created
try:
# Order is reporter, contexts, sampler, integrator
reporter.close() # close to prevent file handle clashes
# clear GPU contexts
# TODO: use cache.empty() calls when openmmtools #690 is resolved
# replace with above
for context in list(sampler.energy_context_cache._lru._data.keys()):
del sampler.energy_context_cache._lru._data[context]
for context in list(sampler.sampler_context_cache._lru._data.keys()):
del sampler.sampler_context_cache._lru._data[context]
# cautiously clear out the global context cache too
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
del openmmtools.cache.global_context_cache._lru._data[context]
# clear GPU context
# Note: use cache.empty() when openmmtools #690 is resolved
for context in list(sampler.energy_context_cache._lru._data.keys()):
del sampler.energy_context_cache._lru._data[context]
for context in list(sampler.sampler_context_cache._lru._data.keys()):
del sampler.sampler_context_cache._lru._data[context]
# cautiously clear out the global context cache too
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
del openmmtools.cache.global_context_cache._lru._data[context]
del sampler.sampler_context_cache, sampler.energy_context_cache
del sampler.sampler_context_cache, sampler.energy_context_cache
if not dry:
del integrator, sampler
# Keep these around in a dry run so we can inspect things
if not dry:
# At this point we know the sampler exists, so we del the integrator
# first since it's associated with the sampler
del integrator, sampler
except UnboundLocalError:
pass
if not dry: # pragma: no-cover
return {