fix: unique world names on scale_out (#303)

If we have back-to-back scale_out - scale_in - scale_out, on the second scale_out, generator might return the same config, creating a world name conflict while configuring multiworld. For that, we added a duplicate count which we increment on each scale_out. After the config is generated, we find the new workers and update their worlds and their peers worlds duplicate count, so we can have unique worlds names between scale_out operations.

Author: raresgaia123

infscale/configs/job.py

@@ -70,7 +70,7 @@ class WorldInfo:
     addr: str = "127.0.0.1"
     backend: Optional[str] = ""
     recover: bool = False
-    duplicate_count: int = 0
+    conflict_count: int = 0
 
 
 @dataclass

infscale/controller/job_context.py

@@ -663,6 +663,7 @@ def __init__(self, ctrl: Controller, job_id: str):
         self._cur_cfg: JobConfig | None = None
         self._new_cfg: JobConfig | None = None
         self._flow_graph_patched = False
+        self._worlds_conflict_count: dict[str, int] = {}
 
         # event to update the config after all agents added ports and ip address
         self.agents_setup_event = asyncio.Event()
@@ -831,6 +832,8 @@ def process_cfg(self) -> None:
 
         self._reconcile_wrk_status(self._cur_cfg, self._new_cfg)
 
+        self._update_worlds_conflict_count(self._cur_cfg, self._new_cfg)
+
         self._new_cfg.reqgen_config = self.ctrl.reqgen_config
 
         agent_resources = self.get_agent_resources_map()
@@ -853,6 +856,33 @@ def process_cfg(self) -> None:
 
         self.job_checker.setup(self._new_cfg)
 
+    def _update_worlds_conflict_count(
+        self, cur_cfg: JobConfig, new_cfg: JobConfig
+    ) -> None:
+        """Update world infos duplicate count."""
+        if cur_cfg:
+            new_workers = JobConfig.get_workers_diff(new_cfg, cur_cfg)
+        else:
+            new_workers = {worker.id for worker in new_cfg.workers}
+
+        for wid, world_list in new_cfg.flow_graph.items():
+            for world_info in world_list:
+                is_peer = any(wrk_id in world_info.peers for wrk_id in new_workers)
+
+                if wid in new_workers or is_peer:
+                    name = world_info.name
+                    self._set_world_conflict_count(name)
+                    world_info.conflict_count = self._worlds_conflict_count[name]
+
+    def _set_world_conflict_count(self, world_name: str) -> None:
+        """Set worlds conflict count."""
+        if world_name in self._worlds_conflict_count:
+            self._worlds_conflict_count[world_name] += 1
+
+            return
+
+        self._worlds_conflict_count[world_name] = 0
+
     def reset_cfg_recover_flags(self) -> None:
         """Reset recover flags on config."""
         self._cur_cfg.reset_recover_flags()
@@ -878,15 +908,19 @@ def _update_recovery_flow_graph(
         recover_flow_graph = cfg.flow_graph[recover_wid]
 
         for world_info in recover_flow_graph:
+            name = world_info.name
+            self._set_world_conflict_count(name)
             world_info.addr = ip
             world_info.recover = True
-            world_info.duplicate_count += 1
+            world_info.conflict_count = self._worlds_conflict_count[name]
 
         for world_list in cfg.flow_graph.values():
             for world_info in world_list:
                 if recover_wid in world_info.peers:
+                    name = world_info.name
+                    self._set_world_conflict_count(name)
                     world_info.recover = True
-                    world_info.duplicate_count += 1
+                    world_info.conflict_count = self._worlds_conflict_count[name]
 
     def _update_recovery_worker_data(
         self, cfg: JobConfig, wrk_id: str, gpu_id: int
@@ -1233,6 +1267,7 @@ def cleanup(self) -> None:
         self._cur_cfg = None
         self._new_cfg = None
         self._flow_graph_patched = False
+        self._worlds_conflict_count = {}
 
     def _release_gpu_resources(self, agent_data: AgentMetaData) -> None:
         resources = self.ctrl.agent_contexts[agent_data.id].resources

infscale/execution/pipeline.py

@@ -458,14 +458,14 @@ def _build_world_infos(self) -> dict[str, WorldInfo]:
                 else:
                     continue
 
-                name, backend, addr, data_port, ctrl_port, recover, duplicate_count = (
+                name, backend, addr, data_port, ctrl_port, recover, conflict_count = (
                     cfg_world_info.name,
                     cfg_world_info.backend,
                     cfg_world_info.addr,
                     cfg_world_info.data_port,
                     cfg_world_info.ctrl_port,
                     cfg_world_info.recover,
-                    cfg_world_info.duplicate_count,
+                    cfg_world_info.conflict_count,
                 )
 
                 world_size = len(cfg_world_info.peers) + 1
@@ -483,8 +483,8 @@ def _build_world_infos(self) -> dict[str, WorldInfo]:
                     "other_id": other_id,
                     "other": other_rank,
                     "recover": recover,
-                    "duplicate_count": duplicate_count,
-                    "multiworld_name": f"{name}-{duplicate_count}",
+                    "conflict_count": conflict_count,
+                    "multiworld_name": f"{name}-{conflict_count}",
                 }
                 world_info = WorldInfo(**data)
                 world_infos[name] = world_info

infscale/execution/world.py

@@ -42,5 +42,5 @@ class WorldInfo:
     other: int  # other peer's rank
 
     recover: bool
-    duplicate_count: int
+    conflict_count: int
     multiworld_name: str