inorbit-edge-executor 3.2.3__tar.gz → 3.2.4__tar.gz

{inorbit_edge_executor-3.2.3 → inorbit_edge_executor-3.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inorbit_edge_executor
-Version: 3.2.3
+Version: 3.2.4
 Summary: InOrbit Edge Missions Executor
 Author-email: Mariano Cereda <mariano.cereda@inorbit.ai>, Hernan Badenes <herchu@inorbit.ai>, Miguel Garcia <mike@inorbit.ai>
 License: MIT

{inorbit_edge_executor-3.2.3 → inorbit_edge_executor-3.2.4}/inorbit_edge_executor/worker.py

@@ -38,6 +38,7 @@ class Worker(Observable):
         self._robot_api_factory = None
         self._behavior_tree = None
         self._shared_memory = shared_memory
+        self._task: asyncio.Task = None
 
     def id(self):
         return self._mission.id
@@ -133,6 +134,10 @@ class Worker(Observable):
             self.set_paused(True)
             self._task.cancel(CANCEL_TASK_PAUSE_MESSAGE)
 
+    def get_task(self) -> asyncio.Task | None:
+        """Get the task associated with the worker. This is used to wait for the task to complete during shutdown."""
+        return self._task
+
     async def resume(self):
         """
         Resumes the worker if it was paused. Returns True if the worker was paused and resumed, False otherwise.

{inorbit_edge_executor-3.2.3 → inorbit_edge_executor-3.2.4}/inorbit_edge_executor/worker_pool.py

@@ -103,12 +103,53 @@ class WorkerPool:
         for worker_state in serialized_workers:
             await self.execute_serialized_worker(worker_state)
 
-    async def shutdown(self):
+    async def shutdown(self, timeout: float = 10.0):
         """
-        Stops the worker pool. It's actually not doing anything more than preventing new jobs
-        to be submitted.
+        Stops the worker pool gracefully by cancelling all running tasks.
+
+        Args:
+            timeout: Maximum time to wait for graceful shutdown before forcing cancellation
         """
+        if not self._running:
+            return
+
         self._running = False
+        logger.info("Starting worker pool shutdown...")
+
+        # Get all running workers and cancel them
+        async with self._mutex:
+            workers_to_cancel = list(self._workers.values())
+
+        if workers_to_cancel:
+            logger.info(f"Cancelling {len(workers_to_cancel)} running workers...")
+
+            # Cancel all workers and collect their tasks
+            cancelled_tasks = []
+            for worker in workers_to_cancel:
+                try:
+                    if worker.cancel():
+                        cancelled_tasks.append(worker.get_task())
+                except Exception as e:
+                    logger.error(f"Error cancelling worker {worker.id()}: {e}")
+
+            if cancelled_tasks:
+                # Wait for tasks to complete cancellation with timeout
+                try:
+                    await asyncio.wait_for(
+                        asyncio.gather(*cancelled_tasks, return_exceptions=True), timeout=timeout
+                    )
+                    logger.info("All workers cancelled gracefully")
+                except asyncio.TimeoutError:
+                    logger.warning(f"Timeout waiting for graceful shutdown after {timeout}s")
+        else:
+            logger.info("No running workers to cancel")
+
+        # Clean up workers dictionary
+        async with self._mutex:
+            self._workers.clear()
+
+        # Shutdown the database connection
+        await self._db.shutdown()
 
     async def notify(self, worker: Worker):
         """Notified when a worker changed its state. Persist it"""

{inorbit_edge_executor-3.2.3 → inorbit_edge_executor-3.2.4}/inorbit_edge_executor.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: inorbit_edge_executor
-Version: 3.2.3
+Version: 3.2.4
 Summary: InOrbit Edge Missions Executor
 Author-email: Mariano Cereda <mariano.cereda@inorbit.ai>, Hernan Badenes <herchu@inorbit.ai>, Miguel Garcia <mike@inorbit.ai>
 License: MIT

{inorbit_edge_executor-3.2.3 → inorbit_edge_executor-3.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "inorbit_edge_executor"
-version = "3.2.3"
+version = "3.2.4"
 description = "InOrbit Edge Missions Executor"
 readme = "README.md"
 authors = [
@@ -47,7 +47,7 @@ target-version = ['py310']
 line-length = 100
 
 [tool.bumpversion]
-current_version = "3.2.3"
+current_version = "3.2.4"
 commit = true
 tag = true
 tag_name = "v{new_version}"