airflow.providers.common.ai.batch.state¶

Run-stable idempotency state for @task.llm_batch, persisted outside XCom.

XCom is cleared at the start of every retry attempt (unlike a deferral resume), so a batch_id pushed to XCom on submit is gone by the time a retry’s execute() runs – the recovery this module exists for cannot rely on XCom at all. State lives as a small JSON file on the same ObjectStoragePath as the batch’s results, at {result_path}/_airflow_batch_state/{key}.json, following the durable-cache pattern in durable/storage.py.

Two invariants that must not be relaxed:

The identity key never includes try_number – that is the entire point of “run-stable”: a retry must compute the same key as the attempt before it, so it can find and re-attach to that attempt’s in-flight batch.
The input fingerprint does include the output schema (see compute_fingerprint()) – changing output_type without changing the prompts must still be treated as a different batch, since the schema is part of what gets sent to the provider.

Attributes¶

SCHEMA_VERSION

Classes¶

BatchStateRecord

The on-disk record at {result_path}/_airflow_batch_state/{key}.json.

Functions¶

`compute_identity_key`(*, dag_id, task_id, run_id, map_index)	Return the run-stable identity key for one task instance (excluding `try_number`).
`key16`(key)	Return the short form of an identity key, used as the `custom_id` prefix and result filename stem.
`compute_output_schema_digest`(output_schema)	Hash the output schema material on its own.
`compute_fingerprint`(*, requests, llm_conn_id, ...)	Hash everything that determines the request content and which account it is billed to.
`write_intent`(result_path, *, key, input_fingerprint, ...)	Phase A: record intent to submit, before any network call.
`write_submitted`(result_path, *, key, ...)	Phase B: overwrite the intent record with the full record, once the provider has accepted the batch.
`read_state`(result_path, key)	Return the recorded state for `key`, or `None` if no record exists.
`delete_state`(result_path, key)	Delete the recorded state for `key` so the next attempt submits a fresh batch.

Module Contents¶

airflow.providers.common.ai.batch.state.SCHEMA_VERSION = 1[source]¶

airflow.providers.common.ai.batch.state.compute_identity_key(*, dag_id, task_id, run_id, map_index)[source]¶

Return the run-stable identity key for one task instance (excluding try_number).

Uses \x00 rather than _ to join the components – a plain _-joined string collides (Dag etl + task load_data and Dag etl_load + task data both yield etl_load_data), which would let one task instance read or overwrite another’s batch state. Mirrors durable/storage.py’s DurableStorage identity hash verbatim, for the same reason.

airflow.providers.common.ai.batch.state.key16(key)[source]¶

Return the short form of an identity key, used as the custom_id prefix and result filename stem.

16 hex characters is 64 bits of the full SHA-256 digest – collision risk across the (at most tens of thousands of requests in) a single batch, or across the handful of concurrent task instances writing under the same result_path, is astronomically below the odds of a provider-side outage; a shorter, more manageable id is worth that trade for something humans read in filenames and provider dashboards.

airflow.providers.common.ai.batch.state.compute_output_schema_digest(output_schema)[source]¶

Hash the output schema material on its own.

Stored alongside the combined input fingerprint so a stale-state error message can say which part changed (prompts vs. the output_type schema) instead of only “something changed”.

airflow.providers.common.ai.batch.state.compute_fingerprint(*, requests, llm_conn_id, model_id, system_prompt, max_tokens, request_params, output_schema)[source]¶

Hash everything that determines the request content and which account it is billed to.

output_schema is caller-supplied (the literal string "str" for an unstructured batch, or an OutputSpec.json_schema dict otherwise) – this module does not know how to derive a schema from output_type itself; that logic lives in batch/output_schema.py, which this module does not import.

Schema is included deliberately: it is sent to the provider as part of every request body (OpenAI’s response_format, Anthropic’s tools[0].input_schema), so a change to it – an added field, an edited description – is a change to the request content, not just to how the response gets parsed. Without this, clearing a task after editing its Pydantic output_type would silently re-attach to a batch whose results were produced under the old schema.

llm_conn_id is included for a different but equally serious reason: two pydanticai connections can point at two different accounts (or even two different providers’ API keys entirely). Without it, switching llm_conn_id and rerunning would silently re-attach to – and return the results of – a batch submitted under a completely different account than the one the current run is configured to use.

class airflow.providers.common.ai.batch.state.BatchStateRecord[source]¶

The on-disk record at {result_path}/_airflow_batch_state/{key}.json.

Written in two phases (see write_intent() / write_submitted()) so a crash between “submit request sent” and “submit response received” leaves a trace the next attempt can act on (see find_orphaned_batch()).

Re-attach eligibility is decided by comparing input_fingerprint, never output_type_ref – the ref is a human-readable label only and can stay identical while the schema underneath it changes.

schema_version: int[source]¶

key: str[source]¶

input_fingerprint: str[source]¶

output_schema_digest: str[source]¶

intent_at: str | None[source]¶

adapter: str | None[source]¶

llm_conn_id: str | None[source]¶

model_id: str | None[source]¶

output_type_ref: str | None[source]¶

batch_id: str | None[source]¶

provider_input_ref: str | None[source]¶

request_count: int | None[source]¶

submitted_at: str | None[source]¶

to_dict()[source]¶

classmethod from_dict(data)[source]¶

airflow.providers.common.ai.batch.state.write_intent(result_path, *, key, input_fingerprint, output_schema_digest, intent_at)[source]¶

Phase A: record intent to submit, before any network call.

batch_id is null; this record exists so a crash between “submit sent” and “submit response received” leaves a trace. Overwritten by write_submitted() once the provider responds. intent_at feeds orphan recovery’s not_before check.

airflow.providers.common.ai.batch.state.write_submitted(result_path, *, key, input_fingerprint, output_schema_digest, intent_at, adapter, llm_conn_id, model_id, output_type_ref, batch_id, provider_input_ref, request_count, submitted_at)[source]¶

Phase B: overwrite the intent record with the full record, once the provider has accepted the batch.

airflow.providers.common.ai.batch.state.read_state(result_path, key)[source]¶

Return the recorded state for key, or None if no record exists.

Only a genuine not-found reads as “no recorded batch”. Any other failure (a transient object-storage error, a corrupt file, a record written by a different schema version) raises LLMBatchStateReadError instead of degrading to None, because “could not read the state” and “there is no state” are different facts and conflating them turns an I/O blip into a duplicate, billable submission.

airflow.providers.common.ai.batch.state.delete_state(result_path, key)[source]¶

Delete the recorded state for key so the next attempt submits a fresh batch.

Only two outcomes call this: a provider-side failed (nothing ran) and a cancellation that lost requests. Every path where the batch is still alive or its results are still worth re-landing leaves the file in place.