query: interrupt query script gracefully on SIGTERM

This now supports gracefully terminating script on SIGTERM.

Also added logs.

Author: skshetry

src/datachain/catalog/catalog.py

@@ -4,6 +4,7 @@
 import os
 import os.path
 import posixpath
+import signal
 import subprocess
 import sys
 import time
@@ -97,6 +98,46 @@ def noop(_: str):
     pass
 
 
+class TerminationSignal(RuntimeError):  # noqa: N818
+    def __init__(self, signal):
+        self.signal = signal
+        super().__init__("Received termination signal", signal)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.signal})"
+
+
+if sys.platform == "win32":
+    SIGINT = signal.CTRL_C_EVENT
+else:
+    SIGINT = signal.SIGINT
+
+
+def shutdown_process(
+    proc: subprocess.Popen,
+    interrupt_timeout: Optional[int] = None,
+    terminate_timeout: Optional[int] = None,
+) -> None:
+    """Shut down the process gracefully with SIGINT -> SIGTERM -> SIGKILL."""
+
+    logger.info("sending interrupt signal to the process %s", proc.pid)
+    proc.send_signal(SIGINT)
+
+    logger.info("waiting for the process %s to finish", proc.pid)
+    try:
+        proc.wait(interrupt_timeout)
+    except subprocess.TimeoutExpired:
+        logger.info(
+            "timed out waiting, sending terminate signal to the process %s", proc.pid
+        )
+        proc.terminate()
+        try:
+            proc.wait(terminate_timeout)
+        except subprocess.TimeoutExpired:
+            logger.info("timed out waiting, killing the process %s", proc.pid)
+            proc.kill()
+
+
 def _process_stream(stream: "IO[bytes]", callback: Callable[[str], None]) -> None:
     buffer = b""
     while byt := stream.read(1):  # Read one byte at a time
@@ -1493,6 +1534,8 @@ def query(
         output_hook: Callable[[str], None] = noop,
         params: Optional[dict[str, str]] = None,
         job_id: Optional[str] = None,
+        interrupt_timeout: Optional[int] = None,
+        terminate_timeout: Optional[int] = None,
     ) -> None:
         cmd = [python_executable, "-c", query_script]
         env = dict(env or os.environ)
@@ -1506,13 +1549,42 @@ def query(
         if capture_output:
             popen_kwargs = {"stdout": subprocess.PIPE, "stderr": subprocess.STDOUT}
 
-        with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
-            if capture_output:
-                args = (proc.stdout, output_hook)
-                thread = Thread(target=_process_stream, args=args, daemon=True)
-                thread.start()
-                thread.join()  # wait for the reader thread
+        def signal_handler(sig: int, _: Any) -> NoReturn:
+            raise TerminationSignal(sig)
+
+        original_sigterm_handler = signal.getsignal(signal.SIGTERM)
+        signal.signal(signal.SIGTERM, signal_handler)
 
+        with subprocess.Popen(cmd, env=env, **popen_kwargs) as proc:  # noqa: S603
+            logger.info("Starting process %s", proc.pid)
+            thread: Optional[Thread] = None
+            try:
+                if capture_output:
+                    args = (proc.stdout, output_hook)
+                    thread = Thread(target=_process_stream, args=args, daemon=True)
+                    thread.start()
+
+                proc.wait()
+            except (KeyboardInterrupt, TerminationSignal) as exc:
+                # NOTE: we ignore `Ctrl-C` signals on CLI so `KeyboardInterrupt`
+                # is not expected to be raised here, but we handle it just in case.
+                # If we don't ignore it in CLI, the child will receive the signal twice.
+                logging.info("Shutting down process %s, received %r", proc.pid, exc)
+                # Rather than forwarding the signal to the child, we try to shut it down
+                # gracefully. This is because we consider the script to be interactive
+                # and special, so we give it time to cleanup before exiting.
+                shutdown_process(proc, interrupt_timeout, terminate_timeout)
+                if proc.returncode:
+                    raise QueryScriptCancelError(
+                        "Query script was canceled by user", return_code=proc.returncode
+                    ) from exc
+            finally:
+                if original_sigterm_handler:
+                    signal.signal(signal.SIGTERM, original_sigterm_handler)
+                if thread:
+                    thread.join()
+
+        logging.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode == QUERY_SCRIPT_CANCELED_EXIT_CODE:
             raise QueryScriptCancelError(
                 "Query script was canceled by user",

src/datachain/cli/commands/query.py

@@ -1,4 +1,5 @@
 import os
+import signal
 import sys
 import traceback
 from typing import TYPE_CHECKING, Optional
@@ -33,6 +34,8 @@ def query(
         params=params,
     )
 
+    # ignore SIGINT in the main process
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
     try:
         catalog.query(
             script_content,

tests/unit/test_query.py

@@ -6,7 +6,10 @@
 
 import pytest
 
-from datachain.catalog.catalog import QUERY_SCRIPT_CANCELED_EXIT_CODE
+from datachain.catalog.catalog import (
+    QUERY_SCRIPT_CANCELED_EXIT_CODE,
+    TerminationSignal,
+)
 from datachain.error import QueryScriptCancelError, QueryScriptRunError
 
 
@@ -63,3 +66,16 @@ def test_non_zero_exitcode(catalog, mock_popen):
         catalog.query("pass")
     assert e.value.return_code == 1
     assert "Query script exited with error code 1" in str(e.value)
+
+
+@pytest.mark.parametrize("side_effect", [KeyboardInterrupt(), TerminationSignal(15)])
+def test_shutdown_process_called(mocker, catalog, mock_popen, side_effect):
+    mock_popen.returncode = -2
+    mock_popen.wait.side_effect = [side_effect]
+    m = mocker.patch("datachain.catalog.catalog.shutdown_process", return_value=-2)
+
+    with pytest.raises(QueryScriptCancelError) as e:
+        catalog.query("pass", interrupt_timeout=0.1, terminate_timeout=0.2)
+    assert e.value.return_code == -2
+    assert "Query script was canceled by user" in str(e.value)
+    m.assert_called_once_with(mock_popen, 0.1, 0.2)