Skip to content

Commit 53d3dad

Browse files
IAlibayatravitz
andauthored
Safely close reporter and clear GPU contexts to avoid UnboundLocalErrors (#1846)
* 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>
1 parent 89d1948 commit 53d3dad

File tree

3 files changed

+74
-37
lines changed

3 files changed

+74
-37
lines changed

news/issue-1845.rst

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
**Added:**
2+
3+
* <news item>
4+
5+
**Changed:**
6+
7+
* <news item>
8+
9+
**Deprecated:**
10+
11+
* <news item>
12+
13+
**Removed:**
14+
15+
* <news item>
16+
17+
**Fixed:**
18+
19+
* Fixed a bug in Protocol termination for the HybridTop and AFE Protocols
20+
which would unnecessarily declare an ``UnboundLocalError``.
21+
22+
**Security:**
23+
24+
* <news item>

src/openfe/protocols/openmm_afe/base_afe_units.py

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1319,24 +1319,31 @@ def run(
13191319
)
13201320

13211321
finally:
1322-
# close reporter when you're done to prevent file handle clashes
1323-
reporter.close()
1324-
1325-
# clear GPU context
1326-
# Note: use cache.empty() when openmmtools #690 is resolved
1327-
for context in list(sampler.energy_context_cache._lru._data.keys()):
1328-
del sampler.energy_context_cache._lru._data[context]
1329-
for context in list(sampler.sampler_context_cache._lru._data.keys()):
1330-
del sampler.sampler_context_cache._lru._data[context]
1331-
# cautiously clear out the global context cache too
1332-
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
1333-
del openmmtools.cache.global_context_cache._lru._data[context]
1334-
1335-
del sampler.sampler_context_cache, sampler.energy_context_cache
1336-
1337-
# Keep these around in a dry run so we can inspect things
1338-
if not dry:
1339-
del integrator, sampler
1322+
# Have to wrap this in a try/except, because we might
1323+
# be in a situation where the reporter or sampler weren't created
1324+
try:
1325+
# Order is reporter, contexts, sampler, integrator
1326+
reporter.close() # close to prevent file handle clashes
1327+
1328+
# clear GPU context
1329+
# Note: use cache.empty() when openmmtools #690 is resolved
1330+
for context in list(sampler.energy_context_cache._lru._data.keys()):
1331+
del sampler.energy_context_cache._lru._data[context]
1332+
for context in list(sampler.sampler_context_cache._lru._data.keys()):
1333+
del sampler.sampler_context_cache._lru._data[context]
1334+
# cautiously clear out the global context cache too
1335+
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
1336+
del openmmtools.cache.global_context_cache._lru._data[context]
1337+
1338+
del sampler.sampler_context_cache, sampler.energy_context_cache
1339+
1340+
# Keep these around in a dry run so we can inspect things
1341+
if not dry:
1342+
# At this point we know the sampler exists, so we del the integrator
1343+
# first since it's associated with the sampler
1344+
del integrator, sampler
1345+
except UnboundLocalError:
1346+
pass
13401347

13411348
if not dry:
13421349
nc = self.shared_basepath / settings["output_settings"].output_filename

src/openfe/protocols/openmm_rfe/hybridtop_units.py

Lines changed: 25 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1227,25 +1227,31 @@ def run(
12271227
dry=dry,
12281228
)
12291229
finally:
1230-
# close reporter when you're done, prevent
1231-
# file handle clashes
1232-
reporter.close()
1233-
1234-
# clear GPU contexts
1235-
# TODO: use cache.empty() calls when openmmtools #690 is resolved
1236-
# replace with above
1237-
for context in list(sampler.energy_context_cache._lru._data.keys()):
1238-
del sampler.energy_context_cache._lru._data[context]
1239-
for context in list(sampler.sampler_context_cache._lru._data.keys()):
1240-
del sampler.sampler_context_cache._lru._data[context]
1241-
# cautiously clear out the global context cache too
1242-
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
1243-
del openmmtools.cache.global_context_cache._lru._data[context]
1244-
1245-
del sampler.sampler_context_cache, sampler.energy_context_cache
1246-
1247-
if not dry:
1248-
del integrator, sampler
1230+
# Have to wrap this in a try/except, because we might
1231+
# be in a situation where the reporter or sampler wasn't created
1232+
try:
1233+
# Order is reporter, contexts, sampler, integrator
1234+
reporter.close() # close to prevent file handle clashes
1235+
1236+
# clear GPU context
1237+
# Note: use cache.empty() when openmmtools #690 is resolved
1238+
for context in list(sampler.energy_context_cache._lru._data.keys()):
1239+
del sampler.energy_context_cache._lru._data[context]
1240+
for context in list(sampler.sampler_context_cache._lru._data.keys()):
1241+
del sampler.sampler_context_cache._lru._data[context]
1242+
# cautiously clear out the global context cache too
1243+
for context in list(openmmtools.cache.global_context_cache._lru._data.keys()):
1244+
del openmmtools.cache.global_context_cache._lru._data[context]
1245+
1246+
del sampler.sampler_context_cache, sampler.energy_context_cache
1247+
1248+
# Keep these around in a dry run so we can inspect things
1249+
if not dry:
1250+
# At this point we know the sampler exists, so we del the integrator
1251+
# first since it's associated with the sampler
1252+
del integrator, sampler
1253+
except UnboundLocalError:
1254+
pass
12491255

12501256
if not dry: # pragma: no-cover
12511257
return {

0 commit comments

Comments
 (0)