This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new af3fa8db653 Let retry policies reschedule a task without using a retry
(#74719) (#74993)
af3fa8db653 is described below
commit af3fa8db653af229b310c5ec085653e3de81d9c2
Author: Rahul Vats <[email protected]>
AuthorDate: Sun Oct 18 18:04:42 2026 +0530
Let retry policies reschedule a task without using a retry (#74719) (#74993)
A retry policy can now return RetryAction.RESCHEDULE for a failure that
means "not yet" rather than "this went wrong", such as a Kubernetes
namespace quota that is full. The task goes to up_for_reschedule, the
same state a sensor in reschedule mode uses, so the try number and
retries are untouched and on_retry_callback does not fire.
Every RESCHEDULE must carry max_reschedules, reschedule_timeout, or
both, and the task runner enforces them: past a bound the failure goes
through the standard retry logic with a retry_reason naming the bound.
ChainRetryPolicy, its common.compat copy and ClassifierRetryPolicy now
copy the deciding decision with dataclasses.replace instead of
rebuilding it field by field, so the bounds survive composition.
Closes #74660
(cherry picked from commit 7c2c9226e0f792947549cf7dd74ba8bd7e43e243)
Co-authored-by: Kaxil Naik <[email protected]>
---
airflow-core/docs/core-concepts/tasks.rst | 89 +++++++++-
.../airflow/example_dags/example_retry_policy.py | 26 ++-
providers/common/ai/docs/retry_policies.rst | 4 +-
.../airflow/providers/common/ai/policies/retry.py | 18 +-
.../ai/tests/unit/common/ai/policies/test_retry.py | 37 +++-
.../providers/common/compat/_retry_policy.py | 5 +-
.../tests/unit/common/compat/test__retry_policy.py | 16 ++
.../src/airflow/sdk/definitions/retry_policy.py | 94 +++++++++-
.../src/airflow/sdk/execution_time/task_runner.py | 63 ++++++-
.../task_sdk/definitions/test_retry_policy.py | 121 ++++++++++++-
.../task_sdk/execution_time/test_task_runner.py | 191 +++++++++++++++++++++
11 files changed, 618 insertions(+), 46 deletions(-)
diff --git a/airflow-core/docs/core-concepts/tasks.rst
b/airflow-core/docs/core-concepts/tasks.rst
index 00f095aa66e..f57de76baa8 100644
--- a/airflow-core/docs/core-concepts/tasks.rst
+++ b/airflow-core/docs/core-concepts/tasks.rst
@@ -81,7 +81,7 @@ The possible states for a Task Instance are:
* ``skipped``: The task was skipped due to branching, LatestOnly, or similar,
or was :ref:`marked as skipped manually <concepts:task-mark-as-skipped>`.
* ``upstream_failed``: An upstream task failed and the :ref:`Trigger Rule
<concepts:trigger-rules>` says we needed it
* ``up_for_retry``: The task failed, but has retry attempts left and will be
rescheduled.
-* ``up_for_reschedule``: The task is a :doc:`Sensor <sensors>` that is in
``reschedule`` mode
+* ``up_for_reschedule``: The task is a :doc:`Sensor <sensors>` in
``reschedule`` mode, or its :ref:`retry policy rescheduled it
<concepts:retry-policies-reschedule>`, and will run again later
* ``deferred``: The task has been :doc:`deferred to a trigger
<../authoring-and-scheduling/deferring>`
* ``awaiting_input``: The task is a :doc:`Human-in-the-loop
<../tutorial/hitl>` task waiting for a human response. It is managed by the
scheduler and uses neither a worker slot nor the triggerer.
* ``removed``: The task has vanished from the Dag since the run started
@@ -218,12 +218,14 @@ The policy runs in the **task worker process** (never the
scheduler) between cat
exception and deciding the task's next state. Each policy decision is logged
in the task
logs as ``Retry policy decision action=<action> reason=<reason>``.
-When a task fails, the policy evaluates the exception and returns one of three
actions:
+When a task fails, the policy evaluates the exception and returns one of four
actions:
* **RETRY** -- retry the task, optionally with a custom delay that overrides
``retry_delay``.
The retry is still subject to the task's ``retries`` count -- a policy can
fail earlier but
cannot extend past the configured maximum.
* **FAIL** -- fail immediately, skipping any remaining retries.
+* **RESCHEDULE** -- run the task again later without using a retry, within
bounds the policy
+ sets. See :ref:`concepts:retry-policies-reschedule`.
* **DEFAULT** -- fall through to the standard retry logic (``retries`` count
and ``retry_delay``).
Rules are evaluated in order; the first matching rule wins.
@@ -254,14 +256,91 @@ Composition with existing parameters
* - Parameter
- Behaviour when ``retry_policy`` is set
* - ``retries``
- - Still the maximum retry count. Policy can fail earlier but not exceed
it.
+ - Still the maximum retry count. Policy can fail earlier but not exceed
it. A
+ reschedule leaves the retry count unchanged.
* - ``retry_delay`` / ``retry_exponential_backoff`` / ``max_retry_delay``
- Used when the policy returns DEFAULT or when
``RetryDecision.retry_delay`` is None.
+ A reschedule without its own delay waits ``retry_delay``, without
backoff.
* - ``on_retry_callback``
- - Fires on all retries, including policy-driven retries.
+ - Fires on retries only, including policy-driven retries.
* - ``AirflowFailException``
- Always takes precedence. The policy is never consulted for this
exception.
+.. _concepts:retry-policies-reschedule:
+
+Rescheduling without using a retry
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. versionadded:: 3.5.0
+
+Some failures mean "not yet" rather than "this went wrong": the system the
task drives is
+out of capacity, such as a Kubernetes namespace whose resource quota is full
or a warehouse
+with no free slots. Retrying spends the task's ``retries`` on waiting, so a
real failure
+later has none left, and ``on_retry_callback`` fires every time.
+
+**RESCHEDULE** puts the task in ``up_for_reschedule`` instead, the same state
a sensor in
+``mode="reschedule"`` uses between pokes. The try number and ``retries`` stay
as they are, and
+it works with ``retries=0``. Like a sensor's reschedule, it counts as waiting
rather than
+failing, so ``on_retry_callback``, ``on_failure_callback``, the
``on_task_instance_failed``
+listener hook, ``email_on_retry`` and the ``ti_failures`` and
``operator_failures`` metrics
+wait until a bound is reached. The failure then goes through the standard
retry logic, and
+they run as they would for any retry or failure.
+
+.. exampleinclude:: /../src/airflow/example_dags/example_retry_policy.py
+ :language: python
+ :start-after: [START retry_policy_reschedule]
+ :end-before: [END retry_policy_reschedule]
+
+A RESCHEDULE always has a bound, so a rule matching a broad exception cannot
loop forever:
+
+* ``max_reschedules`` -- how many times the try may be rescheduled.
+* ``reschedule_timeout`` -- how long after the first reschedule the try may
keep being
+ rescheduled.
+
+Set either or both; a ``RetryRule`` or ``RetryDecision`` with RESCHEDULE and
neither bound
+raises ``ValueError`` when it is created. For a ``RetryRule`` that is when the
Dag is parsed. A
+custom policy creates its decision when the task fails, so there the error is
logged and the
+task falls back to the standard retry logic, as it does for any policy that
raises.
+
+Both bounds count every reschedule of this try, including a sensor's pokes, so
a retry starts
+them afresh. Once a bound is reached, the failure goes through the standard
retry logic, and
+the task's ``retry_reason`` says which bound stopped it, for example
+``Reschedule limit reached (reschedule_timeout=4:00:00): Warehouse busy``. A
RESCHEDULE that
+goes ahead writes its reason to the task log.
+
+Use RESCHEDULE when an external system the task drives is not ready yet. The
policy runs in
+the worker and sees only the exception the task raised, so a disruption of the
Airflow worker
+itself, such as a preempted worker pod, follows the normal retry path.
+
+Telling "out of capacity" apart from a real failure often needs more than the
exception type.
+A custom policy can look at the exception itself; for example, a Kubernetes
API rejection for
+a full namespace quota:
+
+.. code-block:: python
+
+ from datetime import timedelta
+
+ from kubernetes.client.exceptions import ApiException
+
+ from airflow.sdk import RetryDecision, RetryPolicy
+
+
+ class NamespaceQuotaPolicy(RetryPolicy):
+ """Wait for quota to free up instead of spending retries on it."""
+
+ def evaluate(self, exception, try_number, max_tries, context=None):
+ if (
+ isinstance(exception, ApiException)
+ and exception.status == 403
+ and "exceeded quota" in (exception.body or "")
+ ):
+ return RetryDecision.reschedule(
+ timedelta(minutes=10),
+ reschedule_timeout=timedelta(hours=4),
+ reason="Namespace quota full",
+ )
+ return RetryDecision.default()
+
Reusable policies
~~~~~~~~~~~~~~~~~
@@ -305,7 +384,7 @@ Chaining policies
.. versionadded:: 3.4.0
-``ChainRetryPolicy`` consults policies in order. The first RETRY or FAIL wins;
a
+``ChainRetryPolicy`` consults policies in order. The first RETRY, FAIL, or
RESCHEDULE 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.
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 e5ed409ddb5..383c5082a68 100644
--- a/airflow-core/src/airflow/example_dags/example_retry_policy.py
+++ b/airflow-core/src/airflow/example_dags/example_retry_policy.py
@@ -18,8 +18,8 @@
Example DAG demonstrating pluggable retry policies.
A retry policy evaluates the actual exception at failure time and decides
-whether to retry (with a custom delay), fail immediately, or fall through
-to the standard retry logic.
+whether to retry (with a custom delay), fail immediately, reschedule without
+using a retry, or fall through to the standard retry logic.
"""
from __future__ import annotations
@@ -98,3 +98,25 @@ CHAINED_RETRY_POLICY = ChainRetryPolicy(
]
)
# [END retry_policy_chain]
+
+# [START retry_policy_reschedule]
+
+
+class WarehouseBusyError(Exception):
+ """Raised by the task when the warehouse it submits queries to has no free
slots."""
+
+
+# Come back every 10 minutes while the warehouse is busy, for up to 4 hours,
without using
+# any of the task's retries. After 4 hours the standard retry logic takes over.
+WAREHOUSE_RETRY_POLICY = ExceptionRetryPolicy(
+ rules=[
+ RetryRule(
+ exception=WarehouseBusyError,
+ action=RetryAction.RESCHEDULE,
+ retry_delay=timedelta(minutes=10),
+ reschedule_timeout=timedelta(hours=4),
+ reason="Warehouse busy",
+ ),
+ ],
+)
+# [END retry_policy_reschedule]
diff --git a/providers/common/ai/docs/retry_policies.rst
b/providers/common/ai/docs/retry_policies.rst
index 686eb5c7e0f..e4cdd000699 100644
--- a/providers/common/ai/docs/retry_policies.rst
+++ b/providers/common/ai/docs/retry_policies.rst
@@ -350,7 +350,7 @@ The order of events on a failure:
Retry-After header``.
3. If that policy returns DEFAULT, whatever reason it attached, it decided
nothing: the outer ``fallback_rules`` apply, then the task's own retry
- behaviour. Only a RETRY or FAIL from ``fallback_policy`` ends the chain, so
an
+ behaviour. Any other decision from ``fallback_policy`` ends the chain, so an
outer rule such as ``PermissionError -> FAIL`` still holds when both models
are unreachable.
@@ -564,7 +564,7 @@ Both policies share every parameter below except
``categories``,
- None
- ``ClassifierRetryPolicy`` only. A ``RetryPolicy`` to consult when the
classifier is under its bar, reports no confidence, or cannot be
- reached; typically an ``LLMRetryPolicy`` on a text model. Its RETRY or
FAIL
+ reached; typically an ``LLMRetryPolicy`` on a text model. Its decision
is used; a DEFAULT, whatever its reason, means the outer
``fallback_rules``
apply.
* - ``redactor``
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
index cccdfbe4593..a9ca2baacfd 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/policies/retry.py
@@ -39,7 +39,7 @@ from __future__ import annotations
import logging
import traceback
from collections.abc import Mapping
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
from datetime import timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Literal
@@ -501,7 +501,7 @@ class ClassifierRetryPolicy(_ModelRetryPolicy):
:param fallback_policy: A policy to consult when the answer is under its
bar, reports no
confidence, or the model call fails; typically an
:class:`LLMRetryPolicy` on a text
model, so the classifier handles the clear cases and a reasoning model
the rest.
- Its RETRY or FAIL is used, with the reason prefixed by why the
classifier's answer
+ Its decision is used, with the reason prefixed by why the classifier's
answer
was not. A DEFAULT from it counts as no decision, whatever reason it
carries, and this
policy's ``fallback_rules`` and then the task's own retry behaviour
apply. Without
``min_confidence`` the classifier's answer is always acted on, so this
policy is
@@ -617,9 +617,9 @@ class ClassifierRetryPolicy(_ModelRetryPolicy):
why: str,
) -> RetryDecision | None:
"""
- Consult ``fallback_policy`` and return its RETRY or FAIL, or None when
it decided nothing.
+ Consult ``fallback_policy`` and return its decision, or None when it
decided nothing.
- Only RETRY and FAIL are decisions. DEFAULT means the policy had
nothing to add to the
+ Every action but DEFAULT is a decision. DEFAULT means the policy had
nothing to add to the
task's own settings, whatever reason it attached (its own fallback
message, or a matched
rule with ``action=DEFAULT``), so this policy's ``fallback_rules``
still get their say.
A decision comes back with the reason prefixed by why the classifier's
answer was not
@@ -639,11 +639,7 @@ class ClassifierRetryPolicy(_ModelRetryPolicy):
)
return None
prefix = f"escalated ({why})"
- return RetryDecision(
- action=decision.action,
- retry_delay=decision.retry_delay,
- reason=prefix if decision.reason is None else f"{prefix};
{decision.reason}",
- )
+ return replace(decision, reason=prefix if decision.reason is None else
f"{prefix}; {decision.reason}")
def _fall_back(
self,
@@ -665,9 +661,7 @@ class ClassifierRetryPolicy(_ModelRetryPolicy):
prefix = f"classifier answer not applied ({why})"
ruled = self._rules_decision(exception, try_number, max_tries, context)
if ruled.action is not RetryAction.DEFAULT or ruled.reason is not None:
- return RetryDecision(
- action=ruled.action, retry_delay=ruled.retry_delay,
reason=f"{prefix}; {ruled.reason}"
- )
+ return replace(ruled, reason=f"{prefix}; {ruled.reason}")
return RetryDecision(action=RetryAction.DEFAULT, reason=f"{prefix};
task retry settings apply")
def _classify(
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index 303114e2c24..e4c25f3e76e 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -52,6 +52,10 @@ from airflow.sdk._shared.secrets_masker import
reset_secrets_masker
from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision,
RetryPolicy, RetryRule
from airflow.sdk.log import mask_secret
+requires_reschedule = pytest.mark.skipif(
+ not hasattr(RetryAction, "RESCHEDULE"), reason="RESCHEDULE needs Airflow
3.5+"
+)
+
HOOK = "airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook"
@@ -1135,6 +1139,37 @@ class TestOnUncertain:
assert decision.retry_delay == timedelta(seconds=9)
assert decision.reason == "escalated (model_error)"
+ @requires_reschedule
+ def test_reschedule_from_fallback_policy_keeps_its_bounds(self):
+ inner = MagicMock(spec=LLMRetryPolicy)
+ inner.evaluate.return_value = RetryDecision.reschedule(
+ timedelta(minutes=10), max_reschedules=6, reason="quota full"
+ )
+ policy = ClassifierRetryPolicy(llm_conn_id="nonexistent",
min_confidence=0.8, fallback_policy=inner)
+
+ decision = policy.evaluate(RuntimeError("?"), try_number=1,
max_tries=3)
+
+ assert decision == RetryDecision.reschedule(
+ timedelta(minutes=10), max_reschedules=6, reason="escalated
(model_error); quota full"
+ )
+
+ @requires_reschedule
+ def test_reschedule_from_fallback_rules_keeps_its_bounds(self):
+ rule = RetryRule(
+ exception=RuntimeError,
+ action=RetryAction.RESCHEDULE,
+ reschedule_timeout=timedelta(hours=4),
+ reason="quota full",
+ )
+ policy = ClassifierRetryPolicy(llm_conn_id="nonexistent",
min_confidence=0.8, fallback_rules=[rule])
+
+ decision = policy.evaluate(RuntimeError("?"), try_number=1,
max_tries=3)
+
+ assert decision == RetryDecision.reschedule(
+ reschedule_timeout=timedelta(hours=4),
+ reason="classifier answer not applied (model_error); quota full",
+ )
+
@pytest.mark.parametrize(
"inner_decision",
[
@@ -1150,7 +1185,7 @@ class TestOnUncertain:
],
)
def test_any_default_from_fallback_policy_lets_the_outer_rules_run(self,
inner_decision, caplog):
- """Only RETRY or FAIL ends the chain; the reason text on a DEFAULT
does not make it a decision."""
+ """Only a decision other than DEFAULT ends the chain; the reason text
on a DEFAULT does not make it one."""
inner = MagicMock(spec=LLMRetryPolicy)
inner.evaluate.return_value = inner_decision
policy = ClassifierRetryPolicy(
diff --git
a/providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py
b/providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py
index 3ec6f776c30..e0428db8e3f 100644
---
a/providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py
+++
b/providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py
@@ -30,6 +30,7 @@ the compat tests compare the source of both classes and fail
on drift.
from __future__ import annotations
import logging
+from dataclasses import replace
from typing import TYPE_CHECKING
from airflow.sdk.definitions.retry_policy import RetryAction, RetryDecision,
RetryPolicy
@@ -46,7 +47,7 @@ __all__ = ["ChainRetryPolicy"]
class ChainRetryPolicy(RetryPolicy):
"""
- Consult policies in order; the first RETRY or FAIL wins.
+ Consult policies in order; the first decision other than DEFAULT 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
@@ -123,7 +124,7 @@ class ChainRetryPolicy(RetryPolicy):
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)
+ return replace(decision, reason=reason)
# The worker logs this reason as the policy decision; it is not
stored, since no policy
# took a position.
return RetryDecision(action=RetryAction.DEFAULT, reason=f"no policy
decided ({'; '.join(trail)})")
diff --git
a/providers/common/compat/tests/unit/common/compat/test__retry_policy.py
b/providers/common/compat/tests/unit/common/compat/test__retry_policy.py
index 7056de98128..9a12238ca71 100644
--- a/providers/common/compat/tests/unit/common/compat/test__retry_policy.py
+++ b/providers/common/compat/tests/unit/common/compat/test__retry_policy.py
@@ -293,6 +293,22 @@ class TestChainRetryPolicyBackport:
assert decision.retry_delay == timedelta(seconds=7)
assert decision.reason == "Fixed: retry"
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS or not hasattr(RetryAction, "RESCHEDULE"),
+ reason="RESCHEDULE needs Airflow 3.5+",
+ )
+ def test_reschedule_keeps_its_bounds(self):
+ decided = RetryDecision.reschedule(timedelta(minutes=10),
max_reschedules=6, reason="quota full")
+ policy = ChainRetryPolicy([self.fail_403, self._policy(decided)])
+
+ decision = policy.evaluate(ConnectionError("x"), try_number=1,
max_tries=0)
+
+ assert decision == RetryDecision.reschedule(
+ timedelta(minutes=10),
+ max_reschedules=6,
+ reason="Fixed: quota full (after ExceptionRetryPolicy: no
decision)",
+ )
+
def test_nested_chains_compose(self):
inner = ChainRetryPolicy([self.fail_403])
policy = ChainRetryPolicy([inner, self.floor])
diff --git a/task-sdk/src/airflow/sdk/definitions/retry_policy.py
b/task-sdk/src/airflow/sdk/definitions/retry_policy.py
index 63894103bb6..b3d7b5e78ad 100644
--- a/task-sdk/src/airflow/sdk/definitions/retry_policy.py
+++ b/task-sdk/src/airflow/sdk/definitions/retry_policy.py
@@ -19,14 +19,15 @@ Pluggable retry policies for Airflow tasks.
A ``retry_policy`` parameter on any task or operator evaluates the actual
exception at failure time and returns a :class:`RetryDecision` that can
-override the retry/fail decision and the retry delay.
+override the retry/fail decision and the retry delay, or reschedule the task
+without using a retry.
"""
from __future__ import annotations
import abc
import logging
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
from datetime import timedelta
from enum import Enum
from typing import TYPE_CHECKING, Any
@@ -66,6 +67,35 @@ class RetryAction(Enum):
DEFAULT = "default"
"""Fall through to standard retry logic."""
+ RESCHEDULE = "reschedule"
+ """Run the task again later without using a retry.
+
+ For a failure caused by a system outside Airflow that the task drives not
being ready
+ yet, such as a cluster with no free capacity or a full resource quota. The
task goes to
+ ``up_for_reschedule``: the try number and ``retries`` are untouched, and
+ ``on_retry_callback`` does not fire.
+
+ A reschedule must be bounded by ``max_reschedules``,
``reschedule_timeout``, or both.
+ Once a bound is reached, the failure goes through the standard retry logic
instead.
+ """
+
+
+def _validate_reschedule_bounds(
+ action: RetryAction, max_reschedules: int | None, reschedule_timeout:
timedelta | None
+) -> None:
+ if action is not RetryAction.RESCHEDULE:
+ if max_reschedules is not None or reschedule_timeout is not None:
+ raise ValueError(
+ f"max_reschedules and reschedule_timeout only apply to
RetryAction.RESCHEDULE, not {action}."
+ )
+ return
+ if max_reschedules is None and reschedule_timeout is None:
+ raise ValueError("RetryAction.RESCHEDULE needs max_reschedules,
reschedule_timeout, or both.")
+ if max_reschedules is not None and max_reschedules < 1:
+ raise ValueError(f"max_reschedules must be at least 1, got
{max_reschedules}.")
+ if reschedule_timeout is not None and reschedule_timeout <= timedelta(0):
+ raise ValueError(f"reschedule_timeout must be positive, got
{reschedule_timeout}.")
+
@dataclass(frozen=True)
class RetryDecision:
@@ -73,9 +103,16 @@ class RetryDecision:
action: RetryAction = RetryAction.DEFAULT
retry_delay: timedelta | None = None
- """Override delay for this specific retry. ``None`` uses the task's
default."""
+ """Override delay for this specific retry or reschedule. ``None`` uses
the task's ``retry_delay``."""
reason: str | None = None
"""Human-readable reason, logged and stored on the task instance."""
+ max_reschedules: int | None = None
+ """For RESCHEDULE: how many reschedules of this try, from any cause,
before the standard retry logic applies."""
+ reschedule_timeout: timedelta | None = None
+ """For RESCHEDULE: how long after this try's first reschedule the standard
retry logic applies instead."""
+
+ def __post_init__(self) -> None:
+ _validate_reschedule_bounds(self.action, self.max_reschedules,
self.reschedule_timeout)
@classmethod
def fail(cls, reason: str | None = None) -> RetryDecision:
@@ -87,6 +124,24 @@ class RetryDecision:
"""Retry, optionally with a custom delay."""
return cls(action=RetryAction.RETRY, retry_delay=delay, reason=reason)
+ @classmethod
+ def reschedule(
+ cls,
+ delay: timedelta | None = None,
+ *,
+ max_reschedules: int | None = None,
+ reschedule_timeout: timedelta | None = None,
+ reason: str | None = None,
+ ) -> RetryDecision:
+ """Run the task again after *delay* without using a retry, within the
given bounds."""
+ return cls(
+ action=RetryAction.RESCHEDULE,
+ retry_delay=delay,
+ reason=reason,
+ max_reschedules=max_reschedules,
+ reschedule_timeout=reschedule_timeout,
+ )
+
@classmethod
def default(cls) -> RetryDecision:
"""Use the task's standard retry behaviour."""
@@ -170,9 +225,16 @@ class RetryRule:
:param action: What to do when this rule matches.
:param retry_delay: Override the retry delay for this specific error.
+ For RESCHEDULE, the delay before the task runs again.
:param reason: Human-readable reason logged with the decision.
:param match_subclasses: When ``True`` (default), uses ``isinstance``
for matching. When ``False``, only exact type matches count.
+ :param max_reschedules: For RESCHEDULE, how many times a try may be
+ rescheduled before the standard retry logic applies. Every reschedule
+ of the try counts, including a sensor's pokes.
+ :param reschedule_timeout: For RESCHEDULE, how long after the try's first
+ reschedule the standard retry logic applies instead. A RESCHEDULE
+ rule needs ``max_reschedules``, ``reschedule_timeout``, or both.
"""
exception: ExceptionSpec | list[ExceptionSpec]
@@ -180,8 +242,11 @@ class RetryRule:
retry_delay: timedelta | None = None
reason: str | None = None
match_subclasses: bool = True
+ max_reschedules: int | None = None
+ reschedule_timeout: timedelta | None = None
def __post_init__(self) -> None:
+ _validate_reschedule_bounds(self.action, self.max_reschedules,
self.reschedule_timeout)
# Normalise to a list for uniform handling.
specs = self.exception if isinstance(self.exception, list) else
[self.exception]
for spec in specs:
@@ -271,6 +336,10 @@ class RetryRule:
result["retry_delay"] = self.retry_delay.total_seconds()
if self.reason is not None:
result["reason"] = self.reason
+ if self.max_reschedules is not None:
+ result["max_reschedules"] = self.max_reschedules
+ if self.reschedule_timeout is not None:
+ result["reschedule_timeout"] =
self.reschedule_timeout.total_seconds()
return result
@classmethod
@@ -281,6 +350,10 @@ class RetryRule:
retry_delay=timedelta(seconds=data["retry_delay"]) if
"retry_delay" in data else None,
reason=data.get("reason"),
match_subclasses=data.get("match_subclasses", True),
+ max_reschedules=data.get("max_reschedules"),
+ reschedule_timeout=(
+ timedelta(seconds=data["reschedule_timeout"]) if
"reschedule_timeout" in data else None
+ ),
)
@@ -315,6 +388,11 @@ class ExceptionRetryPolicy(RetryPolicy):
rules: list[RetryRule],
default: RetryAction = RetryAction.DEFAULT,
) -> None:
+ if default is RetryAction.RESCHEDULE:
+ raise ValueError(
+ "ExceptionRetryPolicy default cannot be RESCHEDULE, which
needs bounds; "
+ "use a RetryRule with max_reschedules or reschedule_timeout
instead."
+ )
self.rules = rules
self.default = default
@@ -331,6 +409,8 @@ class ExceptionRetryPolicy(RetryPolicy):
action=rule.action,
retry_delay=rule.retry_delay,
reason=rule.reason or f"Matched rule for
{type(exception).__name__}",
+ max_reschedules=rule.max_reschedules,
+ reschedule_timeout=rule.reschedule_timeout,
)
return RetryDecision(action=self.default)
@@ -352,7 +432,7 @@ class ExceptionRetryPolicy(RetryPolicy):
class ChainRetryPolicy(RetryPolicy):
"""
- Consult policies in order; the first RETRY or FAIL wins.
+ Consult policies in order; the first decision other than DEFAULT 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
@@ -379,8 +459,8 @@ class ChainRetryPolicy(RetryPolicy):
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 or a FAIL, and logs it
when no policy
- decided.
+ worker stores it as ``retry_reason`` on a RETRY or a FAIL, and on a
RESCHEDULE past its
+ bound; it logs it on a RESCHEDULE that goes ahead or when no policy
decided.
:param policies: The policies to consult, in order. At least one.
"""
@@ -428,7 +508,7 @@ class ChainRetryPolicy(RetryPolicy):
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)
+ return replace(decision, reason=reason)
# The worker logs this reason as the policy decision; it is not
stored, since no policy
# took a position.
return RetryDecision(action=RetryAction.DEFAULT, reason=f"no policy
decided ({'; '.join(trail)})")
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 88bf9ca01c7..bfa6e493434 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -1808,8 +1808,10 @@ def run(
# not fail-closed: supervisor's final_state would misclassify
# them too, but exiting non-zero would map them to FAILED,
# which is strictly worse than the default. Those need a
- # separate fix in supervisor's final_state.
- if state in (TaskInstanceState.FAILED,
TaskInstanceState.UP_FOR_RETRY):
+ # separate fix in supervisor's final_state. Whenever the task
+ # raised, though, FAILED beats a silent SUCCESS, including when
+ # a retry policy turned the exception into a reschedule.
+ if error is not None or state in (TaskInstanceState.FAILED,
TaskInstanceState.UP_FOR_RETRY):
ti._terminal_state_send_failed = True
# Return the message to make unit tests easier too
@@ -1890,7 +1892,7 @@ def _evaluate_retry_policy(
def _handle_handler_failure(
ti: RuntimeTaskInstance, exception: Exception, log: Logger, context:
Context
-) -> tuple[RetryTask | TaskState, TaskInstanceState, Exception]:
+) -> tuple[RetryTask | RescheduleTask | TaskState, TaskInstanceState,
Exception]:
"""
Decide the outcome for an exception raised by one of the outcome handlers
themselves.
@@ -1932,16 +1934,18 @@ def _handle_current_task_failed(
ti: RuntimeTaskInstance,
exception: BaseException,
log: Logger,
- context: Context | None = None,
-) -> tuple[RetryTask | TaskState, TaskInstanceState]:
+ context: Context,
+) -> tuple[RetryTask | RescheduleTask | TaskState, TaskInstanceState]:
"""
Handle a failed task: evaluate the retry policy (if any) and decide the
next state.
This is the entry point every failure path routes through. When the policy
returns FAIL the task is marked as failed immediately, bypassing the normal
retry-count check. When it returns RETRY with a custom delay, that delay
is
- forwarded in the ``RetryTask`` message. For DEFAULT (or when no policy is
- configured), the standard ``_finalize_task_failure`` logic runs.
+ forwarded in the ``RetryTask`` message. When it returns RESCHEDULE within
its
+ bounds, the task goes to ``up_for_reschedule`` without using a retry; past
its
+ bounds, the standard logic runs with the reason saying why. For DEFAULT
(or when
+ no policy is configured), the standard ``_finalize_task_failure`` logic
runs.
"""
from airflow.sdk.definitions.retry_policy import RetryAction
@@ -1961,9 +1965,54 @@ def _handle_current_task_failed(
return _finalize_task_failure(
ti, log, retry_delay_override=decision.retry_delay,
retry_reason=decision.reason
)
+ if decision is not None and decision.action == RetryAction.RESCHEDULE:
+ if (limit := _reschedule_limit_reached(ti, decision, context)) is None:
+ ti.end_date = datetime.now(tz=UTC)
+ delay = decision.retry_delay if decision.retry_delay is not None
else ti.task.retry_delay
+ reschedule_date = ti.end_date + delay
+ log.info(
+ "Retry policy rescheduled the task, marking task as
UP_FOR_RESCHEDULE",
+ reschedule_date=reschedule_date,
+ reschedule_count=_reschedule_count(ti),
+ max_reschedules=decision.max_reschedules,
+ reschedule_timeout=decision.reschedule_timeout,
+ )
+ return (
+ RescheduleTask(reschedule_date=reschedule_date,
end_date=ti.end_date),
+ TaskInstanceState.UP_FOR_RESCHEDULE,
+ )
+ log.info("Reschedule limit reached, applying the standard retry
logic", limit=limit)
+ reason = f"Reschedule limit reached ({limit})"
+ if decision.reason:
+ reason = f"{reason}: {decision.reason}"
+ return _finalize_task_failure(ti, log, retry_reason=reason)
return _finalize_task_failure(ti, log)
+def _reschedule_count(ti: RuntimeTaskInstance) -> int:
+ """Return how many times this try has been rescheduled so far."""
+ from_server = ti._ti_context_from_server
+ return (from_server.task_reschedule_count or 0) if from_server else 0
+
+
+def _reschedule_limit_reached(
+ ti: RuntimeTaskInstance, decision: RetryDecision, context: Context
+) -> str | None:
+ """
+ Say which bound stops a RESCHEDULE decision, or return None when the task
may reschedule.
+
+ Both bounds count every reschedule of this try, whatever caused it (a
sensor's pokes and a
+ startup reschedule included), so a retry starts them afresh.
+ """
+ if decision.max_reschedules is not None and _reschedule_count(ti) >=
decision.max_reschedules:
+ return f"max_reschedules={decision.max_reschedules}"
+ if decision.reschedule_timeout is not None:
+ first = ti.get_first_reschedule_date(context)
+ if first is not None and datetime.now(tz=UTC) - first >=
decision.reschedule_timeout:
+ return f"reschedule_timeout={decision.reschedule_timeout}"
+ return None
+
+
def _finalize_task_failure(
ti: RuntimeTaskInstance,
log: Logger,
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 1fbc68c0e71..7ff3659a9a3 100644
--- a/task-sdk/tests/task_sdk/definitions/test_retry_policy.py
+++ b/task-sdk/tests/task_sdk/definitions/test_retry_policy.py
@@ -166,6 +166,99 @@ class TestRetryRuleSerialization:
data = rule.serialize()
assert "retry_delay" not in data
assert "reason" not in data
+ assert "max_reschedules" not in data
+ assert "reschedule_timeout" not in data
+
+ def test_reschedule_bounds_roundtrip(self):
+ rule = RetryRule(
+ exception=ConnectionError,
+ action=RetryAction.RESCHEDULE,
+ retry_delay=timedelta(minutes=10),
+ max_reschedules=12,
+ reschedule_timeout=timedelta(hours=4),
+ )
+
+ restored = RetryRule.deserialize(rule.serialize())
+
+ assert restored.action == RetryAction.RESCHEDULE
+ assert restored.retry_delay == timedelta(minutes=10)
+ assert restored.max_reschedules == 12
+ assert restored.reschedule_timeout == timedelta(hours=4)
+
+
+# ---------------------------------------------------------------------------
+# RESCHEDULE bounds
+# ---------------------------------------------------------------------------
+
+
+class TestRescheduleBounds:
+ def test_reschedule_factory_sets_the_bounds(self):
+ decision = RetryDecision.reschedule(
+ timedelta(minutes=10), max_reschedules=6,
reschedule_timeout=timedelta(hours=1), reason="full"
+ )
+
+ assert decision == RetryDecision(
+ action=RetryAction.RESCHEDULE,
+ retry_delay=timedelta(minutes=10),
+ reason="full",
+ max_reschedules=6,
+ reschedule_timeout=timedelta(hours=1),
+ )
+
+ @pytest.mark.parametrize(
+ ("bounds", "match"),
+ [
+ pytest.param({}, "needs max_reschedules, reschedule_timeout, or
both", id="unbounded"),
+ pytest.param({"max_reschedules": 0}, "max_reschedules must be at
least 1", id="zero-reschedules"),
+ pytest.param(
+ {"reschedule_timeout": timedelta(0)}, "reschedule_timeout must
be positive", id="zero-timeout"
+ ),
+ ],
+ )
+ def test_reschedule_needs_a_valid_bound(self, bounds, match):
+ with pytest.raises(ValueError, match=match):
+ RetryDecision.reschedule(**bounds)
+ with pytest.raises(ValueError, match=match):
+ RetryRule(exception=ConnectionError,
action=RetryAction.RESCHEDULE, **bounds)
+
+ @pytest.mark.parametrize("action", [RetryAction.RETRY, RetryAction.FAIL,
RetryAction.DEFAULT])
+ @pytest.mark.parametrize(
+ "bounds",
+ [{"max_reschedules": 3}, {"reschedule_timeout": timedelta(hours=1)}],
+ ids=["count", "timeout"],
+ )
+ def test_bounds_only_apply_to_reschedule(self, action, bounds):
+ with pytest.raises(ValueError, match="only apply to
RetryAction.RESCHEDULE"):
+ RetryDecision(action=action, **bounds)
+ with pytest.raises(ValueError, match="only apply to
RetryAction.RESCHEDULE"):
+ RetryRule(exception=ConnectionError, action=action, **bounds)
+
+ def test_exception_policy_default_cannot_reschedule(self):
+ with pytest.raises(ValueError, match="default cannot be RESCHEDULE"):
+ ExceptionRetryPolicy(rules=[], default=RetryAction.RESCHEDULE)
+
+ def test_exception_policy_passes_the_rule_bounds_on(self):
+ policy = ExceptionRetryPolicy(
+ rules=[
+ RetryRule(
+ exception=ConnectionError,
+ action=RetryAction.RESCHEDULE,
+ retry_delay=timedelta(minutes=10),
+ max_reschedules=6,
+ reschedule_timeout=timedelta(hours=1),
+ reason="cluster full",
+ )
+ ]
+ )
+
+ decision = policy.evaluate(ConnectionError("x"), try_number=1,
max_tries=0)
+
+ assert decision == RetryDecision.reschedule(
+ timedelta(minutes=10),
+ max_reschedules=6,
+ reschedule_timeout=timedelta(hours=1),
+ reason="cluster full",
+ )
# ---------------------------------------------------------------------------
@@ -230,7 +323,7 @@ class TestExceptionRetryPolicy:
class TestChainRetryPolicy:
- """Policies in order; the first RETRY or FAIL wins; any DEFAULT means ask
the next one."""
+ """Policies in order; the first decision other than DEFAULT wins; any
DEFAULT means ask the next one."""
FAIL_403 = ExceptionRetryPolicy(
rules=[RetryRule(exception=PermissionError, action=RetryAction.FAIL,
reason="never retry 403")]
@@ -432,6 +525,18 @@ class TestChainRetryPolicy:
assert decision.retry_delay == timedelta(seconds=7)
assert decision.reason == "Fixed: retry"
+ def test_reschedule_keeps_its_bounds(self):
+ decided = RetryDecision.reschedule(timedelta(minutes=10),
max_reschedules=6, reason="cluster full")
+ policy = ChainRetryPolicy([self.FAIL_403, self._policy(decided)])
+
+ decision = policy.evaluate(ConnectionError("x"), try_number=1,
max_tries=0)
+
+ assert decision == RetryDecision.reschedule(
+ timedelta(minutes=10),
+ max_reschedules=6,
+ reason="Fixed: cluster full (after ExceptionRetryPolicy: no
decision)",
+ )
+
def test_nested_chains_compose(self):
inner = ChainRetryPolicy([self.FAIL_403])
policy = ChainRetryPolicy([inner, self.FLOOR])
@@ -567,7 +672,7 @@ class TestHandleCurrentTaskFailed:
"""Policy FAIL overrides should_retry=True -- task fails
immediately."""
policy =
ExceptionRetryPolicy(rules=[RetryRule(exception=PermissionError,
action=RetryAction.FAIL)])
ti = _make_mock_ti(policy=policy, should_retry=True)
- msg, state = _handle_current_task_failed(ti,
PermissionError("denied"), log)
+ msg, state = _handle_current_task_failed(ti,
PermissionError("denied"), log, {})
assert state == TaskInstanceState.FAILED
def test_retry_with_delay_and_reason(self):
@@ -582,7 +687,7 @@ class TestHandleCurrentTaskFailed:
]
)
ti = _make_mock_ti(policy=policy)
- msg, state = _handle_current_task_failed(ti,
ConnectionError("refused"), log)
+ msg, state = _handle_current_task_failed(ti,
ConnectionError("refused"), log, {})
assert state == TaskInstanceState.UP_FOR_RETRY
assert msg.retry_delay_seconds == 42.0
assert msg.retry_reason == "net error"
@@ -596,14 +701,14 @@ class TestHandleCurrentTaskFailed:
]
)
ti = _make_mock_ti(policy=policy)
- msg, state = _handle_current_task_failed(ti, ValueError("bad"), log)
+ msg, state = _handle_current_task_failed(ti, ValueError("bad"), log,
{})
assert state == TaskInstanceState.UP_FOR_RETRY
assert len(msg.retry_reason) == 500
def test_default_falls_through_to_standard_retry(self):
policy = ExceptionRetryPolicy(rules=[]) # no match -> DEFAULT
ti = _make_mock_ti(policy=policy, should_retry=True)
- msg, state = _handle_current_task_failed(ti, ValueError("bad"), log)
+ msg, state = _handle_current_task_failed(ti, ValueError("bad"), log,
{})
assert state == TaskInstanceState.UP_FOR_RETRY
assert msg.retry_delay_seconds is None # no override
@@ -615,15 +720,15 @@ class TestHandleCurrentTaskFailed:
]
)
ti = _make_mock_ti(policy=policy, should_retry=False)
- msg, state = _handle_current_task_failed(ti, ConnectionError("x"), log)
+ msg, state = _handle_current_task_failed(ti, ConnectionError("x"),
log, {})
assert state == TaskInstanceState.FAILED
def test_no_policy_with_should_retry_true(self):
ti = _make_mock_ti(policy=None, should_retry=True)
- msg, state = _handle_current_task_failed(ti, ValueError("bad"), log)
+ msg, state = _handle_current_task_failed(ti, ValueError("bad"), log,
{})
assert state == TaskInstanceState.UP_FOR_RETRY
def test_no_policy_with_should_retry_false(self):
ti = _make_mock_ti(policy=None, should_retry=False)
- msg, state = _handle_current_task_failed(ti, ValueError("bad"), log)
+ msg, state = _handle_current_task_failed(ti, ValueError("bad"), log,
{})
assert state == TaskInstanceState.FAILED
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index cbce1bb32d1..bf64c087a3e 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -1703,6 +1703,197 @@ def
test_exhausted_logs_about_retry_policy_decision(create_runtime_ti, mock_supe
)
+class _CapacityError(Exception):
+ """Stands in for a cluster rejecting a pod because its resource quota is
full."""
+
+
+class _NoCapacityTask(BaseOperator):
+ def execute(self, context):
+ raise _CapacityError("exceeded quota")
+
+
+def _reschedule_policy(retry_delay: timedelta | None = timedelta(minutes=10),
**bounds):
+ return ExceptionRetryPolicy(
+ rules=[
+ RetryRule(
+ exception=_CapacityError,
+ action=RetryAction.RESCHEDULE,
+ retry_delay=retry_delay,
+ reason="cluster full",
+ **bounds,
+ )
+ ]
+ )
+
+
[email protected](
+ ("bounds", "task_reschedule_count", "first_rescheduled_ago"),
+ [
+ pytest.param({"max_reschedules": 3}, 0, timedelta(0),
id="first-failure"),
+ pytest.param({"max_reschedules": 3}, 2, timedelta(minutes=20),
id="below-max-reschedules"),
+ pytest.param(
+ {"reschedule_timeout": timedelta(hours=1)}, 2,
timedelta(minutes=20), id="within-timeout"
+ ),
+ ],
+)
+def test_retry_policy_reschedule_does_not_use_a_retry(
+ time_machine,
+ create_runtime_ti,
+ mock_supervisor_comms,
+ bounds,
+ task_reschedule_count,
+ first_rescheduled_ago,
+):
+ """RESCHEDULE within its bounds reschedules the try, even with no retries
configured."""
+ now = datetime(2026, 10, 14, 9, 0, tzinfo=UTC)
+ time_machine.move_to(now, tick=False)
+ task = _NoCapacityTask(task_id="no_capacity", retries=0,
retry_policy=_reschedule_policy(**bounds))
+ ti = create_runtime_ti(task=task,
task_reschedule_count=task_reschedule_count)
+ ti._ti_context_from_server.first_task_reschedule_start_date = now -
first_rescheduled_ago
+
+ state, msg, error = run(ti, ti.get_template_context(),
structlog.get_logger())
+
+ assert state == TaskInstanceState.UP_FOR_RESCHEDULE
+ assert msg == RescheduleTask(reschedule_date=now + timedelta(minutes=10),
end_date=now)
+
+
+def test_retry_policy_reschedule_without_delay_uses_the_task_retry_delay(
+ time_machine, create_runtime_ti, mock_supervisor_comms
+):
+ now = datetime(2026, 10, 14, 9, 0, tzinfo=UTC)
+ time_machine.move_to(now, tick=False)
+ task = _NoCapacityTask(
+ task_id="no_capacity",
+ retry_delay=timedelta(minutes=7),
+ retry_policy=_reschedule_policy(retry_delay=None, max_reschedules=3),
+ )
+ ti = create_runtime_ti(task=task)
+
+ state, msg, error = run(ti, ti.get_template_context(),
structlog.get_logger())
+
+ assert state == TaskInstanceState.UP_FOR_RESCHEDULE
+ assert msg.reschedule_date == now + timedelta(minutes=7)
+
+
[email protected](
+ ("bounds", "task_reschedule_count", "first_rescheduled_ago", "limit"),
+ [
+ pytest.param(
+ {"max_reschedules": 3}, 3, timedelta(minutes=30),
"max_reschedules=3", id="max-reschedules"
+ ),
+ pytest.param(
+ {"reschedule_timeout": timedelta(hours=1)},
+ 5,
+ timedelta(hours=2),
+ "reschedule_timeout=1:00:00",
+ id="reschedule-timeout",
+ ),
+ ],
+)
[email protected](
+ ("retries", "expected_state", "expected_msg"),
+ [
+ pytest.param(
+ 1,
+ TaskInstanceState.UP_FOR_RETRY,
+ lambda now, reason: RetryTask(end_date=now, retry_reason=reason),
+ id="retry-left",
+ ),
+ pytest.param(
+ 0,
+ TaskInstanceState.FAILED,
+ lambda now, reason: TaskState(state=TaskInstanceState.FAILED,
end_date=now, retry_reason=reason),
+ id="no-retry-left",
+ ),
+ ],
+)
+def
test_retry_policy_reschedule_past_its_bound_applies_the_standard_retry_logic(
+ time_machine,
+ create_runtime_ti,
+ mock_supervisor_comms,
+ bounds,
+ task_reschedule_count,
+ first_rescheduled_ago,
+ limit,
+ retries,
+ expected_state,
+ expected_msg,
+):
+ """Past its bound the task retries on its own retry_delay, not the
reschedule delay, or fails."""
+ now = datetime(2026, 10, 14, 9, 0, tzinfo=UTC)
+ time_machine.move_to(now, tick=False)
+ task = _NoCapacityTask(task_id="no_capacity", retries=retries,
retry_policy=_reschedule_policy(**bounds))
+ ti = create_runtime_ti(task=task,
task_reschedule_count=task_reschedule_count)
+ ti._ti_context_from_server.first_task_reschedule_start_date = now -
first_rescheduled_ago
+
+ state, msg, error = run(ti, ti.get_template_context(),
structlog.get_logger())
+
+ assert state == expected_state
+ assert msg == expected_msg(now, f"Reschedule limit reached ({limit}):
cluster full")
+
+
+def test_retry_policy_reschedule_runs_no_failure_callbacks_or_metrics(
+ create_runtime_ti, mock_supervisor_comms
+):
+ """A reschedule is waiting, not failing: no retry or failure callback runs
and no failure is counted."""
+ called = []
+ task = _NoCapacityTask(
+ task_id="no_capacity",
+ retries=1,
+ retry_policy=_reschedule_policy(max_reschedules=3),
+ on_retry_callback=lambda context: called.append("on_retry_callback"),
+ on_failure_callback=lambda context:
called.append("on_failure_callback"),
+ )
+ ti = create_runtime_ti(task=task)
+ context = ti.get_template_context()
+ log = structlog.get_logger()
+ stats_backend = mock.MagicMock(spec=StatsLogger)
+
+ with mock.patch("airflow.sdk.execution_time.task_runner.stats",
stats_backend):
+ state, _, error = run(ti, context, log)
+ finalize(ti, state, context, log, error)
+
+ assert state == TaskInstanceState.UP_FOR_RESCHEDULE
+ assert called == []
+ counted = {call.args[0] for call in stats_backend.incr.call_args_list if
call.args}
+ assert counted.isdisjoint({"ti_failures", "operator_failures"})
+
+
+class _SensorStyleReschedule(BaseOperator):
+ def execute(self, context):
+ raise AirflowRescheduleException(reschedule_date=datetime(2026, 10,
14, 10, 0, tzinfo=UTC))
+
+
[email protected](
+ ("task", "should_fail_closed"),
+ [
+ pytest.param(
+ _NoCapacityTask(task_id="no_capacity",
retry_policy=_reschedule_policy(max_reschedules=3)),
+ True,
+ id="policy-reschedule-after-an-exception",
+ ),
+ pytest.param(_SensorStyleReschedule(task_id="poke"), False,
id="task-raised-reschedule"),
+ ],
+)
+def test_run_fails_closed_when_a_policy_reschedule_cannot_be_reported(
+ create_runtime_ti, mock_supervisor_comms, task, should_fail_closed
+):
+ """A task that raised must not exit 0 (and so count as SUCCESS) because
its reschedule was lost."""
+ runtime_ti = create_runtime_ti(task=task)
+
+ def send_side_effect(msg=None, **kwargs):
+ if isinstance(msg, RescheduleTask):
+ raise BrokenPipeError("supervisor IPC broken")
+ return mock.DEFAULT
+
+ mock_supervisor_comms.send.side_effect = send_side_effect
+
+ state, _, _ = run(runtime_ti, runtime_ti.get_template_context(),
structlog.get_logger())
+
+ assert state == TaskInstanceState.UP_FOR_RESCHEDULE
+ assert runtime_ti._terminal_state_send_failed is should_fail_closed
+
+
def test_finalize_emits_endgroup(create_runtime_ti, mock_supervisor_comms):
"""finalize() closes the post-execute log group but does not open it."""
task = BaseOperator(task_id="some_task")
