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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
