Repository: sql-machine-learning/elasticdl

Issue:
Worker optimization with dataset Step 4: support  training-with-evaluation.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/dataset.py b/elasticdl/python/common/dataset.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/common/dataset.py
@@ -0,0 +1,33 @@
+from contextlib import closing
+
+import recordio
+import tensorflow as tf
+
+
+def recordio_dataset(recordio_shards):
+    """
+    Return a RecordIO dataset with records in recordio_shards.
+    recordio_shards is a list of (recordio_file_name, start, end) tuples
+    """
+
+    class _Generator:
+        def __init__(self, recordio_shards):
+            self._shards = recordio_shards
+
+        def gen(self):
+            for s in self._shards:
+                with closing(
+                    recordio.Scanner(s[0], s[1], s[2] - s[1])
+                ) as reader:
+                    while True:
+                        record = reader.record()
+                        if record:
+                            yield record
+                        else:
+                            break
+
+    generator = _Generator(recordio_shards)
+    dataset = tf.data.Dataset.from_generator(
+        generator.gen, (tf.string), (tf.TensorShape([]))
+    )
+    return dataset
diff --git a/elasticdl/python/worker/task_data_service.py b/elasticdl/python/worker/task_data_service.py
--- a/elasticdl/python/worker/task_data_service.py
+++ b/elasticdl/python/worker/task_data_service.py
@@ -6,14 +6,17 @@
 import tensorflow as tf
 
 from elasticdl.proto import elasticdl_pb2
+from elasticdl.python.common.dataset import recordio_dataset
 
 
 class TaskDataService(object):
-    def __init__(self, worker):
+    def __init__(self, worker, training_with_evaluation):
         self._logger = logging.getLogger(__name__)
         self._worker = worker
+        self._training_with_evaluation = training_with_evaluation
         self._lock = threading.Lock()
         self._pending_dataset = True
+        self._pending_eval_tasks = []
         self._reset()
 
     def _reset(self):
@@ -54,6 +57,24 @@ def report_record_done(self, count, err_msg=""):
                 if len(self._pending_tasks_with_counts):
                     self._current_task = self._pending_tasks_with_counts[0][0]
 
+    def get_evaluation_dataset(self):
+        """
+        If there are _pending_eval_tasks, return a RecordIO dataset for
+        an evaluation task and its corresponding model version, task_id.
+        Return None if no _pending_eval_tasks.
+        """
+        if not self._pending_eval_tasks:
+            return None
+        shards = []
+        with self._lock:
+            if self._pending_eval_tasks:
+                task = self._pending_eval_tasks.pop(0)
+                shards.append((task.shard_file_name, task.start, task.end))
+        if shards:
+            return recordio_dataset(shards), task.model_version, task.task_id
+        else:
+            return None
+
     def get_dataset(self):
         """
         Return a RecordIO dataset, or None if no more data.
@@ -90,6 +111,12 @@ def _gen(self):
                     self._logger.info("No more task, stopping")
                 break
             with self._lock:
+                if (
+                    self._training_with_evaluation
+                    and task.type == elasticdl_pb2.EVALUATION
+                ):
+                    self._pending_eval_tasks.append(task)
+                    continue
                 self._record_count += task.end - task.start
                 self._pending_tasks_with_counts.append(
                     (task, self._record_count)
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -70,7 +70,9 @@ def __init__(
             self._stub = elasticdl_pb2_grpc.MasterStub(channel)
         self._max_minibatch_retry_num = max_minibatch_retry_num
         self._model_version = -1
-        self._task_data_service = TaskDataService(self)
+        self._task_data_service = TaskDataService(
+            self, self._job_type == JobType.TRAINING_WITH_EVALUATION
+        )
 
     def get_task(self):
         """
@@ -212,14 +214,16 @@ def _handle_task(self, task):
                 # TODO: Discuss how we separate input_fn for different tasks
                 features, labels = self._input_fn(record_buf)
                 min_model_version = self._process_minibatch(
-                    task, features, labels, min_model_version
+                    task.type, features, labels, min_model_version
                 )
 
-    def _process_minibatch(self, task, features, labels, min_model_version):
+    def _process_minibatch(
+        self, task_type, features, labels, min_model_version
+    ):
         if not self._var_created:
             self._create_variable_and_report(features)
         for _ in range(self._max_minibatch_retry_num):
-            if task.type == elasticdl_pb2.EVALUATION:
+            if task_type == elasticdl_pb2.EVALUATION:
                 if min_model_version == -1:
                     if self._model_version < 0:
                         self.get_model(0, elasticdl_pb2.MINIMUM)
@@ -228,7 +232,7 @@ def _process_minibatch(self, task, features, labels, min_model_version):
                 accepted, _ = self._run_evaluation_task(features, labels)
                 if accepted:
                     break
-            elif task.type == elasticdl_pb2.TRAINING:
+            elif task_type == elasticdl_pb2.TRAINING:
                 # TODO: optimize the logic to avoid unnecessary
                 #       get_model call.
                 self.get_model(
@@ -241,14 +245,14 @@ def _process_minibatch(self, task, features, labels, min_model_version):
                 if accepted:
                     self._logger.info("Loss is %f" % loss.numpy())
                     break
-            elif task.type == elasticdl_pb2.PREDICTION:
+            elif task_type == elasticdl_pb2.PREDICTION:
                 if self._model_version != min_model_version:
                     self.get_model(min_model_version, elasticdl_pb2.FIXED)
                 accepted = self._run_prediction_task(features)
                 if accepted:
                     break
             else:
-                raise RuntimeError("Unrecognized task type, %s" % task.type)
+                raise RuntimeError("Unrecognized task type, %s" % task_type)
         else:
             # Worker got stuck, fail the task.
             # TODO: stop the worker if it fails to make any
@@ -256,6 +260,44 @@ def _process_minibatch(self, task, features, labels, min_model_version):
             raise RuntimeError("Worker got stuck")
         return min_model_version
 
+    def _process_eval_task_if_needed(self):
+        """
+        Check if there are evaluation tasks and process the tasks if any.
+        """
+        eval_info = self._task_data_service.get_evaluation_dataset()
+        if not eval_info:
+            return
+        eval_dataset = eval_info[0]
+        model_version = eval_info[1]
+        task_id = eval_info[2]
+        eval_dataset = self._dataset_fn(eval_dataset)
+        eval_dataset = eval_dataset.batch(self._minibatch_size).prefetch(1)
+        err_msg = ""
+        for data in eval_dataset:
+            data_err_msg = self._process_minibatch_and_report(
+                data[0], data[1], elasticdl_pb2.EVALUATION, model_version
+            )
+            if data_err_msg:
+                err_msg = data_err_msg
+                break
+        del eval_dataset
+        self.report_task_result(task_id, err_msg)
+
+    def _process_minibatch_and_report(
+        self, features, labels, task_type, model_version
+    ):
+        err_msg = ""
+        try:
+            self._process_minibatch(task_type, features, labels, model_version)
+        except RuntimeError as err:
+            err_msg = str(err)
+            traceback.print_exc()
+        except Exception as ex:
+            err_msg = str(ex)
+            traceback.print_exc()
+            raise ex
+        return err_msg
+
     def run_with_dataset(self):
         while True:
             dataset = self._task_data_service.get_dataset()
@@ -264,28 +306,30 @@ def run_with_dataset(self):
             dataset = self._dataset_fn(dataset)
             dataset = dataset.batch(self._minibatch_size).prefetch(1)
             for d in dataset:
+                if self._job_type == JobType.TRAINING_WITH_EVALUATION:
+                    self._process_eval_task_if_needed()
                 task = self._task_data_service.get_current_task()
-                err_msg = ""
-                try:
-                    self._process_minibatch(
-                        task, d[0], d[1], task.model_version
-                    )
-                except RuntimeError as err:
-                    err_msg = str(err)
-                    traceback.print_exc()
-                except Exception as ex:
-                    err_msg = str(ex)
-                    traceback.print_exc()
-                    raise ex
+                err_msg = self._process_minibatch_and_report(
+                    d[0], d[1], task.type, task.model_version
+                )
                 self._task_data_service.report_record_done(
                     self._minibatch_size, err_msg
                 )
+            del dataset
+            # New evaluation tasks may be created after this worker's
+            # training tasks are done, as other workers' may still
+            # have pending training tasks.
+            if self._job_type == JobType.TRAINING_WITH_EVALUATION:
+                self._process_eval_task_if_needed()
 
     def run(self):
         """
         Fetches task from master and performs training or evaluation.
         """
-        if self._job_type == JobType.TRAINING_ONLY:
+        if (
+            self._job_type == JobType.TRAINING_ONLY
+            or self._job_type == JobType.TRAINING_WITH_EVALUATION
+        ):
             return self.run_with_dataset()
         while True:
             task = self.get_task()


Test patch: