kaxil commented on code in PR #74719:
URL: https://github.com/apache/airflow/pull/74719#discussion_r4284033693
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1956,9 +1960,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 | None
+) -> 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 and context is not None:
Review Comment:
Made it required. Every caller already passed one, so
`_handle_current_task_failed` and `_reschedule_limit_reached` now take
`context: Context` and the timeout check always runs (f85ee40f0c).
##########
airflow-core/docs/core-concepts/tasks.rst:
##########
@@ -254,14 +256,83 @@ 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 does not use a retry.
* - ``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 all retries, including policy-driven retries. Does not fire on
a reschedule.
Review Comment:
Yes, intended. It's the same state a sensor's reschedule uses, so it counts
as waiting rather than failing. The docs now list `on_retry_callback`,
`on_failure_callback`, the `on_task_instance_failed` listener hook,
`email_on_retry` and the `ti_failures`/`operator_failures` metrics, and say
they run as usual once a bound is reached.
`test_retry_policy_reschedule_runs_no_failure_callbacks_or_metrics` pins it.
##########
airflow-core/docs/core-concepts/tasks.rst:
##########
@@ -254,14 +256,83 @@ 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 does not use a retry.
* - ``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 all retries, including policy-driven retries. Does not fire on
a reschedule.
* - ``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`` are
untouched and
+``on_retry_callback`` does not fire. It works with ``retries=0``.
+
+.. 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 RESCHEDULE without one raises ``ValueError`` when the
rule or decision
+is created. Both count every reschedule of this try, including a sensor's
pokes, so a retry
Review Comment:
Added. A `RetryRule` fails at parse time; in a custom policy it's at failure
time, logged, and the task falls back to the standard retry logic like any
policy that raises.
##########
airflow-core/docs/core-concepts/tasks.rst:
##########
@@ -254,14 +256,83 @@ 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
Review Comment:
Reworded all four.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -1703,6 +1703,170 @@ def execute(self, context):
)
+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")
+
+
+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(
Review Comment:
`context=None` is gone now that it's required. The finalize test is
`test_retry_policy_reschedule_runs_no_failure_callbacks_or_metrics`: it runs
`run()` and `finalize()` and checks neither callback fires and no failure
metric is counted.
##########
providers/common/compat/src/airflow/providers/common/compat/_retry_policy.py:
##########
@@ -123,7 +124,7 @@ def evaluate(
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)
Review Comment:
Added `test_reschedule_keeps_its_bounds` to the backport tests, skipped
where `RESCHEDULE` doesn't exist.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
