software/src/orchestration/signals.py

1"""Graceful shutdown handling for fault-tolerant translation.23Provides signal handling to allow clean shutdown when the process4receives SIGINT (Ctrl+C) or SIGTERM, ensuring job state is saved.5"""67from __future__ import annotations89import signal10import threading11from types import FrameType12from typing import Any, Callable131415class GracefulShutdown:16    """Context manager for graceful shutdown handling.1718    Captures SIGINT and SIGTERM signals and sets a flag that can be19    checked by worker code to stop gracefully.2021    Example:22        with GracefulShutdown() as shutdown:23            while not shutdown.should_stop:24                # Do work...25                process_next_verse()2627    Attributes:28        should_stop: True when shutdown has been requested.29        signal_received: The signal number that was received, or None.30    """3132    def __init__(self) -> None:33        """Initialize the graceful shutdown handler."""34        self._should_stop = threading.Event()35        self._signal_received: int | None = None36        self._original_handlers: dict[int, Any] = {}37        self._lock = threading.Lock()3839    @property40    def should_stop(self) -> bool:41        """Check if shutdown has been requested.4243        Returns:44            True if SIGINT or SIGTERM was received.45        """46        return self._should_stop.is_set()4748    @property49    def signal_received(self) -> int | None:50        """Get the signal number that triggered shutdown.5152        Returns:53            Signal number (e.g., signal.SIGINT) or None if no signal received.54        """55        return self._signal_received5657    def request_stop(self) -> None:58        """Programmatically request a stop (without signal).5960        Useful for testing or other shutdown triggers.61        """62        self._should_stop.set()6364    def wait_for_stop(self, timeout: float | None = None) -> bool:65        """Wait for shutdown signal.6667        Args:68            timeout: Maximum time to wait in seconds, or None for indefinite.6970        Returns:71            True if stop was signaled, False if timeout occurred.72        """73        return self._should_stop.wait(timeout=timeout)7475    def _signal_handler(self, signum: int, frame: FrameType | None) -> None:76        """Handle shutdown signals.7778        Args:79            signum: Signal number received.80            frame: Current stack frame (unused).81        """82        with self._lock:83            if not self._should_stop.is_set():84                self._signal_received = signum85                self._should_stop.set()8687    def __enter__(self) -> GracefulShutdown:88        """Enter the context manager and install signal handlers.8990        Signal handlers can only be installed from the main thread.91        When used from a non-main thread, the context manager still92        works but without signal handling (should_stop can still be93        set via request_stop()).9495        Returns:96            Self for use in context manager.97        """98        if threading.current_thread() is threading.main_thread():99            # Store original handlers100            self._original_handlers[signal.SIGINT] = signal.signal(101                signal.SIGINT, self._signal_handler102            )103            self._original_handlers[signal.SIGTERM] = signal.signal(104                signal.SIGTERM, self._signal_handler105            )106107        return self108109    def __exit__(110        self,111        exc_type: type | None,112        exc_val: BaseException | None,113        exc_tb: Any | None,114    ) -> bool:115        """Exit the context manager and restore signal handlers.116117        Args:118            exc_type: Exception type if an exception occurred.119            exc_val: Exception value if an exception occurred.120            exc_tb: Exception traceback if an exception occurred.121122        Returns:123            False to propagate any exception.124        """125        # Restore original handlers126        for signum, handler in self._original_handlers.items():127            signal.signal(signum, handler)128129        self._original_handlers.clear()130131        return False132133134class ShutdownCallback:135    """Utility for registering shutdown callbacks.136137    Allows registering cleanup functions to be called when shutdown occurs.138139    Example:140        shutdown = GracefulShutdown()141        callbacks = ShutdownCallback(shutdown)142        callbacks.register(save_state)143        callbacks.register(close_connections)144145        with shutdown:146            while not shutdown.should_stop:147                do_work()148149        callbacks.run_all()  # Called after exiting context150    """151152    def __init__(self, shutdown: GracefulShutdown) -> None:153        """Initialize with a GracefulShutdown instance.154155        Args:156            shutdown: The shutdown handler to monitor.157        """158        self._shutdown = shutdown159        self._callbacks: list[Callable[[], None]] = []160        self._lock = threading.Lock()161162    def register(self, callback: Callable[[], None]) -> None:163        """Register a callback to run on shutdown.164165        Args:166            callback: Function to call during shutdown cleanup.167        """168        with self._lock:169            self._callbacks.append(callback)170171    def run_all(self) -> list[Exception]:172        """Run all registered callbacks.173174        Runs each callback, collecting any exceptions that occur.175        All callbacks are attempted even if some fail.176177        Returns:178            List of exceptions that occurred (empty if all succeeded).179        """180        exceptions = []181        with self._lock:182            for callback in self._callbacks:183                try:184                    callback()185                except Exception as e:186                    exceptions.append(e)187        return exceptions188189190def ignore_signals_in_child() -> None:191    """Configure signal handling for child processes.192193    Child processes in a multiprocessing pool should ignore194    SIGINT/SIGTERM and let the parent handle shutdown coordination.195    """196    signal.signal(signal.SIGINT, signal.SIG_IGN)197    signal.signal(signal.SIGTERM, signal.SIG_IGN)198