diff --git a/airflow-core/docs/core-concepts/tasks.rst b/airflow-core/docs/core-concepts/tasks.rst index 2412dd6d5dc3e..15f110901bd21 100644 --- a/airflow-core/docs/core-concepts/tasks.rst +++ b/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 diff --git a/airflow-core/src/airflow/example_dags/example_retry_policy.py b/airflow-core/src/airflow/example_dags/example_retry_policy.py index 3a6668a29db89..e5ed409ddb568 100644 --- a/airflow-core/src/airflow/example_dags/example_retry_policy.py +++ b/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] diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst index 2c9bbe3f4db80..6ce6213e0c36e 100644 --- a/task-sdk/docs/api.rst +++ b/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 diff --git a/task-sdk/docs/index.rst b/task-sdk/docs/index.rst index 8e0c005b04770..9e15dae85e550 100644 --- a/task-sdk/docs/index.rst +++ b/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` diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py index 6cff747fccc94..aec71b2229054 100644 --- a/task-sdk/src/airflow/sdk/__init__.py +++ b/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", diff --git a/task-sdk/src/airflow/sdk/__init__.pyi b/task-sdk/src/airflow/sdk/__init__.pyi index 37585f5d796eb..6e7a0c8526cfa 100644 --- a/task-sdk/src/airflow/sdk/__init__.pyi +++ b/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", diff --git a/task-sdk/src/airflow/sdk/definitions/retry_policy.py b/task-sdk/src/airflow/sdk/definitions/retry_policy.py index 57d309edf87d0..898047bb0494a 100644 --- a/task-sdk/src/airflow/sdk/definitions/retry_policy.py +++ b/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)})") diff --git a/task-sdk/tests/task_sdk/definitions/test_retry_policy.py b/task-sdk/tests/task_sdk/definitions/test_retry_policy.py index 251a5fa8bc022..1fbc68c0e71ef 100644 --- a/task-sdk/tests/task_sdk/definitions/test_retry_policy.py +++ b/task-sdk/tests/task_sdk/definitions/test_retry_policy.py @@ -24,6 +24,7 @@ from airflow.sdk.api.datamodels._generated import TaskInstanceState from airflow.sdk.definitions.retry_policy import ( + ChainRetryPolicy, ExceptionRetryPolicy, RetryAction, RetryDecision, @@ -228,6 +229,229 @@ def test_serialization_roundtrip_preserves_evaluate_behaviour(self): # --------------------------------------------------------------------------- +class TestChainRetryPolicy: + """Policies in order; the first RETRY or FAIL wins; any DEFAULT means ask the next one.""" + + FAIL_403 = ExceptionRetryPolicy( + rules=[RetryRule(exception=PermissionError, action=RetryAction.FAIL, reason="never retry 403")] + ) + FLOOR = ExceptionRetryPolicy( + rules=[ + RetryRule( + exception=ConnectionError, + action=RetryAction.RETRY, + retry_delay=timedelta(seconds=30), + reason="floor", + ) + ] + ) + + @staticmethod + def _policy(decision): + class Fixed(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + return decision + + return Fixed() + + def test_needs_at_least_one_policy(self): + with pytest.raises(ValueError, match="at least one policy"): + ChainRetryPolicy([]) + + def test_a_single_policy_must_be_wrapped_in_a_sequence(self): + with pytest.raises(TypeError, match="wrap the single policy in a list"): + ChainRetryPolicy(self.FLOOR) # type: ignore[arg-type] + + def test_members_must_be_retry_policies(self): + with pytest.raises(TypeError, match=r"policies\[1\] must be a RetryPolicy, got str"): + ChainRetryPolicy([self.FLOOR, "rules"]) # type: ignore[list-item] + + def test_accepts_any_sequence(self): + policy = ChainRetryPolicy((self.FAIL_403, self.FLOOR)) + + assert policy.policies == [self.FAIL_403, self.FLOOR] + + @pytest.mark.parametrize( + ("first", "exc", "action", "reason"), + [ + pytest.param( + FAIL_403, + PermissionError("403"), + RetryAction.FAIL, + "ExceptionRetryPolicy: never retry 403", + id="fail", + ), + pytest.param( + FLOOR, + ConnectionError("refused"), + RetryAction.RETRY, + "ExceptionRetryPolicy: floor", + id="retry", + ), + ], + ) + def test_first_policy_to_decide_wins_and_later_ones_are_not_consulted(self, first, exc, action, reason): + consulted = [] + + class Recording(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + consulted.append(True) + return RetryDecision.fail(reason="should not run") + + policy = ChainRetryPolicy([first, Recording()]) + + decision = policy.evaluate(exc, try_number=1, max_tries=3) + + assert decision.action == action + assert decision.reason == reason + assert consulted == [] + + def test_default_from_a_policy_moves_to_the_next(self): + policy = ChainRetryPolicy([self.FAIL_403, self.FLOOR]) + + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=30) + assert decision.reason == "ExceptionRetryPolicy: floor (after ExceptionRetryPolicy: no decision)" + + def test_matched_default_rule_passes_control_on(self): + """A rule with action=DEFAULT is not a decision inside a chain, whatever its reason or delay.""" + soft = ExceptionRetryPolicy( + rules=[ + RetryRule( + exception=ConnectionError, + action=RetryAction.DEFAULT, + retry_delay=timedelta(seconds=99), + reason="task settings, please", + ) + ] + ) + policy = ChainRetryPolicy([soft, self.FLOOR]) + + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.retry_delay == timedelta(seconds=30) + assert ( + decision.reason + == "ExceptionRetryPolicy: floor (after ExceptionRetryPolicy: task settings, please)" + ) + + def test_fail_default_ends_the_chain(self): + strict = ExceptionRetryPolicy(rules=[], default=RetryAction.FAIL) + policy = ChainRetryPolicy([strict, self.FLOOR]) + + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "ExceptionRetryPolicy: fail" + + def test_every_policy_abstaining_returns_default_with_the_trail_and_no_delay(self): + soft = ExceptionRetryPolicy( + rules=[ + RetryRule(exception=ValueError, action=RetryAction.DEFAULT, retry_delay=timedelta(seconds=99)) + ] + ) + policy = ChainRetryPolicy([self.FAIL_403, soft]) + + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.DEFAULT + assert decision.retry_delay is None + assert decision.reason == ( + "no policy decided (ExceptionRetryPolicy: no decision; ExceptionRetryPolicy: Matched rule for ValueError)" + ) + + def test_a_policy_that_raises_is_logged_and_skipped(self, caplog): + class Broken(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + raise RuntimeError("policy bug") + + policy = ChainRetryPolicy([Broken(), self.FLOOR]) + + with caplog.at_level("ERROR", logger="airflow.sdk.definitions.retry_policy"): + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.reason == "ExceptionRetryPolicy: floor (after Broken: raised RuntimeError)" + assert "Broken raised while evaluating the retry policy" in caplog.text + assert "policy bug" in caplog.text + + @pytest.mark.parametrize("returned", [None, "retry", {"action": "retry"}]) + def test_a_policy_returning_something_else_is_logged_and_skipped(self, returned, caplog): + policy = ChainRetryPolicy([self._policy(returned), self.FLOOR]) + + with caplog.at_level("ERROR", logger="airflow.sdk.definitions.retry_policy"): + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.reason == "ExceptionRetryPolicy: floor (after Fixed: invalid decision)" + assert "instead of a RetryDecision" in caplog.text + + def test_members_are_called_by_keyword_like_the_worker_does(self): + class KeywordOnly(RetryPolicy): + def evaluate(self, *, exception, try_number, max_tries, context=None): + return RetryDecision.fail(reason="kw-only") + + decision = ChainRetryPolicy([KeywordOnly()]).evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.FAIL + assert decision.reason == "KeywordOnly: kw-only" + + @pytest.mark.parametrize("exc", [KeyboardInterrupt, SystemExit]) + def test_base_exceptions_propagate(self, exc): + class Cancelled(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + raise exc() + + policy = ChainRetryPolicy([Cancelled(), self.FLOOR]) + + with pytest.raises(exc): + policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + def test_every_policy_sees_the_original_exception_and_arguments(self): + seen = [] + + class Recording(RetryPolicy): + def evaluate(self, exception, try_number, max_tries, context=None): + seen.append((exception, try_number, max_tries, context)) + return RetryDecision.default() + + exc = ConnectionError("refused") + ctx = {"params": {}} + ChainRetryPolicy([Recording(), Recording()]).evaluate(exc, try_number=2, max_tries=5, context=ctx) + + assert seen == [(exc, 2, 5, ctx), (exc, 2, 5, ctx)] + + def test_decision_without_a_reason_names_the_action(self): + policy = ChainRetryPolicy([self._policy(RetryDecision.retry(delay=timedelta(seconds=7)))]) + + decision = policy.evaluate(ValueError("x"), try_number=1, max_tries=3) + + assert decision.retry_delay == timedelta(seconds=7) + assert decision.reason == "Fixed: retry" + + def test_nested_chains_compose(self): + inner = ChainRetryPolicy([self.FAIL_403]) + policy = ChainRetryPolicy([inner, self.FLOOR]) + + decision = policy.evaluate(ConnectionError("refused"), try_number=1, max_tries=3) + + assert decision.action == RetryAction.RETRY + assert decision.reason == ( + "ExceptionRetryPolicy: floor (after ChainRetryPolicy: no policy decided (ExceptionRetryPolicy: no decision))" + ) + + def test_runs_through_the_task_runner(self): + ti = _make_mock_ti(policy=ChainRetryPolicy([self.FAIL_403, self.FLOOR])) + + result = _evaluate_retry_policy(ti, ConnectionError("refused"), log) + + assert result.action == RetryAction.RETRY + assert result.retry_delay == timedelta(seconds=30) + + class TestCustomRetryPolicy: def test_context_aware_policy(self): """Policy can use context to make decisions."""