This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b2d1d8183dc Use monotonic clock for OpenAIBatchTrigger polling timeout
(#69534)
b2d1d8183dc is described below
commit b2d1d8183dcea1813a92e2ef5a94aebac4c143e8
Author: Yash jain <[email protected]>
AuthorDate: Sun Aug 2 02:24:37 2026 +0530
Use monotonic clock for OpenAIBatchTrigger polling timeout (#69534)
---
providers/openai/docs/changelog.rst | 15 +++
.../airflow/providers/openai/operators/openai.py | 3 +-
.../airflow/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 61d78b800d0..9c60b5203e7 100644
--- a/providers/openai/docs/changelog.rst
+++ b/providers/openai/docs/changelog.rst
@@ -71,6 +71,21 @@ Doc-only
appropriate section above if needed. Do not delete the lines(!):
+.. 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 9d308bf3987..dfed6e48d51 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 @@ class OpenAITriggerBatchOperator(BaseOperator):
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 17c8361824d..49fd0900cc0 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 @@ from airflow.triggers.base import BaseTrigger, TriggerEvent
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 b8b6c8bac05..9acc25d66c3 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 @@ openai = pytest.importorskip("openai")
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 @@ class TestOpenAIBatchTrigger:
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 @@ class TestOpenAIBatchTrigger:
"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 @@ class TestOpenAIBatchTrigger:
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 @@ class TestOpenAIBatchTrigger:
],
)
@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
@@ -177,7 +250,7 @@ class TestOpenAIBatchTrigger:
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",
