remove process from cache after step_until_terminated

aiidateam · Mar 4, 2021 · 2b6ef68 · 2b6ef68
1 parent 72505e9
commit 2b6ef68
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/plumpy/process_comms.py b/plumpy/process_comms.py
@@ -602,7 +602,10 @@ async def _launch(
             asyncio.ensure_future(proc.step_until_terminated())
             return proc.pid
 
-        await proc.step_until_terminated()
+        try:
+            await proc.step_until_terminated()
+        finally:
+            self._process_cache.pop(proc.pid, None)
 
         return proc.future().result()
 
@@ -638,7 +641,10 @@ async def _continue(
             asyncio.ensure_future(proc.step_until_terminated())
             return proc.pid
 
-        await proc.step_until_terminated()
+        try:
+            await proc.step_until_terminated()
+        finally:
+            self._process_cache.pop(proc.pid, None)
 
         return proc.future().result()