Repository: sql-machine-learning/elasticdl

Issue:
Avoid hang in Allreduce training
If data is distributed unevenly among workers, some workers may have consumed all their data before other workers. This will cause `allreduce` op hang.

One solution is to quit the workers that have already finished training, so FTLib can get a new consensus.



Discussion:
worker 0 has no tasks

```
[2020-07-15 06:29:15,255] [INFO] [task_data_service.py:220:_gen] No tasks for now, maybe more later
[2020-07-15 06:29:20,302] [INFO] [task_data_service.py:220:_gen] No tasks for now, maybe more later
[2020-07-15 06:29:25,349] [INFO] [task_data_service.py:220:_gen] No tasks for now, maybe more later
[2020-07-15 06:29:30,395] [INFO] [task_data_service.py:220:_gen] No tasks for now, maybe more later
[2020-07-15 06:29:35,441] [INFO] [task_data_service.py:220:_gen] No tasks for now, maybe more later
```

worker 1 fails as worker 0 does not have tasks to do

```
ERROR:root:[/pytorch/third_party/gloo/gloo/transport/tcp/unbound_buffer.cc:84] Timed out waiting 60000ms for recv operation to complete
ERROR:root:Application timeout caused pair closure
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/ftlib/impl.py", line 364, in func
    result = ops(*argc, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/ftlib/commlib/pytorch/impl.py", line 72, in barrier
    dist.barrier()
  File "/usr/local/lib/python3.6/dist-packages/torch/distributed/distributed_c10d.py", line 1496, in barrier
    work.wait()
RuntimeError: Application timeout caused pair closure
[2020-07-15 06:30:15,492] [WARNING] [worker.py:771:_broadcast_model_params] Failed to perform barrier operation
[2020-07-15 06:30:15,492] [WARNING] [worker.py:817:_collect_gradients_with_allreduce_robust] Failed to broadcast model parameters
[2020-07-15 06:30:15,492] [WARNING] [worker.py:840:_run_training_task] Failed to perform allreduce operation onthe gradients. Retrying...
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort
[2020-07-15 06:32:55,593] [WARNING] [worker.py:771:_broadcast_model_params] Failed to perform barrier operation
[2020-07-15 06:32:55,593] [WARNING] [worker.py:817:_collect_gradients_with_allreduce_robust] Failed to broadcast model parameters
[2020-07-15 06:32:55,593] [WARNING] [worker.py:840:_run_training_task] Failed to perform allreduce operation onthe gradients. Retrying...
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort
[2020-07-15 06:35:14,858] [WARNING] [worker.py:771:_broadcast_model_params] Failed to perform barrier operation
[2020-07-15 06:35:14,859] [WARNING] [worker.py:817:_collect_gradients_with_allreduce_robust] Failed to broadcast model parameters
[2020-07-15 06:35:14,859] [WARNING] [worker.py:840:_run_training_task] Failed to perform allreduce operation onthe gradients. Retrying...
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort
WARNING:root:Connection reset by peer
WARNING:root:rebuild process returns abort
[2020-07-15 06:36:53,060] [WARNING] [worker.py:771:_broadcast_model_params] Failed to perform barrier operation
[2020-07-15 06:36:53,060] [WARNING] [worker.py:817:_collect_gradients_with_allreduce_robust] Failed to broadcast model parameters
[2020-07-15 06:36:53,060] [WARNING] [worker.py:840:_run_training_task] Failed to perform allreduce operation onthe gradients. Retrying...
WARNING:root:Socket Timeout
WARNING:root:rebuild process returns abort

```

Resolving patch:
diff --git a/elasticdl/python/master/k8s_instance_manager.py b/elasticdl/python/master/k8s_instance_manager.py
--- a/elasticdl/python/master/k8s_instance_manager.py
+++ b/elasticdl/python/master/k8s_instance_manager.py
@@ -386,6 +386,13 @@ def _event_cb(self, event):
             # tolerance.
             self._start_ps(ps_id)
 
+    def get_alive_workers(self):
+        alive_workers = []
+        for pod_name, phase in self._worker_pods_phase.values():
+            if phase == PodStatus.RUNNING:
+                alive_workers.append(pod_name)
+        return alive_workers
+
     @property
     def ps_addrs(self):
         return self._ps_addrs
diff --git a/elasticdl/python/master/master.py b/elasticdl/python/master/master.py
--- a/elasticdl/python/master/master.py
+++ b/elasticdl/python/master/master.py
@@ -160,12 +160,12 @@ def __init__(self, args):
         self.task_d.add_deferred_callback_create_train_end_task()
         self.evaluation_service = self._create_evaluation_service(args)
 
-        # Initialize master service
-        self.master_servicer, self.server = self._create_master_service(args)
-
         # Initialize instance manager
         self.instance_manager = self._create_instance_manager(args)
 
+        # Initialize master service
+        self.master_servicer, self.server = self._create_master_service(args)
+
         self._should_stop = False
         self._exit_code = 0
         threading.Thread(
@@ -376,6 +376,7 @@ def _create_master_service(self, args):
             args.minibatch_size,
             self.task_d,
             evaluation_service=self.evaluation_service,
+            master=self,
         )
         elasticdl_pb2_grpc.add_MasterServicer_to_server(
             master_servicer, server
diff --git a/elasticdl/python/master/servicer.py b/elasticdl/python/master/servicer.py
--- a/elasticdl/python/master/servicer.py
+++ b/elasticdl/python/master/servicer.py
@@ -19,13 +19,14 @@
 
 from elasticdl.proto import elasticdl_pb2, elasticdl_pb2_grpc
 from elasticdl.python.common.log_utils import default_logger as logger
+from elasticdl_client.common.constants import DistributionStrategy
 
 
 class MasterServicer(elasticdl_pb2_grpc.MasterServicer):
     """Master service implementation"""
 
     def __init__(
-        self, minibatch_size, task_d, evaluation_service,
+        self, minibatch_size, task_d, evaluation_service, master,
     ):
         # TODO: group params together into a single object.
         self._task_d = task_d
@@ -34,6 +35,7 @@ def __init__(
         self._version = 0
 
         self._evaluation_service = evaluation_service
+        self._master = master
         self._task_complete_times = {
             elasticdl_pb2.EVALUATION: [],
             elasticdl_pb2.TRAINING: [],
@@ -78,7 +80,16 @@ def get_task(self, request, _):
             # we are trying to pop and invoke the callback.
             # Then the master tells the worker to wait
             # in case of new tasks later.
-            res.type = elasticdl_pb2.WAIT
+            if (
+                self._master.distribution_strategy
+                == DistributionStrategy.ALLREDUCE
+            ):
+                # If there is no more task, master only send wait task to
+                # the last worker and other workers exit.
+                if len(self._master.instance_manager.get_alive_workers()) == 1:
+                    res.type = elasticdl_pb2.WAIT
+            else:
+                res.type = elasticdl_pb2.WAIT
         with self._lock:
             self._worker_liveness_time[request.worker_id] = time.time()
         return res


Test patch:
diff --git a/elasticdl/python/tests/evaluation_service_test.py b/elasticdl/python/tests/evaluation_service_test.py
--- a/elasticdl/python/tests/evaluation_service_test.py
+++ b/elasticdl/python/tests/evaluation_service_test.py
@@ -105,7 +105,9 @@ def testEvaluationService(self):
             None, task_d, 10, 20, 0, False, _eval_metrics_fn,
         )
 
-        _ = MasterServicer(2, task_d, evaluation_service=evaluation_service,)
+        _ = MasterServicer(
+            2, task_d, evaluation_service=evaluation_service, master=None
+        )
 
         # No checkpoint available
         self.assertFalse(evaluation_service.try_to_create_new_job())
@@ -130,7 +132,9 @@ def testEvaluationOnly(self):
         )
         task_d.set_evaluation_service(evaluation_service)
 
-        _ = MasterServicer(2, task_d, evaluation_service=evaluation_service,)
+        _ = MasterServicer(
+            2, task_d, evaluation_service=evaluation_service, master=None,
+        )
 
         self.assertEqual(8, len(task_d._eval_todo))
         for i in range(8):
diff --git a/elasticdl/python/tests/servicer_test.py b/elasticdl/python/tests/servicer_test.py
--- a/elasticdl/python/tests/servicer_test.py
+++ b/elasticdl/python/tests/servicer_test.py
@@ -51,22 +51,23 @@ def optimizer(lr=0.1):
 
 class ServicerTest(unittest.TestCase):
     def testGetEmptyTask(self):
-        master = MasterServicer(
+        master_servicer = MasterServicer(
             3,
             _TaskDispatcher({}, {}, {}, records_per_task=3, num_epochs=2),
             evaluation_service=None,
+            master=None,
         )
 
         req = elasticdl_pb2.GetTaskRequest()
 
         # No task yet, make sure the returned versions are as expected.
         req.worker_id = 1
-        task = master.get_task(req, None)
+        task = master_servicer.get_task(req, None)
         self.assertEqual("", task.shard_name)
         self.assertEqual(0, task.model_version)
 
-        master._version = 1
-        task = master.get_task(req, None)
+        master_servicer._version = 1
+        task = master_servicer.get_task(req, None)
         self.assertEqual("", task.shard_name)
         self.assertEqual(1, task.model_version)
 
@@ -78,7 +79,9 @@ def testReportTaskResult(self):
             records_per_task=3,
             num_epochs=2,
         )
-        master = MasterServicer(3, task_d, evaluation_service=None,)
+        master = MasterServicer(
+            3, task_d, evaluation_service=None, master=None
+        )
 
         # task to number of runs.
         tasks = defaultdict(int)
diff --git a/elasticdl/python/tests/test_utils.py b/elasticdl/python/tests/test_utils.py
--- a/elasticdl/python/tests/test_utils.py
+++ b/elasticdl/python/tests/test_utils.py
@@ -411,7 +411,7 @@ def distributed_train_and_evaluate(
     task_d.set_evaluation_service(evaluation_service)
 
     master = MasterServicer(
-        batch_size, task_d, evaluation_service=evaluation_service,
+        batch_size, task_d, evaluation_service=evaluation_service, master=None,
     )
     callbacks = [
         callback_class(master, worker) for callback_class in callback_classes