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 runner takes care of scheduling and canceling the tasks based on some diffs between old and new config. Also, config runner 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_runner.py

@@ -30,44 +30,54 @@ class ConfigRunner:
     def __init__(self):
         """Initialize config runner instance."""
         self._loop = asyncio.get_event_loop()
-        self._task: asyncio.Task | None = None
         self._event = asyncio.Event()
+        self._world_tasks: dict[str, asyncio.Task] = {}
         self._spec: ServeConfig = None
         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:
+    async def handle_new_spec(self, spec: ServeConfig) -> None:
         """Handle new spec."""
-        self._cancel_cur_cfg = self._should_cancel_current(spec)
+        new_worlds_to_configure = ServeConfig.get_worlds_to_configure(self._spec, spec)
+        worlds_to_cancel = new_worlds_to_configure & self._curr_worlds_to_configure
         self._spec = spec
-        new_world_infos = self._build_world_infos()
 
+        if len(worlds_to_cancel):
+            # since _worlds_to_add are used in pipeline, we need to update them
+            # by removing the ones that will be cancelled.
+            self._worlds_to_add = [
+                world_info
+                for world_info in self._worlds_to_add
+                if world_info.name not in worlds_to_cancel
+            ]
+            await self._cancel_world_configuration(worlds_to_cancel)
+
+        # after cancelling affected worlds, wait for configuration to finish
+        # then process next config.
+        await self._event.wait()
+
+        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."""
-        if self._spec is None:
-            return False
+        self._curr_worlds_to_configure = new - cur
 
-        new_worlds_to_configure = ServeConfig.get_worlds_to_configure(self._spec, spec)
+        # block configuration until current config is processed
+        self._event.clear()
 
-        # 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)
+    def unblock_next_config(self) -> None:
+        """Set task event and unblock next config process."""
+        self._event.set()
 
-    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 reset_curr_worlds_to_configure(self) -> None:
+        """Reset current worlds to configure."""
+        self._curr_worlds_to_configure = set()
 
     def set_world_infos(self) -> None:
         """Set new world infos."""
@@ -98,23 +108,28 @@ def get_words_to_remove(
         """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
-        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()
+    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()
             try:
-                await self._task
+                await task
             except asyncio.CancelledError:
-                pass
-
-        # block again for new run
-        self._event.clear()
-        self._task = self._loop.create_task(self._run(coro_factory))
+                raise
 
     def _build_world_infos(self) -> dict[str, WorldInfo]:
         world_infos: dict[str, WorldInfo] = {}
@@ -170,13 +185,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,32 @@ 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:
+            # logger.warning(f"[{self.rank}] channel setup cancelled")
+            # 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:

infscale/execution/pipeline.py

@@ -104,6 +104,8 @@ async def _configure_multiworld(self, world_info: WorldInfo) -> None:
                 port=port,
                 device=self.device,
             )
+        except asyncio.CancelledError:
+            logger.warning(f"multiworld configuration cancelled for {world_info.name}")
         except Exception as e:
             logger.error(f"failed to initialize a multiworld {name}: {e}")
             condition = self._status != WorkerStatus.UPDATING
@@ -130,9 +132,12 @@ def _set_n_send_worker_status(self, status: WorkerStatus) -> None:
         self.wcomm.send(msg)
 
     async def _configure_control_channel(self, world_info: WorldInfo) -> None:
-        await world_info.channel.setup()
+        try:
+            await world_info.channel.setup()
 
-        await world_info.channel.wait_readiness()
+            await world_info.channel.wait_readiness()
+        except asyncio.CancelledError:
+            logger.warning(f"channel configuration cancelled for {world_info.name} ")
 
     def _reset_multiworld(self, world_info: WorldInfo) -> None:
         self.world_manager.remove_world(world_info.multiworld_name)
@@ -180,7 +185,9 @@ async def _configure(self) -> None:
         tasks = []
         # 1. set up control channel
         for world_info in self.config_runner._worlds_to_add:
-            task = self._configure_control_channel(world_info)
+            task = self.config_runner.schedule_world_cfg(
+                world_info, self._configure_control_channel
+            )
             tasks.append(task)
 
         # TODO: this doesn't handle partial success
@@ -190,7 +197,9 @@ async def _configure(self) -> None:
         tasks = []
         # 2. set up multiworld
         for world_info in self.config_runner._worlds_to_add:
-            task = self._configure_multiworld(world_info)
+            task = self.config_runner.schedule_world_cfg(
+                world_info, self._configure_multiworld
+            )
             tasks.append(task)
 
         # TODO: this doesn't handle partial success
@@ -220,6 +229,9 @@ async def _configure(self) -> None:
 
         worker_status = WorkerStatus.RUNNING if is_first_run else WorkerStatus.UPDATED
 
+        # config is done, do cleanup in config runner
+        self.config_runner.reset_curr_worlds_to_configure()
+        self.config_runner.unblock_next_config()
         self._set_n_send_worker_status(worker_status)
 
         self.cfg_event.set()
@@ -423,14 +435,14 @@ async def _handle_config(self, spec: ServeConfig) -> None:
 
         await self._cleanup_recovered_worlds()
 
-        self.config_runner.handle_new_spec(spec)
-
         self._inspector.configure(self.spec)
 
         self._initialize_once()
-
         # (re)configure the pipeline
-        await self.config_runner.schedule(self._configure)
+        await self.config_runner.handle_new_spec(spec)
+        # run configure as a separate task since we need to unblock receiving
+        # a new config to be processed when current configuration is finished
+        self._configure_task = asyncio.create_task(self._configure())
 
     def _build_world_infos(self) -> dict[str, WorldInfo]:
         world_infos: dict[str, WorldInfo] = {}