Add ChainRetryPolicy to compose retry policies in order

airflow-core/docs/core-concepts/tasks.rst

@@ -265,6 +265,40 @@ share one policy. Per-index variation is not supported on ``.expand()``, but the
 policy's ``evaluate()`` method receives the exception, ``try_number``, and full
 context, so per-index branching can be done inside the policy if needed.
 
+Chaining policies
+~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 3.4.0
+
+``ChainRetryPolicy`` consults policies in order. The first RETRY or FAIL wins; a
+policy that returns DEFAULT has nothing to add and the next one is asked. When every
+policy returns DEFAULT, the task's own ``retries`` and ``retry_delay`` apply. Every
+policy sees the original task exception.
+
+.. exampleinclude:: /../src/airflow/example_dags/example_retry_policy.py
+    :language: python
+    :start-after: [START retry_policy_chain]
+    :end-before: [END retry_policy_chain]
+
+This is how to put a cheap, deterministic policy in front of a slow or costly one, or
+to give a policy that can fail on its own (one that calls a model, for instance) a
+deterministic floor behind it. The policies need not know about each other, and they can
+come from different packages.
+
+Two consequences of "DEFAULT means next" to keep in mind:
+
+* A ``RetryRule`` with ``action=RetryAction.DEFAULT`` passes control on rather than
+  ending the chain. Its ``retry_delay`` is dropped, as it is on any DEFAULT.
+* An ``ExceptionRetryPolicy`` with ``default=RetryAction.FAIL`` fails every exception
+  its rules do not match, so it ends the chain wherever it sits.
+
+A policy that raises an ordinary exception, or returns something other than a
+``RetryDecision``, is logged and treated as DEFAULT, so one broken policy does not take the
+rules after it down with it. The winning decision's reason names the policy that decided and
+then what the earlier ones said (``HTTPStatusRetryPolicy: HTTP 503 (after ExceptionRetryPolicy:
+no decision)``). On a RETRY that string is the task's ``retry_reason``; on FAIL, or when no
+policy decided, it appears in the task log as the ``Retry policy decision`` line.
+
 Custom retry policies
 ~~~~~~~~~~~~~~~~~~~~~
 
@@ -291,7 +325,7 @@ number, max tries, and the full Airflow context (``dag_run``, ``params``, etc.):
                 status = exception.response.status_code
                 if status == 429:  # rate limited -- honour Retry-After header
                     retry_after = int(exception.response.headers.get("Retry-After", 60))
-                    return RetryDecision.retry(retry_delay=timedelta(seconds=retry_after))
+                    return RetryDecision.retry(delay=timedelta(seconds=retry_after))
                 if 500 <= status < 600:  # server error -- worth retrying
                     return RetryDecision.retry()
                 if 400 <= status < 500:  # client error -- not retryable

airflow-core/src/airflow/example_dags/example_retry_policy.py

@@ -27,7 +27,7 @@
 from datetime import timedelta
 
 # [START retry_policy_definition]
-from airflow.sdk import DAG, ExceptionRetryPolicy, RetryAction, RetryRule, task
+from airflow.sdk import DAG, ChainRetryPolicy, ExceptionRetryPolicy, RetryAction, RetryRule, task
 
 API_RETRY_POLICY = ExceptionRetryPolicy(
     rules=[
@@ -79,3 +79,22 @@ def call_external_api():
     ],
 )
 # [END retry_policy_reusable]
+
+# [START retry_policy_chain]
+# Cheap, deterministic rules first, so a slower policy (a model-backed one from a provider,
+# in practice; a second rule set stands in for it here) only sees what they did not settle.
+# The last rung is the floor for when that policy has no answer.
+CHAINED_RETRY_POLICY = ChainRetryPolicy(
+    [
+        ExceptionRetryPolicy(
+            rules=[RetryRule(exception="google.auth.exceptions.RefreshError", action=RetryAction.FAIL)],
+        ),
+        STANDARD_RETRY_POLICY,
+        ExceptionRetryPolicy(
+            rules=[
+                RetryRule(exception=TimeoutError, retry_delay=timedelta(minutes=2), reason="timeout floor")
+            ],
+        ),
+    ]
+)
+# [END retry_policy_chain]

task-sdk/docs/api.rst

@@ -156,6 +156,8 @@ Tasks page in the core docs for usage and design rationale.
 
 .. autoapiclass:: airflow.sdk.ExceptionRetryPolicy
 
+.. autoapiclass:: airflow.sdk.ChainRetryPolicy
+
 .. autoapiclass:: airflow.sdk.RetryRule
 
 .. autoapiclass:: airflow.sdk.RetryDecision

task-sdk/docs/index.rst

@@ -86,6 +86,7 @@ Why use ``airflow.sdk``?
 - :class:`airflow.sdk.BaseOperator`
 - :class:`airflow.sdk.BaseOperatorLink`
 - :class:`airflow.sdk.BaseSensorOperator`
+- :class:`airflow.sdk.ChainRetryPolicy`
 - :class:`airflow.sdk.Connection`
 - :class:`airflow.sdk.Context`
 - :class:`airflow.sdk.DAG`

task-sdk/src/airflow/sdk/__init__.py

@@ -40,6 +40,7 @@
     "BaseXCom",
     "BranchMixIn",
     "ChainMapper",
+    "ChainRetryPolicy",
     "Connection",
     "Context",
     "CronDataIntervalTimetable",
@@ -196,6 +197,7 @@
         YearWindow,
     )
     from airflow.sdk.definitions.retry_policy import (
+        ChainRetryPolicy,
         ExceptionRetryPolicy,
         RetryAction,
         RetryDecision,
@@ -250,6 +252,7 @@
     "BaseXCom": ".bases.xcom",
     "BranchMixIn": ".bases.branch",
     "ChainMapper": ".definitions.partition_mappers.chain",
+    "ChainRetryPolicy": ".definitions.retry_policy",
     "Connection": ".definitions.connection",
     "Context": ".definitions.context",
     "CronDataIntervalTimetable": ".definitions.timetables.interval",

task-sdk/src/airflow/sdk/__init__.pyi

@@ -102,6 +102,7 @@ from airflow.sdk.definitions.partition_mappers.window import (
     YearWindow,
 )
 from airflow.sdk.definitions.retry_policy import (
+    ChainRetryPolicy as ChainRetryPolicy,
     ExceptionRetryPolicy as ExceptionRetryPolicy,
     RetryAction as RetryAction,
     RetryDecision as RetryDecision,
@@ -156,6 +157,7 @@ __all__ = [
     "BaseXCom",
     "BranchMixIn",
     "ChainMapper",
+    "ChainRetryPolicy",
     "Connection",
     "Context",
     "CronDataIntervalTimetable",

task-sdk/src/airflow/sdk/definitions/retry_policy.py

@@ -32,11 +32,14 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from airflow.sdk.definitions.context import Context
 
 log = logging.getLogger(__name__)
 
 __all__ = [
+    "ChainRetryPolicy",
     "ExceptionRetryPolicy",
     "RetryAction",
     "RetryDecision",
@@ -126,8 +129,10 @@ def serialize(self) -> dict[str, Any]:
         Serialize this policy for DAG serialization.
 
         Must return a JSON-serializable dict.  The default implementation
-        raises :class:`NotImplementedError`; built-in policies like
-        :class:`ExceptionRetryPolicy` provide their own implementation.
+        raises :class:`NotImplementedError`; policies whose configuration is plain
+        data, like :class:`ExceptionRetryPolicy`, provide their own implementation.
+        Nothing in Airflow calls this: a Dag's ``retry_policy`` is re-created by
+        parsing the Dag file on the worker.
         """
         raise NotImplementedError(f"{type(self).__name__} must implement serialize()")
 
@@ -342,3 +347,85 @@ def deserialize(cls, data: dict[str, Any]) -> ExceptionRetryPolicy:
         rules = [RetryRule.deserialize(r) for r in data["rules"]]
         default = RetryAction(data.get("default", "default"))
         return cls(rules=rules, default=default)
+
+
+class ChainRetryPolicy(RetryPolicy):
+    """
+    Consult policies in order; the first RETRY or FAIL wins.
+
+    Every policy sees the original task exception. A policy that returns DEFAULT has nothing
+    to add, whatever reason it attached, and the next one is asked. When every policy returns
+    DEFAULT the chain does too, and the task's own ``retries`` and ``retry_delay`` apply.
+
+    This makes the order of a fallback explicit and lets any policies compose, including ones
+    from different packages::
+
+        retry_policy = ChainRetryPolicy(
+            [
+                # Known failures first, so no model call is spent on them.
+                ExceptionRetryPolicy(rules=[RetryRule(exception=PermissionError, action=RetryAction.FAIL)]),
+                # A policy that consults a model (airflow.providers.common.ai.policies.retry).
+                LLMRetryPolicy(llm_conn_id="pydanticai_default"),
+                # The floor when the model is unreachable.
+                ExceptionRetryPolicy(
+                    rules=[RetryRule(exception=ConnectionError, retry_delay=timedelta(seconds=30))]
+                ),
+            ]
+        )
+
+    A policy that raises an ordinary exception, or returns something other than a
+    :class:`RetryDecision`, is logged and treated as DEFAULT. ``BaseException`` propagates.
+
+    The winning decision's reason names the policy that decided, then what every earlier policy
+    said: ``HTTPStatusRetryPolicy: HTTP 404 (after ExceptionRetryPolicy: no decision)``. The
+    worker stores it as ``retry_reason`` on a RETRY and logs it otherwise.
+
+    :param policies: The policies to consult, in order. At least one.
+    """
+
+    def __init__(self, policies: Sequence[RetryPolicy]) -> None:
+        if isinstance(policies, RetryPolicy):
+            raise TypeError(
+                "ChainRetryPolicy takes a sequence of policies; wrap the single policy in a list."
+            )
+        policies = list(policies)
+        if not policies:
+            raise ValueError("ChainRetryPolicy needs at least one policy.")
+        for position, policy in enumerate(policies):
+            if not isinstance(policy, RetryPolicy):
+                raise TypeError(
+                    f"ChainRetryPolicy policies[{position}] must be a RetryPolicy, got {type(policy).__name__}."
+                )
+        self.policies = policies
+
+    def evaluate(
+        self,
+        exception: BaseException,
+        try_number: int,
+        max_tries: int,
+        context: Context | None = None,
+    ) -> RetryDecision:
+        trail: list[str] = []
+        for policy in self.policies:
+            name = type(policy).__name__
+            try:
+                decision = policy.evaluate(
+                    exception=exception, try_number=try_number, max_tries=max_tries, context=context
+                )
+            except Exception as exc:
+                log.exception("%s raised while evaluating the retry policy; treated as no decision", name)
+                trail.append(f"{name}: raised {type(exc).__name__}")
+                continue
+            if not isinstance(decision, RetryDecision):
+                log.error("%s returned %r instead of a RetryDecision; treated as no decision", name, decision)
+                trail.append(f"{name}: invalid decision")
+                continue
+            if decision.action is RetryAction.DEFAULT:
+                trail.append(f"{name}: {decision.reason or 'no decision'}")
+                continue
+            reason = f"{name}: {decision.reason or decision.action.value}"
+            if trail:
+                reason = f"{reason} (after {'; '.join(trail)})"
+            return RetryDecision(action=decision.action, retry_delay=decision.retry_delay, reason=reason)
+        # The worker logs this reason as the policy decision; it is not stored, since nothing is retried by it.
+        return RetryDecision(action=RetryAction.DEFAULT, reason=f"no policy decided ({'; '.join(trail)})")