process_group: wait for futher_thread join before creating new one

Author: dwancn

torchft/process_group.py

@@ -609,8 +609,12 @@ def configure(self, store_addr: str, rank: int, world_size: int) -> None:
         if self._rx is not None:
             self._rx.close()
         if self._future_queue is not None:
+            # wait for the future thread to exit and then close the queue
             self._future_queue.put(_QUEUE_CLOSE)
-            assert self._future_queue is not None
+            assert self._future_thread is not None
+            self._future_thread.join(timeout=10.0)
+            if self._future_thread.is_alive():
+                raise RuntimeError("future thread did not exit")
             self._future_queue.close()
 
         ctx = mp.get_context("spawn")

torchft/process_group_test.py

@@ -208,6 +208,37 @@ def test_baby_gloo_timeout(self) -> None:
         with self.assertRaisesRegex(TimeoutError, "timed out after 0.01 seconds"):
             a.configure(store_addr, 0, 2)
 
+    def test_reconfigure_baby_process_group(self) -> None:
+        store = TCPStore(
+            host_name="localhost", port=0, is_master=True, wait_for_workers=False
+        )
+        store_addr = f"localhost:{store.port}/prefix"
+
+        a = ProcessGroupBabyGloo()
+        a.configure(store_addr, 0, 1)
+        futher_thread_1 = a._future_thread
+        futher_queue_1 = a._future_queue
+        p_1 = a._p
+
+        store_addr = f"localhost:{store.port}/prefix2"
+        a.configure(store_addr, 0, 1)
+        futher_thread_2 = a._future_thread
+        futher_queue_2 = a._future_queue
+        p_2 = a._p
+
+        self.assertNotEqual(futher_thread_1, futher_thread_2)
+        self.assertNotEqual(futher_queue_1, futher_queue_2)
+        self.assertNotEqual(p_1, p_2)
+
+        self.assertFalse(futher_thread_1.is_alive())
+        self.assertTrue(futher_queue_1._closed)
+        self.assertFalse(p_1.is_alive())
+
+        self.assertTrue(futher_thread_2.is_alive())
+        self.assertFalse(futher_queue_2._closed)
+        self.assertTrue(p_2.is_alive())
+
+
     def test_dummy(self) -> None:
         pg = ProcessGroupDummy(0, 1)
         m = nn.Linear(3, 4)