Fix rendezvous error due to EtcdStore get method not waiting in some cases (#137056)

alenawang · tarat44 · pytorchmergebot · commit 18525e185e21 · 2024-10-02T01:45:00.000Z
Fixes #132950 This fixes an issue in `torch/distributed/elastic/rendezvous/etcd_store.py` where the [get method](https://github.com/pytorch/pytorch/blob/v2.4.0/torch/distributed/elastic/rendezvous/etcd_store.py#L60) does not wait as expected when no keys have been written under the store prefix yet (and therefore the store prefix key does not exist). This was because the `_try_wait_get` method would error out immediately [here](https://github.com/alenawang/pytorch/blob/main/torch/distributed/elastic/rendezvous/etcd_store.py#L179) if the prefix was not found instead of continuing to the etcd watch. This was causing upstream issues where distributed jobs using etcd-v2 could not get past the initial rendezvous at all (details in issue #132950). We added a test demonstrating this issue and the fix. Without the fix the test fails with `etcd.EtcdKeyNotFound: Key not found : /torch/elastic/store` instead of waiting for the first key to be written; with the fix the test waits properly. Co-authored-by: tarat44 <32471142+tarat44@users.noreply.github.com> Pull Request resolved: #137056 Approved by: https://github.com/fduwjj Co-authored-by: tarat44 <32471142+tarat44@users.noreply.github.com>
diff --git a/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py b/test/distributed/elastic/rendezvous/etcd_rendezvous_backend_test.py
@@ -7,6 +7,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import subprocess
+import threading
+import time
 from base64 import b64encode
 from typing import cast, ClassVar
 from unittest import TestCase
@@ -18,6 +20,7 @@
     RendezvousConnectionError,
     RendezvousParameters,
 )
+from torch.distributed.elastic.rendezvous.api import RendezvousStoreInfo
 from torch.distributed.elastic.rendezvous.etcd_rendezvous_backend import (
     create_backend,
     EtcdRendezvousBackend,
@@ -146,3 +149,32 @@ def test_create_backend_raises_error_if_read_timeout_is_invalid(self) -> None:
                     ValueError, r"^The read timeout must be a positive integer.$"
                 ):
                     create_backend(self._params)
+
+    def test_get_waits_for_store_prefix_key(self) -> None:
+        def store_get(store, result_dict):
+            start_time = time.perf_counter()
+            result_dict["get_result"] = store.get(
+                RendezvousStoreInfo.MASTER_ADDR_KEY
+            ).decode(encoding="UTF-8")
+            end_time = time.perf_counter()
+            result_dict["time"] = end_time - start_time
+
+        def store_set(store):
+            time.sleep(2)
+            store.set(RendezvousStoreInfo.MASTER_ADDR_KEY, b"foo")
+
+        backend, store = create_backend(self._params)
+        backend.set_state(b"dummy_state")
+        result_dict = {}
+
+        get_thread = threading.Thread(target=store_get, args=(store, result_dict))
+        set_thread = threading.Thread(target=store_set, args=(store,))
+
+        get_thread.start()
+        set_thread.start()
+
+        get_thread.join()
+        set_thread.join()
+
+        assert result_dict["get_result"] == "foo"
+        assert result_dict["time"] >= 2
diff --git a/torch/distributed/elastic/rendezvous/etcd_store.py b/torch/distributed/elastic/rendezvous/etcd_store.py
@@ -176,27 +176,32 @@ def _try_wait_get(self, b64_keys, override_timeout=None):
 
         while True:
             # Read whole directory (of keys), filter only the ones waited for
-            all_nodes = self.client.get(key=self.prefix)
-            req_nodes = {
-                node.key: node.value
-                for node in all_nodes.children
-                if node.key in b64_keys
-            }
-
-            if len(req_nodes) == len(b64_keys):
-                # All keys are available
-                return req_nodes
+            all_nodes = None
+            try:
+                all_nodes = self.client.get(key=self.prefix)
+                req_nodes = {
+                    node.key: node.value
+                    for node in all_nodes.children
+                    if node.key in b64_keys
+                }
+
+                if len(req_nodes) == len(b64_keys):
+                    # All keys are available
+                    return req_nodes
+            except etcd.EtcdKeyNotFound:
+                pass
 
             watch_timeout = deadline - time.time()
             if watch_timeout <= 0:
                 return None
 
             try:
+                index = all_nodes.etcd_index + 1 if all_nodes else 0
                 self.client.watch(
                     key=self.prefix,
                     recursive=True,
                     timeout=watch_timeout,
-                    index=all_nodes.etcd_index + 1,
+                    index=index,
                 )
             except etcd.EtcdWatchTimedOut:
                 if time.time() >= deadline:
