feat: handle individual world configuration

Since having one config for the entire configuration couldn't handle partial worlds update, we implemented per world configuration tasks. Config manager takes care of scheduling and canceling the tasks based on some diffs between old and new config. Also, config manager will orchestrate processing configs based on what it needs to cancel. Updated the code with proper error handlers for cancelling tasks so it properly propagates to parent task. With this, world configuration can be tracked individually, making easier to implement retry on failure in the future.

Author: raresgaia123

infscale/configs/job.py

@@ -157,11 +157,16 @@ def get_worlds_to_configure(
         """Compare two specs and return new and updated worlds."""
         helper = ServeConfigHelper()
 
-        curr_worlds = helper._get_worlds(curr_spec)
         new_worlds = helper._get_worlds(new_spec)
+        new_world_names = set(new_worlds.keys())
+
+        # if current spec is not available,
+        # return worlds from the new spec.
+        if curr_spec is None:
+            return new_world_names
 
+        curr_worlds = helper._get_worlds(curr_spec)
         curr_world_names = set(curr_worlds.keys())
-        new_world_names = set(new_worlds.keys())
 
         deploy_worlds = new_world_names - curr_world_names
 

infscale/execution/config_manager.py

@@ -30,88 +30,126 @@ class ConfigManager:
     def __init__(self):
         """Initialize config manager instance."""
         self._loop = asyncio.get_event_loop()
-        self._task: asyncio.Task | None = None
-        self._event = asyncio.Event()
-        self._spec: ServeConfig = None
-        self._event.set()
-        self._curr_worlds_to_configure: set[str] = set()
-        self._cancel_cur_cfg = False
-        self._world_infos: dict[str, WorldInfo] = {}
-
-    def handle_new_spec(self, spec: ServeConfig) -> None:
+        # semaphore event for back-to-back configs
+        self._config_event = asyncio.Event()
+        self._config_event.set()
+        self._world_tasks: dict[str, asyncio.Task] = {}
+        self._new_spec: ServeConfig = None
+        self._curr_spec: ServeConfig = None
+        self._curr_world_infos: dict[str, WorldInfo] = {}
+        self._new_world_infos: dict[str, WorldInfo] = {}
+        self.worlds_to_cancel = set()
+
+    async def handle_new_spec(self, spec: ServeConfig) -> None:
         """Handle new spec."""
-        self._cancel_cur_cfg = self._should_cancel_current(spec)
-        self._spec = spec
+        new_worlds_to_configure = ServeConfig.get_worlds_to_configure(
+            self._curr_spec, spec
+        )
 
-    def _should_cancel_current(self, spec: ServeConfig) -> bool:
-        """Decide if current configuration should be cancelled."""
-        if self._spec is None:
-            return False
+        # on the first run, both new and cur will be empty sets
+        new = self._new_world_infos.keys()
+        cur = self._curr_world_infos.keys()
+        curr_worlds_to_configure = new - cur
 
-        new_worlds_to_configure = ServeConfig.get_worlds_to_configure(self._spec, spec)
+        self.worlds_to_cancel = new_worlds_to_configure & curr_worlds_to_configure
 
-        # cancel if the new config affects worlds currently being configured
-        # TODO: if there's a overlap between new worlds and curr worlds we cancel
-        # current configuration. This needs to be fixed, to cancel only the worlds that
-        # are affected (eg new_worlds & curr_worlds)
-        return not new_worlds_to_configure.isdisjoint(self._curr_worlds_to_configure)
+        self._curr_world_infos = (
+            self._build_world_infos(self._curr_spec) if self._curr_spec else {}
+        )
+        self._new_world_infos = self._build_world_infos(spec)
+        self._new_spec = spec
 
-    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
+        if len(self.worlds_to_cancel):
+            await self._cancel_world_configuration(self.worlds_to_cancel)
 
-    def set_world_infos(self, worlds: list[WorldInfo]) -> None:
-        """Set new world infos."""
-        for world_info in worlds:
-            self._world_infos[world_info.name] = world_info
+        # wait for current configuration to finish
+        await self._config_event.wait()
 
-    def get_world_infos(self) -> dict[str, WorldInfo]:
-        "Get world infos."
-        return self._world_infos
+        # do cleanup after each successful configuration
+        self._curr_spec = self._new_spec
+        self._new_spec = None
+        self.worlds_to_cancel = set()
+
+        # block handling new spec after doing cleanup for the current one
+        self._config_event.clear()
+
+    def unblock_next_config(self) -> None:
+        """Set task event and unblock next config process."""
+        self._config_event.set()
+
+    def update_world_infos(self, worlds_names: set[str]) -> None:
+        """Update world infos."""
+        for world_name in worlds_names:
+            world_info = self._new_world_infos[world_name]
+            self._curr_world_infos[world_info.name] = world_info
+
+    def get_curr_world_infos(self) -> dict[str, WorldInfo]:
+        "Get current world infos."
+        return self._curr_world_infos
 
     def is_first_run(self) -> bool:
         "Return boolean if is first run or not."
-        return not self._world_infos
+        return not self._curr_world_infos
 
     def remove_world_info(self, world_name: str) -> None:
         """Remove world info by name."""
-        del self._world_infos[world_name]
+        del self._curr_world_infos[world_name]
 
-    def get_worlds_to_add_and_remove(self) -> tuple[list[WorldInfo], list[WorldInfo]]:
+    def get_worlds_to_add_and_remove(self) -> tuple[set[str], set[str]]:
         """Return a list of world infos to add and to remove."""
-        new_world_infos = self._build_world_infos()
-
-        new = new_world_infos.keys()
-        cur = self._world_infos.keys()
+        new = self._new_world_infos.keys()
+        cur = self._curr_world_infos.keys()
 
-        worlds_to_add = [new_world_infos[name] for name in new - cur]
-        worlds_to_remove = [new_world_infos[name] for name in cur - new]
+        worlds_to_add = new - cur
+        worlds_to_remove = cur - new
 
         return worlds_to_add, worlds_to_remove
 
-    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
-        if not self._cancel_cur_cfg:
-            await self._event.wait()
-
-        # cancel current if running
-        if self._task and not self._task.done():
-            self._task.cancel()
-            try:
-                await self._task
-            except asyncio.CancelledError:
-                pass
-
-        # block again for new run
-        self._event.clear()
-        self._task = self._loop.create_task(self._run(coro_factory))
-
-    def _build_world_infos(self) -> dict[str, WorldInfo]:
+    def get_new_world_info(self, world_name) -> WorldInfo:
+        """Return new world info based on world name."""
+        return self._new_world_infos[world_name]
+
+    def get_world_infos_by_name(
+        self, world_names: set[str], to_remove: bool = False
+    ) -> list[WorldInfo]:
+        """Return a list of world infos to add based on world names."""
+        world_infos = []
+
+        target = self._new_world_infos
+
+        if to_remove:
+            target = self._curr_world_infos
+
+        for world_name in world_names:
+            world_infos.append(target[world_name])
+
+        return world_infos
+
+    async def _cancel_world_configuration(self, world_names: set[str]):
+        """Cancel only worlds that are impacted by new spec."""
+        coroutines = [self._cancel_world(w) for w in world_names]
+        await asyncio.gather(*coroutines, return_exceptions=True)
+
+    def schedule_world_cfg(
+        self, world_info: WorldInfo, coro_factory: Callable[[], Awaitable[None]]
+    ):
+        """Schedule configuration for a single world."""
+        task = self._loop.create_task(self._run_world(world_info, coro_factory))
+        self._world_tasks[world_info.name] = task
+        return task
+
+    async def _cancel_world(self, world_name: str):
+        """Cancel an in-progress world config task."""
+        task = self._world_tasks.pop(world_name, None)
+        if task and not task.done():
+            task.cancel()
+            raise asyncio.CancelledError
+
+    def _build_world_infos(self, spec: ServeConfig) -> dict[str, WorldInfo]:
         world_infos: dict[str, WorldInfo] = {}
 
-        my_id = self._spec.stage.id
-        for k, v in self._spec.flow_graph.items():
+        my_id = spec.stage.id
+        for k, v in 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
@@ -161,13 +199,13 @@ def _build_world_infos(self) -> dict[str, WorldInfo]:
 
         return world_infos
 
-    async def _run(self, coro_factory: Callable[[], Awaitable[None]]):
-        """Run coroutine factory."""
+    async def _run_world(
+        self, world_info: WorldInfo, coro_factory: Callable[[], Awaitable[None]]
+    ):
+        """Run and cleanup world configuration."""
         try:
-            await coro_factory()
+            await coro_factory(world_info)
         except asyncio.CancelledError:
-            pass
+            raise
         finally:
-            # reset class attributes and events
-            self._event.set()
-            self._curr_worlds_to_configure = set()
+            self._world_tasks.pop(world_info.name, None)

infscale/execution/control.py

@@ -162,10 +162,31 @@ async def setup(self) -> None:
         if self.rank == 0:
             self._server_task = asyncio.create_task(self._setup_server(setup_done))
         else:
-            _ = asyncio.create_task(self._setup_client(setup_done))
+            client_task = asyncio.create_task(self._setup_client(setup_done))
 
         # wait until setting up either server or client is done
-        await setup_done.wait()
+        try:
+            await setup_done.wait()
+        except asyncio.CancelledError as e:
+            # since both _setup_server and _setup_client are spawned as separate tasks
+            # and the setup itself is a task, we need to handle parent task cancellation
+            # on the awaited line, since cancellation only propagates through awaited calls
+            # here, await setup_done.wait() is the propagation point from parent task to child tasks
+            # so we need to cancel child tasks whenever CancelledError is received
+            if self._server_task and not self._server_task.done():
+                self._server_task.cancel()
+                try:
+                    await self._server_task
+                except asyncio.CancelledError:
+                    pass
+
+            if client_task and not client_task.done():
+                client_task.cancel()
+                try:
+                    await client_task
+                except asyncio.CancelledError:
+                    pass
+            raise
 
     def cleanup(self) -> None:
         if self._server_task is not None: