Repository: sql-machine-learning/elasticdl

Issue:
Worker supports get_model_frequency argument in async-SGD
According to [design doc](https://github.com/sql-machine-learning/elasticdl/blob/develop/elasticdl/doc/async_sgd_design.md#change-in-worker)

To support SSP strategy, the worker pulls the model from PS in every staleness_threshold minibatch step. Also, the worker needs to update the local model with the computed gradients. model pull/updates do not include embedding variables, as we directly access the embedding vectors in the embedding service.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/args.py b/elasticdl/python/common/args.py
--- a/elasticdl/python/common/args.py
+++ b/elasticdl/python/common/args.py
@@ -210,7 +210,7 @@ def add_train_params(parser):
         "--grads_to_wait",
         type=int,
         help="Number of gradients to wait before updating model",
-        default=2,
+        default=1,
     )
     parser.add_argument(
         "--training_data_dir",
@@ -465,6 +465,16 @@ def parse_master_args(master_args=None):
             "checkpoint_filename_for_init is required for running "
             "prediction job"
         )
+    if not args.use_async and args.get_model_steps > 1:
+        args.get_model_steps = 1
+        logger.warning(
+            "get_model_steps is set to 1 when using synchronous SGD."
+        )
+    if args.use_async and args.grads_to_wait > 1:
+        args.grads_to_wait = 1
+        logger.warning(
+            "grads_to_wait is set to 1 when using asynchronous SGD."
+        )
 
     return args
 
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
@@ -104,6 +104,9 @@ def __init__(
             self, self._job_type == JobType.TRAINING_WITH_EVALUATION
         )
         self._get_model_steps = get_model_steps
+        if self._get_model_steps > 1:
+            self._opt = self._opt_fn()
+            self._non_embed_grads = None
 
     # TODO: Multiple tests are currently using this function to initialize
     # self._model, where the initialization should be done via constructor.
@@ -138,6 +141,14 @@ def _reset_embedding(self):
         for layer in self._embedding_layers:
             layer.reset()
 
+    def _update_local_model(self):
+        if not self._non_embed_grads:
+            return
+        self._opt.apply_gradients(
+            zip(self._non_embed_grads, self._non_embed_vars)
+        )
+        self._non_embed_grads = None
+
     def get_task(self):
         """
         get task from master
@@ -362,6 +373,9 @@ def predict_process(self, features):
     def _run_training_task(self, features, labels):
         loss, grads = self.training_process(features, labels)
         accepted, min_model_version = self.report_gradient(grads)
+        if accepted and self._get_model_steps > 1:
+            non_embed_vars_n = len(self._non_embed_vars)
+            self._non_embed_grads = grads[:non_embed_vars_n]
         self._reset_embedding()
         return accepted, min_model_version, loss
 
@@ -375,7 +389,12 @@ def _run_prediction_task(self, features):
         return self.report_prediction_outputs(predictions)
 
     def _process_minibatch(
-        self, task_type, features, labels, min_model_version
+        self,
+        task_type,
+        features,
+        labels,
+        min_model_version,
+        train_with_local_model=False,
     ):
         if self._need_embedding_layer_check or not self._var_created:
             self._run_model_call_before_training(features)
@@ -392,10 +411,11 @@ def _process_minibatch(
             elif task_type == elasticdl_pb2.TRAINING:
                 # TODO: optimize the logic to avoid unnecessary
                 #       get_model call.
-                self.get_model(
-                    max(self._model_version, min_model_version),
-                    elasticdl_pb2.MINIMUM,
-                )
+                if not train_with_local_model:
+                    self.get_model(
+                        max(self._model_version, min_model_version),
+                        elasticdl_pb2.MINIMUM,
+                    )
                 accepted, min_model_version, loss = self._run_training_task(
                     features, labels
                 )
@@ -420,10 +440,14 @@ def _process_minibatch(
     def _process_eval_task_if_needed(self):
         """
         Check if there are evaluation tasks and process the tasks if any.
+
+        Return:
+            A python bool indicating whether worker processed some evaluation
+            tasks.
         """
         eval_info = self._task_data_service.get_evaluation_dataset()
         if not eval_info:
-            return
+            return False
         (eval_dataset, model_version, task_id) = eval_info
         eval_dataset = self._dataset_fn(eval_dataset, Mode.EVALUATION)
         eval_dataset = eval_dataset.batch(self._minibatch_size).prefetch(1)
@@ -437,9 +461,14 @@ def _process_eval_task_if_needed(self):
                 break
         del eval_dataset
         self.report_task_result(task_id, err_msg)
+        return True
 
     def _process_minibatch_and_report(
-        self, dataset_batch, task_type, model_version
+        self,
+        dataset_batch,
+        task_type,
+        model_version,
+        train_with_local_model=False,
     ):
         err_msg = ""
         try:
@@ -449,7 +478,13 @@ def _process_minibatch_and_report(
             else:
                 features = dataset_batch[0]
                 labels = dataset_batch[1]
-            self._process_minibatch(task_type, features, labels, model_version)
+            self._process_minibatch(
+                task_type,
+                features,
+                labels,
+                model_version,
+                train_with_local_model,
+            )
         except RuntimeError as err:
             err_msg = str(err)
             traceback.print_exc()
@@ -469,6 +504,23 @@ def run(self):
             mode = Mode.EVALUATION
         else:
             mode = Mode.TRAINING
+
+        # The worker needs to get model from PS if
+        # `train_with_local_model=False`. This happens when:
+        #     processing first minibatch
+        #     last minibatch is evaluation task
+        #     last minibatch is training task and failed
+        #     local_update_count >= worker._get_model_steps
+        # Otherwise, worker trains with local model, i.e.
+        # `train_with_local_model=True`
+        train_with_local_model = False
+
+        # Initialize `local_update_count=get_model_steps` in order to set
+        # `train_with_local_model` to False inside for-loop for the first
+        # minibatch.
+        local_update_count = self._get_model_steps
+        last_training_minibatch_failed = False
+        last_minibatch_do_evaluation = False
         while True:
             dataset = self._task_data_service.get_dataset()
             if not dataset:
@@ -477,11 +529,34 @@ def run(self):
             dataset = dataset.batch(self._minibatch_size).prefetch(1)
             for dataset_batch in dataset:
                 if self._job_type == JobType.TRAINING_WITH_EVALUATION:
-                    self._process_eval_task_if_needed()
+                    if self._process_eval_task_if_needed():
+                        last_minibatch_do_evaluation = True
                 task = self._task_data_service.get_current_task()
+                if task.type == elasticdl_pb2.TRAINING:
+                    if (
+                        last_minibatch_do_evaluation
+                        or last_training_minibatch_failed
+                        or local_update_count >= self._get_model_steps
+                    ):
+                        local_update_count = 0
+                        train_with_local_model = False
+                    else:
+                        train_with_local_model = True
                 err_msg = self._process_minibatch_and_report(
-                    dataset_batch, task.type, task.model_version
+                    dataset_batch,
+                    task.type,
+                    task.model_version,
+                    train_with_local_model,
                 )
+                if task.type == elasticdl_pb2.TRAINING:
+                    last_minibatch_do_evaluation = False
+                    local_update_count += 1
+                    if err_msg:
+                        last_training_minibatch_failed = True
+                    else:
+                        last_training_minibatch_failed = False
+                        if local_update_count < self._get_model_steps:
+                            self._update_local_model()
                 self._task_data_service.report_record_done(
                     self._minibatch_size, err_msg
                 )
@@ -490,4 +565,5 @@ def run(self):
             # 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()
+                if self._process_eval_task_if_needed():
+                    last_minibatch_do_evaluation = True


Test patch:
diff --git a/elasticdl/python/tests/in_process_master.py b/elasticdl/python/tests/in_process_master.py
--- a/elasticdl/python/tests/in_process_master.py
+++ b/elasticdl/python/tests/in_process_master.py
@@ -1,9 +1,11 @@
 """In process master for unittests"""
+from elasticdl.python.tests import test_call_back
 
 
 class InProcessMaster(object):
-    def __init__(self, master):
+    def __init__(self, master, callbacks=[]):
         self._m = master
+        self._callbacks = callbacks
 
     def GetTask(self, req):
         return self._m.GetTask(req, None)
@@ -15,9 +17,17 @@ def ReportVariable(self, req):
         return self._m.ReportVariable(req, None)
 
     def ReportGradient(self, req):
+        for callback in self._callbacks:
+            if test_call_back.ON_REPORT_GRADIENT_BEGIN in callback.call_times:
+                callback()
         return self._m.ReportGradient(req, None)
 
     def ReportEvaluationMetrics(self, req):
+        for callback in self._callbacks:
+            if test_call_back.ON_REPORT_EVALUATION_METRICS_BEGIN in (
+                callback.call_times
+            ):
+                callback()
         return self._m.ReportEvaluationMetrics(req, None)
 
     def ReportTaskResult(self, req):
diff --git a/elasticdl/python/tests/test_call_back.py b/elasticdl/python/tests/test_call_back.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/tests/test_call_back.py
@@ -0,0 +1,19 @@
+"""Callback for unittests."""
+from abc import ABC, abstractmethod
+
+ON_REPORT_GRADIENT_BEGIN = "on_report_gradient_begin"
+ON_REPORT_EVALUATION_METRICS_BEGIN = "on_report_evaluation_metrics_begin"
+
+
+class BaseCallback(ABC):
+    """Baseclass of callbacks used for unittests."""
+
+    def __init__(self, master, worker, unittest_inst, call_times):
+        self._master = master
+        self._worker = worker
+        self._unittest_inst = unittest_inst
+        self.call_times = call_times
+
+    @abstractmethod
+    def __call__(self):
+        pass
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -17,8 +17,10 @@
 from elasticdl.python.master.evaluation_service import EvaluationService
 from elasticdl.python.master.servicer import MasterServicer
 from elasticdl.python.master.task_dispatcher import _TaskDispatcher
+from elasticdl.python.tests import test_call_back
 from elasticdl.python.tests.in_process_master import InProcessMaster
 from elasticdl.python.tests.mock_kv_store import MockKvStore
+from elasticdl.python.tests.test_call_back import BaseCallback
 from elasticdl.python.worker.worker import Worker
 
 _model_zoo_path = os.path.dirname(os.path.realpath(__file__))
@@ -59,27 +61,80 @@ def create_recordio_file(size):
     return temp_file.name
 
 
+class CheckRetrainCallback(BaseCallback):
+    """Checks the retry functionality of workers.
+
+    When workers report Gradient or evaluation metrics, this callback
+    adds 1 to master's model version. The master rejects the report request and
+    workers retry.
+    """
+
+    def __init__(self, master, worker, unittest_inst):
+        super(CheckRetrainCallback, self).__init__(
+            master,
+            worker,
+            unittest_inst,
+            call_times=[
+                test_call_back.ON_REPORT_GRADIENT_BEGIN,
+                test_call_back.ON_REPORT_EVALUATION_METRICS_BEGIN,
+            ],
+        )
+
+    def __call__(self):
+        if 2 < self._master._version < 80:
+            self._master._version += 1
+
+
+class CheckWorkerModelCallback(BaseCallback):
+    """Checks worker model parameters.
+
+    Before master updating model parameters, master and workers should have
+    same model parameters if `master._grad_n=0`.
+    """
+
+    def __init__(self, master, worker, unittest_inst):
+        super(CheckWorkerModelCallback, self).__init__(
+            master,
+            worker,
+            unittest_inst,
+            call_times=[test_call_back.ON_REPORT_GRADIENT_BEGIN],
+        )
+
+    def __call__(self):
+        if self._master._grad_n:
+            return
+        self._unittest_inst.assertTrue(
+            len(self._worker._non_embed_vars) == len(self._master._model)
+        )
+        for var in self._worker._non_embed_vars:
+            self._unittest_inst.assertTrue(
+                np.isclose(
+                    var.numpy(), self._master._model[var.name].numpy()
+                ).all()
+            )
+
+
 class WorkerTest(unittest.TestCase):
-    def distributed_train_and_evaluate(self, training=True):
+    # TODO (yunjian.lmh): Function `distributed_train_and_evaluate` is used in
+    # `worker_test` and `example_test`, which can be extracted as a reusable
+    # function.
+    def distributed_train_and_evaluate(
+        self,
+        training=True,
+        callback_classes=[],
+        use_async=False,
+        grads_to_wait=2,
+        get_model_steps=1,
+    ):
         """
         Run distributed training and evaluation with a local master.
         grpc calls are mocked by local master call.
         """
 
-        class _Master(InProcessMaster):
-            def ReportGradient(self, req):
-                if 2 < self._m._version < 80:
-                    # For testing of retrain when gradient not accepted.
-                    # Increase master version to reject the gradient.
-                    self._m._version += 1
-                return self._m.ReportGradient(req, None)
-
-            def ReportEvaluationMetrics(self, req):
-                if 2 < self._m._version < 80:
-                    # Testing of evaluation retries. Increase the master
-                    # version so the evaluation metrics will not be accepted.
-                    self._m._version += 1
-                return self._m.ReportEvaluationMetrics(req, None)
+        if use_async and grads_to_wait > 1:
+            raise ValueError(
+                "grads_to_wait should be 1 when using asynchronous SGD."
+            )
 
         job_type = (
             JobType.TRAINING_ONLY if training else JobType.EVALUATION_ONLY
@@ -92,6 +147,7 @@ def ReportEvaluationMetrics(self, req):
             _model_zoo_path,
             model_def="test_module.custom_model",
             channel=None,
+            get_model_steps=get_model_steps,
         )
 
         shards = {create_recordio_file(128): (0, 128)}
@@ -116,7 +172,7 @@ def ReportEvaluationMetrics(self, req):
         else:
             evaluation_service = None
         master = MasterServicer(
-            2,
+            grads_to_wait,
             batch_size,
             worker._opt_fn(),
             task_d,
@@ -124,8 +180,13 @@ def ReportEvaluationMetrics(self, req):
             checkpoint_filename_for_init="",
             checkpoint_service=None,
             evaluation_service=evaluation_service,
+            use_async=use_async,
         )
-        worker._stub = _Master(master)
+        callbacks = [
+            callback_class(master, worker, self)
+            for callback_class in callback_classes
+        ]
+        worker._stub = InProcessMaster(master, callbacks)
 
         for var in worker._model.trainable_variables:
             master.set_model_var(var.name, var.numpy())
@@ -139,10 +200,24 @@ def ReportEvaluationMetrics(self, req):
         self.assertTrue(not task.shard_name)
 
     def test_distributed_train_tf_example(self):
-        self.distributed_train_and_evaluate(training=True)
+        self.distributed_train_and_evaluate(
+            training=True, callback_classes=[CheckRetrainCallback]
+        )
 
     def test_distributed_evaluate_tf_example(self):
-        self.distributed_train_and_evaluate(training=False)
+        self.distributed_train_and_evaluate(
+            training=False, callback_classes=[CheckRetrainCallback]
+        )
+
+    def test_distributed_train_get_model_steps(self):
+        # When `use_async=True`, `grads_to_wait` should be set to 1
+        self.distributed_train_and_evaluate(
+            training=True,
+            callback_classes=[CheckWorkerModelCallback],
+            use_async=True,
+            get_model_steps=4,
+            grads_to_wait=1,
+        )
 
     def test_embedding_layer(self):
         worker = Worker(