refactor: pipeline world infos

Refactored pipeline code and moved world info related stuff into config_runner. With this, config_runner will be responsible of computing worlds to add and to remove and it will be used for when we refactor the code to handle per world configuration.

Author: raresgaia123

infscale/execution/config_runner.py

@@ -20,6 +20,8 @@
 from typing import Awaitable, Callable
 
 from infscale.configs.job import ServeConfig
+from infscale.execution.control import Channel as CtrlCh
+from infscale.execution.world import WorldInfo
 
 
 class ConfigRunner:
@@ -34,11 +36,21 @@ def __init__(self):
         self._event.set()  # initially no configure running
         self._curr_worlds_to_configure: set[str] = set()
         self._cancel_cur_cfg = False
+        self._worlds_to_add: list[WorldInfo]
+        self._worlds_to_remove: list[WorldInfo]
+        self._world_infos: dict[str, WorldInfo] = {}
 
     def handle_new_spec(self, spec: ServeConfig) -> None:
         """Handle new spec."""
         self._cancel_cur_cfg = self._should_cancel_current(spec)
         self._spec = spec
+        new_world_infos = self._build_world_infos()
+
+        new = new_world_infos.keys()
+        cur = self._world_infos.keys()
+
+        self._worlds_to_add = self.get_words_to_add(new_world_infos, new, cur)
+        self._worlds_to_remove = self.get_words_to_remove(self._world_infos, new, cur)
 
     def _should_cancel_current(self, spec: ServeConfig) -> bool:
         """Decide if current configuration should be cancelled."""
@@ -57,6 +69,35 @@ def set_worlds_to_configure(self, world_names: set[str]) -> None:
         """Set the world names currently being configured."""
         self._curr_worlds_to_configure = world_names
 
+    def set_world_infos(self) -> None:
+        """Set new world infos."""
+        for world_info in self._worlds_to_add:
+            self._world_infos[world_info.name] = world_info
+
+    def get_world_infos(self) -> dict[str, WorldInfo]:
+        "Get world infos."
+        return self._world_infos
+
+    def is_first_run(self) -> bool:
+        "Return boolean if is first run or not."
+        return not self._world_infos
+
+    def remove_world_info(self, world_name: str) -> None:
+        """Remove world info by name."""
+        del self._world_infos[world_name]
+
+    def get_words_to_add(
+        self, world_infos: list[WorldInfo], new: set[str], cur: set[str]
+    ) -> list[WorldInfo]:
+        """Return a list of world infos to add."""
+        return [world_infos[name] for name in new - cur]
+
+    def get_words_to_remove(
+        self, world_infos: list[WorldInfo], new: set[str], cur: set[str]
+    ) -> list[WorldInfo]:
+        """Return a list of world infos to remove."""
+        return [world_infos[name] for name in cur - new]
+
     async def schedule(self, coro_factory: Callable[[], Awaitable[None]]):
         """Cancel any in-progress configure and schedule a new one."""
         # wait for current to finish if we do not want to cancel
@@ -75,6 +116,60 @@ async def schedule(self, coro_factory: Callable[[], Awaitable[None]]):
         self._event.clear()
         self._task = self._loop.create_task(self._run(coro_factory))
 
+    def _build_world_infos(self) -> dict[str, WorldInfo]:
+        world_infos: dict[str, WorldInfo] = {}
+
+        my_id = self._spec.stage.id
+        for k, v in self._spec.flow_graph.items():
+            for cfg_world_info in v:
+                # NOTE: no. of peers is always 1 for now
+                assert len(cfg_world_info.peers) == 1
+
+                if my_id == k:
+                    my_rank = 0
+                    other_rank = 1
+                    other_id = cfg_world_info.peers[0]
+                elif my_id in cfg_world_info.peers:
+                    # NOTE: this is always 1 for now
+                    my_rank = cfg_world_info.peers.index(my_id) + 1
+                    other_rank = 0
+                    other_id = k
+                else:
+                    continue
+
+                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.conflict_count,
+                )
+
+                world_size = len(cfg_world_info.peers) + 1
+                ctrl_ch = CtrlCh(my_rank, world_size, addr, ctrl_port)
+
+                data = {
+                    "name": name,
+                    "size": world_size,
+                    "addr": addr,
+                    "port": data_port,
+                    "backend": backend,
+                    "channel": ctrl_ch,
+                    "my_id": my_id,
+                    "me": my_rank,
+                    "other_id": other_id,
+                    "other": other_rank,
+                    "recover": recover,
+                    "conflict_count": conflict_count,
+                    "multiworld_name": f"{name}-{conflict_count}",
+                }
+                world_info = WorldInfo(**data)
+                world_infos[name] = world_info
+
+        return world_infos
+
     async def _run(self, coro_factory: Callable[[], Awaitable[None]]):
         """Run coroutine factory."""
         try:

infscale/execution/pipeline.py

@@ -68,7 +68,6 @@ def __init__(
         self.wcomm = wcomm
         self.spec: ServeConfig = None
         self.device = None
-        self.world_infos: dict[str, WorldInfo] = {}
         self.cfg_event = asyncio.Event()
         self._micro_batch_size = 1
         self._initialized = False
@@ -118,7 +117,9 @@ def _set_worker_status(self, status: WorkerStatus) -> None:
         """Set worker status in pipeline and channel."""
         self._status = status
 
-        for world_info in self.world_infos.values():
+        world_infos = self.config_runner.get_world_infos()
+
+        for world_info in world_infos.values():
             world_info.channel.set_worker_status(status)
 
     def _set_n_send_worker_status(self, status: WorkerStatus) -> None:
@@ -143,52 +144,42 @@ def _reset_control_channel(self, world_info: WorldInfo) -> None:
 
     async def _cleanup_recovered_worlds(self) -> None:
         """Clean up world infos for recovered worlds."""
+        world_infos = self.config_runner.get_world_infos()
+
         # if I'm the recovered worker, return
-        if len(self.world_infos) == 0:
+        if len(world_infos) == 0:
             return
 
         recover_worlds = [
             world_info
             for world_list in self.spec.flow_graph.values()
             for world_info in world_list
-            if world_info.recover and world_info.name in self.world_infos
+            if world_info.recover and world_info.name in world_infos
         ]
 
         # no worlds to recover
         if len(recover_worlds) == 0:
             return
 
         for world_info in recover_worlds:
-            wi = self.world_infos.get(world_info.name, None)
+            wi = world_infos.get(world_info.name, None)
 
             await self.router.cleanup_world(wi)
             self._reset_control_channel(wi)
             self._reset_multiworld(wi)
 
-            del self.world_infos[wi.name]
+            self.config_runner.remove_world_info(wi.name)
 
     async def _configure(self) -> None:
         """(Re)configure multiworld, control channel and router."""
-        await self._cleanup_recovered_worlds()
-
-        is_first_run = not self.world_infos
+        is_first_run = self.config_runner.is_first_run()
 
         if not is_first_run:
             self._set_worker_status(WorkerStatus.UPDATING)
 
-        new_world_infos = self._build_world_infos()
-        new = new_world_infos.keys()
-        cur = self.world_infos.keys()
-
-        worlds_to_add = [new_world_infos[name] for name in new - cur]
-        worlds_to_remove = [self.world_infos[name] for name in cur - new]
-
-        self.config_runner.set_worlds_to_configure(new - cur)
-
-        # handle new worlds
         tasks = []
         # 1. set up control channel
-        for world_info in worlds_to_add:
+        for world_info in self.config_runner._worlds_to_add:
             task = self._configure_control_channel(world_info)
             tasks.append(task)
 
@@ -198,7 +189,7 @@ async def _configure(self) -> None:
 
         tasks = []
         # 2. set up multiworld
-        for world_info in worlds_to_add:
+        for world_info in self.config_runner._worlds_to_add:
             task = self._configure_multiworld(world_info)
             tasks.append(task)
 
@@ -207,23 +198,25 @@ async def _configure(self) -> None:
         await asyncio.gather(*tasks)
 
         # update world_info for added worlds
-        for world_info in worlds_to_add:
-            self.world_infos[world_info.name] = world_info
+        self.config_runner.set_world_infos()
 
         # configure router with worlds to add and remove
         await self.router.configure(
-            self.spec, self.device, worlds_to_add, worlds_to_remove
+            self.spec,
+            self.device,
+            self.config_runner._worlds_to_add,
+            self.config_runner._worlds_to_remove,
         )
 
         # handle unnecessary world
         # remove is executed in the reverse order of add
-        for world_info in worlds_to_remove:
+        for world_info in self.config_runner._worlds_to_remove:
             # 1. remove unnecessary world from control channel
             self._reset_control_channel(world_info)
             # 2. remove unnecessary world from multiworld
             self._reset_multiworld(world_info)
 
-            del self.world_infos[world_info.name]
+            self.config_runner.remove_world_info(world_info.name)
 
         worker_status = WorkerStatus.RUNNING if is_first_run else WorkerStatus.UPDATED
 
@@ -426,10 +419,12 @@ async def _handle_config(self, spec: ServeConfig) -> None:
         if spec is None:
             return
 
-        self.config_runner.handle_new_spec(spec)
-
         self._configure_variables(spec)
 
+        await self._cleanup_recovered_worlds()
+
+        self.config_runner.handle_new_spec(spec)
+
         self._inspector.configure(self.spec)
 
         self._initialize_once()