Status

Motivation

As part of the effort to make Airflow AI-native, this AIP introduces a pluggable extension point for retry decisions. The core abstraction (RetryPolicy) is provider-agnostic and useful on its own for declarative exception handling. The LLMRetryPolicy in the Common AI provider builds on this to bring intelligent, LLM-powered error classification to any Airflow task, a concrete example of how AI capabilities integrate into Airflow's existing task lifecycle rather than requiring new operators or wrappers.

Airflow's retry mechanism applies uniform behavior regardless of error type. When a task fails, the system retries based on a static count (retries) and delay (retry_delay) with no awareness of what went wrong. This creates three recurring problems:

1. Wasted retries on permanent failures. Authentication errors, schema validation errors, and permission denials are retried the same number of times as transient network issues. Each retry is guaranteed to fail, wasting compute, delaying DAG completion, and cluttering logs.
2. Suboptimal timing for rate-limited APIs. When a task hits a rate limit with a "retry after 60 seconds" header, Airflow either retries too early (triggering the same 429) or too late (using the default 5-minute delay when 60 seconds would suffice). There is no mechanism to adapt the retry delay to the error.
3. No separation between retry logic and task logic. Users who want error-aware retries must wrap their task code in try/except blocks, manually raise AirflowFailException for non-retryable errors, and sleep for custom durations. This mixes retry policy with business logic and is not reusable across tasks.

Current workarounds and their limitations

Today, users have three options:

• AirflowFailException: Raise it inside a task to prevent retries. This works but requires modifying every task's code and doesn't allow custom delays.
• on_retry_callback: Called on retry, but cannot prevent the retry or change the delay. It's fire-and-forget.
• External provider: The airflow-provider-smart-retry wraps tasks in an LLMSmartRetryOperator. This requires replacing the task's operator, couples retry logic to a specific LLM backend (Ollama), and breaks the standard DAG authoring experience.

None of these approaches provide a clean, reusable way to configure per-exception retry behavior as a parameter on any task.

Quick start

from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy, RetryRule, RetryAction
from datetime import timedelta

@task(
    retries=5,
    retry_policy=ExceptionRetryPolicy(rules=[
        RetryRule(exception="requests.exceptions.HTTPError", action=RetryAction.RETRY,
                  retry_delay=timedelta(minutes=5), reason="Rate limit"),
        RetryRule(exception="google.auth.exceptions.RefreshError", action=RetryAction.FAIL,
                  reason="Auth failure, not retryable"),
    ]),
)
def call_api():
    ...

# For AI-based retries
llm_policy = LLMRetryPolicy(
    llm_conn_id="pydanticai_default",
    fallback_rules=[
        RetryRule(exception=ConnectionError, action=RetryAction.RETRY, retry_delay=timedelta(seconds=10)),
        RetryRule(exception=PermissionError, action=RetryAction.FAIL),
    ],
)

Relationship to Other AIPs

• AIP-72 (Task SDK): RetryPolicy is defined in the Task SDK and evaluated in the task runner process, consistent with AIP-72's model of user code executing in workers.
• AIP-99 (AI Provider Integration): The Common AI provider ships an LLMRetryPolicy that uses PydanticAIHook for LLM-based error classification. This is an optional enhancement, not a core dependency.
• AIP-103 (Task State Management): RetryDecision reasons could be exposed via Task State for observability, though this is out of scope for the initial implementation.

Considerations

What change do you propose?

Add a retry_policy parameter to BaseOperator that accepts a RetryPolicy object. The policy evaluates the exception at failure time and returns a RetryDecision that can override the retry/fail decision and the retry delay. The existing flat parameters (retries, retry_delay, retry_exponential_backoff, max_retry_delay) remain unchanged and serve as defaults when no policy is configured or when the policy defers to standard behavior.

Core types

New module in the Task SDK: airflow.sdk.definitions.retry_policy

class RetryAction(Enum):
    """What should happen after a failure."""
    RETRY = "retry"      # Retry the task
    FAIL = "fail"        # Fail immediately, skip remaining retries
    DEFAULT = "default"  # Fall through to standard retry logic


@dataclass(frozen=True)
class RetryDecision:
    """The result of evaluating a RetryPolicy."""
    action: RetryAction = RetryAction.DEFAULT
    retry_delay: timedelta | None = None   # Override delay; None = use task's default
    reason: str | None = None              # Logged and stored for observability

    @classmethod
    def fail(cls, reason=None) -> RetryDecision:
        return cls(action=RetryAction.FAIL, reason=reason)

    @classmethod
    def retry(cls, delay=None, reason=None) -> RetryDecision:
        return cls(action=RetryAction.RETRY, retry_delay=delay, reason=reason)

    @classmethod
    def default(cls) -> RetryDecision:
        return cls(action=RetryAction.DEFAULT)


class RetryPolicy(abc.ABC):
    """
    Base class for retry policies.

    The evaluate() method runs in the task worker process and has full
    access to the exception object and the Airflow context.
    The scheduler never calls this.
    """

    @abc.abstractmethod
    def evaluate(
        self,
        exception: BaseException,
        try_number: int,
        max_tries: int,
        context: Context | None = None,
    ) -> RetryDecision:
        """Evaluate whether and how to retry given the failure."""
        ...

Declarative policy for common cases

ExceptionRetryPolicy maps exception types to retry behaviors without writing code:

@dataclass
class RetryRule:
    """A single exception-to-behavior mapping."""
    exception: type[BaseException] | str | list  # Class, dotted path, or list of either
    action: RetryAction = RetryAction.RETRY
    retry_delay: timedelta | None = None
    reason: str | None = None
    match_subclasses: bool = True  # True = isinstance; False = exact type


class ExceptionRetryPolicy(RetryPolicy):
    def __init__(self, rules: list[RetryRule], default: RetryAction = RetryAction.DEFAULT):
        self.rules = rules
        self.default = default

    def evaluate(self, exception, try_number, max_tries, context=None) -> RetryDecision:
        for rule in self.rules:
            if rule.matches(exception):
                return RetryDecision(
                    action=rule.action,
                    retry_delay=rule.retry_delay,
                    reason=rule.reason or f"Matched rule for {type(exception).__name__}",
                )
        return RetryDecision(action=self.default)

Exception string validation

RetryRule validates exception strings at definition time:

• Strings without a dot (e.g., "ValueError" instead of "builtins.ValueError") raise ValueError immediately, catching typos early.
• Dotted paths that can't be resolved at parse time produce a warning (they may resolve on the worker, which can have different packages installed).
• By default, rules use isinstance matching, so a rule for OSError also matches ConnectionError (a subclass). Set match_subclasses=False for exact type matching.

A rule can match multiple exception types at once:

# Multiple exceptions sharing the same behaviour
RetryRule(
    exception=[ConnectionError, TimeoutError, "requests.exceptions.HTTPError"],
    action=RetryAction.RETRY,
    retry_delay=timedelta(seconds=30),
    reason="Transient network error",
)

User-facing API

The retry_policy parameter works on any task or operator:

from airflow.sdk.definitions.retry_policy import ExceptionRetryPolicy, RetryRule, RetryAction
from datetime import timedelta

@task(
    retries=5,
    retry_delay=timedelta(minutes=1),
    retry_policy=ExceptionRetryPolicy(rules=[
        RetryRule(
            exception="requests.exceptions.HTTPError",
            action=RetryAction.RETRY,
            retry_delay=timedelta(minutes=5),
            reason="Likely rate limit, backing off",
        ),
        RetryRule(
            exception="google.auth.exceptions.RefreshError",
            action=RetryAction.FAIL,
            reason="Auth failure, not retryable",
        ),
        RetryRule(
            exception=ConnectionError,
            action=RetryAction.RETRY,
            retry_delay=timedelta(seconds=30),
        ),
    ]),
)
def call_external_api():
    response = requests.get("https://api.example.com/data")
    response.raise_for_status()
    return response.json()

Works on operators, via default_args, and with mapped tasks:

# Operator
op = MyOperator(task_id="x", retries=3, retry_policy=my_policy)

# Shared across a DAG via default_args
with DAG("dag", default_args={"retry_policy": my_policy}):
    ...

# Mapped tasks -- policy applies per-instance
@task.partial(retry_policy=my_policy).expand(input=[1, 2, 3])
def my_mapped_task(input):
    ...

Composition with existing parameters

Evaluation priority:

1. AirflowFailException / AirflowSensorTimeout: Caught in a separate handler; policy never sees these. Explicit "do not retry" from task code always wins.
2. retry_policy.evaluate(): Runs for all other exceptions if a policy is configured. Can override to FAIL, RETRY with custom delay, or DEFAULT.
3. Standard retry logic: should_retry based on try_number <= max_tries. Used when no policy is configured or policy returns DEFAULT.

Execution flow

The context parameter passed to evaluate() is the standard Airflow template context dict, including: dag_run, params, task_instance, ds, logical_date, var (variables), conn (connections), and all other template reference fields.

The policy evaluates in the task worker process, between catching the exception and communicating state to the API server. Each decision is logged in the task logs as Retry policy decision action=<action> reason=<reason>.

def _evaluate_retry_policy(
    ti: RuntimeTaskInstance,
    exception: BaseException,
    log: Logger,
    context: Context | None = None,
) -> RetryDecision | None:
    """Evaluate the task's retry policy. Returns None if no policy configured."""
    policy = getattr(ti.task, "retry_policy", None)
    if policy is None:
        return None
    try:
        decision = policy.evaluate(
            exception=exception,
            try_number=ti.try_number,
            max_tries=ti._ti_context_from_server.max_tries,
            context=context,
        )
        if decision.reason:
            log.info("Retry policy decision: %s (%s)", decision.action.value, decision.reason)
        return decision
    except Exception:
        log.exception("Retry policy evaluation failed, using default behavior")
        return None

Communicating the retry delay override to the scheduler

The policy runs in the worker. The retry delay is computed in the scheduler via next_retry_datetime(). The override must cross the process boundary.

Protocol change: Add retry_delay_seconds and retry_reason to TIRetryStatePayload:

class TIRetryStatePayload(StrictBaseModel):
    state: Literal[IntermediateTIState.UP_FOR_RETRY]
    end_date: UtcDateTime
    rendered_map_index: str | None = None
    retry_delay_seconds: float | None = None   # NEW: policy-overridden delay
    retry_reason: str | None = None            # NEW: human-readable reason

Data model change: Add two nullable columns to TaskInstance:

retry_delay_override: Mapped[float | None] = mapped_column(Float, nullable=True)
retry_reason: Mapped[str | None] = mapped_column(String(500), nullable=True)

Scheduler change: next_retry_datetime() checks the override first:

def next_retry_datetime(self):
    if self.retry_delay_override is not None:
        return self.end_date + timedelta(seconds=self.retry_delay_override)
    # Existing exponential backoff calculation (unchanged)
    ...

Execution API versioning: New version file with Cadwyn version change. Old workers that don't send the new fields omit them (defaults to None). Old servers ignore them via Cadwyn stripping. Full backward compatibility in both directions.

DAG serialization

The serialized DAG stores a boolean flag has_retry_policy (matching the pattern used for callbacks like has_on_retry_callback). The full policy object is not serialized. The scheduler only needs to know whether a policy exists (for UI display); it never evaluates the policy. The API server has no use for the policy object either. The worker parses the DAG file directly and has the live policy object, which is where evaluate() runs.

This means:

• No custom serialization/deserialization code for policy classes
• No requirement for custom RetryPolicy subclasses to be importable on the API server
• Older Airflow versions that encounter has_retry_policy in a serialized DAG simply ignore it

Security considerations

• Exception content as LLM input: When using LLMRetryPolicy, exception messages are sent to the configured LLM. Users handling sensitive data can use local models (Ollama/vLLM) so error logs never leave the infrastructure.
• Policy execution scope: The policy runs in the same process as the task, with the same permissions. It cannot escalate privileges.
• Timeout protection: LLMRetryPolicy enforces a configurable timeout (default 30 seconds) on LLM calls. If the provider is degraded, the policy falls back to fallback_rules within the timeout. Connection failures fall back in under 1 second.

Integration with AIP-99 (AI Provider Integration)

The Common AI provider (apache-airflow-providers-common-ai) ships an LLMRetryPolicy as an optional, higher-level policy built on the core RetryPolicy abstraction.

LLMRetryPolicy

class ErrorClassification(BaseModel):
    """Structured LLM output for error classification."""
    category: str           # rate_limit, auth, network, data, transient, permanent
    should_retry: bool
    suggested_delay_seconds: int = 0
    reasoning: str


class LLMRetryPolicy(RetryPolicy):
    """
    Uses PydanticAIHook to call any LLM provider for error classification.
    Falls back to fallback_rules when LLM call fails.
    Enforces a configurable timeout (default 30s).
    """
    def __init__(self, llm_conn_id, model_id=None, instructions=None,
                 fallback_rules=None, timeout=30.0):
        self.llm_conn_id = llm_conn_id
        self.model_id = model_id
        self.instructions = instructions or self._default_instructions()
        self.fallback_rules = fallback_rules
        self.timeout = timeout

    def evaluate(self, exception, try_number, max_tries, context=None):
        try:
            return self._classify_with_timeout(exception, try_number, max_tries)
        except Exception:
            if self.fallback_rules:
                return ExceptionRetryPolicy(rules=self.fallback_rules).evaluate(
                    exception, try_number, max_tries, context)
            return RetryDecision.default()

    def _classify_with_timeout(self, exception, try_number, max_tries):
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(self._classify, exception, try_number, max_tries)
            return future.result(timeout=self.timeout)

    def _classify(self, exception, try_number, max_tries):
        from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
        hook = PydanticAIHook(llm_conn_id=self.llm_conn_id, model_id=self.model_id)
        agent = hook.create_agent(output_type=ErrorClassification,
                                  instructions=self.instructions)
        result = agent.run_sync(
            f"Classify this error (attempt {try_number}/{max_tries}):\n\n{exception}")
        c = result.output
        if not c.should_retry:
            return RetryDecision.fail(reason=f"{c.category}: {c.reasoning}")
        delay = (timedelta(seconds=c.suggested_delay_seconds)
                 if c.suggested_delay_seconds > 0 else None)
        return RetryDecision.retry(delay=delay,
                                   reason=f"{c.category}: {c.reasoning}")

Usage:

from airflow.providers.common.ai.policies.retry import LLMRetryPolicy

# Any LLM provider (OpenAI, Anthropic, Bedrock, Vertex, Ollama, etc.)
@task(retries=5, retry_policy=LLMRetryPolicy(llm_conn_id="my_llm"))
def call_flaky_api(): ...

# Local-only (Ollama -- error logs never leave infrastructure)
@task(retries=3, retry_policy=LLMRetryPolicy(
    llm_conn_id="ollama_local", model_id="ollama:llama3.2"))
def sensitive_task(): ...

# LLM with declarative fallback
@task(retries=5, retry_policy=LLMRetryPolicy(
    llm_conn_id="my_llm",
    timeout=15.0,
    fallback_rules=[
        RetryRule(exception=ConnectionError, retry_delay=timedelta(seconds=10)),
    ]))
def call_api(): ...

Which users are affected?

DAG authors: Gain a new optional parameter (retry_policy) on all tasks and operators. No existing DAGs are affected.

Operator authors: No changes required. RetryPolicy is evaluated by the task runner, not by operators.

Platform operators: A database migration adds two nullable columns (retry_delay_override, retry_reason) to the task_instance table. On PostgreSQL and MySQL 8+, nullable columns with no defaults are metadata-only operations (no table rewrite).

Migration effort

None. The feature is entirely opt-in. Existing DAGs with retries and retry_delay continue to work identically. The new columns have no effect on tasks without a retry_policy.

Downsides

• Additional latency on failure path: When using LLMRetryPolicy, the LLM call adds 1-3 seconds to the failure handling path. This happens only on failures, not on the happy path. For the declarative ExceptionRetryPolicy, overhead is negligible.
• LLM reliability: If the LLM call fails during retry evaluation, the policy falls back to fallback_rules or default behavior within the configured timeout (default 30s). Connection failures fall back in under 1 second.
• New DB columns: Two nullable columns on task_instance. The retry_reason column (VARCHAR 500) adds minimal storage overhead.

Out of scope

• Retry budgets per DAG or DAG run. A system-level cap on total retries across all tasks in a DAG run is a separate feature.
• Automatic retry policy suggestion. Analyzing historical failures to recommend policies could be a follow-up feature.
• UI for configuring retry policies. Policies are defined in code (DAG files), not via the UI.
• Retry delay as a callable. Making retry_delay accept a callable was considered but rejected in favor of the policy pattern, which is more composable.

What defines this AIP as "done"?

1. RetryPolicy, RetryDecision, RetryAction, RetryRule, and ExceptionRetryPolicy classes available in airflow.sdk.definitions.retry_policy.
2. retry_policy parameter accepted on BaseOperator and the @task decorator.
3. Policy evaluation wired into task_runner.py exception handling, with proper fallback on policy errors.
4. TIRetryStatePayload accepts retry_delay_seconds and retry_reason.
5. TaskInstance model has retry_delay_override and retry_reason columns with corresponding migration.
6. next_retry_datetime() respects retry_delay_override.
7. Execution API version change ensures backward compatibility with older workers/servers.
8. has_retry_policy boolean flag in serialized DAGs (matching the callback pattern).
9. LLMRetryPolicy ships in apache-airflow-providers-common-ai with timeout and fallback support.
10. Documentation covers the feature with examples for common use cases.
11. Unit and integration tests cover: rule matching, policy evaluation in task_runner, delay override in scheduler, LLM policy with fallback.