fix: worker info update

Signed-off-by: Hao Lin <linhaomails@gmail.com>

Author: andylin-hao

rlinf/scheduler/manager/worker_manager.py

@@ -335,15 +335,8 @@ def get_worker_info(self, worker_address: WorkerAddress) -> WorkerInfo:
                 return node._worker_info
         return None
 
-    def update_worker_info(
-        self, worker_address: WorkerAddress, worker_info: WorkerInfo
-    ):
-        """Update worker information by its address.
-
-        Args:
-            worker_address (WorkerAddress): The address of the worker to update.
-            worker_info (WorkerInfo): The new worker information.
-        """
+    def update_worker_info(self, worker_address: WorkerAddress, worker_info: WorkerInfo):
+        """Update worker information by its address."""
         for root in self._root_workers:
             node = WorkerNode.find_node(root, worker_address)
             if node is not None:
@@ -386,18 +379,21 @@ def get_group_ranks(self, root_group_name: str) -> list[int]:
         return []
 
     def set_group_workers(self, root_group_name: str, workers: list[WorkerInfo]):
-        """Replace all workers of a root group with provided worker infos."""
-        root_node = None
-        for root in self._root_workers:
-            if (
-                root._worker_address is not None
-                and root._worker_address.get_name() == root_group_name
-            ):
-                root_node = root
-                break
-        if root_node is None:
-            root_node = WorkerNode(WorkerAddress(root_group_name, []))
-            self._root_workers.append(root_node)
-        root_node._nodes = []
+        """Replace all workers of a root group with provided worker infos.
+
+        Rebuilds membership via register/unregister APIs to preserve the WorkerNode
+        tree invariants instead of mutating internal tree state directly.
+        """
+        for rank in list(self.get_group_ranks(root_group_name)):
+            self.unregister_worker(WorkerAddress(root_group_name, rank))
+
         for worker_info in sorted(workers, key=lambda info: info.rank):
-            root_node.add_child(worker_info.rank, worker_info)
+            if worker_info.address.root_group_name != root_group_name:
+                raise ValueError(
+                    f"Worker {worker_info.address.get_name()} does not belong to group {root_group_name}."
+                )
+            existing = self.get_worker_info(worker_info.address)
+            if existing is None:
+                self.register_worker(worker_info.address, worker_info)
+            else:
+                self.update_worker_info(worker_info.address, worker_info)

rlinf/scheduler/worker/worker.py

@@ -1274,6 +1274,7 @@ def update_state(
         master_addr: Optional[str] = None,
         root_group_names: Optional[list[str]] = None,
         rebuild_collective: bool = False,
+        sync_manager: bool = True,
     ) -> "WorkerInfo":
         """Refresh worker state after membership/rank changes."""
         if world_size is not None:
@@ -1314,5 +1315,4 @@ def update_state(
 
             self._collective = Collective(self)
 
-        self._manager_proxy.update_worker_info(self._worker_address, self._worker_info)
         return self._worker_info

rlinf/scheduler/worker/worker_group.py

@@ -242,7 +242,7 @@ def scale_up(
 
         self._placement_strategy = placement_strategy
         self._create_workers_from_placements(placements)
-        self.update_state(reindex_ranks=True)
+        self.update_state()
         self._is_ready()
         return self
 
@@ -271,7 +271,7 @@ def scale_down(
                 remaining_workers.append(worker_info)
 
         self._workers = remaining_workers
-        self.update_state(reindex_ranks=True)
+        self.update_state()
         return self
 
     def _unregister_worker(self, rank: int):
@@ -283,40 +283,39 @@ def _unregister_worker(self, rank: int):
         )
         worker_manager.unregister_worker(worker_address)
 
-    def update_state(self, reindex_ranks: bool = False):
+    def update_state(self):
         """Refresh worker-group state after membership changes."""
         from ..manager import CollectiveManager, WorkerManager
 
         worker_manager = WorkerManager.get_proxy()
-        if reindex_ranks:
-            self._sort_workers()
-            world_size = len(self._workers)
-            if world_size > 0:
-                first_worker_info = worker_manager.get_worker_info(
-                    WorkerAddress(self._worker_group_name, self._workers[0].rank)
-                )
-                assert first_worker_info is not None, (
-                    f"Cannot get worker info for rank {self._workers[0].rank} in group {self._worker_group_name}."
-                )
-                master_addr = first_worker_info.node_ip
-                refs = []
-                for new_rank, worker_info in enumerate(self._workers):
-                    refs.append(
-                        worker_info.worker.update_state.remote(
-                            new_rank=new_rank,
-                            world_size=world_size,
-                            master_addr=master_addr,
-                            rebuild_collective=True,
-                        )
+        self._sort_workers()
+        world_size = len(self._workers)
+        if world_size > 0:
+            first_worker_info = worker_manager.get_worker_info(
+                WorkerAddress(self._worker_group_name, self._workers[0].rank)
+            )
+            assert first_worker_info is not None, (
+                f"Cannot get worker info for rank {self._workers[0].rank} in group {self._worker_group_name}."
+            )
+            master_addr = first_worker_info.node_ip
+            refs = []
+            for new_rank, worker_info in enumerate(self._workers):
+                refs.append(
+                    worker_info.worker.update_state.remote(
+                        new_rank=new_rank,
+                        world_size=world_size,
+                        master_addr=master_addr,
+                        rebuild_collective=True,
                     )
-                updated_infos = ray.get(refs)
-                worker_manager.set_group_workers(self._worker_group_name, updated_infos)
-                self._workers = [
-                    WorkerGroup.WorkerRank(rank=rank, worker=worker_info.worker)
-                    for rank, worker_info in enumerate(self._workers)
-                ]
-            self._world_size = world_size
-            self._group_size = world_size
+                )
+            updated_infos = ray.get(refs)
+            worker_manager.set_group_workers(self._worker_group_name, updated_infos)
+            self._workers = [
+                WorkerGroup.WorkerRank(rank=rank, worker=worker_info.worker)
+                for rank, worker_info in enumerate(self._workers)
+            ]
+        self._world_size = world_size
+        self._group_size = world_size
 
         coll_manager = CollectiveManager.get_proxy()
         related_groups = coll_manager.get_related_worker_groups(self._worker_group_name)

tests/unit_tests/test_worker.py

@@ -198,11 +198,11 @@ def test_worker_group_scale_up_and_down(self, cluster: Cluster):
 
         worker_group.scale_down([1])
         infos = worker_group.get_env_info().wait()
-        assert sorted(info["rank"] for info in infos) == [0, 2]
+        assert sorted(info["rank"] for info in infos) == [0, 1]
         assert all(info["world_size"] == 2 for info in infos)
 
-        subset_infos = worker_group.execute_on(0, 2).get_env_info().wait()
-        assert sorted(info["rank"] for info in subset_infos) == [0, 2]
+        subset_infos = worker_group.execute_on(0, 1).get_env_info().wait()
+        assert sorted(info["rank"] for info in subset_infos) == [0, 1]
 
 
 class TestLoadUserExtensions: