UN-2609 Added graceful shutdown support for tool containers with signal handling and configurable timeout

src/unstract/sdk/tool/entrypoint.py

@@ -1,11 +1,22 @@
+import logging
+import signal
+
 from unstract.sdk.tool.base import BaseTool
 from unstract.sdk.tool.executor import ToolExecutor
 from unstract.sdk.tool.parser import ToolArgsParser
 
+logger = logging.getLogger(__name__)
 
 class ToolEntrypoint:
     """Class that contains methods for the entrypoint for a tool."""
 
+    @staticmethod
+    def _signal_handler(signum: int, frame: Any) -> None:
+        """Handle SIGTERM and SIGINT signals."""
+        signal_name = "SIGTERM" if signum == signal.SIGTERM else "SIGINT"
+        logger.info(f"Received {signal_name} signal")
+        # Don't exit - let K8s send SIGKILL if tool doesn't finish
+
     @staticmethod
     def launch(tool: BaseTool, args: list[str]) -> None:
         """Entrypoint function for a tool.
@@ -17,6 +28,10 @@ def launch(tool: BaseTool, args: list[str]) -> None:
             tool (AbstractTool): Tool to execute
             args (List[str]): Arguments passed to a tool
         """
+        # Register signal handlers for graceful shutdown
+        signal.signal(signal.SIGTERM, ToolEntrypoint._signal_handler)
+        signal.signal(signal.SIGINT, ToolEntrypoint._signal_handler)
+
         parsed_args = ToolArgsParser.parse_args(args)
         executor = ToolExecutor(tool=tool)
         executor.execute(parsed_args)