This is an automated email from the ASF dual-hosted git repository. Lee-W pushed a commit to branch openai-batch-termination-reason in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 95cd0514bc24d3a3cc435e888e3866d4c2cf0a6d Author: Wei Lee <[email protected]> AuthorDate: Mon Sep 28 11:44:05 2026 +0900 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 | 15 ++---- .../airflow/providers/openai/operators/openai.py | 52 +++++++++++++++++- .../tests/unit/openai/operators/test_openai.py | 63 ++++++++++++++++++++++ 3 files changed, 118 insertions(+), 12 deletions(-) diff --git a/providers/openai/docs/changelog.rst b/providers/openai/docs/changelog.rst index a5f332b3d19..6c2fe382826 100644 --- a/providers/openai/docs/changelog.rst +++ b/providers/openai/docs/changelog.rst @@ -52,16 +52,11 @@ Changelog setting. .. note:: - A deferred ``OpenAITriggerBatchOperator`` that times out 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. + 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``. Breaking changes ~~~~~~~~~~~~~~~~ diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index d1332a55086..f58258639f4 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/providers/openai/src/airflow/providers/openai/operators/openai.py @@ -364,8 +364,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. - Applies in both deferrable and non-deferrable mode. Defaults to 24 hours, which is the SLA for - OpenAI Batch API. + Applies in both deferrable and non-deferrable mode: 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 hours, 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. :param metadata: Optional. A set of key-value pairs that can be attached to the batch. (templated) @@ -469,14 +477,54 @@ class OpenAITriggerBatchOperator(BaseOperator): ``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 66613601cb1..1c40f2951e5 100644 --- a/providers/openai/tests/unit/openai/operators/test_openai.py +++ b/providers/openai/tests/unit/openai/operators/test_openai.py @@ -974,3 +974,66 @@ class TestOpenAITriggerBatchOperatorExecuteComplete: 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): + """The resumed task is a fresh operator instance, so ``self.batch_id`` is ``None`` here. + 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()
