Skip to content

Commit 9cfb17d

Browse files
UN-2813 [FIX] Implement graceful shutdown for signal handler
Improve signal handler to properly shut down backend instead of immediate exit that can leave tasks in inconsistent state. **Problem:** Signal handler immediately called sys.exit(0) without giving backend a chance to finish current tasks or clean up resources: ```python def _signal_handler(self, signum, frame): logger.info("Initiating worker shutdown") sys.exit(0) # ❌ Immediate exit, no cleanup ``` This can cause: - Tasks left in inconsistent states - Database connections not closed - Incomplete work lost - No proper cleanup of backend resources **Solution:** Implement graceful shutdown that: 1. Sets shutdown_requested flag 2. Calls backend.stop(graceful=True) if available 3. Logs errors with full traceback 4. Allows backend to finish current tasks ```python def _signal_handler(self, signum, frame): self.shutdown_requested = True logger.info("Initiating graceful worker shutdown") if self.backend and hasattr(self.backend, "stop"): try: self.backend.stop(graceful=True) except Exception as e: logger.error(f"Error during graceful shutdown: {e}", exc_info=True) sys.exit(0) ``` **Changes:** 1. ✅ Added backend.stop(graceful=True) call before exit 2. ✅ Proper try/except with exc_info logging 3. ✅ Check for backend availability with hasattr 4. ✅ More descriptive log messages 5. ✅ shutdown_requested flag still set for future use **Impact:** - ✅ Tasks can complete gracefully before shutdown - ✅ Backend resources properly cleaned up - ✅ Better error visibility during shutdown - ✅ Safer production deployments **Note:** Backend implementations should support stop() method for proper graceful shutdown: - Stop accepting new tasks - Wait for in-progress tasks (with timeout) - Clean up connections and resources - Exit cleanly Resolves: CodeRabbit comment #2404866807 Related: #1555
1 parent 3ee8100 commit 9cfb17d

File tree

1 file changed

+10
-4
lines changed
  • task-backend/src/unstract/task_backend

1 file changed

+10
-4
lines changed

task-backend/src/unstract/task_backend/worker.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -223,13 +223,19 @@ def _get_worker_config(self) -> dict:
223223
return base_config
224224

225225
def _signal_handler(self, signum, frame) -> None:
226-
"""Handle shutdown signals."""
226+
"""Handle shutdown signals gracefully."""
227227
logger.info(f"Received shutdown signal: {signum}")
228228
self.shutdown_requested = True
229229

230-
# Note: The backend's run_worker() method should handle
231-
# graceful shutdown. For now, we'll just exit.
232-
logger.info("Initiating worker shutdown")
230+
logger.info("Initiating graceful worker shutdown")
231+
if self.backend and hasattr(self.backend, "stop"):
232+
try:
233+
# Give the backend a chance to finish current tasks
234+
logger.info("Requesting backend graceful shutdown")
235+
self.backend.stop(graceful=True)
236+
except Exception as e:
237+
logger.error(f"Error during graceful shutdown: {e}", exc_info=True)
238+
233239
sys.exit(0)
234240

235241

0 commit comments

Comments
 (0)