Use monotonic clock for OpenAIBatchTrigger polling timeout

providers/openai/docs/changelog.rst

@@ -37,6 +37,21 @@ explicitly, or install the SDK extra directly:
 
     pip install 'openai[datalib]'
 
+.. Bug fix
+
+``OpenAIBatchTrigger`` now measures its polling timeout with
+:func:`time.monotonic` instead of :func:`time.time`, so a batch task's
+timeout is no longer affected by wall-clock adjustments (NTP corrections,
+DST, VM pause/resume) that happen while the trigger is deferred.
+
+To make that possible the trigger's preferred constructor argument changed
+from ``end_time`` (an absolute wall-clock deadline) to ``timeout`` (a
+duration in seconds). ``OpenAITriggerBatchOperator`` has been updated to
+pass ``timeout``. ``OpenAIBatchTrigger`` still accepts the legacy
+``end_time`` argument so that triggers serialized by the previous version
+of the operator continue to run after an upgrade; direct users of the
+trigger should switch to ``timeout``.
+
 1.8.0
 .....
 

providers/openai/src/airflow/providers/openai/operators/openai.py

@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import time
 from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Literal
@@ -194,7 +193,7 @@ def execute(self, context: Context) -> str | None:
                         conn_id=self.conn_id,
                         batch_id=self.batch_id,
                         poll_interval=60,
-                        end_time=time.time() + self.timeout,
+                        timeout=self.timeout,
                     ),
                     method_name="execute_complete",
                 )

providers/openai/src/airflow/providers/openai/triggers/openai.py

@@ -26,44 +26,87 @@
 
 
 class OpenAIBatchTrigger(BaseTrigger):
-    """Triggers OpenAI Batch API."""
+    """
+    Triggers OpenAI Batch API.
+
+    :param conn_id: The OpenAI connection ID to use.
+    :param batch_id: The ID of the OpenAI batch to wait on.
+    :param poll_interval: Seconds between batch status polls.
+    :param timeout: Total seconds to wait for the batch to reach a terminal
+        state before giving up. This is the preferred way to bound the wait
+        because the trigger measures elapsed time with :func:`time.monotonic`,
+        which is not affected by wall-clock jumps (NTP corrections, DST,
+        container clock skew, VM pause/resume). Either ``timeout`` or
+        ``end_time`` must be provided.
+    :param end_time: Deprecated. Absolute wall-clock deadline (``time.time()``
+        based) after which the trigger reports a timeout. Kept for backward
+        compatibility with triggers that were serialized by the previous
+        version of the operator and are still in flight during an upgrade.
+        Prefer ``timeout`` for new code.
+    """
 
     def __init__(
         self,
         conn_id: str,
         batch_id: str,
         poll_interval: float,
-        end_time: float,
+        end_time: float | None = None,
+        timeout: float | None = None,
     ) -> None:
+        if timeout is None and end_time is None:
+            raise ValueError("OpenAIBatchTrigger requires either 'timeout' or 'end_time'.")
+        if timeout is not None and end_time is not None:
+            raise ValueError("OpenAIBatchTrigger accepts either 'timeout' or 'end_time', not both.")
         super().__init__()
         self.conn_id = conn_id
         self.poll_interval = poll_interval
         self.batch_id = batch_id
         self.end_time = end_time
+        self.timeout = timeout
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serialize OpenAIBatchTrigger arguments and class path."""
-        return (
-            "airflow.providers.openai.triggers.openai.OpenAIBatchTrigger",
-            {
-                "conn_id": self.conn_id,
-                "batch_id": self.batch_id,
-                "poll_interval": self.poll_interval,
-                "end_time": self.end_time,
-            },
-        )
+        """
+        Serialize OpenAIBatchTrigger arguments and class path.
+
+        The trigger stores exactly the argument it was constructed with
+        (``timeout`` or ``end_time``) so that a rolling upgrade never rewrites
+        the schema of an in-flight deferred trigger.
+        """
+        kwargs: dict[str, Any] = {
+            "conn_id": self.conn_id,
+            "batch_id": self.batch_id,
+            "poll_interval": self.poll_interval,
+        }
+        if self.timeout is not None:
+            kwargs["timeout"] = self.timeout
+        else:
+            kwargs["end_time"] = self.end_time
+        return ("airflow.providers.openai.triggers.openai.OpenAIBatchTrigger", kwargs)
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Make connection to OpenAI Client, and poll the status of batch."""
+        # Measure elapsed time with time.monotonic() so the timeout is robust
+        # against wall-clock adjustments (NTP, DST, VM pause/resume, etc.).
+        # For legacy ``end_time`` callers we derive a best-effort remaining
+        # duration from the wall clock exactly once, then track the rest with
+        # the monotonic clock.
+        if self.timeout is not None:
+            timeout = self.timeout
+        else:
+            timeout = max(0.0, self.end_time - time.time())  # type: ignore[operator]
+        start_monotonic = time.monotonic()
         hook = OpenAIHook(conn_id=self.conn_id)
         try:
             while (batch := hook.get_batch(self.batch_id)) and BatchStatus.is_in_progress(batch.status):
-                if self.end_time < time.time():
+                elapsed = time.monotonic() - start_monotonic
+                if elapsed >= timeout:
                     yield TriggerEvent(
                         {
                             "status": "error",
-                            "message": f"Batch {self.batch_id} has not reached a terminal status after "
-                            f"{time.time() - self.end_time} seconds.",
+                            "message": (
+                                f"Batch {self.batch_id} has not reached a terminal status after "
+                                f"{elapsed:.0f} seconds."
+                            ),
                             "batch_id": self.batch_id,
                         }
                     )

providers/openai/tests/unit/openai/triggers/test_openai.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+import itertools
 import time
 from typing import Literal
 from unittest import mock
@@ -34,7 +35,8 @@
 class TestOpenAIBatchTrigger:
     BATCH_ID = "batch_id"
     CONN_ID = "openai_default"
-    END_TIME = time.time() + 24 * 60 * 60
+    TIMEOUT = 24 * 60 * 60
+    LEGACY_END_TIME = time.time() + 24 * 60 * 60
     POLL_INTERVAL = 3.0
 
     def mock_get_batch(
@@ -60,23 +62,61 @@ def mock_get_batch(
             status=status,
         )
 
-    def test_serialization(self):
-        """Assert TestOpenAIBatchTrigger correctly serializes its arguments and class path."""
+    def test_serialization_with_timeout(self):
+        """Trigger constructed with ``timeout`` round-trips through ``serialize``."""
         trigger = OpenAIBatchTrigger(
             conn_id=self.CONN_ID,
             batch_id=self.BATCH_ID,
             poll_interval=self.POLL_INTERVAL,
-            end_time=self.END_TIME,
+            timeout=self.TIMEOUT,
         )
         class_path, kwargs = trigger.serialize()
         assert class_path == "airflow.providers.openai.triggers.openai.OpenAIBatchTrigger"
         assert kwargs == {
             "conn_id": self.CONN_ID,
             "batch_id": self.BATCH_ID,
             "poll_interval": self.POLL_INTERVAL,
-            "end_time": self.END_TIME,
+            "timeout": self.TIMEOUT,
         }
 
+    def test_serialization_with_legacy_end_time(self):
+        """A trigger constructed with the legacy ``end_time`` re-serializes with ``end_time``.
+
+        This preserves on-disk compatibility with triggers that were serialized by the
+        pre-fix operator and are still in flight during a rolling upgrade.
+        """
+        trigger = OpenAIBatchTrigger(
+            conn_id=self.CONN_ID,
+            batch_id=self.BATCH_ID,
+            poll_interval=self.POLL_INTERVAL,
+            end_time=self.LEGACY_END_TIME,
+        )
+        _, kwargs = trigger.serialize()
+        assert kwargs == {
+            "conn_id": self.CONN_ID,
+            "batch_id": self.BATCH_ID,
+            "poll_interval": self.POLL_INTERVAL,
+            "end_time": self.LEGACY_END_TIME,
+        }
+
+    def test_requires_one_of_timeout_or_end_time(self):
+        with pytest.raises(ValueError, match="requires either 'timeout' or 'end_time'"):
+            OpenAIBatchTrigger(
+                conn_id=self.CONN_ID,
+                batch_id=self.BATCH_ID,
+                poll_interval=self.POLL_INTERVAL,
+            )
+
+    def test_rejects_both_timeout_and_end_time(self):
+        with pytest.raises(ValueError, match="not both"):
+            OpenAIBatchTrigger(
+                conn_id=self.CONN_ID,
+                batch_id=self.BATCH_ID,
+                poll_interval=self.POLL_INTERVAL,
+                timeout=self.TIMEOUT,
+                end_time=self.LEGACY_END_TIME,
+            )
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         ("mock_batch_status", "mock_status", "mock_message"),
@@ -102,7 +142,7 @@ async def test_openai_batch_for_terminal_status(
             conn_id=self.CONN_ID,
             batch_id=self.BATCH_ID,
             poll_interval=self.POLL_INTERVAL,
-            end_time=self.END_TIME,
+            timeout=self.TIMEOUT,
         )
         expected_result = {
             "status": mock_status,
@@ -124,26 +164,59 @@ async def test_openai_batch_for_terminal_status(
         ],
     )
     @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch")
-    @mock.patch("time.time")
-    async def test_openai_batch_for_timeout(self, mock_check_time, mock_batch, mock_batch_status):
-        """Assert that run trigger messages in case of batch is still running after timeout"""
-        MOCK_TIME = 1724068066.6468632
+    @mock.patch("airflow.providers.openai.triggers.openai.time.monotonic")
+    async def test_openai_batch_for_timeout(self, mock_monotonic, mock_batch, mock_batch_status):
+        """Trigger reports a timeout error once the monotonic elapsed time exceeds ``timeout``.
+
+        ``time.monotonic`` is patched with an ever-increasing counter rather than a
+        fixed list: the asyncio event loop also calls ``time.monotonic`` internally,
+        so a finite ``side_effect`` would be exhausted by the loop and raise
+        ``StopIteration``. The trigger reads the clock twice with no ``await`` in
+        between, so its measured elapsed time is exactly one counter step (100s),
+        which exceeds the 10s timeout.
+        """
+        mock_monotonic.side_effect = itertools.count(start=0.0, step=100.0).__next__
         mock_batch.return_value = self.mock_get_batch(mock_batch_status)
-        mock_check_time.return_value = MOCK_TIME + 1
         trigger = OpenAIBatchTrigger(
             conn_id=self.CONN_ID,
             batch_id=self.BATCH_ID,
             poll_interval=self.POLL_INTERVAL,
-            end_time=MOCK_TIME,
+            timeout=10.0,
         )
-        expected_result = {
-            "status": "error",
-            "message": f"Batch {self.BATCH_ID} has not reached a terminal status after {mock_check_time.return_value - MOCK_TIME} seconds.",
-            "batch_id": self.BATCH_ID,
-        }
         task = asyncio.create_task(trigger.run().__anext__())
         await asyncio.sleep(0.1)
-        assert TriggerEvent(expected_result) == task.result()
+        event = task.result()
+        assert event.payload["status"] == "error"
+        assert f"Batch {self.BATCH_ID} has not reached a terminal status after" in event.payload["message"]
+        asyncio.get_event_loop().stop()
+
+    @pytest.mark.asyncio
+    @mock.patch("airflow.providers.openai.hooks.openai.OpenAIHook.get_batch")
+    @mock.patch("airflow.providers.openai.triggers.openai.time.monotonic")
+    @mock.patch("airflow.providers.openai.triggers.openai.time.time")
+    async def test_timeout_uses_monotonic_not_wall_clock(self, mock_wall, mock_monotonic, mock_batch):
+        """Regression: the polling timeout is decided by the monotonic clock only.
+
+        The pre-fix trigger compared ``self.end_time`` against ``time.time()`` inside
+        the loop, so a wall-clock jump (NTP correction, DST, VM pause/resume) could
+        extend or shorten the timeout. A ``timeout``-constructed trigger must never
+        consult the wall clock for its timeout decision; ``time.time`` is asserted
+        unused. ``time.monotonic`` uses an ever-increasing counter for the same
+        event-loop reason as ``test_openai_batch_for_timeout``.
+        """
+        mock_monotonic.side_effect = itertools.count(start=0.0, step=100.0).__next__
+        mock_batch.return_value = self.mock_get_batch(str(BatchStatus.IN_PROGRESS))
+        trigger = OpenAIBatchTrigger(
+            conn_id=self.CONN_ID,
+            batch_id=self.BATCH_ID,
+            poll_interval=self.POLL_INTERVAL,
+            timeout=10.0,
+        )
+        task = asyncio.create_task(trigger.run().__anext__())
+        await asyncio.sleep(0.1)
+        event = task.result()
+        assert event.payload["status"] == "error"
+        mock_wall.assert_not_called()
         asyncio.get_event_loop().stop()
 
     @pytest.mark.asyncio
@@ -155,7 +228,7 @@ async def test_openai_batch_for_unexpected_error(self, mock_batch):
             conn_id=self.CONN_ID,
             batch_id=self.BATCH_ID,
             poll_interval=self.POLL_INTERVAL,
-            end_time=self.END_TIME,
+            timeout=self.TIMEOUT,
         )
         expected_result = {
             "status": "error",