Skip signal handling if not main thread

Such as when called together with progress logging

Author: FlorentinD

graphdatascience/query_runner/protocol/project_protocols.py

@@ -1,13 +1,14 @@
 import signal
 from abc import ABC, abstractmethod
+from logging import DEBUG, getLogger
 from typing import Any, Optional
 
 from pandas import DataFrame
 from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import retry_unless_signal
+from graphdatascience.query_runner.protocol.retry_utils import before_log, retry_unless_signal
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
@@ -125,7 +126,10 @@ def is_not_done(result: DataFrame) -> bool:
             status: str = result.squeeze()["status"]
             return status != Status.DONE.name
 
+        logger = getLogger()
+
         @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]),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
         )

graphdatascience/query_runner/protocol/retry_utils.py

@@ -1,4 +1,6 @@
+import logging
 import signal
+import typing
 from types import FrameType
 from typing import Optional
 
@@ -12,10 +14,30 @@ def __init__(self, signals: list[signal.Signals]) -> None:
         self.signal_received = False
 
         def receive_signal(sig: int, frame: Optional[FrameType]) -> None:
+            logging.debug(f"Received signal {sig}. Interrupting retry.")
             self.signal_received = True
 
-        for sig in signals:
-            signal.signal(sig, receive_signal)
+        try:
+            for sig in signals:
+                signal.signal(sig, receive_signal)
+        except ValueError as e:
+            # signal.signal() can raise ValueError if this is not called in the main thread (such as when an algorithm is called in a ThreadPool)
+            logging.debug(f"Cannot set signal handler for retries {e}")
 
     def __call__(self, retry_state: RetryCallState) -> bool:
         return not self.signal_received
+
+
+def before_log(
+    fn_name: str,
+    logger: logging.Logger,
+    log_level: int,
+    sec_format: str = "%0.3f",
+) -> typing.Callable[[RetryCallState], None]:
+    def log_it(retry_state: RetryCallState) -> None:
+        logger.log(
+            log_level,
+            f"Retry of '{fn_name}', " f"attempt: {retry_state.attempt_number}",
+        )
+
+    return log_it

graphdatascience/query_runner/protocol/write_protocols.py

@@ -4,11 +4,11 @@
 from typing import Any, Optional
 
 from pandas import DataFrame
-from tenacity import after_log, retry, retry_if_result, wait_incrementing
+from tenacity import retry, retry_if_result, wait_incrementing
 
 from graphdatascience import QueryRunner
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.query_runner.protocol.retry_utils import retry_unless_signal
+from graphdatascience.query_runner.protocol.retry_utils import before_log, retry_unless_signal
 from graphdatascience.query_runner.protocol.status import Status
 from graphdatascience.session.dbms.protocol_version import ProtocolVersion
 
@@ -128,7 +128,11 @@ def is_not_completed(result: DataFrame) -> bool:
         @retry(
             retry=retry_if_result(is_not_completed) and retry_unless_signal([signal.SIGTERM, signal.SIGINT]),
             wait=wait_incrementing(start=0.2, increment=0.2, max=2),
-            after=after_log(logger, logging.WARN),
+            before=before_log(
+                f"Write-Back (graph: `{parameters['graphName']}`, jobId: `{parameters['jobId']}`)",
+                logger,
+                logging.DEBUG,
+            ),
         )
         def write_fn() -> DataFrame:
             return query_runner.call_procedure(