chore: update_workforce_stop_logic (#3447)

Author: fengju0213

camel/societies/workforce/worker.py

@@ -45,6 +45,7 @@ def __init__(
     ) -> None:
         super().__init__(description, node_id=node_id)
         self._active_task_ids: Set[str] = set()
+        self._running_tasks: Set[asyncio.Task] = set()
 
     def __repr__(self):
         return f"Worker node {self.node_id} ({self.description})"
@@ -113,15 +114,12 @@ async def _listen_to_channel(self):
         self._running = True
         logger.info(f"{self} started.")
 
-        # Keep track of running task coroutines
-        running_tasks: Set[asyncio.Task] = set()
-
         while self._running:
             try:
                 # Clean up completed tasks
-                completed_tasks = [t for t in running_tasks if t.done()]
+                completed_tasks = [t for t in self._running_tasks if t.done()]
                 for completed_task in completed_tasks:
-                    running_tasks.remove(completed_task)
+                    self._running_tasks.discard(completed_task)
                     # Check for exceptions in completed tasks
                     try:
                         await completed_task
@@ -138,11 +136,11 @@ async def _listen_to_channel(self):
                     task_coroutine = asyncio.create_task(
                         self._process_single_task(task)
                     )
-                    running_tasks.add(task_coroutine)
+                    self._running_tasks.add(task_coroutine)
 
                 except asyncio.TimeoutError:
                     # No tasks available, continue loop
-                    if not running_tasks:
+                    if not self._running_tasks:
                         # No tasks running and none available, short sleep
                         await asyncio.sleep(0.1)
                     continue
@@ -155,12 +153,12 @@ async def _listen_to_channel(self):
                 continue
 
         # Wait for all remaining tasks to complete when stopping
-        if running_tasks:
+        if self._running_tasks:
             logger.info(
-                f"{self} stopping, waiting for {len(running_tasks)} "
+                f"{self} stopping, waiting for {len(self._running_tasks)} "
                 f"tasks to complete..."
             )
-            await asyncio.gather(*running_tasks, return_exceptions=True)
+            await asyncio.gather(*self._running_tasks, return_exceptions=True)
 
         logger.info(f"{self} stopped.")
 
@@ -171,6 +169,20 @@ async def start(self):
 
     @check_if_running(True)
     def stop(self):
-        r"""Stop the worker."""
+        r"""Forcefully stop the worker.
+
+        Cancels all running tasks immediately and sets the stop flag.
+        The worker will exit after completing cancellation.
+        """
+        # First cancel all running tasks to interrupt ongoing work
+        tasks_to_cancel = list(self._running_tasks)
+        for task in tasks_to_cancel:
+            if not task.done():
+                task.cancel()
+
+        # Clear the running tasks set since they're all cancelled
+        self._running_tasks.clear()
+
+        # Set stop flag to exit the listen loop
         self._running = False
         return

camel/societies/workforce/workforce.py

@@ -1802,6 +1802,95 @@ def stop_gracefully(self) -> None:
                 f"(event-loop not yet started)."
             )
 
+    async def _async_stop_immediately(self) -> None:
+        r"""Force-stop the workforce without waiting for in-flight tasks.
+
+        Note: This method does not wait for child nodes to fully stop.
+        Child nodes will receive stop signals but may still be cleaning up
+        when this method returns.
+        """
+        # Guard against redundant calls
+        if not self._running and self._state == WorkforceState.STOPPED:
+            logger.debug(
+                f"Workforce {self.node_id} already stopped, skipping."
+            )
+            return
+
+        self._stop_requested = True
+        self._pause_event.set()
+        logger.info(f"Workforce {self.node_id} force stop requested.")
+
+        # Remove pending tasks and clear channel postings to avoid new work
+        self._pending_tasks.clear()
+
+        # Try to clear in-flight tasks from channel
+        if self._channel:
+            try:
+                in_flight = await self._channel.get_in_flight_tasks(
+                    self.node_id
+                )
+                for task in in_flight:
+                    await self._channel.remove_task(task.id)
+                # Only reset counter if channel cleanup succeeded
+                self._in_flight_tasks = 0
+            except Exception as e:
+                logger.error(
+                    f"Failed to clear in-flight tasks during force stop: {e}",
+                    exc_info=True,
+                )
+                # Still reset counter for consistency
+                self._in_flight_tasks = 0
+
+        # Stop child nodes and cancel their listening tasks immediately
+        for child in self._children:
+            if child._running:
+                child.stop()
+        for listening_task in self._child_listening_tasks:
+            if not listening_task.done():
+                listening_task.cancel()
+
+        self._running = False
+        self._state = WorkforceState.STOPPED
+
+    def stop_immediately(self) -> None:
+        r"""Force-stop without waiting for current tasks to finish.
+
+        Note: This method does not wait for child nodes to fully stop.
+        Child nodes will receive stop signals but may still be cleaning up
+        when this method returns.
+        """
+        # Guard against redundant calls
+        if not self._running and self._state == WorkforceState.STOPPED:
+            logger.debug(
+                f"Workforce {self.node_id} already stopped, skipping."
+            )
+            return
+
+        if self._loop and not self._loop.is_closed():
+            self._submit_coro_to_loop(self._async_stop_immediately())
+        else:
+            # No running loop; best-effort synchronous cleanup
+            self._stop_requested = True
+            self._pause_event.set()
+            self._pending_tasks.clear()
+            self._in_flight_tasks = 0
+
+            # Stop children
+            for child in self._children:
+                if child._running:
+                    child.stop()
+            # Cancel listening tasks
+            for task in self._child_listening_tasks:
+                if not task.done():
+                    task.cancel()
+
+            self._running = False
+            self._state = WorkforceState.STOPPED
+            logger.info(
+                f"Workforce {self.node_id} force stopped "
+                f"(event-loop not yet started)."
+            )
+
     async def _async_skip_gracefully(self) -> None:
         r"""Async implementation of skip_gracefully to run on the event
         loop.
@@ -5059,54 +5148,34 @@ async def start(self) -> None:
 
     @check_if_running(True)
     def stop(self) -> None:
-        r"""Stop all the child nodes under it. The node itself will be stopped
-        by its parent node.
-        """
-        # Stop all child nodes first
-        for child in self._children:
-            if child._running:
-                child.stop()
+        r"""Forcefully stop the workforce and its children immediately.
 
-        # Cancel child listening tasks
-        if self._child_listening_tasks:
-            try:
-                loop = asyncio.get_running_loop()
-                if loop and not loop.is_closed():
-                    # Create graceful cleanup task
-                    async def cleanup():
-                        await asyncio.sleep(0.1)  # Brief grace period
-                        for task in self._child_listening_tasks:
-                            if not task.done():
-                                task.cancel()
-
-                        # Handle both asyncio.Task and concurrent.futures.
-                        # Future
-                        awaitables = []
-                        for task in self._child_listening_tasks:
-                            if isinstance(task, concurrent.futures.Future):
-                                # Convert Future to awaitable
-                                awaitables.append(asyncio.wrap_future(task))
-                            else:
-                                # Already an asyncio.Task
-                                awaitables.append(task)
-
-                        await asyncio.gather(
-                            *awaitables,
-                            return_exceptions=True,
-                        )
-
-                    self._cleanup_task = loop.create_task(cleanup())
-                else:
-                    # No active loop, cancel immediately
-                    for task in self._child_listening_tasks:
-                        task.cancel()
-            except (RuntimeError, Exception) as e:
-                # Fallback: cancel immediately
-                logger.debug(f"Exception during task cleanup: {e}")
-                for task in self._child_listening_tasks:
-                    task.cancel()
-
-        self._running = False
+        This is now an immediate stop (was previously a graceful lifecycle
+        cleanup). It cancels child listeners, clears pending/in-flight tasks,
+        and sets state to STOPPED without waiting for active work to finish.
+        """
+        if self._loop and not self._loop.is_closed():
+            self._submit_coro_to_loop(self._async_stop_immediately())
+        else:
+            # No running loop; perform synchronous best-effort force stop
+            self._stop_requested = True
+            self._pause_event.set()
+            self._pending_tasks.clear()
+            self._in_flight_tasks = 0
+            # Stop children
+            for child in self._children:
+                if child._running:
+                    child.stop()
+            # Cancel listening tasks
+            for listening_task in self._child_listening_tasks:
+                if not listening_task.done():
+                    listening_task.cancel()
+            self._running = False
+            self._state = WorkforceState.STOPPED
+            logger.info(
+                f"Workforce {self.node_id} force stopped "
+                f"(event-loop not yet started)."
+            )
 
     def clone(self, with_memory: bool = False) -> 'Workforce':
         r"""Creates a new instance of Workforce with the same configuration.