handle graceful shutdown

Author: v1r3n

src/conductor/client/automator/async_task_runner.py

@@ -110,6 +110,7 @@ def __init__(
 
         # Semaphore will be created in run() within the event loop
         self._semaphore = None
+        self._shutdown = False  # Flag to indicate graceful shutdown
 
     async def run(self) -> None:
         """Main async loop - runs continuously in single event loop."""
@@ -150,12 +151,49 @@ async def run(self) -> None:
         )
 
         try:
-            while True:
+            while not self._shutdown:
                 await self.run_once()
         finally:
-            # Cleanup async client on exit
-            if self.async_api_client:
+            # Cleanup resources on exit
+            await self._cleanup()
+
+    async def stop(self) -> None:
+        """Signal the runner to stop gracefully."""
+        self._shutdown = True
+
+    async def _cleanup(self) -> None:
+        """Clean up async resources."""
+        logger.debug("Cleaning up AsyncTaskRunner resources...")
+
+        # Cancel any running tasks (EAFP style)
+        try:
+            for task in list(self._running_tasks):
+                if not task.done():
+                    task.cancel()
+        except AttributeError:
+            pass  # No tasks to cancel
+
+        # Close async HTTP client
+        if self.async_api_client:
+            try:
                 await self.async_api_client.close()
+                logger.debug("Async API client closed successfully")
+            except (IOError, OSError) as e:
+                logger.warning(f"Error closing async client: {e}")
+
+        # Clear event listeners
+        self.event_dispatcher = None
+
+        logger.debug("AsyncTaskRunner cleanup completed")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit - ensures cleanup."""
+        await self._cleanup()
+        return False  # Don't suppress exceptions
 
     async def __async_register_task_definition(self) -> None:
         """
@@ -384,10 +422,8 @@ async def __async_register_task_definition(self) -> None:
     async def run_once(self) -> None:
         """Execute one iteration of the polling loop (async version)."""
         try:
-            # Cleanup completed tasks
-            self.__cleanup_completed_tasks()
-
-            # Check if we can accept more tasks
+            # No need for manual cleanup - tasks remove themselves via add_done_callback
+            # Just check capacity directly
             current_capacity = len(self._running_tasks)
             if current_capacity >= self._max_workers:
                 # At capacity - sleep briefly then return
@@ -435,10 +471,6 @@ async def run_once(self) -> None:
         except Exception as e:
             logger.error("Error in run_once: %s", traceback.format_exc())
 
-    def __cleanup_completed_tasks(self) -> None:
-        """Remove completed task futures from tracking set (same as TaskRunner)."""
-        self._running_tasks = {f for f in self._running_tasks if not f.done()}
-
     async def __async_batch_poll(self, count: int) -> list:
         """Async batch poll for multiple tasks (async version of TaskRunner.__batch_poll_tasks)."""
         task_definition_name = self.worker.get_task_definition_name()
@@ -533,7 +565,8 @@ async def __async_batch_poll(self, count: int) -> list:
 
     async def __async_execute_and_update_task(self, task: Task) -> None:
         """Execute task and update result (async version - runs in event loop, not thread pool)."""
-        # Acquire semaphore to limit concurrency
+        # Acquire semaphore for entire task lifecycle (execution + update)
+        # This ensures we never exceed thread_count tasks in any stage of processing
         async with self._semaphore:
             try:
                 task_result = await self.__async_execute_task(task)

src/conductor/client/automator/task_handler.py

@@ -347,10 +347,22 @@ def __stop_process(self, process: Process):
         try:
             logger.debug("Terminating process: %s", process.pid)
             process.terminate()
+            # Wait for graceful termination
+            process.join(timeout=5.0)
+            if process.is_alive():
+                logger.warning("Process %s did not terminate within 5 seconds, killing...", process.pid)
+                process.kill()
+                process.join(timeout=1.0)  # Wait after kill too
+                if process.is_alive():
+                    logger.error("Process %s could not be killed", process.pid)
         except Exception as e:
             logger.debug("Failed to terminate process: %s, reason: %s", process.pid, e)
-            process.kill()
-            logger.debug("Killed process: %s", process.pid)
+            try:
+                process.kill()
+                process.join(timeout=1.0)
+                logger.debug("Killed process: %s", process.pid)
+            except Exception as kill_error:
+                logger.error("Failed to kill process: %s, reason: %s", process.pid, kill_error)
 
 
 # Setup centralized logging queue

src/conductor/client/automator/task_runner.py

@@ -5,6 +5,7 @@
 import time
 import traceback
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Optional, Any
 
 from conductor.client.configuration.configuration import Configuration
 from conductor.client.configuration.settings.metrics_settings import MetricsSettings
@@ -90,6 +91,7 @@ def __init__(
         self._max_workers = max_workers
         self._last_poll_time = 0  # Track last poll to avoid excessive polling when queue is empty
         self._consecutive_empty_polls = 0  # Track empty polls to implement backoff
+        self._shutdown = False  # Flag to indicate graceful shutdown
 
     def run(self) -> None:
         if self.configuration is not None:
@@ -114,8 +116,53 @@ def run(self) -> None:
             self.worker.get_polling_interval_in_seconds()
         )
 
-        while True:
-            self.run_once()
+        try:
+            while not self._shutdown:
+                self.run_once()
+        finally:
+            # Cleanup resources on exit
+            self._cleanup()
+
+    def stop(self) -> None:
+        """Signal the runner to stop gracefully."""
+        self._shutdown = True
+
+    def _cleanup(self) -> None:
+        """Clean up resources - called on exit."""
+        logger.debug("Cleaning up TaskRunner resources...")
+
+        # Shutdown ThreadPoolExecutor (EAFP style - more Pythonic)
+        try:
+            self._executor.shutdown(wait=True, cancel_futures=True)
+            logger.debug("ThreadPoolExecutor shut down successfully")
+        except AttributeError:
+            pass  # No executor to shutdown
+        except (RuntimeError, ValueError) as e:
+            logger.warning(f"Error shutting down executor: {e}")
+
+        # Close HTTP client (EAFP style)
+        try:
+            rest_client = self.task_client.api_client.rest_client
+            rest_client.close()
+            logger.debug("HTTP client closed successfully")
+        except AttributeError:
+            pass  # No client to close or no close method
+        except (IOError, OSError) as e:
+            logger.warning(f"Error closing HTTP client: {e}")
+
+        # Clear event listeners
+        self.event_dispatcher = None
+
+        logger.debug("TaskRunner cleanup completed")
+
+    def __enter__(self):
+        """Context manager entry - returns self for 'with' statement usage."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures cleanup even if exception occurs."""
+        self._cleanup()
+        return False  # Don't suppress exceptions
 
     def __register_task_definition(self) -> None:
         """
@@ -350,8 +397,11 @@ def run_once(self) -> None:
             self.__cleanup_completed_tasks()
 
             # Check if we can accept more tasks (based on thread_count)
-            # Account for pending async tasks in capacity calculation
-            pending_async_count = len(getattr(self.worker, '_pending_async_tasks', {}))
+            # Account for pending async tasks in capacity calculation (thread-safe)
+            pending_async_count = 0
+            if hasattr(self.worker, '_pending_tasks_lock') and hasattr(self.worker, '_pending_async_tasks'):
+                with self.worker._pending_tasks_lock:
+                    pending_async_count = len(self.worker._pending_async_tasks)
             current_capacity = len(self._running_tasks) + pending_async_count
             if current_capacity >= self._max_workers:
                 # At capacity - sleep briefly then return to check again
@@ -397,9 +447,11 @@ def run_once(self) -> None:
             logger.error("Error in run_once: %s", traceback.format_exc())
 
     def __cleanup_completed_tasks(self) -> None:
-        """Remove completed task futures from tracking set"""
-        # Fast path: use difference_update for better performance
-        self._running_tasks = {f for f in self._running_tasks if not f.done()}
+        """Remove completed task futures from tracking set (thread-safe)"""
+        # Avoid recreating the set - modify in place to prevent race conditions
+        completed = [f for f in self._running_tasks if f.done()]
+        for f in completed:
+            self._running_tasks.discard(f)
 
     def __check_completed_async_tasks(self) -> None:
         """Check for completed async tasks and update Conductor"""

src/conductor/client/http/api_client.py

@@ -765,7 +765,9 @@ def __refresh_auth_token(self) -> None:
             return
         # Initial token generation - apply backoff if there were previous failures
         token = self.__get_new_token(skip_backoff=False)
-        self.configuration.update_token(token)
+        # If token is None and auth was disabled (404), skip update
+        if token is not None or self.configuration.authentication_settings is not None:
+            self.configuration.update_token(token)
 
     def force_refresh_auth_token(self) -> bool:
         """
@@ -780,6 +782,10 @@ def force_refresh_auth_token(self) -> bool:
         if token:
             self.configuration.update_token(token)
             return True
+        # Check if auth was disabled during token refresh (404 response)
+        if self.configuration.authentication_settings is None:
+            logger.info('Authentication was disabled (no auth endpoint found)')
+            return False
         return False
 
     def __force_refresh_auth_token(self) -> bool:
@@ -857,6 +863,28 @@ def __get_new_token(self, skip_backoff: bool = False) -> str:
             )
             return None
 
+        except rest.ApiException as ae:
+            # Check if it's a 404 - indicates no authentication endpoint (Conductor OSS)
+            if ae.is_not_found():
+                logger.info(
+                    'Authentication endpoint /token not found (404). '
+                    'Running in open mode without authentication (Conductor OSS).'
+                )
+                # Disable authentication to prevent future attempts
+                self.configuration.authentication_settings = None
+                self.configuration.AUTH_TOKEN = None
+                # Reset failure counter since this is not a failure
+                self._token_refresh_failures = 0
+                return None
+            else:
+                # Other API errors
+                self._token_refresh_failures += 1
+                logger.error(
+                    f'API error when getting token (attempt {self._token_refresh_failures}): '
+                    f'{ae.status} - {ae.reason}'
+                )
+                return None
+
         except Exception as e:
             # Other errors (network, etc)
             self._token_refresh_failures += 1

src/conductor/client/worker/worker.py

@@ -330,6 +330,8 @@ def __init__(self,
 
         # Track pending async tasks: {task_id -> (future, task, submit_time)}
         self._pending_async_tasks = {}
+        # Add thread lock for safe concurrent access to _pending_async_tasks
+        self._pending_tasks_lock = threading.Lock()
 
     def execute(self, task: Task) -> TaskResult:
         task_input = {}
@@ -367,15 +369,17 @@ def execute(self, task: Task) -> TaskResult:
                 # This allows high concurrency for async I/O-bound workloads
                 future = self._background_loop.submit_coroutine(task_output)
 
-                # Store future for later retrieval
+                # Store future for later retrieval (thread-safe)
                 submit_time = time.time()
-                self._pending_async_tasks[task.task_id] = (future, task, submit_time)
+                with self._pending_tasks_lock:
+                    self._pending_async_tasks[task.task_id] = (future, task, submit_time)
+                    pending_count = len(self._pending_async_tasks)
 
                 logger.debug(
                     "Submitted async task: %s (task_id=%s, pending_count=%d, submit_time=%s)",
                     task.task_def_name,
                     task.task_id,
-                    len(self._pending_async_tasks),
+                    pending_count,
                     submit_time
                 )
 
@@ -455,11 +459,16 @@ def check_completed_async_tasks(self) -> list:
         completed_results = []
         tasks_to_remove = []
 
-        pending_count = len(self._pending_async_tasks)
+        # Create snapshot of pending tasks to avoid iteration during modification
+        with self._pending_tasks_lock:
+            tasks_snapshot = list(self._pending_async_tasks.items())
+            pending_count = len(self._pending_async_tasks)
+
         if pending_count > 0:
             logger.debug(f"Checking {pending_count} pending async tasks")
 
-        for task_id, (future, task, submit_time) in list(self._pending_async_tasks.items()):
+        # Process snapshot outside of lock
+        for task_id, (future, task, submit_time) in tasks_snapshot:
             if future.done():  # Non-blocking check
                 done_time = time.time()
                 actual_duration = done_time - submit_time
@@ -531,9 +540,11 @@ def check_completed_async_tasks(self) -> list:
                     completed_results.append((task_id, task_result, submit_time, task))
                     tasks_to_remove.append(task_id)
 
-        # Remove completed tasks
-        for task_id in tasks_to_remove:
-            del self._pending_async_tasks[task_id]
+        # Remove completed tasks (thread-safe)
+        with self._pending_tasks_lock:
+            for task_id in tasks_to_remove:
+                # Use pop to avoid KeyError if task was already removed
+                self._pending_async_tasks.pop(task_id, None)
 
         return completed_results