Safely close reporter and clear GPU contexts to avoid UnboundLocalErrors

news/issue-1845.rst

@@ -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>

src/openfe/protocols/openmm_afe/base_afe_units.py

@@ -1319,24 +1319,31 @@ def run(
             )
 
         finally:
-            # close reporter when you're done to prevent file handle clashes
-            reporter.close()
-
-            # 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
-
-            # Keep these around in a dry run so we can inspect things
-            if not dry:
-                del integrator, sampler
+            # 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]
+
+                del sampler.sampler_context_cache, sampler.energy_context_cache
+
+                # 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

src/openfe/protocols/openmm_rfe/hybridtop_units.py

@@ -1227,25 +1227,31 @@ def run(
                 dry=dry,
             )
         finally:
-            # close reporter when you're done, prevent
-            # file handle clashes
-            reporter.close()
-
-            # 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]
-
-            del sampler.sampler_context_cache, sampler.energy_context_cache
-
-            if not dry:
-                del integrator, sampler
+            # 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 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
+
+                # 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 {