more on finalize

python/mxnet/__init__.py

@@ -10,12 +10,26 @@
 
 from .context import Context, current_context, cpu, gpu
 from .base import MXNetError
+from . import base
 from . import ndarray
 from . import symbol
 from . import kvstore as kv
 from . import io
 # use mx.nd as short for mx.ndarray
 from . import ndarray as nd
 from . import random
+import atexit
 
 __version__ = "0.1.0"
+
+def finalize():
+    """Stop all the components in mxnet.
+
+    There is no need to call this function.
+    This function will be automatically called at module exit.
+    """
+    # pylint: disable=protected-access
+    base.check_call(base._LIB.MXFinalize())
+    kv._cleanup()
+
+atexit.register(finalize)

python/mxnet/base.py

@@ -8,7 +8,6 @@
 import ctypes
 import platform
 import numpy as np
-import atexit
 
 __all__ = ['MXNetError']
 #----------------------------
@@ -180,10 +179,3 @@ def ctypes2numpy_shared(cptr, shape):
     dbuffer = (mx_float * size).from_address(ctypes.addressof(cptr.contents))
     return np.frombuffer(dbuffer, dtype=np.float32).reshape(shape)
 
-
-def stop_all():
-    """Stop All the components in mxnet."""
-    check_call(_LIB.MXFinalize())
-
-
-atexit.register(stop_all)

python/mxnet/kvstore.py

@@ -6,9 +6,8 @@
 from .ndarray import NDArray
 from .base import _LIB
 from .base import check_call, c_array, NDArrayHandle
-import atexit
 
-__all__ = ['start', 'init', 'push', 'pull', 'stop', 'set_updater']
+__all__ = ['start', 'init', 'push', 'pull', 'set_updater']
 
 def _ctype_key_value(keys, vals):
     """
@@ -213,11 +212,9 @@ def set_updater(updater):
     _updater_func = _updater_proto(_updater_wrapper(updater))
     check_call(_LIB.MXKVStoreSetUpdater(_updater_func))
 
-def stop():
-    """ Stop the kvstore """
+def _cleanup():
+    """ cleanup callbacks """
     # need to clear _updater_func before _LIB
     global _updater_func
     _updater_func = None
 
-
-atexit.register(stop)

src/global.cc

@@ -0,0 +1,21 @@
+/*!
+ *  Copyright (c) 2015 by Contributors
+ * \file global.cc
+ * \brief Implementation of project global related functions.
+ */
+#include <mxnet/base.h>
+#include <mxnet/engine.h>
+#include <mxnet/storage.h>
+#include <mxnet/resource.h>
+#include <mxnet/kvstore.h>
+
+namespace mxnet {
+// finalize the mxnet modules
+void Finalize() {
+  ResourceManager::Get()->Finalize();
+  KVStore::Get()->Finalize();
+  Engine::Get()->WaitForAll();
+  Engine::Get()->Finalize();
+  Storage::Get()->Finalize();
+}
+}  // namespace mxnet

tests/python/unittest/test_kvstore.py

@@ -12,10 +12,6 @@ def init_kv():
     # list
     mx.kv.init(keys, [mx.nd.zeros(shape)] * len(keys))
 
-def stop_kv():
-    """stop kv """
-    mx.kv.stop()
-
 def check_diff_to_scalar(A, x):
     """ assert A == x"""
     assert(np.sum(np.abs((A - x).asnumpy())) == 0)
@@ -30,7 +26,6 @@ def test_single_kv_pair():
     mx.kv.pull(3, out = val)
     check_diff_to_scalar(val, 1)
 
-    stop_kv()
 
 def test_list_kv_pair():
     """list key-value pair push & pull"""
@@ -43,7 +38,6 @@ def test_list_kv_pair():
     for v in val:
         check_diff_to_scalar(v, 4)
 
-    stop_kv()
 
 def test_aggregator():
     """aggregate value on muliple devices"""
@@ -72,7 +66,6 @@ def test_aggregator():
         for v in vv:
             check_diff_to_scalar(v, num_devs * 2.0)
 
-    stop_kv()
 
 def updater(key, recv, local):
     """use updater: +="""
@@ -110,7 +103,6 @@ def test_updater(dev = 'cpu'):
         for v in vv:
             check_diff_to_scalar(v, num_devs * num_push)
 
-    stop_kv()
 
 if __name__ == '__main__':
     test_single_kv_pair()