Lee-W commented on code in PR #72149:
URL: https://github.com/apache/airflow/pull/72149#discussion_r4090968653


##########
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:
   Dropped the branch.
   `execute_complete` now reads `event["batch_id"]` the same way the success 
path does and calls `_cancel_batch_quietly` unconditionally on a timeout.



-- 
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]

Reply via email to