Fix #13521 (#13537)

zhreshold · web-flow · commit f6b4665995f8 · 2018-12-05T13:39:17.000-08:00
* fix pool release

* fix
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()
