Allow interruption also if run with progress logging

Author: FlorentinD

changelog.md

@@ -20,6 +20,7 @@
 * Display progress bar for remote projection and open-ended tasks.
 * Improve progress bar by showing the description of the currently running task.
 * Allow passing the optional graph filter also as type `str` to `gds.graph.list()` instead of only `Graph`.
+* Listen and to SIGINT and SIGTERM and interrupt projection and write-backs based on GDS Session. Note this only works if the query runs in the main thread.
 
 
 ## Other changes

graphdatascience/query_runner/progress/query_progress_logger.py

@@ -82,6 +82,7 @@ def _log(
 
         if pbar is not None:
             self._finish_pbar(pbar)
+        # TODO show as cancelled if future was interrupted
 
     def _init_pbar(self, task_with_progress: TaskWithProgress) -> tqdm:  # type: ignore
         root_task_name = task_with_progress.task_name

graphdatascience/query_runner/protocol/project_protocols.py

@@ -1,4 +1,3 @@
-import signal
 from abc import ABC, abstractmethod
 from logging import DEBUG, getLogger
 from typing import Any, Optional
@@ -8,8 +7,9 @@
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log, retry_unless_signal
+from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -27,6 +27,7 @@ def run_projection(
         query_runner: QueryRunner,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -61,6 +62,7 @@ def run_projection(
         query_runner: QueryRunner,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -89,6 +91,7 @@ def run_projection(
         query_runner: QueryRunner,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -118,6 +121,7 @@ def run_projection(
         query_runner: QueryRunner,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -130,10 +134,11 @@ def is_not_done(result: DataFrame) -> bool:
 
         @retry(
             before=before_log(f"Projection (graph: `{params['graph_name']}`)", logger, DEBUG),
-            retry=retry_if_result(is_not_done) and retry_unless_signal([signal.SIGTERM, signal.SIGINT]),
+            retry=retry_if_result(is_not_done),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
         )
         def project_fn() -> DataFrame:
+            terminationFlag.assert_running()
             return query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name(endpoint), params, yields, database, logging, False
             )

graphdatascience/query_runner/protocol/write_protocols.py

@@ -1,5 +1,4 @@
 import logging
-import signal
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 
@@ -8,8 +7,9 @@
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import before_log, retry_unless_signal
+from graphdatascience.query_runner.protocol.retry_utils import before_log
 from graphdatascience.query_runner.protocol.status import Status
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
 
@@ -28,7 +28,11 @@ def write_back_params(
 
     @abstractmethod
     def run_write_back(
-        self, query_runner: QueryRunner, parameters: CallParameters, yields: Optional[list[str]]
+        self,
+        query_runner: QueryRunner,
+        parameters: CallParameters,
+        yields: Optional[list[str]],
+        terminationFlag: TerminationFlag,
     ) -> DataFrame:
         """Executes the write-back procedure"""
         pass
@@ -59,7 +63,11 @@ def write_back_params(
         )
 
     def run_write_back(
-        self, query_runner: QueryRunner, parameters: CallParameters, yields: Optional[list[str]]
+        self,
+        query_runner: QueryRunner,
+        parameters: CallParameters,
+        yields: Optional[list[str]],
+        terminationFlag: TerminationFlag,
     ) -> DataFrame:
         return query_runner.call_procedure(
             ProtocolVersion.V1.versioned_procedure_name("gds.arrow.write"),
@@ -93,7 +101,11 @@ def write_back_params(
         )
 
     def run_write_back(
-        self, query_runner: QueryRunner, parameters: CallParameters, yields: Optional[list[str]]
+        self,
+        query_runner: QueryRunner,
+        parameters: CallParameters,
+        yields: Optional[list[str]],
+        terminationFlag: TerminationFlag,
     ) -> DataFrame:
         return query_runner.call_procedure(
             ProtocolVersion.V2.versioned_procedure_name("gds.arrow.write"),
@@ -117,7 +129,11 @@ def write_back_params(
         return RemoteWriteBackV2().write_back_params(graph_name, job_id, config, arrow_config, database)
 
     def run_write_back(
-        self, query_runner: QueryRunner, parameters: CallParameters, yields: Optional[list[str]]
+        self,
+        query_runner: QueryRunner,
+        parameters: CallParameters,
+        yields: Optional[list[str]],
+        terminationFlag: TerminationFlag,
     ) -> DataFrame:
         def is_not_completed(result: DataFrame) -> bool:
             status: str = result.squeeze()["status"]
@@ -126,7 +142,7 @@ def is_not_completed(result: DataFrame) -> bool:
         logger = logging.getLogger()
 
         @retry(
-            retry=retry_if_result(is_not_completed) and retry_unless_signal([signal.SIGTERM, signal.SIGINT]),
+            retry=retry_if_result(is_not_completed),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
             before=before_log(
                 f"Write-Back (graph: `{parameters['graphName']}`, jobId: `{parameters['jobId']}`)",
@@ -135,6 +151,7 @@ def is_not_completed(result: DataFrame) -> bool:
             ),
         )
         def write_fn() -> DataFrame:
+            terminationFlag.assert_running()
             return query_runner.call_procedure(
                 ProtocolVersion.V3.versioned_procedure_name("gds.arrow.write"),
                 parameters,

graphdatascience/query_runner/session_query_runner.py

@@ -8,6 +8,7 @@
 
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
 from graphdatascience.query_runner.progress.query_progress_logger import QueryProgressLogger
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.server_version.server_version import ServerVersion
 
 from ..call_parameters import CallParameters
@@ -69,10 +70,12 @@ def call_procedure(
             params = CallParameters()
 
         if SessionQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME in endpoint:
-            return self._remote_projection(endpoint, params, yields, database, logging)
+            terminationFlag = TerminationFlag.create()
+            return self._remote_projection(endpoint, params, terminationFlag, yields, database, logging)
 
         elif ".write" in endpoint and self.is_remote_projected_graph(params["graph_name"]):
-            return self._remote_write_back(endpoint, params, yields, database, logging, custom_error)
+            terminationFlag = TerminationFlag.create()
+            return self._remote_write_back(endpoint, params, terminationFlag, yields, database, logging, custom_error)
 
         return self._gds_query_runner.call_procedure(endpoint, params, yields, database, logging, custom_error)
 
@@ -126,6 +129,7 @@ def _remote_projection(
         self,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -144,7 +148,7 @@ def _remote_projection(
 
             def run_projection() -> DataFrame:
                 return project_protocol.run_projection(
-                    self._db_query_runner, endpoint, project_params, yields, database, logging
+                    self._db_query_runner, endpoint, project_params, terminationFlag, yields, database, logging
                 )
 
             if self._resolve_show_progress(logging):
@@ -159,6 +163,7 @@ def _remote_write_back(
         self,
         endpoint: str,
         params: CallParameters,
+        terminationFlag: TerminationFlag,
         yields: Optional[list[str]] = None,
         database: Optional[str] = None,
         logging: bool = False,
@@ -180,6 +185,7 @@ def _remote_write_back(
         gds_write_result = self._gds_query_runner.call_procedure(
             endpoint, params, yields, database, logging, custom_error
         )
+        terminationFlag.assert_running()
 
         self._inject_arrow_config(db_arrow_config)
 
@@ -193,7 +199,7 @@ def _remote_write_back(
         write_back_start = time.time()
 
         def run_write_back() -> DataFrame:
-            return write_protocol.run_write_back(self._db_query_runner, write_back_params, yields)
+            return write_protocol.run_write_back(self._db_query_runner, write_back_params, yields, terminationFlag)
 
         try:
             if self._resolve_show_progress(logging):

graphdatascience/query_runner/termination_flag.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+import signal
+import threading
+from abc import ABC, abstractmethod
+from types import FrameType
+from typing import Optional
+
+
+class TerminationFlag(ABC):
+    @abstractmethod
+    def is_set(self) -> bool:
+        pass
+
+    @abstractmethod
+    def set(self) -> None:
+        pass
+
+    def assert_running(self) -> None:
+        if self.is_set():
+            raise RuntimeError("Query has been terminated")
+
+    @staticmethod
+    def create(signals: Optional[list[signal.Signals]] = None) -> TerminationFlag:
+        if signals is None:
+            signals = [signal.SIGINT, signal.SIGTERM]
+
+        if threading.current_thread() == threading.main_thread():
+            return TerminationFlagImpl(signals)
+        else:
+            logging.debug("Cannot set terminationFlag for query runner in non-main thread")
+            return TerminationFlagNoop()
+
+
+class TerminationFlagImpl(TerminationFlag):
+    def __init__(self, signals: list[signal.Signals]) -> None:
+        self._event = threading.Event()
+
+        def receive_signal(sig: int, frame: Optional[FrameType]) -> None:
+            logging.debug(f"Received signal {sig}. Interrupting query.")
+            self._event.set()
+
+        for sig in signals:
+            signal.signal(sig, receive_signal)
+
+    def is_set(self) -> bool:
+        return self._event.is_set()
+
+    def set(self) -> None:
+        self._event.set()
+
+    def assert_running(self) -> None:
+        if self.is_set():
+            raise RuntimeError("Query has been terminated")
+
+
+class TerminationFlagNoop(TerminationFlag):
+    def __init__(self) -> None:
+        pass
+
+    def is_set(self) -> bool:
+        return False
+
+    def set(self) -> None:
+        pass
+
+    def assert_running(self) -> None:
+        pass