Develop a reader to read records from recordio files (#2502)

* Fix the example of TensorFlow 1.x

* Develop a reader to read records from recordio files

* The master does not add the worker host if it is in the next group

* Add a docstring to find which block the shard is in

Author: workingloong

elasticai_api/io/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2021 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

elasticai_api/io/recordio_reader.py

@@ -0,0 +1,66 @@
+# Copyright 2020 The ElasticDL Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections
+from contextlib import closing
+
+import recordio
+
+RecordBlock = collections.namedtuple("RecordBlock", ("name", "start", "end"))
+
+
+class RecordIOReader(object):
+    def __init__(self, recordio_files):
+        self._init_file_record_count(recordio_files)
+
+    def _init_file_record_count(self, recordio_files):
+        self._data_blocks = []
+        start = 0
+        for file_path in recordio_files:
+            with closing(recordio.Index(file_path)) as rio:
+                num_records = rio.num_records()
+                end = start + num_records
+                self._data_blocks.append(RecordBlock(file_path, start, end))
+                start = end
+
+    def read_records(self, start, end):
+        target_files = self._get_record_file(start, end)
+        for file_path, start, count in target_files:
+            with closing(recordio.Scanner(file_path, start, count)) as reader:
+                while True:
+                    record = reader.record()
+                    if record:
+                        yield record
+                    else:
+                        break
+
+    def _get_record_file(self, start, end):
+        """The block ranges in data_blocks are sorted in
+        increasing order. For example,
+        blocks are [[0,100),[100, 200),[200,300)]. So we
+        can find which block the shard is in by sequential search.
+        """
+        target_files = []
+        for block in self._data_blocks:
+            if start < block.end:
+                if end < block.end:
+                    target_files.append(
+                        (block.name, start - block.start, end - start)
+                    )
+                    break
+                else:
+                    target_files.append(
+                        (block.name, start - block.start, block.end - start)
+                    )
+                    start = block.end
+        return target_files

elasticai_api/tensorflow/controller.py

@@ -37,11 +37,7 @@
 
 
 def create_elastic_controller(
-    batch_size,
-    num_epochs=None,
-    dataset_size=None,
-    shuffle=False,
-    training_data=None,
+    batch_size, num_epochs=None, dataset_size=None, shuffle=False,
 ):
     """Create an elastic AllReduce controller with data shard service.
     Users can use the `controller.data_shard_service` to get data
@@ -64,7 +60,6 @@ def create_elastic_controller(
         num_epochs=num_epochs,
         dataset_size=dataset_size,
         shuffle=shuffle,
-        training_data=training_data,
     )
     if _IS_TF2:
         controller = TensorFlowV2AllReduceController(

model_zoo/mnist/mnist_train_tfv1.py

@@ -12,11 +12,11 @@
 # limitations under the License.
 
 import argparse
-from contextlib import closing
+import os
 
-import recordio
 import tensorflow as tf
 
+from elasticai_api.io.recordio_reader import RecordIOReader
 from elasticai_api.tensorflow.controller import create_elastic_controller
 from elasticai_api.tensorflow.optimizer import (
     AdjustBackwardPassesPerStepHook,
@@ -27,31 +27,35 @@
 layers = tf.layers
 
 
-def get_dataset_gen(data_shard_service):
+def get_dataset_gen(data_shard_service, data_reader):
     def gen():
         while True:
             shard = data_shard_service.fetch_shard()
             if not shard:
                 raise StopIteration("No data")
-            with closing(
-                recordio.Scanner(
-                    shard.name, shard.start, shard.end - shard.start,
-                )
-            ) as reader:
-                for i in range(shard.start, shard.end):
-                    record = reader.record()
-                    if record:
-                        yield record
+            count = 0
+            for record in data_reader.read_records(shard.start, shard.end):
+                count += 1
+                yield record
 
     return gen
 
 
-def create_dataset(data_shard_service):
-    gen = get_dataset_gen(data_shard_service)
+def create_dataset(data_shard_service, training_data_dir):
+    data_files = _get_data_files(training_data_dir)
+    data_reader = RecordIOReader(data_files)
+    gen = get_dataset_gen(data_shard_service, data_reader)
     dataset = tf.data.Dataset.from_generator(gen, tf.string)
     return dataset
 
 
+def _get_data_files(data_dir):
+    data_files = []
+    for filename in os.listdir(data_dir):
+        data_files.append(os.path.join(data_dir, filename))
+    return data_files
+
+
 def conv_model(feature, target, mode):
     """2-layer convolution model."""
     # Convert the target to a one-hot tensor of shape (batch_size, 10) and
@@ -109,9 +113,11 @@ def train(args):
     allreduce_controller = create_elastic_controller(
         batch_size=args.batch_size,
         num_epochs=args.num_epochs,
-        training_data=args.training_data,
+        dataset_size=50000,
+    )
+    dataset = create_dataset(
+        allreduce_controller.data_shard_service, args.training_data
     )
-    dataset = create_dataset(allreduce_controller.data_shard_service)
     dataset = feed(dataset)
     dataset = dataset.batch(args.batch_size).prefetch(1)
     dataset_it = dataset.make_one_shot_iterator()