kaxil commented on code in PR #73027:
URL: https://github.com/apache/airflow/pull/73027#discussion_r4042333959
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1931,9 +1932,17 @@ def _finalize_task_failure(
if retry_reason is not None:
retry_kwargs["retry_reason"] = retry_reason[:500]
return RetryTask(**retry_kwargs), TaskInstanceState.UP_FOR_RETRY
+ if retry_reason is not None and ti._ti_context_from_server is not None:
+ # max_tries is the retry count, not the attempt count -- total
attempts is max_tries + 1.
+ total_attempts = ti._ti_context_from_server.max_tries + 1
+ suffix = f"; retries exhausted ({ti.try_number} of {total_attempts})"
Review Comment:
Ran this at head: a task with `retry_policy=` and no `retries=` persists
`"rate limit; retries exhausted (1 of 1)"`. `[core] default_task_retries` is 0,
so `max_tries` is 0, `should_retry` is False on the first failure, and this
branch fires on attempt 1. The count is right and matches core's own "Starting
attempt 1 of 1", but "exhausted" is the wrong word when no budget ever existed,
and the thing the user needs to do (set `retries=N`) is exactly what the
wording hides. Skipping the suffix when `max_tries == 0` covers it.
Worth deciding now whether the worker should bake this sentence into the
stored value at all, because nothing reads the column yet (`retry_reason` greps
empty across `core_api/` and `ui/`) and once a UI renders it the format is
fixed. `max_tries` is cumulative, `ti.max_tries = ti.try_number + task.retries`
on clear, so after a couple of clears the denominator reads like a retry count
nobody configured. `try_number > max_tries` is derivable from the same row, so
a renderer could say this in the viewer's locale and the column could hold just
the policy's own text.
Either way this branch logs nothing, so the task log's last word is `Retry
policy decision action=retry` and then FAILED, with nothing saying why the
retry did not happen. A `log.info` here would close that in the channel that
works today.
##########
airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py:
##########
@@ -52,3 +52,13 @@ class AddCallbackRunEndpoint(VersionChange):
instructions_to_migrate_to_previous_version = (
endpoint("/callbacks/{callback_id}/run", ["PATCH"]).didnt_exist,
)
+
+
+class AddTerminalStateRetryReasonField(VersionChange):
+ """Add the `retry_reason` field to TITerminalStatePayload for failed
retry-policy decisions."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(TITerminalStatePayload).field("retry_reason").didnt_exist,
Review Comment:
Correcting my own ask from last round: I measured this at head and the
instruction does not gate anything. With `Airflow-API-Version` pinned, a body
carrying `retry_reason` is accepted at 2025-04-11, 2026-04-06, 2026-06-30 and
2026-10-30 alike, while a bogus field is rejected at all four with
`extra_forbidden` on the `_terminal_` union tag. That tag in the error path is
the tell: `TIStateUpdate` is a discriminated union and cadwyn does not reach
through it to swap in the versioned copies, so the route keeps validating
against the head models. The shipped `rendered_map_index` gate at 2025-04-28 is
inert in exactly the same way, so this is not something you introduced.
I do not think it blocks and I am not asking you to fix cadwyn. Keep the
class, it is the documented recipe. Two things worth knowing: a client
generated from a pinned old version inherits a contract claiming `retry_reason`
existed on `TITerminalStatePayload` back to 2025-04-11, which shipped in 3.0.0,
and the per-version test that `AGENTS.md` step 4 asks for would fail today if
you wrote it, which is why this stays invisible. I will file the framework
issue separately.
This does not generalise to the supervisor half, which I checked separately
and which does bind.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -1196,6 +1196,65 @@ def execute(self, context):
assert counted.count("operator_failures") == 1
+def test_retry_policy_fail_persists_reason(create_runtime_ti,
mock_supervisor_comms):
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(
+ task_id="fail_with_reason",
+ retry_policy=ExceptionRetryPolicy(
+ rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL,
reason="do not retry")]
+ ),
+ )
+ ti = create_runtime_ti(task=task, should_retry=True)
+
+ state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.retry_reason == "do not retry"
+
+
+def
test_retry_policy_retry_exhausted_persists_combined_reason(create_runtime_ti,
mock_supervisor_comms):
+ """A policy-chosen RETRY that hits an exhausted budget still fails, with
both reasons recorded."""
+
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(
+ task_id="retry_exhausted",
+ retry_policy=ExceptionRetryPolicy(
+ rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY,
reason="rate limit")]
+ ),
+ )
+ ti = create_runtime_ti(task=task, try_number=3, max_tries=2,
should_retry=False)
+
+ state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.retry_reason == "rate limit; retries exhausted (3 of 3)"
+
+
+def test_plain_retries_exhausted_has_no_reason(create_runtime_ti,
mock_supervisor_comms):
+ """Without a retry policy, exhausting the retry budget must not synthesize
a reason."""
+
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(task_id="plain_exhausted")
+ ti = create_runtime_ti(task=task, try_number=3, max_tries=2,
should_retry=False)
+
+ state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.retry_reason is None
Review Comment:
`RetryAction.RETRY`'s docstring still says "When all retries are exhausted,
RETRY behaves identically to DEFAULT". This test and the one above it now
assert those two halves diverging: RETRY-exhausted stores a reason,
DEFAULT-exhausted stores NULL. Since that docstring is the rendered public docs
for the enum, worth updating it here. `FAIL` and `DEFAULT`'s docstrings are
both still accurate.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -1196,6 +1196,65 @@ def execute(self, context):
assert counted.count("operator_failures") == 1
+def test_retry_policy_fail_persists_reason(create_runtime_ti,
mock_supervisor_comms):
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(
+ task_id="fail_with_reason",
+ retry_policy=ExceptionRetryPolicy(
+ rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL,
reason="do not retry")]
+ ),
+ )
+ ti = create_runtime_ti(task=task, should_retry=True)
+
+ state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.retry_reason == "do not retry"
+
+
+def
test_retry_policy_retry_exhausted_persists_combined_reason(create_runtime_ti,
mock_supervisor_comms):
+ """A policy-chosen RETRY that hits an exhausted budget still fails, with
both reasons recorded."""
+
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(
+ task_id="retry_exhausted",
+ retry_policy=ExceptionRetryPolicy(
+ rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY,
reason="rate limit")]
+ ),
+ )
+ ti = create_runtime_ti(task=task, try_number=3, max_tries=2,
should_retry=False)
+
+ state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock())
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(msg, TaskState)
+ assert msg.retry_reason == "rate limit; retries exhausted (3 of 3)"
Review Comment:
I reverted last round's truncation back to `retry_reason[:500]` and this
test plus the other new ones all still passed, so nothing guards the fix. With
a 600-char reason the old form produces 528 chars, the route cuts it to 500,
and the suffix is amputated entirely, which is the defect the fix was for.
Widening the reason here to `"z" * 600` and adding `assert
len(msg.retry_reason) == 500` plus `assert msg.retry_reason.endswith("; retries
exhausted (3 of 3)")` closes it. The `endswith` is the half that fails on the
old code.
Related and cheap in the same pass: `_handle_current_task_failed`'s FAIL
branch passes `decision.reason` with no bound, so it is the one site that
leaves truncation entirely to the route. `decision.reason[:500]` there would
make the rule uniform.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -1196,6 +1196,65 @@ def execute(self, context):
assert counted.count("operator_failures") == 1
+def test_retry_policy_fail_persists_reason(create_runtime_ti,
mock_supervisor_comms):
+ class _AlwaysFails(BaseOperator):
+ def execute(self, context):
+ raise RuntimeError("boom")
+
+ task = _AlwaysFails(
+ task_id="fail_with_reason",
+ retry_policy=ExceptionRetryPolicy(
+ rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL,
reason="do not retry")]
+ ),
+ )
+ ti = create_runtime_ti(task=task, should_retry=True)
Review Comment:
`should_retry=True` with the operator's `retries` left at its default is
another triple the server will not produce, the same shape as the one from last
round: server-side it is `max_tries != 0 and try_number <= max_tries`, so with
`max_tries=0` this can only be False. The assertion still holds, since the FAIL
branch returns before `should_retry` is read, so it is not false coverage. But
the fixture reads as "budget available, policy overrides it" without setting
that up, which is the most valuable thing to pin about `RetryAction.FAIL`.
`retries=2` on the operator plus `max_tries=2` here would deliver it.
Same idea on 1232 and 1249: adding `retries=2` to the operator lets the
fixture derive `max_tries` and `should_retry` itself instead of overriding all
three axes, so they cannot drift back out of the reachable space.
##########
task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py:
##########
@@ -34,3 +35,11 @@ class AddArgBindingsToSupervisorTIRunContext(VersionChange):
description = __doc__
instructions_to_migrate_to_previous_version =
(schema(TIRunContext).field("arg_bindings").didnt_exist,)
+
+
+class AddRetryReasonToTaskState(VersionChange):
+ """Add `retry_reason` to `TaskState`."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version =
(schema(TaskState).field("retry_reason").didnt_exist,)
Review Comment:
This gate does work, unlike the execution-API one. I measured it:
downgrading a head `TaskState` to 2026-06-16 strips `retry_reason`, 2026-10-30
keeps it, and upgrading a 2026-06-16 body yields `None`, so java-sdk's pin
stays coherent. It just is not asserted anywhere.
`AddArgBindingsToSupervisorTIRunContext` right above has
`TestRealBundleArgBindingsDowngrade` driving the real migrator, and a
two-assert sibling would pin this one. The upgrade direction is the one foreign
runtimes are actually on.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -670,7 +670,9 @@ def _create_ti_state_update_query_and_update_state(
query = query.values(state=updated_state, next_method=None,
next_kwargs=None)
if updated_state == TaskInstanceState.FAILED:
- # This is the only case needs extra handling for
TITerminalStatePayload
+ if isinstance(ti_patch_payload, TITerminalStatePayload) and
ti_patch_payload.retry_reason:
+ failed_retry_reason: str | None =
ti_patch_payload.retry_reason[:500]
Review Comment:
Small one alongside the truncation thread on this line: the `str | None`
annotation is wider than the guard above it, which already narrows to a truthy
`str`. Inlining `query =
query.values(retry_reason=ti_patch_payload.retry_reason[:500])` drops the local
and the annotation together. The hunk also dropped the `# This is the only case
needs extra handling for TITerminalStatePayload` comment, which was the one
line explaining why the branch exists.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -670,7 +670,9 @@ def _create_ti_state_update_query_and_update_state(
query = query.values(state=updated_state, next_method=None,
next_kwargs=None)
if updated_state == TaskInstanceState.FAILED:
- # This is the only case needs extra handling for
TITerminalStatePayload
+ if isinstance(ti_patch_payload, TITerminalStatePayload) and
ti_patch_payload.retry_reason:
+ failed_retry_reason: str | None =
ti_patch_payload.retry_reason[:500]
+ query = query.values(retry_reason=failed_retry_reason)
Review Comment:
Rebase hazard rather than anything wrong here.
`providers/common/ai/docs/retry_policies.rst` on main (from #73158, merged
after you branched) says "``category`` and ``reasoning`` are only recorded on a
RETRY ... On a FAIL they are not written anywhere". That text is not on your
branch, so it survives the merge and then contradicts this line, since
`LLMRetryPolicy` does set the reason on FAIL. Worth updating that paragraph
after rebasing, including the part about a terminally-failed row keeping the
value rather than having it cleared when the next attempt starts.
--
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]
