From 26e4d233a8824869d7c50f8f7236c8b2e4dd2d1b Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Tue, 18 Aug 2026 15:03:32 +0800 Subject: [PATCH 1/2] Distinguish OpenAI batch timeout and cancellation in deferrable tasks A deferred OpenAITriggerBatchOperator reported every non-success outcome as the same OpenAIBatchJobException, so a task that merely ran out of patience looked identical to a batch that failed, and neither could be handled separately. The synchronous path has always raised OpenAIBatchTimeout for the same condition. The trigger now reports why it stopped, and the resuming task picks the matching exception from that field rather than from the message text. --- providers/openai/docs/changelog.rst | 12 +++++ .../airflow/providers/openai/exceptions.py | 11 +++++ .../airflow/providers/openai/hooks/openai.py | 30 ++++++++++-- .../providers/openai/operators/openai.py | 21 +++++++-- .../providers/openai/triggers/openai.py | 17 ++++++- .../tests/unit/openai/hooks/test_openai.py | 19 ++++++++ .../unit/openai/operators/test_openai.py | 41 +++++++++++++++- .../tests/unit/openai/test_exceptions.py | 15 +++++- .../tests/unit/openai/triggers/test_openai.py | 47 ++++++++++++++++--- 9 files changed, 195 insertions(+), 18 deletions(-) diff --git a/providers/openai/docs/changelog.rst b/providers/openai/docs/changelog.rst index cc2a0bd901fe3..710e1235cfa30 100644 --- a/providers/openai/docs/changelog.rst +++ b/providers/openai/docs/changelog.rst @@ -20,6 +20,18 @@ 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. + 1.8.2 ..... diff --git a/providers/openai/src/airflow/providers/openai/exceptions.py b/providers/openai/src/airflow/providers/openai/exceptions.py index 85f015880c548..4b2b2c7db0d21 100644 --- a/providers/openai/src/airflow/providers/openai/exceptions.py +++ b/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.""" diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py b/providers/openai/src/airflow/providers/openai/hooks/openai.py index 97dcc4ceec8dc..0a8dba11af03e 100644 --- a/providers/openai/src/airflow/providers/openai/hooks/openai.py +++ b/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( diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index dfed6e48d51a9..bfb02f658bd21 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/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: @@ -149,6 +152,13 @@ class OpenAITriggerBatchOperator(BaseOperator): :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,11 +217,14 @@ 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). """ event = validate_execute_complete_event(event) if event["status"] != "success": - raise OpenAIBatchJobException(event["message"]) + raise build_batch_error(event["message"], event.get("termination_reason")) self.log.info("%s completed successfully.", self.task_id) return event["batch_id"] diff --git a/providers/openai/src/airflow/providers/openai/triggers/openai.py b/providers/openai/src/airflow/providers/openai/triggers/openai.py index 49fd0900cc022..767aaa0c4bc41 100644 --- a/providers/openai/src/airflow/providers/openai/triggers/openai.py +++ b/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,7 +144,8 @@ 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, } ) @@ -148,9 +153,17 @@ async def run(self) -> AsyncIterator[TriggerEvent]: 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, + } + ) diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py b/providers/openai/tests/unit/openai/hooks/test_openai.py index d132a6a3d9582..155e45c46c68e 100644 --- a/providers/openai/tests/unit/openai/hooks/test_openai.py +++ b/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 diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index 0ec78bb182c51..27665bf253291 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -23,7 +23,12 @@ from openai.types.responses import Response from airflow.providers.common.compat.sdk import Context, TaskDeferred -from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAITriggerEventError +from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, + OpenAIBatchJobException, + OpenAIBatchTimeout, + OpenAITriggerEventError, +) from airflow.providers.openai.hooks.openai import OpenAIHook from airflow.providers.openai.operators.openai import ( OpenAIEmbeddingOperator, @@ -187,3 +192,37 @@ def test_failed_event_raises(self, event): def test_invalid_event_raises_instead_of_succeeding(self, event): with pytest.raises(OpenAITriggerEventError): self._operator().execute_complete(Context(), event) + + @pytest.mark.parametrize( + ("termination_reason", "expected_exc"), + [ + pytest.param("timeout", OpenAIBatchTimeout, id="timeout"), + pytest.param("cancelled", OpenAIBatchCancelled, id="cancelled"), + pytest.param("failed", OpenAIBatchJobException, id="failed"), + pytest.param("expired", OpenAIBatchJobException, id="expired"), + pytest.param("unexpected_status", OpenAIBatchJobException, id="unexpected-status"), + pytest.param("polling_error", OpenAIBatchJobException, id="polling-error"), + ], + ) + def test_execute_complete_raises_exception_matching_termination_reason( + self, termination_reason, expected_exc + ): + event = { + "status": "error", + "termination_reason": termination_reason, + "message": "boom", + "batch_id": BATCH_ID, + } + with pytest.raises(expected_exc, match="boom"): + self._operator().execute_complete(Context(), event) + + @pytest.mark.parametrize("status", ["error", "cancelled"]) + def test_execute_complete_missing_termination_reason_falls_back(self, status): + """A trigger serialized before ``termination_reason`` existed sends an event without + that key; ``execute_complete`` must fall back to ``OpenAIBatchJobException`` exactly, + not raise ``KeyError``. + """ + event = {"status": status, "message": "boom", "batch_id": BATCH_ID} + with pytest.raises(OpenAIBatchJobException, match="boom") as exc_info: + self._operator().execute_complete(Context(), event) + assert type(exc_info.value) is OpenAIBatchJobException diff --git a/providers/openai/tests/unit/openai/test_exceptions.py b/providers/openai/tests/unit/openai/test_exceptions.py index fabaad35343f0..b9c6a14e44667 100644 --- a/providers/openai/tests/unit/openai/test_exceptions.py +++ b/providers/openai/tests/unit/openai/test_exceptions.py @@ -21,7 +21,11 @@ import pytest -from airflow.providers.openai.exceptions import OpenAIBatchJobException, OpenAIBatchTimeout +from airflow.providers.openai.exceptions import ( + OpenAIBatchCancelled, + OpenAIBatchJobException, + OpenAIBatchTimeout, +) from airflow.providers.openai.hooks.openai import OpenAIHook @@ -30,6 +34,7 @@ [ OpenAIBatchTimeout, OpenAIBatchJobException, + OpenAIBatchCancelled, ], ) def test_wait_for_batch_raise_exception(exception_class): @@ -38,3 +43,11 @@ def test_wait_for_batch_raise_exception(exception_class): hook = mock_hook_instance with pytest.raises(exception_class): hook.wait_for_batch(batch_id="batch_id") + + +def test_batch_cancelled_is_subclass_of_batch_job_exception(): + """Cancellation is deliberately a subclass, not a sibling, of the generic batch failure + exception: existing ``except OpenAIBatchJobException`` handlers must keep working + unchanged after cancellation gets its own exception type. + """ + assert issubclass(OpenAIBatchCancelled, OpenAIBatchJobException) diff --git a/providers/openai/tests/unit/openai/triggers/test_openai.py b/providers/openai/tests/unit/openai/triggers/test_openai.py index 0a2c4322a5ed6..d72620ed3c501 100644 --- a/providers/openai/tests/unit/openai/triggers/test_openai.py +++ b/providers/openai/tests/unit/openai/triggers/test_openai.py @@ -119,22 +119,28 @@ def test_rejects_both_timeout_and_end_time(self): @pytest.mark.asyncio @pytest.mark.parametrize( - ("mock_batch_status", "mock_status", "mock_message"), + ("mock_batch_status", "mock_status", "mock_termination_reason", "mock_message"), [ - (str(BatchStatus.COMPLETED), "success", "Batch batch_id has completed successfully."), - (str(BatchStatus.CANCELLING), "cancelled", "Batch batch_id has been cancelled."), - (str(BatchStatus.CANCELLED), "cancelled", "Batch batch_id has been cancelled."), - (str(BatchStatus.FAILED), "error", "Batch failed:\nbatch_id"), + ( + str(BatchStatus.COMPLETED), + "success", + "completed", + "Batch batch_id has completed successfully.", + ), + (str(BatchStatus.CANCELLING), "cancelled", "cancelled", "Batch batch_id has been cancelled."), + (str(BatchStatus.CANCELLED), "cancelled", "cancelled", "Batch batch_id has been cancelled."), + (str(BatchStatus.FAILED), "error", "failed", "Batch failed:\nbatch_id"), ( str(BatchStatus.EXPIRED), "error", - "Batch couldn't be completed within the hour time window :\nbatch_id", + "expired", + "Batch couldn't be completed within its completion window:\nbatch_id", ), ], ) @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch") async def test_openai_batch_for_terminal_status( - self, mock_batch, mock_batch_status, mock_status, mock_message + self, mock_batch, mock_batch_status, mock_status, mock_termination_reason, mock_message ): """Assert that run trigger messages in case of job finished""" mock_batch.return_value = self.mock_get_batch(mock_batch_status) @@ -146,6 +152,7 @@ async def test_openai_batch_for_terminal_status( ) expected_result = { "status": mock_status, + "termination_reason": mock_termination_reason, "message": mock_message, "batch_id": self.BATCH_ID, } @@ -187,6 +194,7 @@ async def test_openai_batch_for_timeout(self, mock_monotonic, mock_batch, mock_b await asyncio.sleep(0.1) event = task.result() assert event.payload["status"] == "error" + assert event.payload["termination_reason"] == "timeout" assert f"Batch {self.BATCH_ID} has not reached a terminal status after" in event.payload["message"] asyncio.get_event_loop().stop() @@ -235,6 +243,7 @@ async def test_openai_batch_yields_single_terminal_event(self, mock_batch): TriggerEvent( { "status": "success", + "termination_reason": "completed", "message": f"Batch {self.BATCH_ID} has completed successfully.", "batch_id": self.BATCH_ID, } @@ -254,6 +263,7 @@ async def test_openai_batch_for_unexpected_error(self, mock_batch): ) expected_result = { "status": "error", + "termination_reason": "polling_error", "message": "'float' object has no attribute 'status'", "batch_id": self.BATCH_ID, } @@ -261,3 +271,26 @@ async def test_openai_batch_for_unexpected_error(self, mock_batch): await asyncio.sleep(0.1) assert TriggerEvent(expected_result) == task.result() asyncio.get_event_loop().stop() + + @pytest.mark.asyncio + @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch") + async def test_openai_batch_for_unexpected_status(self, mock_batch): + """A batch status outside the known terminal set falls into the `unexpected_status` branch.""" + mock_batch.return_value = self.mock_get_batch("validating") + mock_batch.return_value.status = "some_future_status" + trigger = OpenAIBatchTrigger( + conn_id=self.CONN_ID, + batch_id=self.BATCH_ID, + poll_interval=self.POLL_INTERVAL, + timeout=self.TIMEOUT, + ) + expected_result = { + "status": "error", + "termination_reason": "unexpected_status", + "message": f"Batch {self.BATCH_ID} has failed.", + "batch_id": self.BATCH_ID, + } + task = asyncio.create_task(trigger.run().__anext__()) + await asyncio.sleep(0.1) + assert TriggerEvent(expected_result) == task.result() + asyncio.get_event_loop().stop() From 0725f39d7fb3fe3f282a1cad9d269028e46f5fdc Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Tue, 18 Aug 2026 21:24:13 +0800 Subject: [PATCH 2/2] Cancel the OpenAI batch when a deferred task times out The synchronous path cancels the batch before raising a timeout, but the deferrable path only failed the task and left the batch running and billing on OpenAI's side. A deferred timeout now requests cancellation with the batch id carried by the trigger event. Cancellation is asynchronous, so the batch reports "cancelling" for a while before it settles; a cancel that fails is logged and never masks the timeout. Deliberately not cancelling on a polling_error: batches.cancel is irreversible, and the failure mode there is unknown and often a transient, Airflow-side error rather than a real batch problem. OpenAI's batch has a 24-hour completion window that bounds it on its own, so leaving it alone costs a bounded amount; wrongly cancelling it is unrecoverable data loss. Anthropic cancels on its equivalent error branch, but its session has no such natural end point, so that precedent does not apply here. --- providers/openai/docs/changelog.rst | 7 ++ .../providers/openai/operators/openai.py | 51 +++++++++++- .../unit/openai/operators/test_openai.py | 78 +++++++++++++++++++ 3 files changed, 135 insertions(+), 1 deletion(-) diff --git a/providers/openai/docs/changelog.rst b/providers/openai/docs/changelog.rst index 710e1235cfa30..742450896dd34 100644 --- a/providers/openai/docs/changelog.rst +++ b/providers/openai/docs/changelog.rst @@ -32,6 +32,13 @@ Changelog 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 ..... diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index bfb02f658bd21..57b0c0c365f15 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -148,7 +148,16 @@ 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. @@ -221,14 +230,54 @@ def execute_complete(self, context: Context, event: Any = None) -> str: ``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": + 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: diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py b/providers/openai/tests/unit/openai/operators/test_openai.py index 27665bf253291..556cecb3e7698 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -226,3 +226,81 @@ def test_execute_complete_missing_termination_reason_falls_back(self, status): with pytest.raises(OpenAIBatchJobException, match="boom") as exc_info: self._operator().execute_complete(Context(), event) assert type(exc_info.value) is OpenAIBatchJobException + + def test_timeout_requests_cancellation_using_event_batch_id(self): + """This is a regression lock: the resumed task is a fresh operator instance, so + ``self.batch_id`` is ``None`` here (``execute``'s assignment never happened on this + object). Cancellation must use ``event["batch_id"]``; if this test is made to pass by + reading ``self.batch_id`` instead, it should fail again as soon as that read returns + ``None`` for a real resumed task. + """ + operator = self._operator() + assert operator.batch_id is None + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = { + "status": "error", + "termination_reason": "timeout", + "message": "boom", + "batch_id": BATCH_ID, + } + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_called_once_with(BATCH_ID) + + def test_cancel_failure_does_not_mask_timeout(self): + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + mock_hook_instance.cancel_batch.side_effect = RuntimeError("cancel failed") + operator.hook = mock_hook_instance + event = { + "status": "error", + "termination_reason": "timeout", + "message": "boom", + "batch_id": BATCH_ID, + } + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_called_once_with(BATCH_ID) + + @pytest.mark.parametrize( + "termination_reason", + [ + "failed", + "cancelled", + "expired", + "polling_error", + "unexpected_status", + None, # a trigger serialized before `termination_reason` existed sends no such key + ], + ) + def test_non_timeout_termination_reasons_do_not_cancel(self, termination_reason): + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = {"status": "error", "message": "boom", "batch_id": BATCH_ID} + if termination_reason is not None: + event["termination_reason"] = termination_reason + + with pytest.raises(OpenAIBatchJobException): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_not_called() + + def test_timeout_with_no_batch_id_skips_cancellation(self): + """A timeout event that (hypothetically) carries no ``batch_id`` must not crash and + must not call ``cancel_batch`` — but the original timeout error still has to surface. + """ + operator = self._operator() + mock_hook_instance = Mock(spec=OpenAIHook) + operator.hook = mock_hook_instance + event = {"status": "error", "termination_reason": "timeout", "message": "boom"} + + with pytest.raises(OpenAIBatchTimeout): + operator.execute_complete(Context(), event) + + mock_hook_instance.cancel_batch.assert_not_called()