diff --git a/distarray/cleanup.py b/distarray/cleanup.py index cd18e1e0..06d444ab 100644 --- a/distarray/cleanup.py +++ b/distarray/cleanup.py @@ -8,13 +8,15 @@ from distarray.ipython_utils import IPythonClient -def cleanup(view, module_name, context_name): +def cleanup(view, module_name, prefix): """ Delete Context object with the given name from the given module""" - def _cleanup(module_name, context_name): + def _cleanup(module_name, prefix): ns = __import__(module_name) - delattr(ns, context_name) + for name in vars(ns).copy(): + if name.startswith(prefix): + delattr(ns, name) - view.apply_async(_cleanup, module_name, context_name) + view.apply_sync(_cleanup, module_name, prefix) def cleanup_all(module_name, prefix): diff --git a/distarray/context.py b/distarray/context.py index 41b6656d..270ba293 100644 --- a/distarray/context.py +++ b/distarray/context.py @@ -168,7 +168,7 @@ def delete_key(self, key): def cleanup(self): """ Delete keys that this context created from all the engines. """ - cleanup.cleanup(view=self.view, module_name='__main__', context_name=self.context_key) + cleanup.cleanup(view=self.view, module_name='__main__', prefix=self.context_key) def close(self): self.cleanup()