remove _fork (#353)

* remove _fork

* No rely on latest pr

* remove _fork

* No rely on latest pr

* add unittest for record file

* add comments

* remove _fork

* No rely on latest pr

* add unittest for record file

* add comments

* mark test as serial

* disable record

* fix test case

Author: zhreshold

gluonnlp/data/dataloader.py

@@ -19,13 +19,33 @@
 """DataLoader. An extension of Gluon data loader that allows multi-shard sampling."""
 __all__ = ['ShardedDataLoader']
 
+import sys
 from mxnet import context
-from mxnet.gluon.data.dataloader import DataLoader, _MultiWorkerIter, _as_in_context
+from mxnet.gluon.data.dataloader import DataLoader
+from mxnet.gluon.data.dataloader import _MultiWorkerIter, _as_in_context
+from mxnet.recordio import MXRecordIO
+
+def _recursive_fork_recordio(obj, depth, max_depth=1000):
+    """Recursively find instance of MXRecordIO and reset file handler.
+    This is required for MXRecordIO which holds a C pointer to a opened file after fork.
+    """
+    if depth >= max_depth:
+        return
+    if isinstance(obj, MXRecordIO):
+        obj.close()
+        obj.open()  # re-obtain file hanlder in new process
+    elif (hasattr(obj, '__dict__')):
+        for _, v in obj.__dict__.items():
+            _recursive_fork_recordio(v, depth + 1, max_depth)
 
 
 def worker_loop(dataset, key_queue, data_queue, batchify_fn):
     """Worker loop for multiprocessing DataLoader."""
-    dataset._fork()
+    # re-fork a new recordio handler in new process if applicable
+    limit = sys.getrecursionlimit()
+    max_recursion_depth = min(limit - 5, max(10, limit // 2))
+    _recursive_fork_recordio(dataset, 0, max_recursion_depth)
+
     while True:
         idx, samples = key_queue.get()
         if idx is None:

tests/unittest/train/test_dataloader.py

@@ -1,7 +1,9 @@
 import numpy as np
+import os
 import mxnet as mx
 from gluonnlp.data import FixedBucketSampler, ShardedDataLoader
 from mxnet import gluon
+from mxnet.gluon.utils import download
 import pytest
 
 
@@ -36,3 +38,25 @@ def test_sharded_data_loader():
                     assert mx.test_utils.almost_equal(seqs[j][1].asnumpy(),
                                                       Y[(i*num_shards+j)*2-num_shards:
                                                         (i*num_shards+j+1)*2-num_shards])
+
+def test_sharded_data_loader_record_file():
+    # test record file
+    url_format = 'https://apache-mxnet.s3-accelerate.amazonaws.com/gluon/dataset/pikachu/{}'
+    filename = 'val.rec'
+    idx_filename = 'val.idx'
+    download(url_format.format(filename), path=os.path.join('tests', 'data', filename))
+    download(url_format.format(idx_filename), path=os.path.join('tests', 'data', idx_filename))
+    rec_dataset = gluon.data.vision.ImageRecordDataset(os.path.join('tests', 'data', filename))
+
+    num_workers = 2
+    num_shards = 4
+    X = np.random.uniform(size=(100, 20))
+    Y = np.random.uniform(size=(100,))
+    batch_sampler = FixedBucketSampler(lengths=[X.shape[1]] * X.shape[0],
+                                       batch_size=2,
+                                       num_buckets=1,
+                                       shuffle=False,
+                                       num_shards=num_shards)
+    loader = ShardedDataLoader(rec_dataset, batch_sampler=batch_sampler, num_workers=num_workers)
+    for i, seqs in enumerate(loader):
+        assert len(seqs) == num_shards