refactor: rewrite update process management logic

Author: specture724

checkpoint_engine/ps.py

@@ -931,6 +931,15 @@ def init_process_group(
         )
         logger.info(f"[rank{self._rank}] init process group successfully.")
 
+    def store_based_barrier(self, store: dist.TCPStore):
+        dist.distributed_c10d._store_based_barrier(
+            rank=self._rank,
+            store=store,
+            group_name="parameter_server_barrier",
+            rendezvous_count=self._world_size,
+            timeout=timedelta(minutes=5),
+        )
+
     def update(
         self,
         checkpoint_name: str,
@@ -958,34 +967,27 @@ def update(
                 timeout=timedelta(minutes=10),
                 is_master=self._rank == 0,
             )
-            # if both ranks is None or [], it will use fully broadcast to update to all ranks
-            if not ranks:
-                if self._auto_pg and not dist.is_initialized():
-                    self.init_process_group()
-                ranks_group = dist.new_group()
-                self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks)
-            else:
-                if self._auto_pg:
+            if self._auto_pg:
+                if not ranks:
+                    if not dist.is_initialized():
+                        self.init_process_group()
+                else:
                     if dist.is_initialized():
                         dist.destroy_process_group()
                         # HACK: wait 2s to ensure destroy is finished
                         time.sleep(2)
                     self.init_process_group()
-                ranks_group = dist.new_group(ranks)
-                logger.info(
-                    f"[rank{self._rank}] default pg: {dist.group.WORLD}, ranks group: {ranks_group}"
-                )
+
+            # if both ranks is None or [], it will use fully broadcast to update to all ranks
+            ranks_group = dist.new_group(ranks if ranks else None)
+            if not ranks:
+                self._update_per_bucket(checkpoint_name, req_func, ranks_group)
+            else:
                 self._update_per_bucket(checkpoint_name, req_func, ranks_group, ranks)
 
-            dist.distributed_c10d._store_based_barrier(
-                rank=self._rank,
-                store=manager_store,
-                group_name="manager_store_barrier",
-                rendezvous_count=self._world_size,
-                timeout=timedelta(minutes=5),
-            )
+            self.store_based_barrier(manager_store)
+
             dist.destroy_process_group(ranks_group)
-            del ranks_group
             if self._auto_pg:
                 dist.destroy_process_group()