process_group: add PG timeouts + automatically assign manager port

Author: d4l3k

torchft/manager.py

@@ -46,7 +46,7 @@
     from torchft.process_group import ProcessGroup
 
 MANAGER_ADDR_KEY: str = "manager_addr"
-MANAGER_DEFAULT_PORT: int = int(os.environ.get("TORCHFT_MANAGER_PORT", 29511))
+MANAGER_PORT_ENV: str = "TORCHFT_MANAGER_PORT"
 REPLICA_ID_KEY: str = "replica_id"
 
 T = TypeVar("T")
@@ -84,7 +84,6 @@ def __init__(
         load_state_dict: Callable[[T], None],
         state_dict: Callable[[], T],
         min_replica_size: int,
-        port: int = MANAGER_DEFAULT_PORT,
         use_async_quorum: bool = True,
         timeout: timedelta = timedelta(seconds=60),
         rank: Optional[int] = None,
@@ -94,6 +93,7 @@ def __init__(
         store_port: Optional[int] = None,
         lighthouse_addr: Optional[str] = None,
         replica_id: Optional[str] = None,
+        port: Optional[int] = None,
     ) -> None:
         """
         Args:
@@ -147,6 +147,13 @@ def _manager_state_dict() -> Dict[str, T]:
 
         if rank == 0:
             hostname = socket.gethostname()
+
+            if port is None:
+                port_str = os.environ.get(MANAGER_PORT_ENV)
+                if port_str is None:
+                    port = 0
+                else:
+                    port = int(port_str)
             addr = f"http://{hostname}:{port}"
             bind = f"[::]:{port}"
             lighthouse_addr = lighthouse_addr or os.environ["TORCHFT_LIGHTHOUSE"]
@@ -163,7 +170,7 @@ def _manager_state_dict() -> Dict[str, T]:
                 world_size=world_size,
             )
 
-            self._store.set(MANAGER_ADDR_KEY, addr)
+            self._store.set(MANAGER_ADDR_KEY, self._manager.address())
             self._store.set(REPLICA_ID_KEY, replica_id)
 
         addr = self._store.get(MANAGER_ADDR_KEY).decode("utf-8")

torchft/process_group.py

@@ -187,7 +187,6 @@ def __repr__(self) -> str:
 
 
 class ProcessGroupWrapper(ProcessGroup):
-    PG_CLASS: Type[BaseProcessGroup]  # pyre-fixme[13]: never initialized
     """
     This is a wrapper around any ProcessGroup with a reconfiguration method.
     """
@@ -209,9 +208,10 @@ def configure(self, store_addr: str, rank: int, world_size: int) -> None:
 
         store = create_store_client(store_addr)
 
-        # TODO: set global timeout
-        # pyre-fixme[20]: expects argument options
-        self._pg = self.PG_CLASS(store, rank, world_size)
+        self._pg = self._create_pg(store, rank, world_size)
+
+    def _create_pg(self, store: Store, rank: int, world_size: int) -> BaseProcessGroup:
+        raise NotImplementedError("not implemented")
 
     def allreduce(self, tensors: List[torch.Tensor], opts: object) -> Work:
         return self.parent.allreduce(tensors, opts)
@@ -244,9 +244,12 @@ class ProcessGroupGloo(ProcessGroupWrapper):
     This is a reconfigurable version of ProcessGroupGloo.
     """
 
-    PG_CLASS: Type[BaseProcessGroup] = (
-        BaseProcessGroupGloo  # pyre-fixme[16]: no attribute 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:
+        return BaseProcessGroupGloo(store, rank, world_size, self._timeout)
 
     def getBackendName(self) -> str:
         return "torchft-gloo"
@@ -263,9 +266,12 @@ class ProcessGroupNCCL(ProcessGroupWrapper):
     abort when reconfiguring, we need to ensure this is safe.
     """
 
-    PG_CLASS: Type[BaseProcessGroup] = (
-        BaseProcessGroupNCCL  # pyre-fixme[16]: no attribute ProcessGroupNCCL
-    )
+    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:
+        return BaseProcessGroupNCCL(store, rank, world_size, self._timeout)
 
     def getBackendName(self) -> str:
         return "torchft-nccl"