Test Cases AIP-105 Pluggable Retry Policies

Feature: AIP-105 Pluggable Retry Policies Target release: Airflow 3.3.0 (core) + apache-airflow-providers-common-ai (LLM backend) AIP: https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-105%3A+Pluggable+Retry+Policies Author: Kaxil Naik

Run all tests using breeze from main against SQLite (default), PostgreSQL, and MySQL.

Note on existing unit coverage — these layers are already covered in isolation; the cases below focus on migration, full-stack/integration, serialization round-trips, Cadwyn compat, and the LLM backend:

• task-sdk/.../definitions/test_retry_policy.py — RetryRule matching (class / dotted-string / list / match_subclasses), ExceptionRetryPolicy first-match-wins, serialize/deserialize round-trips, _evaluate_retry_policy / _handle_current_task_failed (FAIL bypasses count, RETRY with delay+reason, reason truncation, retries-exhausted → DEFAULT).
• task-sdk/.../execution_time/test_task_runner.py — worker failure routing.
• providers/common/ai/.../policies/test_retry.py — LLMRetryPolicy classify decisions and fallback behaviour.
• airflow-core/.../models/test_taskinstance.py, test_serialized_objects.py, test_dag_serialization.py, execution_api/versions/head/test_task_instances.py — TI columns, serialization, Execution API payload.

How the pieces fit together

A retry policy is an object you attach to a task with retry_policy=.... When the task fails, the policy's evaluate(exception, try_number, max_tries, context) runs in the worker process (it has the live exception object) and returns a RetryDecision:

Two backends ship:

1. ExceptionRetryPolicy (core, airflow.sdk) — declarative RetryRule list, first match wins. No AI dependency.
2. LLMRetryPolicy (apache-airflow-providers-common-ai) — classifies the error with an LLM via PydanticAIHook, with a declarative fallback_rules safety net.

Decision flow back to the scheduler: the worker never writes scheduling state directly. A RETRY decision is carried in the RetryTask Execution-API payload as retry_delay_seconds + retry_reason; the API route writes them to two TI columns, retry_delay_override (Float) and retry_reason (String(500)). TaskInstance.next_retry_datetime() reads retry_delay_override and lets it take precedence over the static delay / exponential backoff. The columns are cleared when the TI next enters RUNNING (ti_run), and a per-try snapshot is preserved in task_instance_history.

Serialization: the policy object itself is not serialized into the DAG. Only a has_retry_policy: bool flag is stored (mirroring has_on_retry_callback). The worker re-parses the DAG file and reads the live ti.task.retry_policy. The scheduler never instantiates or runs the policy.

Reference QA DAG (drop into breeze; exercises F-* below deterministically)

from __future__ import annotations

from datetime import timedelta

from airflow.sdk import DAG, ExceptionRetryPolicy, RetryAction, RetryRule, task

QA_POLICY = ExceptionRetryPolicy(
    rules=[
        RetryRule(
            exception=PermissionError,
            action=RetryAction.FAIL,
            reason="Auth failure, not retryable",
        ),
        RetryRule(
            exception=ConnectionError,
            action=RetryAction.RETRY,
            retry_delay=timedelta(seconds=30),
            reason="Transient, backing off 30s",
        ),
        # list form + dotted-string path, sharing one behaviour
        RetryRule(
            exception=[TimeoutError, "builtins.OSError"],
            action=RetryAction.RETRY,
            retry_delay=timedelta(seconds=10),
        ),
    ],
    default=RetryAction.DEFAULT,
)

with DAG("aip105_retry_policy_qa", schedule=None, catchup=False, tags=["qa", "retry_policy"]):

    @task(retries=3, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def fail_fast_on_auth():
        raise PermissionError("403 Forbidden")  # FAIL → no retry despite retries=3

    @task(retries=3, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def retry_with_override():
        raise ConnectionError("connection reset")  # RETRY → 30s override, not 5m

    @task(retries=3, retry_delay=timedelta(minutes=5), retry_policy=QA_POLICY)
    def default_path():
        raise ValueError("unmatched error")  # no rule → DEFAULT → standard 5m retry

    @task(retries=0, retry_policy=QA_POLICY)
    def retry_exhausted():
        raise ConnectionError("no retries left")  # RETRY but retries=0 → behaves as DEFAULT → FAILED

    fail_fast_on_auth()
    retry_with_override()
    default_path()
    retry_exhausted()

1. Migration testing

Migration 0113_3_3_0_add_retry_policy_fields_to_ti.py (revision b8f3e4a1d2c9, down-revision fde9ed84d07b) adds two nullable columns — retry_delay_override (Float) and retry_reason (String(500)) — to both task_instance and task_instance_history.

2. Functional — ExceptionRetryPolicy (declarative backend)

Full-stack via a running worker. Trigger aip105_retry_policy_qa (reference DAG above) unless noted.

3. Functional — fail-fast exceptions that bypass the policy

The policy is consulted for AirflowException, AirflowTaskTimeout, AirflowRuntimeError, SystemExit, and generic BaseException. It is never consulted for AirflowFailException, AirflowSensorTimeout, or AirflowTaskTerminated — those always fail immediately.

4. Delay-override & reason persistence lifecycle

5. Custom RetryPolicy subclasses & mapped tasks

6. Serialization

The policy object is not serialized; only has_retry_policy: bool is stored. The worker re-parses the DAG to obtain the live policy.

7. LLM backend — LLMRetryPolicy (apache-airflow-providers-common-ai)

Requires Airflow 3.3+ and apache-airflow-providers-common-ai[<llm-extra>]. Example DAG: example_llm_retry_policy.py. For deterministic CI runs, mock PydanticAIHook; for an E2E pass, use a local model (Ollama) or a low-cost hosted model.

Connection setup (per the example): pydanticai_default, conn_type='pydanticai', password=<API key>, extra='{"model": "<provider:model>"}'.

8. Cadwyn / API version compatibility

Execution-API version v2026_06_16 (AddRetryPolicyFields) adds retry_delay_seconds and retry_reason to TIRetryStatePayload.

9. Safety / negative cases

ID	Test Case	Expected Result	Pass Criteria	Execution Result
N-01 [P0]	Scheduler never instantiates or runs the policy	Parse a DAG whose policy object would error if evaluate were called server-side.	Scheduler healthy; policy only ever runs on the worker.	✅ Pass broken/failing policies never affected scheduler; only worker evaluates (all runs scheduled fine)
N-02 [P1]	A non-callable / malformed retry_policy is rejected at definition time	Pass a non-RetryPolicy object.	Clear error at task definition / parse, not at runtime.	Fail malformed retry_policy is NOT rejected at definition time (discrepancy)
N-03 [P1]	Backwards compatibility: no retry_policy == pre-3.3 behaviour	Run existing DAGs unchanged on 3.3.	Identical retry behaviour; columns stay NULL.	✅ Pass no_policy UP_FOR_RETRY via standard retries, no decision log

10. Documentation review