Distinguish OpenAI batch timeout and cancellation in deferrable tasks

providers/openai/docs/changelog.rst

@@ -20,6 +20,25 @@
 Changelog
 ---------
 
+.. note::
+    A deferred ``OpenAITriggerBatchOperator`` that runs out of patience now raises
+    ``OpenAIBatchTimeout``, matching the exception the non-deferrable path has always raised
+    for the same condition. Previously every non-success outcome of a deferred batch raised
+    the same ``OpenAIBatchJobException``, so a timeout and a genuine batch failure could not
+    be told apart or handled separately.
+
+    A cancelled batch now raises ``OpenAIBatchCancelled``, a subclass of
+    ``OpenAIBatchJobException``, so existing code that catches ``OpenAIBatchJobException``
+    keeps working unchanged while callers that want to distinguish cancellation can catch the
+    subclass specifically.
+
+.. note::
+    A deferred ``OpenAITriggerBatchOperator`` that times out now requests cancellation of the
+    batch, matching the non-deferrable path. Previously a deferred timeout only failed the
+    task and left the batch running (and billing) on OpenAI's side. Cancellation on OpenAI's
+    side is asynchronous, so the batch reports ``cancelling`` for a while before it settles as
+    ``cancelled``.
+
 1.8.2
 .....
 

providers/openai/src/airflow/providers/openai/exceptions.py

@@ -24,6 +24,17 @@ class OpenAIBatchJobException(AirflowException):
     """Raise when OpenAI Batch Job fails to start AFTER processing the request."""
 
 
+class OpenAIBatchCancelled(OpenAIBatchJobException):
+    """
+    Raise when an OpenAI Batch Job was cancelled.
+
+    Cancellation is a decision, not a failure, so it gets its own subclass: callers
+    that want to distinguish "someone cancelled this batch" from "the batch failed"
+    can catch this specifically, while existing handlers written against
+    ``OpenAIBatchJobException`` keep working unchanged.
+    """
+
+
 class OpenAIBatchTimeout(AirflowException):
     """Raise when OpenAI Batch Job times out."""
 

providers/openai/src/airflow/providers/openai/hooks/openai.py

@@ -54,8 +54,9 @@
     from openai.types.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.common.compat.module_loading import import_string
-from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import AirflowException, BaseHook
 from airflow.providers.openai.exceptions import (
+    OpenAIBatchCancelled,
     OpenAIBatchJobException,
     OpenAIBatchTimeout,
     OpenAITriggerEventError,
@@ -95,6 +96,29 @@ def is_in_progress(cls, status: str) -> bool:
 #: Statuses the provider's trigger emits in its terminal event.
 TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "cancelled"})
 
+#: Maps the trigger's ``termination_reason`` field to the exception ``execute_complete``
+#: should raise. Keyed on the reason field, never on the message text, so that a
+#: rewording of the trigger's message never silently changes which exception a
+#: downstream task can catch.
+_TERMINATION_REASON_EXCEPTIONS: dict[str, type[AirflowException]] = {
+    "timeout": OpenAIBatchTimeout,
+    "cancelled": OpenAIBatchCancelled,
+}
+
+
+def build_batch_error(message: str, termination_reason: str | None) -> AirflowException:
+    """
+    Build (but do not raise) the exception matching a trigger event's termination reason.
+
+    ``termination_reason`` is ``None`` when the event was produced by a trigger
+    serialized before this field existed (a rolling upgrade in flight); that case
+    falls back to ``OpenAIBatchJobException``, matching today's behavior.
+    """
+    if termination_reason is None:
+        return OpenAIBatchJobException(message)
+    exception_class = _TERMINATION_REASON_EXCEPTIONS.get(termination_reason, OpenAIBatchJobException)
+    return exception_class(message)
+
 
 def validate_execute_complete_event(event: dict[str, Any] | None = None) -> dict[str, Any]:
     """
@@ -670,10 +694,10 @@ def wait_for_batch(self, batch_id: str, wait_seconds: float = 3, timeout: float
             if batch.status == BatchStatus.FAILED:
                 raise OpenAIBatchJobException(f"Batch failed - \n{batch_id}")
             if batch.status in (BatchStatus.CANCELLED, BatchStatus.CANCELLING):
-                raise OpenAIBatchJobException(f"Batch failed - batch was cancelled:\n{batch_id}")
+                raise OpenAIBatchCancelled(f"Batch failed - batch was cancelled:\n{batch_id}")
             if batch.status == BatchStatus.EXPIRED:
                 raise OpenAIBatchJobException(
-                    f"Batch failed - batch couldn't be completed within the hour time window :\n{batch_id}"
+                    f"Batch failed - batch couldn't be completed within its completion window:\n{batch_id}"
                 )
 
             raise OpenAIBatchJobException(

providers/openai/src/airflow/providers/openai/operators/openai.py

@@ -22,8 +22,11 @@
 from typing import TYPE_CHECKING, Any, Literal
 
 from airflow.providers.common.compat.sdk import BaseOperator, conf
-from airflow.providers.openai.exceptions import OpenAIBatchJobException
-from airflow.providers.openai.hooks.openai import OpenAIHook, validate_execute_complete_event
+from airflow.providers.openai.hooks.openai import (
+    OpenAIHook,
+    build_batch_error,
+    validate_execute_complete_event,
+)
 from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger
 
 if TYPE_CHECKING:
@@ -145,10 +148,26 @@ class OpenAITriggerBatchOperator(BaseOperator):
     :param wait_seconds: Optional. Number of seconds between checks. Only used when ``deferrable`` is False.
         Defaults to 3 seconds.
     :param timeout: Optional. The amount of time, in seconds, to wait for the request to complete.
-        Only used when ``deferrable`` is False. Defaults to 24 hour, which is the SLA for OpenAI Batch API.
+        Used in both modes: in the synchronous path it bounds ``wait_for_batch``; in the
+        deferrable path it bounds the trigger's poll loop. When the deferrable path times out,
+        the operator requests cancellation of the batch using the batch id carried by the
+        trigger event, mirroring the synchronous path. Cancellation on OpenAI's side is
+        asynchronous — the batch reports ``cancelling`` for up to 10 minutes before it settles
+        as ``cancelled`` — so this only *requests* cancellation, it does not wait for it. If
+        ``execution_timeout`` is set shorter than ``timeout``, the scheduler's deferral timeout
+        fires first: the task is failed with ``TaskDeferralTimeout`` before the trigger ever
+        times out, ``execute_complete`` is never called, and this cancellation path does not
+        run. Defaults to 24 hour, which is the SLA for OpenAI Batch API.
     :param wait_for_completion: Optional. Whether to wait for the batch to complete. If set to False, the operator
         will return immediately after triggering the batch. Defaults to True.
 
+    When ``deferrable`` is True and the batch does not reach a terminal state, ``execute_complete``
+    raises :class:`~airflow.providers.openai.exceptions.OpenAIBatchTimeout`, matching the exception
+    raised by the synchronous path for the same condition. A cancelled batch raises
+    :class:`~airflow.providers.openai.exceptions.OpenAIBatchCancelled` (a subclass of
+    :class:`~airflow.providers.openai.exceptions.OpenAIBatchJobException`), and any other failure
+    raises :class:`~airflow.providers.openai.exceptions.OpenAIBatchJobException`.
+
     .. seealso::
         For more information on how to use this operator, please take a look at the guide:
         :ref:`howto/operator:OpenAITriggerBatchOperator`
@@ -207,15 +226,58 @@ def execute_complete(self, context: Context, event: Any = None) -> str:
         Invoke this callback when the trigger fires; return immediately.
 
         Relies on trigger to throw an exception, otherwise it assumes execution was
-        successful.
+        successful. The exception raised depends on the event's ``termination_reason``:
+        ``OpenAIBatchTimeout`` for a timeout, ``OpenAIBatchCancelled`` for a cancellation,
+        and ``OpenAIBatchJobException`` for any other failure (including events from a
+        trigger serialized before ``termination_reason`` existed).
+
+        On a timeout, cancellation of the batch is requested before the timeout is raised
+        (see :meth:`_cancel_batch_quietly`). No other termination reason triggers
+        cancellation: a ``polling_error`` may be a transient, Airflow-side failure rather than
+        a real batch problem, and cancellation is irreversible, so it is left alone to run to
+        its own 24-hour completion window instead.
         """
         event = validate_execute_complete_event(event)
         if event["status"] != "success":
-            raise OpenAIBatchJobException(event["message"])
+            if event.get("termination_reason") == "timeout":
+                batch_id = event.get("batch_id")
+                if batch_id:
+                    self.log.warning(
+                        "%s timed out waiting for batch %s; requesting cancellation.",
+                        self.task_id,
+                        batch_id,
+                    )
+                    self._cancel_batch_quietly(batch_id)
+                else:
+                    self.log.warning(
+                        "%s timed out but the trigger event carried no batch_id; "
+                        "skipping cancellation request.",
+                        self.task_id,
+                    )
+            raise build_batch_error(event["message"], event.get("termination_reason"))
 
         self.log.info("%s completed successfully.", self.task_id)
         return event["batch_id"]
 
+    def _cancel_batch_quietly(self, batch_id: str) -> None:
+        """
+        Best-effort request to cancel a batch; never raises.
+
+        Called from ``execute_complete`` after a deferred timeout, using the batch id carried
+        by the trigger event rather than ``self.batch_id`` — this method runs on a resumed task
+        instance, a fresh operator object on which ``execute``'s assignment to ``self.batch_id``
+        never happened, so ``self.batch_id`` is ``None`` here.
+
+        Cancellation on OpenAI's side is asynchronous: the batch reports ``cancelling`` for up
+        to 10 minutes before it settles as ``cancelled``, so this only requests cancellation. A
+        failure to cancel is logged, not raised, so it never masks the timeout that is the
+        task's real failure reason.
+        """
+        try:
+            self.hook.cancel_batch(batch_id)
+        except Exception as e:
+            self.log.warning("Failed to request cancellation of batch %s: %s", batch_id, e)
+
     def on_kill(self) -> None:
         """Cancel the batch if task is cancelled."""
         if self.batch_id:

providers/openai/src/airflow/providers/openai/triggers/openai.py

@@ -103,6 +103,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                     yield TriggerEvent(
                         {
                             "status": "error",
+                            "termination_reason": "timeout",
                             "message": (
                                 f"Batch {self.batch_id} has not reached a terminal status after "
                                 f"{elapsed:.0f} seconds."
@@ -116,6 +117,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 yield TriggerEvent(
                     {
                         "status": "success",
+                        "termination_reason": "completed",
                         "message": f"Batch {self.batch_id} has completed successfully.",
                         "batch_id": self.batch_id,
                     }
@@ -124,6 +126,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 yield TriggerEvent(
                     {
                         "status": "cancelled",
+                        "termination_reason": "cancelled",
                         "message": f"Batch {self.batch_id} has been cancelled.",
                         "batch_id": self.batch_id,
                     }
@@ -132,6 +135,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 yield TriggerEvent(
                     {
                         "status": "error",
+                        "termination_reason": "failed",
                         "message": f"Batch failed:\n{self.batch_id}",
                         "batch_id": self.batch_id,
                     }
@@ -140,17 +144,26 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
                 yield TriggerEvent(
                     {
                         "status": "error",
-                        "message": f"Batch couldn't be completed within the hour time window :\n{self.batch_id}",
+                        "termination_reason": "expired",
+                        "message": f"Batch couldn't be completed within its completion window:\n{self.batch_id}",
                         "batch_id": self.batch_id,
                     }
                 )
             else:
                 yield TriggerEvent(
                     {
                         "status": "error",
+                        "termination_reason": "unexpected_status",
                         "message": f"Batch {self.batch_id} has failed.",
                         "batch_id": self.batch_id,
                     }
                 )
         except Exception as e:
-            yield TriggerEvent({"status": "error", "message": str(e), "batch_id": self.batch_id})
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "termination_reason": "polling_error",
+                    "message": str(e),
+                    "batch_id": self.batch_id,
+                }
+            )

providers/openai/tests/unit/openai/hooks/test_openai.py

@@ -39,6 +39,7 @@
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.models import Connection
 from airflow.providers.openai.exceptions import (
+    OpenAIBatchCancelled,
     OpenAIBatchJobException,
     OpenAIBatchTimeout,
     OpenAITriggerEventError,
@@ -644,6 +645,24 @@ def test_wait_for_in_progress_batch_timeout(mock_openai_hook, mock_wip_batch):
     assert mock_openai_hook.conn.batches.cancel.call_count == 1
 
 
+@pytest.mark.parametrize("status", ["cancelled", "cancelling"])
+def test_wait_for_cancelled_batch_raises_exact_cancelled_type(mock_openai_hook, status):
+    """``OpenAIBatchCancelled`` is a subclass of ``OpenAIBatchJobException``, so asserting
+    only the base class would stay green even if this raised the wrong (base) type. Assert
+    the exact type to prove the exception was actually narrowed.
+    """
+    mock_openai_hook.conn.batches.retrieve.return_value = create_batch(status)
+    with pytest.raises(OpenAIBatchCancelled):
+        mock_openai_hook.wait_for_batch(batch_id=BATCH_ID)
+
+
+def test_wait_for_expired_batch_message_does_not_mention_hour_window(mock_openai_hook):
+    mock_openai_hook.conn.batches.retrieve.return_value = create_batch("expired")
+    with pytest.raises(OpenAIBatchJobException, match="completion window") as exc_info:
+        mock_openai_hook.wait_for_batch(batch_id=BATCH_ID)
+    assert "hour time window" not in str(exc_info.value)
+
+
 def test_openai_hook_test_connection(mock_openai_hook):
     result, message = mock_openai_hook.test_connection()
     assert result is True