Fixes cleanup issues and improves cleanup logic.

distarray/cleanup.py

@@ -4,32 +4,49 @@
 #  Distributed under the terms of the BSD License.  See COPYING.rst.
 # ---------------------------------------------------------------------------
 
+from __future__ import print_function
+
 from distarray.ipython_utils import IPythonClient
 
 
-def cleanup(view, prefix):
+def engine_cleanup(module_name, prefix):
+    """ Remove variables with ``prefix`` prefix from the namespace of the
+    module with ``module_name``. 
+    """
+    mod = __import__(module_name)
+    ns = mod.__dict__
+    keys = tuple(ns.keys())
+    deleted = 0
+    for k in keys:
+        if k.startswith(prefix):
+            del ns[k]
+            deleted += 1
+    count = 0
+    for k in ns:
+        if k.startswith(prefix):
+            count += 1
+    return (deleted, count)
+
+
+def cleanup(view, module_name, prefix):
     """ Delete keys with prefix from client's engines. """
-    # Delete keys only from this context.
-    def engine_cleanup(prefix):
-        glb = globals()
-        global_keys = list(globals())
-        for gk in global_keys:
-            if gk.startswith(prefix):
-                del glb[gk]
-    view.apply_async(engine_cleanup, prefix)
-
-def cleanup_all(prefix):
+    remaining = view.apply_async(engine_cleanup, module_name, prefix).get_dict()
+    return remaining
+
+
+def cleanup_all(module_name, prefix):
     """ Connects to all engines and runs ``cleanup()`` on them. """
     try:
         c = IPythonClient()
     except IOError:  # If we can't create a client, return silently.
         return
     try:
         v = c[:]
-        cleanup(v, prefix)
+        cleanup(v, module_name, prefix)
     finally:
         c.close()
 
+
 def get_local_keys(view, prefix):
     """ Returns a dictionary of keyname -> target_list mapping for all names
         that start with ``prefix`` on engines in ``view``.

distarray/context.py

@@ -21,9 +21,8 @@
 from distarray.client_map import Distribution
 from distarray.ipython_utils import IPythonClient
 
+
 DISTARRAY_BASE_NAME = '__distarray__'
-atexit.register(cleanup.cleanup_all, DISTARRAY_BASE_NAME)
-atexit.register(cleanup.clear_all)
 
 
 class Context(object):
@@ -39,8 +38,14 @@ class Context(object):
 
     '''
 
+    _CLEANUP = None
+
     def __init__(self, client=None, targets=None):
 
+        if not Context._CLEANUP:
+            Context._CLEANUP = (atexit.register(cleanup.clear_all),
+                                atexit.register(cleanup.cleanup_all, '__main__', DISTARRAY_BASE_NAME))
+
         if client is None:
             self.client = IPythonClient()
             self.owns_client = True
@@ -160,7 +165,7 @@ def delete_key(self, key):
 
     def cleanup(self):
         """ Delete keys that this context created from all the engines. """
-        cleanup.cleanup(view=self.view, prefix=self._key_prefix())
+        cleanup.cleanup(view=self.view, module_name='__main__', prefix=self._key_prefix())
 
     def close(self):
         self.cleanup()