This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4b9c3b54a7b Validate trigger events in Openai deferrable tasks (#69506)
4b9c3b54a7b is described below
commit 4b9c3b54a7bab8022075294b1a9b2066c88350f2
Author: Takayoshi Makabe <[email protected]>
AuthorDate: Sun Aug 2 02:58:54 2026 +0900
Validate trigger events in Openai deferrable tasks (#69506)
* Add validate trigger events deferrable tasks
* Record the cancelled-batch behaviour change for openai users
A deferred batch that was cancelled previously completed successfully
with no results, so operators upgrading need to know their tasks will
start failing where they used to go green.
---------
Co-authored-by: Jarek Potiuk <[email protected]>
---
providers/openai/docs/changelog.rst | 15 ++++++++
.../src/airflow/providers/openai/exceptions.py | 4 +++
.../src/airflow/providers/openai/hooks/openai.py | 26 +++++++++++++-
.../airflow/providers/openai/operators/openai.py | 5 +--
.../airflow/providers/openai/triggers/openai.py | 16 ++++-----
.../openai/tests/unit/openai/hooks/test_openai.py | 39 +++++++++++++++++++--
.../tests/unit/openai/operators/test_openai.py | 40 ++++++++++++++++++++++
.../tests/unit/openai/triggers/test_openai.py | 22 ++++++++++++
8 files changed, 154 insertions(+), 13 deletions(-)
diff --git a/providers/openai/docs/changelog.rst
b/providers/openai/docs/changelog.rst
index 3d304eb8a94..61d78b800d0 100644
--- a/providers/openai/docs/changelog.rst
+++ b/providers/openai/docs/changelog.rst
@@ -20,6 +20,21 @@
Changelog
---------
+.. note::
+ ``OpenAITriggerBatchOperator`` now fails the task when a deferred batch
ends in a
+ non-success state. Previously only ``status="error"`` raised, so a batch
that was
+ **cancelled** while the task was deferred completed **successfully** with
no results,
+ and an unrecognized or missing trigger event either succeeded silently or
crashed with
+ ``TypeError: 'NoneType' object is not subscriptable``.
+
+ This aligns the deferrable path with the non-deferrable one, which has
raised for
+ ``CANCELLED``/``CANCELLING`` since ``OpenAIHook.wait_for_batch`` was fixed.
+
+ Dags whose batches are cancelled will now fail where they previously
reported success.
+ That is the intended correction — the earlier green runs produced no batch
output — but
+ if you were relying on cancellation being treated as success, handle it
explicitly, for
+ example with a trigger rule or by catching ``OpenAIBatchJobException``
downstream.
+
1.8.1
.....
diff --git a/providers/openai/src/airflow/providers/openai/exceptions.py
b/providers/openai/src/airflow/providers/openai/exceptions.py
index 05eee6b24fb..85f015880c5 100644
--- a/providers/openai/src/airflow/providers/openai/exceptions.py
+++ b/providers/openai/src/airflow/providers/openai/exceptions.py
@@ -26,3 +26,7 @@ class OpenAIBatchJobException(AirflowException):
class OpenAIBatchTimeout(AirflowException):
"""Raise when OpenAI Batch Job times out."""
+
+
+class OpenAITriggerEventError(AirflowException):
+ """Raise when a deferred task resumes with a missing or malformed trigger
event."""
diff --git a/providers/openai/src/airflow/providers/openai/hooks/openai.py
b/providers/openai/src/airflow/providers/openai/hooks/openai.py
index 8315a907ffa..97dcc4ceec8 100644
--- a/providers/openai/src/airflow/providers/openai/hooks/openai.py
+++ b/providers/openai/src/airflow/providers/openai/hooks/openai.py
@@ -55,7 +55,11 @@ if TYPE_CHECKING:
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.openai.exceptions import OpenAIBatchJobException,
OpenAIBatchTimeout
+from airflow.providers.openai.exceptions import (
+ OpenAIBatchJobException,
+ OpenAIBatchTimeout,
+ OpenAITriggerEventError,
+)
#: The OpenAI Assistants API (``beta.assistants``/``beta.threads``) is
deprecated by OpenAI. The hook
#: methods wrapping it warn and point at the Responses and Conversations APIs
(``create_response`` /
@@ -88,6 +92,26 @@ class BatchStatus(str, Enum):
return status in (cls.VALIDATING, cls.IN_PROGRESS, cls.FINALIZING)
+#: Statuses the provider's trigger emits in its terminal event.
+TRIGGER_EVENT_STATUSES = frozenset({"success", "error", "cancelled"})
+
+
+def validate_execute_complete_event(event: dict[str, Any] | None = None) ->
dict[str, Any]:
+ """
+ Validate the event a deferred task resumes with, returning it if
well-formed.
+
+ The event crosses the triggerer/worker boundary through the metadata DB,
so a
+ resuming task can receive ``None`` or a status its handler does not
recognize
+ (version skew, a custom trigger). Both must fail loudly instead of crashing
+ opaquely or being misread as an outcome.
+ """
+ if event is None:
+ raise OpenAITriggerEventError("Trigger error: event is None")
+ if event.get("status") not in TRIGGER_EVENT_STATUSES:
+ raise OpenAITriggerEventError(f"Unexpected trigger event status
{event.get('status')!r}: {event!r}")
+ return event
+
+
class OpenAIHook(BaseHook):
"""
Use OpenAI SDK to interact with OpenAI APIs.
diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py
b/providers/openai/src/airflow/providers/openai/operators/openai.py
index 0f040a03b55..9d308bf3987 100644
--- a/providers/openai/src/airflow/providers/openai/operators/openai.py
+++ b/providers/openai/src/airflow/providers/openai/operators/openai.py
@@ -24,7 +24,7 @@ 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
+from airflow.providers.openai.hooks.openai import OpenAIHook,
validate_execute_complete_event
from airflow.providers.openai.triggers.openai import OpenAIBatchTrigger
if TYPE_CHECKING:
@@ -210,7 +210,8 @@ class OpenAITriggerBatchOperator(BaseOperator):
Relies on trigger to throw an exception, otherwise it assumes
execution was
successful.
"""
- if event["status"] == "error":
+ event = validate_execute_complete_event(event)
+ if event["status"] != "success":
raise OpenAIBatchJobException(event["message"])
self.log.info("%s completed successfully.", self.task_id)
diff --git a/providers/openai/src/airflow/providers/openai/triggers/openai.py
b/providers/openai/src/airflow/providers/openai/triggers/openai.py
index 5800e948490..17c8361824d 100644
--- a/providers/openai/src/airflow/providers/openai/triggers/openai.py
+++ b/providers/openai/src/airflow/providers/openai/triggers/openai.py
@@ -101,13 +101,13 @@ class OpenAIBatchTrigger(BaseTrigger):
"batch_id": self.batch_id,
}
)
-
- yield TriggerEvent(
- {
- "status": "error",
- "message": f"Batch {self.batch_id} has failed.",
- "batch_id": self.batch_id,
- }
- )
+ else:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "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})
diff --git a/providers/openai/tests/unit/openai/hooks/test_openai.py
b/providers/openai/tests/unit/openai/hooks/test_openai.py
index 5e5882e3548..d132a6a3d95 100644
--- a/providers/openai/tests/unit/openai/hooks/test_openai.py
+++ b/providers/openai/tests/unit/openai/hooks/test_openai.py
@@ -38,8 +38,12 @@ from openai.types.vector_stores import VectorStoreFile,
VectorStoreFileBatch, Ve
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
-from airflow.providers.openai.exceptions import OpenAIBatchJobException,
OpenAIBatchTimeout
-from airflow.providers.openai.hooks.openai import OpenAIHook
+from airflow.providers.openai.exceptions import (
+ OpenAIBatchJobException,
+ OpenAIBatchTimeout,
+ OpenAITriggerEventError,
+)
+from airflow.providers.openai.hooks.openai import OpenAIHook,
validate_execute_complete_event
ASSISTANT_ID = "test_assistant_abc123"
ASSISTANT_NAME = "Test Assistant"
@@ -872,3 +876,34 @@ def
test_get_conn_workload_identity_custom_missing_token_provider():
conn = _make_workload_identity_conn("wi_custom_missing",
{"workload_identity_provider": "custom"})
with pytest.raises(ValueError, match="Missing required 'token_provider'"):
OpenAIHook(conn_id=conn.conn_id).get_conn()
+
+
+class TestValidateTriggerEvent:
+ @pytest.mark.parametrize(
+ ("event", "match"),
+ [
+ pytest.param(None, "event is None", id="none"),
+ pytest.param({}, "Unexpected trigger event status None",
id="missing-status"),
+ pytest.param(
+ {"status": "expired", "batch_id": BATCH_ID},
+ "Unexpected trigger event status",
+ id="unknown-status",
+ ),
+ ],
+ )
+ def test_invalid_event_raises(self, event, match):
+ with pytest.raises(OpenAITriggerEventError, match=match):
+ validate_execute_complete_event(event)
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param({"status": "success", "batch_id": BATCH_ID},
id="success"),
+ pytest.param({"status": "error", "batch_id": BATCH_ID, "message":
"boom"}, id="error"),
+ pytest.param(
+ {"status": "cancelled", "batch_id": BATCH_ID, "message":
"cancelled"}, id="cancelled"
+ ),
+ ],
+ )
+ def test_valid_event_is_returned(self, event):
+ assert validate_execute_complete_event(event) is event
diff --git a/providers/openai/tests/unit/openai/operators/test_openai.py
b/providers/openai/tests/unit/openai/operators/test_openai.py
index 954306d42a5..0ec78bb182c 100644
--- a/providers/openai/tests/unit/openai/operators/test_openai.py
+++ b/providers/openai/tests/unit/openai/operators/test_openai.py
@@ -23,6 +23,7 @@ from openai.types.batch import Batch
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.hooks.openai import OpenAIHook
from airflow.providers.openai.operators.openai import (
OpenAIEmbeddingOperator,
@@ -147,3 +148,42 @@ def
test_openai_trigger_batch_operator_with_deferred(mock_batch, wait_for_comple
else:
batch_id = operator.execute(context)
assert batch_id == BATCH_ID
+
+
+class TestOpenAITriggerBatchOperatorExecuteComplete:
+ def _operator(self):
+ return OpenAITriggerBatchOperator(
+ task_id=TASK_ID,
+ conn_id=CONN_ID,
+ file_id=FILE_ID,
+ endpoint=BATCH_ENDPOINT,
+ )
+
+ def test_success_returns_batch_id(self):
+ event = {"status": "success", "message": "done", "batch_id": BATCH_ID}
+ assert self._operator().execute_complete(Context(), event) == BATCH_ID
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param({"status": "error", "message": "boom", "batch_id":
BATCH_ID}, id="error"),
+ pytest.param(
+ {"status": "cancelled", "message": "Batch has been
cancelled.", "batch_id": BATCH_ID},
+ id="cancelled",
+ ),
+ ],
+ )
+ def test_failed_event_raises(self, event):
+ with pytest.raises(OpenAIBatchJobException, match=event["message"]):
+ self._operator().execute_complete(Context(), event)
+
+ @pytest.mark.parametrize(
+ "event",
+ [
+ pytest.param(None, id="none"),
+ pytest.param({"status": "expired", "batch_id": BATCH_ID},
id="unknown-status"),
+ ],
+ )
+ def test_invalid_event_raises_instead_of_succeeding(self, event):
+ with pytest.raises(OpenAITriggerEventError):
+ self._operator().execute_complete(Context(), event)
diff --git a/providers/openai/tests/unit/openai/triggers/test_openai.py
b/providers/openai/tests/unit/openai/triggers/test_openai.py
index 1c1b661f1d8..b8b6c8bac05 100644
--- a/providers/openai/tests/unit/openai/triggers/test_openai.py
+++ b/providers/openai/tests/unit/openai/triggers/test_openai.py
@@ -146,6 +146,28 @@ class TestOpenAIBatchTrigger:
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_yields_single_terminal_event(self, mock_batch):
+ """A terminal batch must produce exactly one event, not a success
followed by an error."""
+ mock_batch.return_value =
self.mock_get_batch(str(BatchStatus.COMPLETED))
+ trigger = OpenAIBatchTrigger(
+ conn_id=self.CONN_ID,
+ batch_id=self.BATCH_ID,
+ poll_interval=self.POLL_INTERVAL,
+ end_time=self.END_TIME,
+ )
+ events = [event async for event in trigger.run()]
+ assert events == [
+ TriggerEvent(
+ {
+ "status": "success",
+ "message": f"Batch {self.BATCH_ID} has completed
successfully.",
+ "batch_id": self.BATCH_ID,
+ }
+ )
+ ]
+
@pytest.mark.asyncio
@mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch")
async def test_openai_batch_for_unexpected_error(self, mock_batch):
