Fix apache#13521 (apache#13537)

* fix pool release * fix
zhaoyao73 · Dec 5, 2018 · f6b4665 · f6b4665
1 parent f6f8401
commit f6b4665
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/python/mxnet/gluon/data/dataloader.py b/python/mxnet/gluon/data/dataloader.py
@@ -556,3 +556,9 @@ def same_process_iter():
 
     def __len__(self):
         return len(self._batch_sampler)
+
+    def __del__(self):
+        if self._worker_pool:
+            # manually terminate due to a bug that pool is not automatically terminated on linux
+            assert isinstance(self._worker_pool, multiprocessing.pool.Pool)
+            self._worker_pool.terminate()
diff --git a/tests/python/unittest/test_gluon_data.py b/tests/python/unittest/test_gluon_data.py
@@ -244,6 +244,17 @@ def test_multi_worker_forked_data_loader():
         for i, data in enumerate(loader):
             pass
 
+@with_seed()
+def test_multi_worker_dataloader_release_pool():
+    # will trigger too many open file if pool is not released properly
+    for _ in range(100):
+        A = np.random.rand(999, 2000)
+        D = mx.gluon.data.DataLoader(A, batch_size=8, num_workers=8)
+        the_iter = iter(D)
+        next(the_iter)
+        del the_iter
+        del D
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()