airflow.providers.common.ai.batch.polling

The poll-until-terminal decision logic shared by the trigger and the operator’s sync path.

Both loops do the same thing with different I/O primitives (await asyncio.to_thread(...) vs a direct call, asyncio.sleep vs time.sleep). The decision of what to do with each poll result lives here, once, so the two cannot drift: the caller feeds in either a BatchState or the exception a poll raised, and gets back “keep polling” or a finished event dict, plus whether the batch should be cancelled first.

Attributes

MAX_CONSECUTIVE_POLL_FAILURES

Classes

PollOutcome

What the poll loop should do after one status check.

BatchPoller

Tracks consecutive failures and the wall-clock deadline across polls.

Functions

terminal_event(batch_id, state)

Reshape a terminal BatchState into the shared event dict.

Module Contents

airflow.providers.common.ai.batch.polling.MAX_CONSECUTIVE_POLL_FAILURES = 5[source]
class airflow.providers.common.ai.batch.polling.PollOutcome[source]

What the poll loop should do after one status check.

event is None while the batch is still running. Otherwise it is the dict the trigger yields as its TriggerEvent payload (and the operator’s sync loop returns), with status in success/failed/expired/ cancelled/timeout/error. cancel says the caller should cancel the batch before emitting the event (our own deadline passed with cancel_on_timeout=True); the caller then reports the result of that cancel through BatchPoller.finish_timeout().

event: dict[str, Any] | None[source]
cancel: bool = False[source]
class airflow.providers.common.ai.batch.polling.BatchPoller(*, batch_id, end_time, timeout, cancel_on_timeout)[source]

Tracks consecutive failures and the wall-clock deadline across polls.

Parameters:
  • batch_id (str) – The provider batch id, echoed into every event.

  • end_time (float) – Wall-clock deadline in epoch seconds.

  • timeout (int) – The configured budget in seconds, for the timeout message.

  • cancel_on_timeout (bool) – Whether the deadline should cancel the batch.

batch_id[source]
end_time[source]
timeout[source]
cancel_on_timeout[source]
consecutive_failures = 0[source]
property deadline_iso: str[source]
on_state(state, *, now)[source]

Decide after a successful status check.

on_error(exc, *, now)[source]

Decide after a status check raised.

Persistent failures past the deadline are a real timeout and honor cancel_on_timeout. Persistent failures inside the deadline give up with "error" and leave the batch alone: its health is unknown, so cancelling it would be presumptuous.

finish_timeout(*, cancelled, cancel_error=None)[source]

Build the final timeout event once the caller has (not) cancelled the batch.

The message names the budget, the deadline, and what happened to the batch, so an on-call reader can tell “cancelled, a retry resubmits” from “still running and billing, a retry re-attaches”.

airflow.providers.common.ai.batch.polling.terminal_event(batch_id, state)[source]

Reshape a terminal BatchState into the shared event dict.

Was this entry helpful?