fix: pending requests not forwarded on worker termination (#316)

When a new config is received that is supposed to terminate some workers, we blindly close connection between these workers and their receivers, which blocks the worker to be terminated forward remaining requests to it's receiver. In order to fix this, code was updated in a way that only the senders of the workers to be terminated will close their connection and the receivers of the workers to be terminated will close the connection when a broken world exception is raised, meaning that the worker is terminated and all requests were forwarded. Also, refactored the code so the cleanup is done in router only.

Author: raresgaia123

infscale/execution/pipeline.py

@@ -134,14 +134,6 @@ async def _configure_control_channel(self, world_info: WorldInfo) -> None:
 
         await world_info.channel.wait_readiness()
 
-    def _reset_multiworld(self, world_info: WorldInfo) -> None:
-        self.world_manager.remove_world(world_info.multiworld_name)
-        logger.info(f"remove world {world_info.multiworld_name} from multiworld")
-
-    def _reset_control_channel(self, world_info: WorldInfo) -> None:
-        world_info.channel.cleanup()
-        logger.info(f"remove world {world_info.name} from control channel")
-
     async def _cleanup_recovered_worlds(self) -> None:
         """Clean up world infos for recovered worlds."""
         world_infos = self.config_manager.get_world_infos()
@@ -165,8 +157,6 @@ async def _cleanup_recovered_worlds(self) -> None:
             wi = world_infos.get(world_info.name, None)
 
             await self.router.cleanup_world(wi)
-            self._reset_control_channel(wi)
-            self._reset_multiworld(wi)
 
             self.config_manager.remove_world_info(wi.name)
 
@@ -217,10 +207,11 @@ async def _configure(self) -> None:
         # handle unnecessary world
         # remove is executed in the reverse order of add
         for world_info in 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)
+            # cleanup of control channel and multiworld was moved into router
+            # since we need to do async world cleanup based on certain scenarios
+            # sender can do the cleanup when new config is processed to stop
+            # sending requests to failed / removed worker
+            # received needs to keep waiting for requests until an exception is raised
 
             self.config_manager.remove_world_info(world_info.name)
 

infscale/execution/router.py

@@ -153,6 +153,7 @@ async def configure(
             if world_info.me == 0:
                 continue
 
+            # do cleanup for send task and world
             await self.cleanup_world(world_info)
 
     async def _recv(self, world_info: WorldInfo, cancellable: asyncio.Event) -> None:
@@ -201,7 +202,7 @@ async def wait_on_term_ready(self) -> None:
     async def cleanup_world(
         self, world_info: WorldInfo, cancel_task: bool = True
     ) -> None:
-        """Cleanup state info on world."""
+        """Cleanup state info, channel and multiworld for world."""
         # reset tx q related to a given world info
         self._cleanup_tx_q(world_info)
 
@@ -217,6 +218,9 @@ async def cleanup_world(
         except Exception as e:
             logger.warning(f"failed to cancel task for world {name}: {e}")
 
+        world_info.channel.cleanup()
+        self.world_manager.remove_world(world_info.multiworld_name)
+
     def _find_tx_q(self, world_info: WorldInfo) -> asyncio.Queue:
         for _, v in self.__tx_qs.items():
             for wi, q in v: