Repository: sql-machine-learning/elasticdl

Issue:
Fail to fetch shard using multi-process in Python.
```python
class ImageDataset(IterableDataset):
    def __init__(self, data_shard_service, shuffle=False):
        self.data_shard_service = data_shard_service
        self._shuffle = shuffle

    def __iter__(self):
        while True:
            index = self.data_shard_service.fetch_shard()
            image, label = read_images(index)
            yield image, label

dataset = ImageDataset(
      allreduce_controller.data_shard_service, shuffle=True
  )

data_loader = DataLoader(
      dataset=dataset,
      batch_size=args.batch_size,
      num_workers=2
  )
```

![image](https://user-images.githubusercontent.com/18071380/103133964-8ef96c80-46e8-11eb-99a8-ce16197fcadd.png)




Discussion:

Resolving patch:
diff --git a/elasticai_api/common/data_shard_service.py b/elasticai_api/common/data_shard_service.py
--- a/elasticai_api/common/data_shard_service.py
+++ b/elasticai_api/common/data_shard_service.py
@@ -11,8 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import random
 import threading
+import time
 from collections import deque
+from multiprocessing import SimpleQueue
 
 from elasticai_api.common.constants import TaskExecCounterKey
 from elasticai_api.common.master_client import build_master_client
@@ -134,3 +137,51 @@ def fetch_shard(self):
             return None
 
         return task.shard
+
+
+class RecordIndexService(DataShardService):
+    def __init__(
+        self,
+        master_client,
+        batch_size,
+        num_epochs=None,
+        dataset_size=None,
+        task_type=elasticai_api_pb2.TRAINING,
+        shuffle=False,
+    ):
+        super(RecordIndexService, self).__init__(
+            master_client, batch_size, num_epochs, dataset_size, task_type
+        )
+        self._shuffle = shuffle
+        self._shard_queue = SimpleQueue()
+        threading.Thread(
+            target=self._get_shard_indices,
+            name="fetch_shard_indices",
+            daemon=True,
+        ).start()
+
+    def _get_shard_indices(self):
+        while True:
+            if self._shard_queue.empty():
+                task = self.get_task(self._task_type)
+                if not task.shard or task.type != self._task_type:
+                    break
+                ids = list(range(task.shard.start, task.shard.end))
+                if self._shuffle:
+                    random.shuffle(ids)
+                for i in ids:
+                    self._shard_queue.put(i)
+            else:
+                time.sleep(1)
+
+    def fetch_record_index(self):
+        """Fetch an index of the record. The function get an index
+        from a queue because there may be multiple sub-process to call
+        the function.
+        """
+        for _ in range(30):
+            if not self._shard_queue.empty():
+                return self._shard_queue.get()
+            else:
+                time.sleep(1)
+        raise StopIteration
diff --git a/elasticai_api/pytorch/controller.py b/elasticai_api/pytorch/controller.py
--- a/elasticai_api/pytorch/controller.py
+++ b/elasticai_api/pytorch/controller.py
@@ -21,7 +21,7 @@
     AllReduceController,
 )
 from elasticai_api.common.constants import WorkerEnv
-from elasticai_api.common.data_shard_service import DataShardService
+from elasticai_api.common.data_shard_service import RecordIndexService
 from elasticai_api.common.master_client import build_master_client
 from elasticai_api.util.log_utils import default_logger as logger
 
@@ -38,15 +38,14 @@
     hvd = None
 
 
-def create_elastic_controller(batch_size, num_epochs=None, dataset_size=None):
-    """Create an elastic AllReduce controller with data shard service.
+def create_elastic_controller(
+    batch_size, num_epochs=None, dataset_size=None, shuffle=False
+):
+    """Create an elastic AllReduce controller with record index service.
     Users can use the `controller.data_shard_service` to get data
     shards like:
     ```python
-    while True:
-        shard = controller.data_shard_service.fetch_shard()
-        for i in range(shard.start, shard.end):
-            yield i
+    index = controller.data_shard_service.fetch_record_index()
     ```
 
     Users also can use the controller to do an elastic training.
@@ -80,11 +79,17 @@ def train_one_batch(model, optimizer, data, target):
         dataset_size: The total size of dataset.
     """
     master_client = build_master_client()
-    data_shard_service = DataShardService(
-        master_client, batch_size, num_epochs, dataset_size
+    record_index_service = RecordIndexService(
+        master_client=master_client,
+        batch_size=batch_size,
+        num_epochs=num_epochs,
+        dataset_size=dataset_size,
+        shuffle=shuffle,
     )
 
-    controller = PyTorchAllReduceController(master_client, data_shard_service)
+    controller = PyTorchAllReduceController(
+        master_client, record_index_service
+    )
     controller.init_horovod_locally()
     return controller
 


Test patch: