Worker can report training data to the master if using RecordIO (#2494)

* Worker can report training data to the master if using recordio

* Pre-commit

* Fix by comments

Author: workingloong

elasticai_api/common/data_shard_service.py

@@ -54,6 +54,7 @@ def __init__(
         shuffle_shards=False,
         task_type=elasticai_api_pb2.TRAINING,
         num_minibatches_per_shard=0,
+        training_data=None,
     ):
         self._mc = master_client
         self._batch_size = batch_size
@@ -63,6 +64,7 @@ def __init__(
         self._shuffle_shards = shuffle_shards
         self._task_type = task_type
         self._num_minibatches_per_shard = num_minibatches_per_shard
+        self._training_data = training_data
         self._lock = threading.Lock()
         self._failed_record_count = 0
         self._reported_record_count = 0
@@ -79,6 +81,7 @@ def _report_training_params(self):
                 shuffle=self._shuffle,
                 shuffle_shards=self._shuffle_shards,
                 num_minibatches_per_shard=self._num_minibatches_per_shard,
+                training_data=self._training_data,
             )
 
     def get_minibatch_count_per_epoch(self):
@@ -167,6 +170,7 @@ def __init__(
         dataset_size=None,
         task_type=elasticai_api_pb2.TRAINING,
         shuffle=False,
+        training_data=None,
     ):
         super(RecordIndexService, self).__init__(
             master_client=master_client,
@@ -175,6 +179,7 @@ def __init__(
             dataset_size=dataset_size,
             shuffle=shuffle,
             task_type=task_type,
+            training_data=training_data,
         )
         self._shard_queue = SimpleQueue()
         threading.Thread(

elasticai_api/common/master_client.py

@@ -118,6 +118,7 @@ def report_training_params(
         shuffle=False,
         shuffle_shards=False,
         num_minibatches_per_shard=0,
+        training_data=None,
     ):
         request = elasticai_api_pb2.ReportTrainingParamsRequest()
         request.batch_size = batch_size
@@ -127,5 +128,7 @@ def report_training_params(
             request.num_epochs = num_epochs
         if dataset_size is not None:
             request.dataset_size = dataset_size
+        if training_data is not None:
+            request.training_data = training_data
         request.num_minibatches_per_shard = num_minibatches_per_shard
         return self._stub.report_training_params(request)

elasticai_api/proto/elasticai_api.proto

@@ -60,6 +60,7 @@ message ReportTrainingParamsRequest {
   bool shuffle = 4;
   bool shuffle_shards = 5;
   int32 num_minibatches_per_shard = 6;
+  string training_data = 7;
 }
 
 message GetTaskRequest {

elasticdl/python/master/servicer.py

@@ -174,6 +174,7 @@ def report_training_params(self, request, _):
             request.shuffle,
             request.shuffle_shards,
             request.num_minibatches_per_shard,
+            request.training_data,
         )
         return empty_pb2.Empty()
 

elasticdl/python/master/task_manager.py

@@ -13,6 +13,7 @@
 
 """TaskQueue Implementation"""
 
+import os
 import random
 import threading
 import time
@@ -229,17 +230,20 @@ def set_training_params(
         shuffle,
         shuffle_shards,
         num_minibatches_per_shard,
+        training_data=None,
     ):
         logger.info(
             "Set training parameters: "
             "batch_size={}, num_epochs={}, dataset_size={}, shuffle={}, "
-            "shuffle_shards={}, num_minibatches_per_shard={}".format(
+            "shuffle_shards={}, num_minibatches_per_shard={}, "
+            "training_data={}".format(
                 batch_size,
                 num_epochs,
                 dataset_size,
                 shuffle,
                 shuffle_shards,
                 num_minibatches_per_shard,
+                training_data,
             )
         )
 
@@ -270,9 +274,20 @@ def set_training_params(
             self._dataset_size = (
                 dataset_size if dataset_size > 0 else self._dataset_size
             )
-            self._training_shards = self._create_shards_by_dataset_size(
-                dataset_size
-            )
+            if (
+                not dataset_size
+                and training_data
+                and os.path.isdir(training_data)
+            ):
+                # The training_data is a directory only using RecordIO.
+                self._create_training_tasks(training_data, None)
+            elif dataset_size > 0:
+                self._training_shards = self._create_shards_by_dataset_size(
+                    dataset_size
+                )
+            else:
+                logger.warning("No data to create shards")
+
             if self._training_shards:
                 logger.info("Starting epoch %d", self._epoch)
                 self.create_tasks(elasticai_api_pb2.TRAINING)

elasticdl/python/tests/task_manager_test.py

@@ -11,12 +11,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import tempfile
 import threading
 import time
 import unittest
 
 from elasticai_api.proto import elasticai_api_pb2
-from elasticdl.python.tests.test_utils import create_task_manager
+from elasticdl.python.tests.test_utils import (
+    DatasetName,
+    create_recordio_file,
+    create_task_manager,
+)
 
 
 class TaskManagerTest(unittest.TestCase):
@@ -176,6 +181,20 @@ def test_set_training_params(self):
         )
         self.assertEqual(len(task_manager._todo), 4)
 
+        task_manager = create_task_manager([], [])
+        num_records = 128
+        with tempfile.TemporaryDirectory() as temp_dir_name:
+            shard_name = create_recordio_file(
+                num_records, DatasetName.TEST_MODULE, 1, temp_dir=temp_dir_name
+            )
+
+            task_manager.set_training_params(
+                1, 1, 0, False, False, 3, temp_dir_name
+            )
+            self.assertEqual(
+                task_manager._training_shards, [(shard_name, 0, num_records)]
+            )
+
 
 if __name__ == "__main__":
     unittest.main()