From c1ac1292f0f0ca897a2b0e5cdd9e449762cdc814 Mon Sep 17 00:00:00 2001 From: Yash Jain <23dec512@lnmiit.ac.in> Date: Wed, 8 Jul 2026 22:49:35 +0530 Subject: [PATCH] Use monotonic clock for OpenAIBatchTrigger polling timeout --- providers/openai/docs/changelog.rst | 15 +++ .../providers/openai/operators/openai.py | 3 +- .../providers/openai/triggers/openai.py | 73 +++++++++--- .../tests/unit/openai/triggers/test_openai.py | 111 +++++++++++++++--- 4 files changed, 166 insertions(+), 36 deletions(-) diff --git a/providers/openai/docs/changelog.rst b/providers/openai/docs/changelog.rst index 2dfae17aec05b..1302acc458187 100644 --- a/providers/openai/docs/changelog.rst +++ b/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 ..... diff --git a/providers/openai/src/airflow/providers/openai/operators/openai.py b/providers/openai/src/airflow/providers/openai/operators/openai.py index 0f040a03b5505..c9038fef5e4c7 100644 --- a/providers/openai/src/airflow/providers/openai/operators/openai.py +++ b/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", ) diff --git a/providers/openai/src/airflow/providers/openai/triggers/openai.py b/providers/openai/src/airflow/providers/openai/triggers/openai.py index 5800e948490d3..4d4fcac52f8fd 100644 --- a/providers/openai/src/airflow/providers/openai/triggers/openai.py +++ b/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, } ) diff --git a/providers/openai/tests/unit/openai/triggers/test_openai.py b/providers/openai/tests/unit/openai/triggers/test_openai.py index 1c1b661f1d825..0f5cbbf46d4fa 100644 --- a/providers/openai/tests/unit/openai/triggers/test_openai.py +++ b/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,13 +62,13 @@ 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" @@ -74,9 +76,47 @@ def test_serialization(self): "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",