Repository: sql-machine-learning/elasticdl

Issue:
Master locks the rendezvous server when the worker initializes Horovod.



Discussion:

Resolving patch:
diff --git a/elasticdl/python/master/rendezvous_server.py b/elasticdl/python/master/rendezvous_server.py
--- a/elasticdl/python/master/rendezvous_server.py
+++ b/elasticdl/python/master/rendezvous_server.py
@@ -10,6 +10,8 @@
 # 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 time
+from threading import Lock
 
 try:
     from horovod.runner.common.util.hosts import (
@@ -33,6 +35,10 @@ def __init__(self, server_host):
         self._worker_hosts = []
         self._rendezvous_server = RendezvousServer(verbose=True)
         self._rendezvous_port = None
+        self._next_worker_hosts = None
+        self._ready_worker_hosts = set()
+        self._rendezvous_completed = True
+        self._lock = Lock()
 
     def start(self):
         self._rendezvous_port = self._rendezvous_server.start()
@@ -44,13 +50,17 @@ def set_worker_hosts(self, worker_hosts):
         Args:
             worker_hosts: List of host string.
         """
-        if sorted(worker_hosts) == sorted(self._worker_hosts):
-            return
 
-        self._rendezvous_id += 1
-        self._worker_hosts = worker_hosts
+        if sorted(worker_hosts) != sorted(self._worker_hosts):
+            self._next_worker_hosts = worker_hosts
+
+    def _init_rendezvous_server(self):
+        self._worker_hosts = self._next_worker_hosts
+        self._next_worker_hosts = None
         host_alloc_plan = self._get_host_plan()
         self._rendezvous_server.init(host_alloc_plan)
+        self._rendezvous_id += 1
+        self._rendezvous_completed = False
 
     def _get_host_plan(self):
         hosts = []
@@ -68,10 +78,24 @@ def get_rendezvous_port(self):
         return self._rendezvous_port
 
     def get_worker_host_rank(self, host):
-        # -1 if host not in worker_hosts list.
-        if host not in self._worker_hosts:
-            return -1
-        return self._worker_hosts.index(host)
+        with self._lock:
+            if self._next_worker_hosts and self._rendezvous_completed:
+                time.sleep(2)  # Wait 2s for workers to complete rendezvous.
+                self._init_rendezvous_server()
+
+            # -1 if host not in worker_hosts list.
+            if host not in self._worker_hosts:
+                return -1
+
+            if not self._rendezvous_completed:
+                self._ready_worker_hosts.add(host)
+                # If all active workers in the rendezvous are ready,
+                # the server can start to set hosts for the next rendezvous
+                if self._ready_worker_hosts == set(self._worker_hosts):
+                    self._rendezvous_completed = True
+                    self._ready_worker_hosts = set()
+
+            return self._worker_hosts.index(host)
 
     def get_size(self):
         return len(self._worker_hosts)
diff --git a/elasticdl/python/worker/allreduce_trainer.py b/elasticdl/python/worker/allreduce_trainer.py
--- a/elasticdl/python/worker/allreduce_trainer.py
+++ b/elasticdl/python/worker/allreduce_trainer.py
@@ -129,7 +129,7 @@ def _broadcast_model(self):
         broadcast_variables(self._optimizer.variables(), root_rank=0)
 
     def _run_model_call_locally(self, features, labels):
-        """Call `self._model.call` locally to Create variables of the model
+        """Call `self._model.call` locally to create variables of the model
         and optimizer. Because we should have variables before broadcasting.
         """
         with tf.GradientTape() as tape:


Test patch:
diff --git a/elasticdl/python/tests/horovod_rendezvous_test.py b/elasticdl/python/tests/horovod_rendezvous_test.py
--- a/elasticdl/python/tests/horovod_rendezvous_test.py
+++ b/elasticdl/python/tests/horovod_rendezvous_test.py
@@ -36,10 +36,16 @@ def test_get_host_plan(self):
     def test_set_worker_hosts(self):
         worker_hosts = ["127.0.0.2", "127.0.0.3"]
         self.rendezvous_server.set_worker_hosts(worker_hosts)
+        rank_0 = self.rendezvous_server.get_worker_host_rank("127.0.0.2")
+        rank_1 = self.rendezvous_server.get_worker_host_rank("127.0.0.3")
+        self.assertEqual(rank_0, 0)
+        self.assertEqual(rank_1, 1)
+        self.assertEqual(self.rendezvous_server._rendezvous_completed, True)
         self.assertEqual(self.rendezvous_server._rendezvous_id, 1)
 
         new_worker_hosts = ["127.0.0.1", "127.0.0.3"]
         self.rendezvous_server.set_worker_hosts(new_worker_hosts)
+        self.rendezvous_server._init_rendezvous_server()
         self.assertEqual(self.rendezvous_server._rendezvous_id, 2)
 
     def test_get_attr(self):