D-43840 Support abort signal in container python plugin

digitalai/release/integration/base_task.py

@@ -1,4 +1,5 @@
 import sys
+import threading
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional
 
@@ -20,7 +21,8 @@ def __init__(self):
         self.release_context = None
         self.release_server_url = None
         self.input_properties = None
-        self.output_context = None
+        self.output_context = OutputContext(0, "", {}, [])
+        self._abort_event = threading.Event()
 
     def execute_task(self) -> None:
         """
@@ -29,7 +31,6 @@ def execute_task(self) -> None:
         is raised, the task's exit code is also set to 1.
         """
         try:
-            self.output_context = OutputContext(0, "", {}, [])
             self.execute()
         except AbortException:
             dai_logger.info("Abort requested")
@@ -51,10 +52,31 @@ def execute(self) -> None:
 
     def abort(self) -> None:
         """
-        Sets the task's exit code to 1 and exits the program with a status code of 1.
+        Marks the task as aborted (exit code 1, "Abort requested") and raises
+        AbortException. When called from within execute(), this is caught by
+        execute_task()'s existing AbortException handling, stopping the task
+        immediately. When called from the SDK's signal handler (outside of
+        execute_task()), the caller is responsible for catching this exception
+        and performing its own cleanup/exit.
         """
         self.set_exit_code(1)
-        sys.exit(1)
+        self.set_error_message("Abort requested")
+        raise AbortException()
+
+    def is_abort_requested(self) -> bool:
+        """
+        Returns True once the SDK has signaled that an abort was requested
+        (e.g. the process received SIGABRT/SIGTERM). Tasks may poll this from
+        within execute() to cooperatively stop long-running work promptly.
+        """
+        return self._abort_event.is_set()
+
+    def _request_abort(self) -> None:
+        """
+        Internal SDK hook: signals is_abort_requested() to start returning True.
+        Not intended to be called by task authors.
+        """
+        self._abort_event.set()
 
     def get_output_context(self) -> OutputContext:
         """

digitalai/release/integration/wrapper.py

@@ -7,13 +7,15 @@
 import os
 import signal
 import sys
+import threading
 import time
 from typing import Any, Dict, Optional, Tuple
 
 import requests
 
 from digitalai.release.integration import k8s, watcher
 from .base_task import BaseTask
+from .exceptions import AbortException
 from .input_context import InputContext
 from .job_data_encryptor import AESJobDataEncryptor, NoOpJobDataEncryptor
 from .logger import dai_logger
@@ -45,6 +47,10 @@
 HTTP_CONNECT_TIMEOUT = float(os.getenv('HTTP_CONNECT_TIMEOUT', '10'))
 HTTP_READ_TIMEOUT = float(os.getenv('HTTP_READ_TIMEOUT', '60'))
 
+# How often the main thread polls the worker thread (running the task's
+# execute_task()) via join(), so it stays responsive to SIGABRT/SIGTERM.
+WORKER_POLL_INTERVAL_SECONDS = 0.5
+
 # A single Session reused across all callback requests so the underlying urllib3
 # connection pool is shared (instead of opening a fresh connection per call).
 _http_session: requests.Session = requests.Session()
@@ -63,24 +69,88 @@ def get_encryptor():
 # Initialize the global task object
 dai_task_object: Optional[BaseTask] = None
 
+# Single-writer guard so the output context is flushed exactly once, whether
+# the task finishes normally (worker thread's completion path) or is aborted
+# (main thread's signal handler). Reset at the start of each execute_task()
+# call so daemon-mode iterations each get a fresh guard.
+_report_lock = threading.Lock()
+_result_reported = False
+
+
+def _reset_report_guard() -> None:
+    """Clears the single-writer flush guard ahead of a new task execution."""
+    global _result_reported
+    with _report_lock:
+        _result_reported = False
+
+
+def _report_once() -> bool:
+    """
+    Flushes the current dai_task_object's output context exactly once. The
+    first caller (either the worker thread's normal completion path or the
+    main thread's abort handler) wins; subsequent calls are no-ops.
+
+    Returns True if this call performed the flush, False if it was already
+    reported by another caller.
+    """
+    global _result_reported
+    with _report_lock:
+        if _result_reported:
+            return False
+        _result_reported = True
+
+    if dai_task_object is not None and dai_task_object.get_output_context() is not None:
+        update_output_context(dai_task_object.get_output_context())
+    else:
+        dai_logger.error("No output context available to report task result")
+        update_output_context(OutputContext(1, "Task produced no output context", {}, []))
+    return True
+
 
 def abort_handler(signum, frame):
     """
-    This function handles the abort request by calling the abort method on the global task object, if it exists.
-    If the task object does not exist, it logs a message and exits with a status code of 1.
+    Handles SIGTERM/SIGABRT by requesting a cooperative abort and running the
+    task's abort() cleanup hook (if a task object exists), then flushing the
+    (aborted) output context exactly once via _report_once(), and finally
+    terminating the process immediately with os._exit(1).
+
+    os._exit() (rather than sys.exit()) is used deliberately: it bypasses all
+    Python finally/atexit machinery, which matters in daemon mode where run()'s
+    finally block would otherwise restart a blocking input-context watcher and
+    prevent the process from actually exiting.
+
+    The worker thread that may still be running the task's execute() is a
+    daemon thread and is not waited on; process exit here is sufficient to
+    end it.
     """
-    dai_logger.info("Received SIGTERM to gracefully stop the process")
+    signal_name = signal.Signals(signum).name
+    dai_logger.info(f"Received {signal_name} to gracefully stop the process")
     global dai_task_object
 
     if dai_task_object:
-        dai_task_object.abort()
+        dai_task_object._request_abort()
+        try:
+            dai_task_object.abort()
+        except AbortException:
+            # Expected: the default abort() raises this after marking the
+            # output context as aborted.
+            pass
+        except SystemExit:
+            # A task's overridden abort() might still call sys.exit()/exit();
+            # ignore it so we can still flush and exit deliberately below.
+            pass
+        except Exception:
+            dai_logger.error("Unexpected error occurred in abort()", exc_info=True)
     else:
         dai_logger.info("Abort requested")
-        sys.exit(1)
+
+    _report_once()
+    os._exit(1)
 
 
 # Register abort handler
 signal.signal(signal.SIGTERM, abort_handler)
+signal.signal(signal.SIGABRT, abort_handler)
 
 
 def _post_callback(url: str, encrypted_json) -> requests.Response:
@@ -224,26 +294,35 @@ def should_retry_callback_request(encrypted_data) -> bool:
 
 def execute_task(task_object: BaseTask):
     """
-    Execute the given task object by setting it as the global task object and starting its execution.
-    If an exception is raised during execution, log the error. Finally, update the output context file
-    using the output context of the task object.
+    Execute the given task object on a dedicated daemon worker thread, while the
+    main thread waits responsively (polling worker.join with a short timeout)
+    so it stays free to handle SIGABRT/SIGTERM even if the task is blocked in
+    an uninterruptible call.
+
+    The worker thread is a daemon thread: if the process exits (e.g. via the
+    signal handler's os._exit(1) during an abort), it is not waited on.
+
+    After the worker finishes on its own (the non-aborted path), the output
+    context is flushed exactly once via _report_once() -- a no-op if the
+    abort handler already reported it first.
     """
     global dai_task_object
-    try:
-        dai_task_object = task_object
-        dai_logger.info("Starting task execution")
-        dai_task_object.execute_task()
-    except Exception:
-        dai_logger.error("Unexpected error occurred.", exc_info=True)
-    finally:
-        # Guard against a task object that was never constructed or that failed
-        # before producing an output context, so the finally block does not raise
-        # a second exception that masks the original one.
-        if dai_task_object is not None and dai_task_object.get_output_context() is not None:
-            update_output_context(dai_task_object.get_output_context())
-        else:
-            dai_logger.error("No output context available to report task result")
-            update_output_context(OutputContext(1, "Task produced no output context", {}, []))
+    _reset_report_guard()
+    dai_task_object = task_object
+    dai_logger.info("Starting task execution")
+
+    def _run_task():
+        try:
+            task_object.execute_task()
+        except Exception:
+            dai_logger.error("Unexpected error occurred.", exc_info=True)
+
+    worker = threading.Thread(target=_run_task, name="dai-task-worker", daemon=True)
+    worker.start()
+    while worker.is_alive():
+        worker.join(timeout=WORKER_POLL_INTERVAL_SECONDS)
+
+    _report_once()
 
 
 def find_class_file(root_dir, class_name):

docs/classes/base_task.md

@@ -9,7 +9,7 @@
 | **Class** | `BaseTask` |
 | **Module** | `digitalai.release.integration.base_task` |
 | **Source** | [`digitalai/release/integration/base_task.py`](../../digitalai/release/integration/base_task.py) |
-| **Methods** | 18 |
+| **Methods** | 20 |
 
 Subclass `BaseTask` to build a container-based integration task. Implement the
 abstract [`execute`](#execute) method with your task logic; the Release runtime
@@ -46,7 +46,8 @@ class Hello(BaseTask):
 |--------|---------|-------------|
 | [`execute`](#execute) | `None` | Abstract method holding the task's main logic; implemented by subclasses. |
 | [`execute_task`](#execute_task) | `None` | Runs `execute` with error handling and exit-code management; called by the runtime. |
-| [`abort`](#abort) | `None` | Sets the exit code to 1 and exits the program. |
+| [`abort`](#abort) | `None` | Marks the task as aborted (exit code 1, "Abort requested") and raises `AbortException`. |
+| [`is_abort_requested`](#is_abort_requested) | `bool` | Returns `True` once the SDK has signaled that an abort was requested. |
 | [`get_input_properties`](#get_input_properties) | `dict[str, Any]` | Returns the task's input properties. |
 | [`set_output_property`](#set_output_property) | `None` | Sets a named output property of the task. |
 | [`get_output_properties`](#get_output_properties) | `dict[str, Any]` | Returns the output properties of the task. |
@@ -102,12 +103,44 @@ _No parameters._
 
 ### `abort`
 
-Sets the task's exit code to `1` and exits the program with a status code of `1`.
+Marks the task as aborted by setting the exit code to `1` and the error message to
+`"Abort requested"`, then raises `AbortException`. The default implementation is
+the user-overridable hook invoked by the SDK when the process receives `SIGABRT`
+or `SIGTERM` (e.g. when a Release job is aborted). Override it to perform custom
+cleanup; call `super().abort()` (or raise `AbortException` yourself) to preserve
+the exit code / error message contract, or call `self.set_exit_code(...)` /
+`self.set_error_message(...)` directly if you need different semantics.
+
+For long-running tasks, prefer polling [`is_abort_requested`](#is_abort_requested)
+from within `execute` to stop cooperatively and promptly, rather than relying
+solely on `abort` being called.
 
 _No parameters._
 
 **Returns:** _None._
 
+**Raises:** `AbortException`.
+
+[↑ Method index](#method-index)
+
+### `is_abort_requested`
+
+Returns `True` once the SDK has signaled that an abort was requested (the process
+received `SIGABRT`/`SIGTERM`). Tasks may poll this from within `execute` to
+cooperatively stop long-running work promptly instead of waiting to be
+forcibly terminated.
+
+```python
+class MyTask(BaseTask):
+    def execute(self) -> None:
+        while not self.is_abort_requested():
+            do_one_unit_of_work()
+```
+
+_No parameters._
+
+**Returns:** `bool` — `True` if an abort has been requested, `False` otherwise.
+
 [↑ Method index](#method-index)
 
 ### `get_input_properties`

docs/examples.md

@@ -104,7 +104,8 @@ class Deploy(BaseTask):
 
 Raising an exception from `execute` fails the task: `execute_task` records the
 message and sets a non-zero exit code. To stop immediately, call `abort`, which
-sets the exit code to `1` and exits the process.
+marks the task as aborted (exit code `1`, "Abort requested") and raises
+`AbortException`, unwinding out of `execute`.
 
 **Classes:** [`BaseTask`](classes/base_task.md)
 
@@ -126,6 +127,23 @@ class GuardedTask(BaseTask):
         self.add_comment("Proceeding")
 ```
 
+For long-running tasks, the SDK also handles `SIGABRT`/`SIGTERM` (e.g. when a
+Release job is aborted) by running `abort` on your task automatically — you
+don't need to install a signal handler yourself. If your `execute` performs
+work in a loop or in chunks, poll `is_abort_requested` to stop promptly instead
+of waiting for the whole task to finish:
+
+```python
+class LongRunningTask(BaseTask):
+
+    def execute(self) -> None:
+        for item in self.input_properties.get("items", []):
+            if self.is_abort_requested():
+                self.add_comment("Stopping early: abort requested")
+                return
+            process(item)
+```
+
 [↑ Contents](#contents)
 
 ## Call the Release API with ApiBaseTask

tests/release/integration/blocking_task.py

@@ -0,0 +1,14 @@
+import time
+
+from digitalai.release.integration import BaseTask
+
+
+class BlockingTask(BaseTask):
+    """
+    Fixture task used by test_wrapper_abort.py to verify that SIGABRT interrupts
+    a task promptly even while it is blocked in an uninterruptible call (i.e.
+    without polling is_abort_requested()).
+    """
+
+    def execute(self) -> None:
+        time.sleep(30)