process_group: set timeout for TCPStore client connect (#145)

Author: d4l3k

torchft/process_group.py

@@ -84,7 +84,7 @@
 T = TypeVar("T")
 
 
-def create_store_client(store_addr: str) -> Store:
+def create_store_client(store_addr: str, timeout: timedelta) -> Store:
     """
     Creates a PrefixStore(TCPStore(...)) client from an address in the format:
 
@@ -100,6 +100,7 @@ def create_store_client(store_addr: str) -> Store:
         port=int(port),
         is_master=False,
         wait_for_workers=False,
+        timeout=timeout,
     )
     store = PrefixStore(prefix, store)
     return store
@@ -350,11 +351,20 @@ def __repr__(self) -> str:
 class ProcessGroupWrapper(ProcessGroup):
     """
     This is a wrapper around any ProcessGroup with a reconfiguration method.
+
+    Args:
+        timeout: timeout for reconfiguration for TCPStore
+        pg: optional ProcessGroup to use, if None a new one will be created
     """
 
-    def __init__(self, pg: Optional[ProcessGroup] = None) -> None:
+    def __init__(
+        self,
+        timeout: timedelta = timedelta(seconds=60),
+        pg: Optional[ProcessGroup] = None,
+    ) -> None:
         super().__init__(0, 1)
         self._pg: Optional[BaseProcessGroup] = pg
+        self._timeout = timeout
 
     def configure(self, store_addr: str, rank: int, world_size: int) -> None:
         pg = self._pg
@@ -365,7 +375,7 @@ def configure(self, store_addr: str, rank: int, world_size: int) -> None:
         # abort if already initialized
         self.abort()
 
-        store = create_store_client(store_addr)
+        store = create_store_client(store_addr, timeout=self._timeout)
 
         self._pg = self._create_pg(store, rank, world_size)
 
@@ -511,10 +521,6 @@ class ProcessGroupGloo(ProcessGroupWrapper):
     This is a reconfigurable version of ProcessGroupGloo.
     """
 
-    def __init__(self, timeout: timedelta = timedelta(seconds=60.0)) -> None:
-        super().__init__()
-        self._timeout = timeout
-
     def _create_pg(self, store: Store, rank: int, world_size: int) -> BaseProcessGroup:
         pg = BaseProcessGroup(store, rank, world_size)
         pg._set_default_backend(ProcessGroup.BackendType.GLOO)
@@ -648,8 +654,7 @@ class ProcessGroupNCCL(ProcessGroupWrapper):
     """
 
     def __init__(self, timeout: timedelta = timedelta(seconds=60.0)) -> None:
-        super().__init__()
-        self._timeout = timeout
+        super().__init__(timeout)
         self._use_abort: bool = torch.cuda.nccl.version() >= (2, 25)
 
     def _opts_hook(self, opts: T) -> T:
@@ -877,7 +882,7 @@ class ErrorSwallowingProcessGroupWrapper(ProcessGroupWrapper):
     """
 
     def __init__(self, pg: ProcessGroup) -> None:
-        super().__init__(pg)
+        super().__init__(pg=pg)
 
         self._error: Optional[Exception] = None
 
@@ -958,7 +963,7 @@ class ManagedProcessGroup(ProcessGroupWrapper):
     """
 
     def __init__(self, manager: "Manager") -> None:
-        super().__init__(manager._pg)
+        super().__init__(pg=manager._pg)
 
         self._manager = manager
 
@@ -1195,7 +1200,11 @@ def _worker(
             if curr_device >= 0 and torch.cuda.is_available():
                 torch.cuda.set_device(curr_device)
 
-            store = create_store_client(store_addr)
+            store = create_store_client(
+                store_addr,
+                # default TCPStore timeout is 5 minutes
+                timeout=timedelta(minutes=5),
+            )
 
             try:
                 pg = cls._create_pg(store, rank, world_size)

torchft/process_group_test.py

@@ -552,6 +552,23 @@ def test_nccl_apis(self) -> None:
 
         torch.cuda.synchronize()
 
+    # pyre-fixme[56]: Pyre was not able to infer the type of the decorator
+    @skipUnless(
+        torch.cuda.is_available(),
+        "needs CUDA",
+    )
+    def test_nccl_init_timeout(self) -> None:
+        store = TCPStore(
+            host_name="localhost", port=0, is_master=True, wait_for_workers=False
+        )
+        store_addr = f"localhost:{store.port}/prefix"
+        del store
+
+        pg = ProcessGroupNCCL(timeout=timedelta(seconds=0.01))
+
+        with self.assertRaisesRegex(RuntimeError, "timed out after 10ms"):
+            pg.configure(store_addr, 0, 2)
+
     def test_baby_gloo_timeout(self) -> None:
         store = TCPStore(
             host_name="localhost", port=0, is_master=True, wait_for_workers=False
@@ -710,7 +727,7 @@ def test_functional_collectives(self) -> None:
 
     def test_process_group_wrapper(self) -> None:
         pg = ProcessGroupDummy(0, 1)
-        wrapper = ProcessGroupWrapper(pg)
+        wrapper = ProcessGroupWrapper(pg=pg)
         self.assertIs(wrapper.parent, pg)
 
         wrapper.configure("addr", 0, 1)