You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kaxil (via GitHub)" <gi...@apache.org> on 2026/09/18 18:08:45 UTC
Re: [PR] Distinguish OpenAI batch timeout and cancellation in deferrable tasks [airflow]
kaxil commented on code in PR #72149:
URL: https://github.com/apache/airflow/pull/72149#discussion_r4049396085
##########
providers/openai/src/airflow/providers/openai/hooks/openai.py:
##########
@@ -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,
Review Comment:
This dict and the trigger's literal termination_reason strings are a producer/consumer contract with no shared type, so a typo on either side wouldn't get caught by mypy. The parametrized tests would catch a mismatch today so it's not urgent, but a small shared `Literal[...]` or `StrEnum` between triggers/openai.py and here would make that a type error instead of relying on test coverage.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -446,15 +464,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:
Review Comment:
This wraps the cancel call so a failure can't mask the timeout, but the other two callers of `cancel_batch` don't: `on_kill` right below (line 523) and `wait_for_batch`'s timeout branch (hooks/openai.py:711) both call it bare. If cancellation happens to fail there, it replaces the intended exception instead of just logging a warning. Might be worth having `on_kill` call `self._cancel_batch_quietly(self.batch_id)` too. Also, the "mirroring the synchronous path" line above is a little generous since the sync path doesn't actually have this protection.
##########
providers/openai/src/airflow/providers/openai/operators/openai.py:
##########
@@ -446,15 +464,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:
Review Comment:
Every `TriggerEvent` `OpenAIBatchTrigger` has ever yielded (back to #41554) sets `batch_id` unconditionally from a required constructor arg, and this operator only ever defers to that one trigger. I don't think this branch is reachable, and there's no test for it. Might be worth dropping it and calling `_cancel_batch_quietly` unconditionally, or leaving a one-line comment that it's precautionary for some future/custom trigger rather than something the current code depends on.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org