process_group pipe fix

Author: H-Huang

torchft/process_group.py

@@ -23,17 +23,17 @@
 from datetime import timedelta
 from multiprocessing.connection import Connection
 from typing import (
-    TYPE_CHECKING,
     Any,
     Callable,
+    cast,
     Dict,
     Generator,
     List,
     Optional,
     Tuple,
+    TYPE_CHECKING,
     TypeVar,
     Union,
-    cast,
 )
 
 import torch
@@ -44,14 +44,14 @@
 # pyre-fixme[21]: no attribute ProcessGroupGloo
 from torch.distributed import (
     DeviceMesh,
+    get_rank,
+    init_device_mesh,
     PrefixStore,
     ProcessGroup as BaseProcessGroup,
     ProcessGroupGloo as BaseProcessGroupGloo,
     ProcessGroupNCCL as BaseProcessGroupNCCL,
     Store,
     TCPStore,
-    get_rank,
-    init_device_mesh,
 )
 from torch.distributed.distributed_c10d import (
     AllgatherOptions,
@@ -970,6 +970,26 @@ def shutdown(self) -> None:
             self._p.kill()
 
     def configure(self, store_addr: str, rank: int, world_size: int) -> None:
+        """
+        Structure
+        +-------------------+
+        |                   |
+        |   Main Process    |  (updates self._futures)
+        |                   |   <---------------
+        +-------------------+                  |
+                | Pipe 1                       |
+                v                              |
+        +-------------------+          +-------------------+
+        |                   |          |                   |
+        |  Worker Process   |   ->     |  Future Thread    |
+        |                   |  Pipe 2  |                   |
+        +-------------------+          +-------------------+
+
+        Main Process: Central controller, maintains self._futures.
+        Worker Process: Handles tasks, communicates with Future Thread.
+        Future Thread: Manages asynchronous tasks, updates self._futures.
+        """
+
         self._world_size = world_size
 
         self.shutdown()
@@ -990,7 +1010,7 @@ def configure(self, store_addr: str, rank: int, world_size: int) -> None:
                 rank,
                 world_size,
                 req_remote,
-                future_remote,
+                future_local,
                 curr_device,
             ),
             daemon=True,
@@ -1003,7 +1023,7 @@ def configure(self, store_addr: str, rank: int, world_size: int) -> None:
         self._futures = {}
         self._future_thread = threading.Thread(
             target=self._future_handler,
-            args=(future_local,),
+            args=(future_remote,),
             daemon=True,
         )
         self._future_thread.start()
@@ -1163,11 +1183,14 @@ def callback(fut: Future[object], metadata: _OpMetadata) -> None:
 
     def _future_handler(self, future_pipe: _MonitoredPipe) -> None:
         try:
-            while True:
+            while not self._future_thread_shutdown_flag.is_set():
                 try:
                     cmd = future_pipe.recv(timedelta(seconds=10))
                 except TimeoutError:
                     continue
+                # except EOFError:
+                #     # Pipe was closed, exit the loop
+                #     break
 
                 op_id, mode, data, event = cast(
                     Tuple[int, str, object, Optional[torch.cuda.Event]], cmd