We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Master locks the rendezvous server when the worker initializes Horovod.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticdl/python/master/rendezvous_server.py`
Content:
```
1 # Copyright 2020 The ElasticDL Authors. All rights reserved.
2 # Licensed under the Apache License, Version 2.0 (the "License");
3 # you may not use this file except in compliance with the License.
4 # You may obtain a copy of the License at
5 #
6 # http://www.apache.org/licenses/LICENSE-2.0
7 #
8 # Unless required by applicable law or agreed to in writing, software
9 # distributed under the License is distributed on an "AS IS" BASIS,
10 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 # See the License for the specific language governing permissions and
12 # limitations under the License.
13 
14 try:
15     from horovod.runner.common.util.hosts import (
16         get_host_assignments,
17         parse_hosts,
18     )
19     from horovod.runner.http.http_server import RendezvousServer
20 
21     _HOROVOD_INSTALLED = True
22 except ImportError:
23     _HOROVOD_INSTALLED = False
24 
25 _WORKER_SLOT_NUMBER = 1
26 _HOST_SEP = ","
27 
28 
29 class HorovodRendezvousServer(object):
30     def __init__(self, server_host):
31         self._rendezvous_host = server_host
32         self._rendezvous_id = 0
33         self._worker_hosts = []
34         self._rendezvous_server = RendezvousServer(verbose=True)
35         self._rendezvous_port = None
36 
37     def start(self):
38         self._rendezvous_port = self._rendezvous_server.start()
39 
40     def set_worker_hosts(self, worker_hosts):
41         """
42         Set worker hosts into RendezvousServer.
43 
44         Args:
45             worker_hosts: List of host string.
46         """
47         if sorted(worker_hosts) == sorted(self._worker_hosts):
48             return
49 
50         self._rendezvous_id += 1
51         self._worker_hosts = worker_hosts
52         host_alloc_plan = self._get_host_plan()
53         self._rendezvous_server.init(host_alloc_plan)
54 
55     def _get_host_plan(self):
56         hosts = []
57         for host in self._worker_hosts:
58             hosts.append(host + ":" + str(_WORKER_SLOT_NUMBER))
59 
60         host_infos = parse_hosts(_HOST_SEP.join(hosts))
61         host_alloc_plan = get_host_assignments(host_infos, len(host_infos))
62         return host_alloc_plan
63 
64     def get_rendezvous_host(self):
65         return self._rendezvous_host
66 
67     def get_rendezvous_port(self):
68         return self._rendezvous_port
69 
70     def get_worker_host_rank(self, host):
71         # -1 if host not in worker_hosts list.
72         if host not in self._worker_hosts:
73             return -1
74         return self._worker_hosts.index(host)
75 
76     def get_size(self):
77         return len(self._worker_hosts)
78 
79     def get_rendezvous_id(self):
80         return self._rendezvous_id
81
```

Path: `elasticdl/python/worker/allreduce_trainer.py`
Content:
```
1 # Copyright 2020 The ElasticDL Authors. All rights reserved.
2 # Licensed under the Apache License, Version 2.0 (the "License");
3 # you may not use this file except in compliance with the License.
4 # You may obtain a copy of the License at
5 #
6 # http://www.apache.org/licenses/LICENSE-2.0
7 #
8 # Unless required by applicable law or agreed to in writing, software
9 # distributed under the License is distributed on an "AS IS" BASIS,
10 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 # See the License for the specific language governing permissions and
12 # limitations under the License.
13 
14 import os
15 import time
16 
17 import tensorflow as tf
18 from tensorflow.python.framework.errors_impl import UnknownError
19 
20 from elasticdl.python.common.constants import HorovodEnv
21 from elasticdl.python.common.log_utils import default_logger as logger
22 
23 try:
24     from horovod.tensorflow.functions import broadcast_variables
25     import horovod.tensorflow as hvd
26 except ImportError:
27     hvd = None
28 
29 
30 # The default maximum number of retries for allreduce operation
31 # if allreduce-based distributed training strategy is used.
32 DEFAULT_MAX_ALLREDUCE_RETRY_NUM = 5
33 
34 
35 class AllReduceTrainer(object):
36     def __init__(self, master_client, master_addr, model, loss_fn, optimizer):
37         if not hvd:
38             raise RuntimeError("Horovod is not installed for AllReduce")
39         self._master_client = master_client
40         self._rendezvous_addr = master_addr
41         self._model = model
42         self._loss = loss_fn
43         self._optimizer = optimizer
44         self._rendezvous_id = None
45         self._need_broadcast = True
46         self._var_created = False
47         self.set_horovod_env()
48 
49     @tf.function
50     def _training_process(self, features, labels):
51         with tf.GradientTape() as tape:
52             outputs = self._model.call(features, training=True)
53             loss = self._loss(labels, outputs)
54             if self._model.losses:
55                 loss += tf.math.add_n(self._model.losses)
56 
57         tape = hvd.DistributedGradientTape(tape)
58         grads = tape.gradient(loss, self._model.trainable_variables)
59         # Take care of the order of grads and vars if worker modifies
60         # `_non_embed_vars` during training.
61         self._optimizer.apply_gradients(
62             zip(grads, self._model.trainable_variables)
63         )
64         return loss
65 
66     def training_process_with_fault_tolerance(self, features, labels):
67         if not self._var_created:
68             self._run_model_call_locally(features, labels)
69 
70         for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM + 1):
71             try:
72                 if self._need_broadcast:
73                     self._broadcast_model()
74                     self._need_broadcast = False
75                 loss = self._training_process(features, labels)
76                 version = self._optimizer.iterations.numpy()
77                 return version, loss
78             except UnknownError as e:
79                 logger.warning(
80                     "Failed to perform allreduce operation on "
81                     "the gradients. Retrying..."
82                 )
83                 # Those error message show that the communication
84                 # to merge gradient fails and we can rebuild the
85                 # communication.
86                 if (
87                     "HorovodAllreduce" in e.message
88                     or "HorovodAllgather" in e.message
89                     or "HorovodBroadcast" in e.message
90                 ):
91                     time.sleep(3)
92                     self.init_horovod_if_needed()
93 
94     def init_horovod_if_needed(self):
95         for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM):
96             rank_response = self._master_client.get_comm_rank()
97             if rank_response.rank_id < 0:
98                 logger.warning(
99                     "The master has not added the worker host into "
100                     "rendezvous yet. Retrying to get rank"
101                 )
102                 time.sleep(5)
103             else:
104                 break
105 
106         # If the rendezvous from master is unequal to self._rendezvous_id,
107         # the worker should rebuild the communication because the master
108         # has updated the communication group.
109         if rank_response.rendezvous_id != self._rendezvous_id:
110             os.environ[HorovodEnv.RENDEZVOUS_PORT] = str(
111                 rank_response.rendezvous_port
112             )
113             os.environ[HorovodEnv.RANK] = str(rank_response.rank_id)
114             os.environ[HorovodEnv.SIZE] = str(rank_response.world_size)
115             hvd.shutdown()
116             hvd.init()
117             self._rendezvous_id = rank_response.rendezvous_id
118             self._need_broadcast = True
119 
120     def set_horovod_env(self):
121         if self._rendezvous_addr:
122             os.environ[HorovodEnv.RENDEZVOUS_ADDR] = self._rendezvous_addr
123         os.environ[HorovodEnv.CONTROLLER] = "gloo"
124         os.environ[HorovodEnv.CPU_OPERATIONS] = "gloo"
125         os.environ[HorovodEnv.HOSTNAME] = "master"
126 
127     def _broadcast_model(self):
128         broadcast_variables(self._model.variables, root_rank=0)
129         broadcast_variables(self._optimizer.variables(), root_rank=0)
130 
131     def _run_model_call_locally(self, features, labels):
132         """Call `self._model.call` locally to Create variables of the model
133         and optimizer. Because we should have variables before broadcasting.
134         """
135         with tf.GradientTape() as tape:
136             outputs = self._model.call(features)
137             loss = self._loss(labels, outputs)
138 
139         grads = tape.gradient(loss, self._model.trainable_variables)
140         self._optimizer.apply_gradients(
141             zip(grads, self._model.trainable_variables)
142         )
143         self._var_created = True
144
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
         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 @@
         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 @@
         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:

{"golden_diff": "diff --git a/elasticdl/python/master/rendezvous_server.py b/elasticdl/python/master/rendezvous_server.py\n--- a/elasticdl/python/master/rendezvous_server.py\n+++ b/elasticdl/python/master/rendezvous_server.py\n@@ -10,6 +10,8 @@\n # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n # See the License for the specific language governing permissions and\n # limitations under the License.\n+import time\n+from threading import Lock\n \n try:\n     from horovod.runner.common.util.hosts import (\n@@ -33,6 +35,10 @@\n         self._worker_hosts = []\n         self._rendezvous_server = RendezvousServer(verbose=True)\n         self._rendezvous_port = None\n+        self._next_worker_hosts = None\n+        self._ready_worker_hosts = set()\n+        self._rendezvous_completed = True\n+        self._lock = Lock()\n \n     def start(self):\n         self._rendezvous_port = self._rendezvous_server.start()\n@@ -44,13 +50,17 @@\n         Args:\n             worker_hosts: List of host string.\n         \"\"\"\n-        if sorted(worker_hosts) == sorted(self._worker_hosts):\n-            return\n \n-        self._rendezvous_id += 1\n-        self._worker_hosts = worker_hosts\n+        if sorted(worker_hosts) != sorted(self._worker_hosts):\n+            self._next_worker_hosts = worker_hosts\n+\n+    def _init_rendezvous_server(self):\n+        self._worker_hosts = self._next_worker_hosts\n+        self._next_worker_hosts = None\n         host_alloc_plan = self._get_host_plan()\n         self._rendezvous_server.init(host_alloc_plan)\n+        self._rendezvous_id += 1\n+        self._rendezvous_completed = False\n \n     def _get_host_plan(self):\n         hosts = []\n@@ -68,10 +78,24 @@\n         return self._rendezvous_port\n \n     def get_worker_host_rank(self, host):\n-        # -1 if host not in worker_hosts list.\n-        if host not in self._worker_hosts:\n-            return -1\n-        return self._worker_hosts.index(host)\n+        with self._lock:\n+            if self._next_worker_hosts and self._rendezvous_completed:\n+                time.sleep(2)  # Wait 2s for workers to complete rendezvous.\n+                self._init_rendezvous_server()\n+\n+            # -1 if host not in worker_hosts list.\n+            if host not in self._worker_hosts:\n+                return -1\n+\n+            if not self._rendezvous_completed:\n+                self._ready_worker_hosts.add(host)\n+                # If all active workers in the rendezvous are ready,\n+                # the server can start to set hosts for the next rendezvous\n+                if self._ready_worker_hosts == set(self._worker_hosts):\n+                    self._rendezvous_completed = True\n+                    self._ready_worker_hosts = set()\n+\n+            return self._worker_hosts.index(host)\n \n     def get_size(self):\n         return len(self._worker_hosts)\ndiff --git a/elasticdl/python/worker/allreduce_trainer.py b/elasticdl/python/worker/allreduce_trainer.py\n--- a/elasticdl/python/worker/allreduce_trainer.py\n+++ b/elasticdl/python/worker/allreduce_trainer.py\n@@ -129,7 +129,7 @@\n         broadcast_variables(self._optimizer.variables(), root_rank=0)\n \n     def _run_model_call_locally(self, features, labels):\n-        \"\"\"Call `self._model.call` locally to Create variables of the model\n+        \"\"\"Call `self._model.call` locally to create variables of the model\n         and optimizer. Because we should have variables before broadcasting.\n         \"\"\"\n         with tf.GradientTape() as tape:\n", "issue": "Master locks the rendezvous server when the worker initializes Horovod.\n\n", "before_files": [{"content": "# Copyright 2020 The ElasticDL Authors. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\ntry:\n    from horovod.runner.common.util.hosts import (\n        get_host_assignments,\n        parse_hosts,\n    )\n    from horovod.runner.http.http_server import RendezvousServer\n\n    _HOROVOD_INSTALLED = True\nexcept ImportError:\n    _HOROVOD_INSTALLED = False\n\n_WORKER_SLOT_NUMBER = 1\n_HOST_SEP = \",\"\n\n\nclass HorovodRendezvousServer(object):\n    def __init__(self, server_host):\n        self._rendezvous_host = server_host\n        self._rendezvous_id = 0\n        self._worker_hosts = []\n        self._rendezvous_server = RendezvousServer(verbose=True)\n        self._rendezvous_port = None\n\n    def start(self):\n        self._rendezvous_port = self._rendezvous_server.start()\n\n    def set_worker_hosts(self, worker_hosts):\n        \"\"\"\n        Set worker hosts into RendezvousServer.\n\n        Args:\n            worker_hosts: List of host string.\n        \"\"\"\n        if sorted(worker_hosts) == sorted(self._worker_hosts):\n            return\n\n        self._rendezvous_id += 1\n        self._worker_hosts = worker_hosts\n        host_alloc_plan = self._get_host_plan()\n        self._rendezvous_server.init(host_alloc_plan)\n\n    def _get_host_plan(self):\n        hosts = []\n        for host in self._worker_hosts:\n            hosts.append(host + \":\" + str(_WORKER_SLOT_NUMBER))\n\n        host_infos = parse_hosts(_HOST_SEP.join(hosts))\n        host_alloc_plan = get_host_assignments(host_infos, len(host_infos))\n        return host_alloc_plan\n\n    def get_rendezvous_host(self):\n        return self._rendezvous_host\n\n    def get_rendezvous_port(self):\n        return self._rendezvous_port\n\n    def get_worker_host_rank(self, host):\n        # -1 if host not in worker_hosts list.\n        if host not in self._worker_hosts:\n            return -1\n        return self._worker_hosts.index(host)\n\n    def get_size(self):\n        return len(self._worker_hosts)\n\n    def get_rendezvous_id(self):\n        return self._rendezvous_id\n", "path": "elasticdl/python/master/rendezvous_server.py"}, {"content": "# Copyright 2020 The ElasticDL Authors. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport time\n\nimport tensorflow as tf\nfrom tensorflow.python.framework.errors_impl import UnknownError\n\nfrom elasticdl.python.common.constants import HorovodEnv\nfrom elasticdl.python.common.log_utils import default_logger as logger\n\ntry:\n    from horovod.tensorflow.functions import broadcast_variables\n    import horovod.tensorflow as hvd\nexcept ImportError:\n    hvd = None\n\n\n# The default maximum number of retries for allreduce operation\n# if allreduce-based distributed training strategy is used.\nDEFAULT_MAX_ALLREDUCE_RETRY_NUM = 5\n\n\nclass AllReduceTrainer(object):\n    def __init__(self, master_client, master_addr, model, loss_fn, optimizer):\n        if not hvd:\n            raise RuntimeError(\"Horovod is not installed for AllReduce\")\n        self._master_client = master_client\n        self._rendezvous_addr = master_addr\n        self._model = model\n        self._loss = loss_fn\n        self._optimizer = optimizer\n        self._rendezvous_id = None\n        self._need_broadcast = True\n        self._var_created = False\n        self.set_horovod_env()\n\n    @tf.function\n    def _training_process(self, features, labels):\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features, training=True)\n            loss = self._loss(labels, outputs)\n            if self._model.losses:\n                loss += tf.math.add_n(self._model.losses)\n\n        tape = hvd.DistributedGradientTape(tape)\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        # Take care of the order of grads and vars if worker modifies\n        # `_non_embed_vars` during training.\n        self._optimizer.apply_gradients(\n            zip(grads, self._model.trainable_variables)\n        )\n        return loss\n\n    def training_process_with_fault_tolerance(self, features, labels):\n        if not self._var_created:\n            self._run_model_call_locally(features, labels)\n\n        for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM + 1):\n            try:\n                if self._need_broadcast:\n                    self._broadcast_model()\n                    self._need_broadcast = False\n                loss = self._training_process(features, labels)\n                version = self._optimizer.iterations.numpy()\n                return version, loss\n            except UnknownError as e:\n                logger.warning(\n                    \"Failed to perform allreduce operation on \"\n                    \"the gradients. Retrying...\"\n                )\n                # Those error message show that the communication\n                # to merge gradient fails and we can rebuild the\n                # communication.\n                if (\n                    \"HorovodAllreduce\" in e.message\n                    or \"HorovodAllgather\" in e.message\n                    or \"HorovodBroadcast\" in e.message\n                ):\n                    time.sleep(3)\n                    self.init_horovod_if_needed()\n\n    def init_horovod_if_needed(self):\n        for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM):\n            rank_response = self._master_client.get_comm_rank()\n            if rank_response.rank_id < 0:\n                logger.warning(\n                    \"The master has not added the worker host into \"\n                    \"rendezvous yet. Retrying to get rank\"\n                )\n                time.sleep(5)\n            else:\n                break\n\n        # If the rendezvous from master is unequal to self._rendezvous_id,\n        # the worker should rebuild the communication because the master\n        # has updated the communication group.\n        if rank_response.rendezvous_id != self._rendezvous_id:\n            os.environ[HorovodEnv.RENDEZVOUS_PORT] = str(\n                rank_response.rendezvous_port\n            )\n            os.environ[HorovodEnv.RANK] = str(rank_response.rank_id)\n            os.environ[HorovodEnv.SIZE] = str(rank_response.world_size)\n            hvd.shutdown()\n            hvd.init()\n            self._rendezvous_id = rank_response.rendezvous_id\n            self._need_broadcast = True\n\n    def set_horovod_env(self):\n        if self._rendezvous_addr:\n            os.environ[HorovodEnv.RENDEZVOUS_ADDR] = self._rendezvous_addr\n        os.environ[HorovodEnv.CONTROLLER] = \"gloo\"\n        os.environ[HorovodEnv.CPU_OPERATIONS] = \"gloo\"\n        os.environ[HorovodEnv.HOSTNAME] = \"master\"\n\n    def _broadcast_model(self):\n        broadcast_variables(self._model.variables, root_rank=0)\n        broadcast_variables(self._optimizer.variables(), root_rank=0)\n\n    def _run_model_call_locally(self, features, labels):\n        \"\"\"Call `self._model.call` locally to Create variables of the model\n        and optimizer. Because we should have variables before broadcasting.\n        \"\"\"\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features)\n            loss = self._loss(labels, outputs)\n\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        self._optimizer.apply_gradients(\n            zip(grads, self._model.trainable_variables)\n        )\n        self._var_created = True\n", "path": "elasticdl/python/worker/allreduce_trainer.py"}], "after_files": [{"content": "# Copyright 2020 The ElasticDL Authors. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport time\nfrom threading import Lock\n\ntry:\n    from horovod.runner.common.util.hosts import (\n        get_host_assignments,\n        parse_hosts,\n    )\n    from horovod.runner.http.http_server import RendezvousServer\n\n    _HOROVOD_INSTALLED = True\nexcept ImportError:\n    _HOROVOD_INSTALLED = False\n\n_WORKER_SLOT_NUMBER = 1\n_HOST_SEP = \",\"\n\n\nclass HorovodRendezvousServer(object):\n    def __init__(self, server_host):\n        self._rendezvous_host = server_host\n        self._rendezvous_id = 0\n        self._worker_hosts = []\n        self._rendezvous_server = RendezvousServer(verbose=True)\n        self._rendezvous_port = None\n        self._next_worker_hosts = None\n        self._ready_worker_hosts = set()\n        self._rendezvous_completed = True\n        self._lock = Lock()\n\n    def start(self):\n        self._rendezvous_port = self._rendezvous_server.start()\n\n    def set_worker_hosts(self, worker_hosts):\n        \"\"\"\n        Set worker hosts into RendezvousServer.\n\n        Args:\n            worker_hosts: List of host string.\n        \"\"\"\n\n        if sorted(worker_hosts) != sorted(self._worker_hosts):\n            self._next_worker_hosts = worker_hosts\n\n    def _init_rendezvous_server(self):\n        self._worker_hosts = self._next_worker_hosts\n        self._next_worker_hosts = None\n        host_alloc_plan = self._get_host_plan()\n        self._rendezvous_server.init(host_alloc_plan)\n        self._rendezvous_id += 1\n        self._rendezvous_completed = False\n\n    def _get_host_plan(self):\n        hosts = []\n        for host in self._worker_hosts:\n            hosts.append(host + \":\" + str(_WORKER_SLOT_NUMBER))\n\n        host_infos = parse_hosts(_HOST_SEP.join(hosts))\n        host_alloc_plan = get_host_assignments(host_infos, len(host_infos))\n        return host_alloc_plan\n\n    def get_rendezvous_host(self):\n        return self._rendezvous_host\n\n    def get_rendezvous_port(self):\n        return self._rendezvous_port\n\n    def get_worker_host_rank(self, host):\n        with self._lock:\n            if self._next_worker_hosts and self._rendezvous_completed:\n                time.sleep(2)  # Wait 2s for workers to complete rendezvous.\n                self._init_rendezvous_server()\n\n            # -1 if host not in worker_hosts list.\n            if host not in self._worker_hosts:\n                return -1\n\n            if not self._rendezvous_completed:\n                self._ready_worker_hosts.add(host)\n                # If all active workers in the rendezvous are ready,\n                # the server can start to set hosts for the next rendezvous\n                if self._ready_worker_hosts == set(self._worker_hosts):\n                    self._rendezvous_completed = True\n                    self._ready_worker_hosts = set()\n\n            return self._worker_hosts.index(host)\n\n    def get_size(self):\n        return len(self._worker_hosts)\n\n    def get_rendezvous_id(self):\n        return self._rendezvous_id\n", "path": "elasticdl/python/master/rendezvous_server.py"}, {"content": "# Copyright 2020 The ElasticDL Authors. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport time\n\nimport tensorflow as tf\nfrom tensorflow.python.framework.errors_impl import UnknownError\n\nfrom elasticdl.python.common.constants import HorovodEnv\nfrom elasticdl.python.common.log_utils import default_logger as logger\n\ntry:\n    from horovod.tensorflow.functions import broadcast_variables\n    import horovod.tensorflow as hvd\nexcept ImportError:\n    hvd = None\n\n\n# The default maximum number of retries for allreduce operation\n# if allreduce-based distributed training strategy is used.\nDEFAULT_MAX_ALLREDUCE_RETRY_NUM = 5\n\n\nclass AllReduceTrainer(object):\n    def __init__(self, master_client, master_addr, model, loss_fn, optimizer):\n        if not hvd:\n            raise RuntimeError(\"Horovod is not installed for AllReduce\")\n        self._master_client = master_client\n        self._rendezvous_addr = master_addr\n        self._model = model\n        self._loss = loss_fn\n        self._optimizer = optimizer\n        self._rendezvous_id = None\n        self._need_broadcast = True\n        self._var_created = False\n        self.set_horovod_env()\n\n    @tf.function\n    def _training_process(self, features, labels):\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features, training=True)\n            loss = self._loss(labels, outputs)\n            if self._model.losses:\n                loss += tf.math.add_n(self._model.losses)\n\n        tape = hvd.DistributedGradientTape(tape)\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        # Take care of the order of grads and vars if worker modifies\n        # `_non_embed_vars` during training.\n        self._optimizer.apply_gradients(\n            zip(grads, self._model.trainable_variables)\n        )\n        return loss\n\n    def training_process_with_fault_tolerance(self, features, labels):\n        if not self._var_created:\n            self._run_model_call_locally(features, labels)\n\n        for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM + 1):\n            try:\n                if self._need_broadcast:\n                    self._broadcast_model()\n                    self._need_broadcast = False\n                loss = self._training_process(features, labels)\n                version = self._optimizer.iterations.numpy()\n                return version, loss\n            except UnknownError as e:\n                logger.warning(\n                    \"Failed to perform allreduce operation on \"\n                    \"the gradients. Retrying...\"\n                )\n                # Those error message show that the communication\n                # to merge gradient fails and we can rebuild the\n                # communication.\n                if (\n                    \"HorovodAllreduce\" in e.message\n                    or \"HorovodAllgather\" in e.message\n                    or \"HorovodBroadcast\" in e.message\n                ):\n                    time.sleep(3)\n                    self.init_horovod_if_needed()\n\n    def init_horovod_if_needed(self):\n        for _ in range(DEFAULT_MAX_ALLREDUCE_RETRY_NUM):\n            rank_response = self._master_client.get_comm_rank()\n            if rank_response.rank_id < 0:\n                logger.warning(\n                    \"The master has not added the worker host into \"\n                    \"rendezvous yet. Retrying to get rank\"\n                )\n                time.sleep(5)\n            else:\n                break\n\n        # If the rendezvous from master is unequal to self._rendezvous_id,\n        # the worker should rebuild the communication because the master\n        # has updated the communication group.\n        if rank_response.rendezvous_id != self._rendezvous_id:\n            os.environ[HorovodEnv.RENDEZVOUS_PORT] = str(\n                rank_response.rendezvous_port\n            )\n            os.environ[HorovodEnv.RANK] = str(rank_response.rank_id)\n            os.environ[HorovodEnv.SIZE] = str(rank_response.world_size)\n            hvd.shutdown()\n            hvd.init()\n            self._rendezvous_id = rank_response.rendezvous_id\n            self._need_broadcast = True\n\n    def set_horovod_env(self):\n        if self._rendezvous_addr:\n            os.environ[HorovodEnv.RENDEZVOUS_ADDR] = self._rendezvous_addr\n        os.environ[HorovodEnv.CONTROLLER] = \"gloo\"\n        os.environ[HorovodEnv.CPU_OPERATIONS] = \"gloo\"\n        os.environ[HorovodEnv.HOSTNAME] = \"master\"\n\n    def _broadcast_model(self):\n        broadcast_variables(self._model.trainable_variables, root_rank=0)\n        broadcast_variables(self._optimizer.variables(), root_rank=0)\n\n    def _run_model_call_locally(self, features, labels):\n        \"\"\"Call `self._model.call` locally to create variables of the model\n        and optimizer. Because we should have variables before broadcasting.\n        \"\"\"\n        with tf.GradientTape() as tape:\n            outputs = self._model.call(features)\n            loss = self._loss(labels, outputs)\n\n        grads = tape.gradient(loss, self._model.trainable_variables)\n        self._optimizer.apply_gradients(\n            zip(grads, self._model.trainable_variables)\n        )\n        self._var_created = True\n", "path": "elasticdl/python/worker/allreduce_trainer.py"}]}