From b93cc1844ece8f0ce7bc5e57a8117931045f2810 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Sat, 12 Sep 2026 15:08:25 +0530 Subject: [PATCH 1/8] Persist retry_reason not just for retries, but even when a task fails --- .../execution_api/datamodels/taskinstance.py | 1 + .../execution_api/routes/task_instances.py | 6 +- .../execution_api/versions/__init__.py | 8 ++- .../execution_api/versions/v2026_10_30.py | 12 +++- .../versions/head/test_task_instances.py | 66 +++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 14 +++- .../airflow/sdk/api/datamodels/_generated.py | 1 + .../src/airflow/sdk/execution_time/comms.py | 1 + .../sdk/execution_time/schema/schema.json | 12 ++++ .../airflow/sdk/execution_time/supervisor.py | 3 + .../airflow/sdk/execution_time/task_runner.py | 9 ++- task-sdk/tests/task_sdk/api/test_client.py | 7 +- .../execution_time/test_supervisor.py | 22 +++++++ .../execution_time/test_task_runner.py | 59 +++++++++++++++++ 14 files changed, 214 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ddf31db971838..ac1cb781349d6 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -85,6 +85,7 @@ class TITerminalStatePayload(StrictBaseModel): end_date: UtcDateTime """When the task completed executing""" rendered_map_index: str | None = None + retry_reason: str | None = None class TISuccessStatePayload(StrictBaseModel): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index a49d8fe30e546..d4b773b028365 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -670,7 +670,11 @@ 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) + if ti is not None: + ti.retry_reason = failed_retry_reason if ti is not None: _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag) elif isinstance(ti_patch_payload, TIRetryStatePayload): diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 70ea9be2be65c..22e762e39ae95 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -55,11 +55,17 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30 import ( AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, + AddTerminalStateRetryReasonField, ) bundle = VersionBundle( HeadVersion(), - Version("2026-10-30", AddArgBindingsToTIRunContext, AddCallbackRunEndpoint), + Version( + "2026-10-30", + AddArgBindingsToTIRunContext, + AddCallbackRunEndpoint, + AddTerminalStateRetryReasonField, + ), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py index 0620053b3e2eb..ad4b8f55b31ca 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -26,7 +26,7 @@ schema, ) -from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext, TITerminalStatePayload class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects): @@ -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, + ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index ccaecce822d6e..1acd717cc819f 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -2446,6 +2446,72 @@ def test_ti_update_state_to_failed_table_check(self, client, session, create_tas assert ti.next_kwargs is None assert ti.duration == 3600.00 + def test_ti_update_state_to_failed_persists_retry_reason(self, client, session, create_task_instance): + ti = create_task_instance( + task_id="test_ti_update_state_to_failed_persists_retry_reason", + state=State.RUNNING, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": TerminalTIState.FAILED, + "end_date": DEFAULT_END_DATE.isoformat(), + "retry_reason": "auth error, do not retry", + }, + ) + + assert response.status_code == 204 + + session.expire_all() + ti = session.get(TaskInstance, ti.id) + assert ti.state == State.FAILED + assert ti.retry_reason == "auth error, do not retry" + + def test_ti_update_state_to_failed_truncates_retry_reason(self, client, session, create_task_instance): + ti = create_task_instance( + task_id="test_ti_update_state_to_failed_truncates_retry_reason", + state=State.RUNNING, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": TerminalTIState.FAILED, + "end_date": DEFAULT_END_DATE.isoformat(), + "retry_reason": "x" * 600, + }, + ) + + assert response.status_code == 204 + + session.expire_all() + ti = session.get(TaskInstance, ti.id) + assert ti.retry_reason == "x" * 500 + + def test_ti_update_state_to_failed_without_retry_reason(self, client, session, create_task_instance): + ti = create_task_instance( + task_id="test_ti_update_state_to_failed_without_retry_reason", + state=State.RUNNING, + ) + session.commit() + + response = client.patch( + f"/execution/task-instances/{ti.id}/state", + json={ + "state": TerminalTIState.FAILED, + "end_date": DEFAULT_END_DATE.isoformat(), + }, + ) + + assert response.status_code == 204 + + session.expire_all() + ti = session.get(TaskInstance, ti.id) + assert ti.retry_reason is None + def test_ti_update_state_not_running(self, client, session, create_task_instance): """Test that a 409 error is returned when attempting to update a TI that is not in RUNNING state.""" ti = create_task_instance( diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index ea220c68fde71..11c170e55021b 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -271,13 +271,23 @@ def start(self, id: uuid.UUID, pid: int, when: datetime) -> TIRunContext: raise return TIRunContext.model_validate_json(resp.read()) - def finish(self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index): + def finish( + self, + id: uuid.UUID, + state: TerminalStateNonSuccess, + when: datetime, + rendered_map_index, + retry_reason: str | None = None, + ): """Tell the API server that this TI has reached a terminal state.""" if state == TaskInstanceState.SUCCESS: raise ValueError("Logic error. SUCCESS state should call the `succeed` function instead") # TODO: handle the naming better. finish sounds wrong as "even" deferred is essentially finishing. body = TITerminalStatePayload( - end_date=when, state=TerminalStateNonSuccess(state), rendered_map_index=rendered_map_index + end_date=when, + state=TerminalStateNonSuccess(state), + rendered_map_index=rendered_map_index, + retry_reason=retry_reason, ) self.client.patch(f"task-instances/{id}/state", content=body.model_dump_json()) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 3635bad9628a0..a5f9df3d95dec 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -724,6 +724,7 @@ class TITerminalStatePayload(BaseModel): state: TerminalStateNonSuccess end_date: Annotated[AwareDatetime, Field(title="End Date")] rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None + retry_reason: Annotated[str | None, Field(title="Retry Reason")] = None class XComArgBinding(BaseModel): diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 4a26f56e297dc..ee7aa11f331f0 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -869,6 +869,7 @@ class TaskState(BaseModel): end_date: datetime | None = None type: Literal["TaskState"] = "TaskState" rendered_map_index: str | None = None + retry_reason: str | None = None class SucceedTask(TISuccessStatePayload): diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 85f8b7cfa4e2b..23640ce3e8921 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4203,6 +4203,18 @@ ], "default": null, "title": "Rendered Map Index" + }, + "retry_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Retry Reason" } }, "required": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 0a4808512856f..9c5feac20105b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1407,6 +1407,7 @@ class ActivitySubprocess(WatchedSubprocess): _terminal_state: str | None = attrs.field(default=None, init=False) _final_state: str | None = attrs.field(default=None, init=False) + _retry_reason: str | None = attrs.field(default=None, init=False) # The terminal-state message currently being processed by `_handle_request`, # captured BEFORE the dedicated API call (succeed / retry / defer / # reschedule). If the API call raises (network blip, server 5xx, etc.), @@ -1567,6 +1568,7 @@ def update_task_state_if_needed(self): state=self.final_state, when=datetime.now(tz=timezone.utc), rendered_map_index=self._rendered_map_index, + retry_reason=self._retry_reason, ) def _send_terminal_state_msg( @@ -1803,6 +1805,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: self._terminal_state = msg.state self._task_end_time_monotonic = time.monotonic() self._rendered_map_index = msg.rendered_map_index + self._retry_reason = msg.retry_reason elif isinstance(msg, SucceedTask): self._task_end_time_monotonic = time.monotonic() self._rendered_map_index = msg.rendered_map_index 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 eeb2788062248..da9e2a0813e7c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1887,6 +1887,7 @@ def _handle_current_task_failed( state=TaskInstanceState.FAILED, end_date=ti.end_date, rendered_map_index=ti.rendered_map_index, + retry_reason=decision.reason, ), TaskInstanceState.FAILED, ) @@ -1931,9 +1932,15 @@ 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 = ti._ti_context_from_server.max_tries + retry_reason = f"{retry_reason}; retries exhausted ({ti.try_number} of {max_tries})" return ( TaskState( - state=TaskInstanceState.FAILED, end_date=end_date, rendered_map_index=ti.rendered_map_index + state=TaskInstanceState.FAILED, + end_date=end_date, + rendered_map_index=ti.rendered_map_index, + retry_reason=retry_reason, ), TaskInstanceState.FAILED, ) diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 684d4e7dcbb17..e31829797e4a6 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -432,6 +432,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert actual_body["end_date"] == "2024-10-31T12:00:00Z" assert actual_body["state"] == state assert actual_body["rendered_map_index"] == "test" + assert actual_body["retry_reason"] == "auth error, do not retry" return httpx.Response( status_code=204, ) @@ -439,7 +440,11 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(transport=httpx.MockTransport(handle_request)) client.task_instances.finish( - ti_id, state=state, when="2024-10-31T12:00:00Z", rendered_map_index="test" + ti_id, + state=state, + when="2024-10-31T12:00:00Z", + rendered_map_index="test", + retry_reason="auth error, do not retry", ) def test_task_instance_heartbeat(self): diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 93422f4426a00..8ae01f8ddba22 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3705,6 +3705,28 @@ def test_update_task_state_no_recovery_without_pending_msg(self, watched_subproc watched_subprocess.client.task_instances.finish.assert_not_called() watched_subprocess.client.task_instances.succeed.assert_not_called() + def test_task_state_retry_reason_forwarded_to_finish(self, watched_subprocess, mocker): + """A TaskState message's retry_reason must reach the deferred finish() call.""" + watched_subprocess, _ = watched_subprocess + watched_subprocess._exit_code = 0 + + msg = TaskState( + state=TaskInstanceState.FAILED, + end_date=timezone.parse("2024-10-31T12:00:00Z"), + retry_reason="auth error, do not retry", + ) + watched_subprocess._handle_request(msg, mocker.Mock(), req_id=1) + + watched_subprocess.update_task_state_if_needed() + + watched_subprocess.client.task_instances.finish.assert_called_once_with( + id=watched_subprocess.id, + state=TaskInstanceState.FAILED, + when=mocker.ANY, + rendered_map_index=None, + retry_reason="auth error, do not retry", + ) + class TestSetSupervisorComms: class DummyComms: 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 cc9fb77e08921..5feac21c28f69 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 @@ -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=2, 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 (2 of 2)" + + +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=2, 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 + + def test_run_downstream_skipped(mocked_parse, create_runtime_ti, mock_supervisor_comms, listener_manager): listener = TestTaskRunnerCallsListeners.CustomListener() listener_manager(listener) From a6f747cd5c9a5723f264951d076ec6d26630f394 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 15 Sep 2026 17:04:51 +0530 Subject: [PATCH 2/8] comments from kaxil --- .../api_fastapi/execution_api/routes/task_instances.py | 2 -- .../sdk/execution_time/schema/versions/__init__.py | 3 ++- .../sdk/execution_time/schema/versions/v2026_10_30.py | 9 +++++++++ task-sdk/src/airflow/sdk/execution_time/task_runner.py | 6 ++++-- task-sdk/tests/task_sdk/api/test_client.py | 7 ++++--- .../tests/task_sdk/execution_time/test_task_runner.py | 6 +++--- 6 files changed, 22 insertions(+), 11 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index d4b773b028365..4464e6fc35480 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -673,8 +673,6 @@ def _create_ti_state_update_query_and_update_state( 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) - if ti is not None: - ti.retry_reason = failed_retry_reason if ti is not None: _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag) elif isinstance(ti_patch_payload, TIRetryStatePayload): diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 7e5ce93f86bdc..06be64d346170 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -39,11 +39,12 @@ def get_bundle() -> VersionBundle: from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( AddArgBindingsToSupervisorTIRunContext, + AddRetryReasonToTaskState, ) return VersionBundle( HeadVersion(), - Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext), + Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext, AddRetryReasonToTaskState), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index e6b93f5dea805..e6f6e920d4520 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -20,6 +20,7 @@ from cadwyn import VersionChange, schema from airflow.sdk.api.datamodels._generated import TIRunContext +from airflow.sdk.execution_time.comms import TaskState class AddArgBindingsToSupervisorTIRunContext(VersionChange): @@ -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,) 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 da9e2a0813e7c..74c9d30fc8b83 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1933,8 +1933,10 @@ def _finalize_task_failure( 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 = ti._ti_context_from_server.max_tries - retry_reason = f"{retry_reason}; retries exhausted ({ti.try_number} of {max_tries})" + # 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})" + retry_reason = f"{retry_reason[: 500 - len(suffix)]}{suffix}" return ( TaskState( state=TaskInstanceState.FAILED, diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index e31829797e4a6..d3c47cebbf685 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -422,7 +422,8 @@ def handle_request(request: httpx.Request) -> httpx.Response: @pytest.mark.parametrize( "state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS] ) - def test_task_instance_finish(self, state): + @pytest.mark.parametrize("retry_reason", [None, "auth error, do not retry"]) + def test_task_instance_finish(self, state, retry_reason): # Simulate a successful response from the server that finishes (moved to terminal state) a task ti_id = uuid6.uuid7() @@ -432,7 +433,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert actual_body["end_date"] == "2024-10-31T12:00:00Z" assert actual_body["state"] == state assert actual_body["rendered_map_index"] == "test" - assert actual_body["retry_reason"] == "auth error, do not retry" + assert actual_body["retry_reason"] == retry_reason return httpx.Response( status_code=204, ) @@ -444,7 +445,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: state=state, when="2024-10-31T12:00:00Z", rendered_map_index="test", - retry_reason="auth error, do not retry", + retry_reason=retry_reason, ) def test_task_instance_heartbeat(self): 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 5feac21c28f69..4c4f50b46d243 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 @@ -1229,13 +1229,13 @@ def execute(self, context): rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY, reason="rate limit")] ), ) - ti = create_runtime_ti(task=task, try_number=2, max_tries=2, should_retry=False) + 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 (2 of 2)" + 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): @@ -1246,7 +1246,7 @@ def execute(self, context): raise RuntimeError("boom") task = _AlwaysFails(task_id="plain_exhausted") - ti = create_runtime_ti(task=task, try_number=2, max_tries=2, should_retry=False) + 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()) From baa26759fd902267ecf4a8dc585a4337ebd61bdd Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 18 Sep 2026 10:16:58 +0530 Subject: [PATCH 3/8] review comments from kaxil --- .../execution_api/routes/task_instances.py | 4 +- .../airflow/sdk/definitions/retry_policy.py | 3 +- .../airflow/sdk/execution_time/task_runner.py | 17 +++++---- .../execution_time/schema/test_migrator.py | 34 +++++++++++++++++ .../execution_time/test_task_runner.py | 38 +++++++++++++++++-- 5 files changed, 82 insertions(+), 14 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 4464e6fc35480..ba3dc458eed7c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -670,9 +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) + query = query.values(retry_reason=ti_patch_payload.retry_reason[:500]) if ti is not None: _handle_fail_fast_for_dag(ti=ti, dag_id=dag_id, session=session, dag_bag=dag_bag) elif isinstance(ti_patch_payload, TIRetryStatePayload): diff --git a/task-sdk/src/airflow/sdk/definitions/retry_policy.py b/task-sdk/src/airflow/sdk/definitions/retry_policy.py index 57d309edf87d0..d2a70a1e01ac4 100644 --- a/task-sdk/src/airflow/sdk/definitions/retry_policy.py +++ b/task-sdk/src/airflow/sdk/definitions/retry_policy.py @@ -53,7 +53,8 @@ class RetryAction(Enum): The retry is still subject to the task's ``retries`` count -- the policy can fail a task earlier but cannot extend past the configured maximum. - When all retries are exhausted, RETRY behaves identically to DEFAULT. + When all retries are exhausted, RETRY fails the task like DEFAULT does, + but still records the policy's reason; DEFAULT records none. """ FAIL = "fail" 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 74c9d30fc8b83..db2fd4fdbd825 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1887,19 +1887,20 @@ def _handle_current_task_failed( state=TaskInstanceState.FAILED, end_date=ti.end_date, rendered_map_index=ti.rendered_map_index, - retry_reason=decision.reason, + retry_reason=decision.reason[:500] if decision.reason is not None else None, ), TaskInstanceState.FAILED, ) if decision is not None and decision.action == RetryAction.RETRY: return _finalize_task_failure( - ti, retry_delay_override=decision.retry_delay, retry_reason=decision.reason + ti, log, retry_delay_override=decision.retry_delay, retry_reason=decision.reason ) - return _finalize_task_failure(ti) + return _finalize_task_failure(ti, log) def _finalize_task_failure( ti: RuntimeTaskInstance, + log: Logger, retry_delay_override: timedelta | None = None, retry_reason: str | None = None, ) -> tuple[RetryTask, TaskInstanceState] | tuple[TaskState, TaskInstanceState]: @@ -1933,10 +1934,12 @@ def _finalize_task_failure( 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})" - retry_reason = f"{retry_reason[: 500 - len(suffix)]}{suffix}" + max_tries = ti._ti_context_from_server.max_tries + if max_tries > 0: + # max_tries is the retry count, not the attempt count -- total attempts is max_tries + 1. + suffix = f"; retries exhausted ({ti.try_number} of {max_tries + 1})" + retry_reason = f"{retry_reason[: 500 - len(suffix)]}{suffix}" + log.info("Retry policy decision", action="fail", reason=retry_reason) return ( TaskState( state=TaskInstanceState.FAILED, diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index cd5f5fff5fb56..19790f81cb5ac 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -470,3 +470,37 @@ def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): assert isinstance(defaulted, LiteralArgBinding) assert defaulted.from_default is True assert defaulted.value_schema.root == {"type": "integer", "format": "int64"} + + +class TestRealBundleRetryReasonUpgrade: + """ + Drive the *real* supervisor bundle through the ``retry_reason`` migration. + + ``TaskState`` flows foreign-runtime -> supervisor, the opposite direction from + ``arg_bindings`` above, so a runtime pinned to an older schema is exercised + through ``upgrade`` rather than ``downgrade``. + """ + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): + from airflow.sdk.execution_time.comms import TaskState + + body = {"type": "TaskState", "state": "failed", "end_date": None, "rendered_map_index": None} + out = real_migrator.upgrade(body, TaskState, "2026-06-16") + assert out["retry_reason"] is None + + def test_upgrade_keeps_retry_reason_at_head(self, real_migrator): + from airflow.sdk.execution_time.comms import TaskState + + body = { + "type": "TaskState", + "state": "failed", + "end_date": None, + "rendered_map_index": None, + "retry_reason": "auth error, do not retry", + } + out = real_migrator.upgrade(body, TaskState, "2026-10-30") + assert out["retry_reason"] == "auth error, do not retry" 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 4c4f50b46d243..5dc5b26432df2 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 @@ -1203,11 +1203,12 @@ def execute(self, context): task = _AlwaysFails( task_id="fail_with_reason", + retries=2, retry_policy=ExceptionRetryPolicy( rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL, reason="do not retry")] ), ) - ti = create_runtime_ti(task=task, should_retry=True) + ti = create_runtime_ti(task=task) state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock()) @@ -1225,11 +1226,12 @@ def execute(self, context): task = _AlwaysFails( task_id="retry_exhausted", + retries=2, 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) + ti = create_runtime_ti(task=task, try_number=3) state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock()) @@ -1238,6 +1240,34 @@ def execute(self, context): assert msg.retry_reason == "rate limit; retries exhausted (3 of 3)" +def test_retry_policy_retry_exhausted_reason_is_truncated_with_suffix_kept( + create_runtime_ti, mock_supervisor_comms +): + """A long reason is truncated to 500 chars total, with the exhausted-suffix always kept.""" + + class _AlwaysFails(BaseOperator): + def execute(self, context): + raise RuntimeError("boom") + + long_reason = "z" * 600 + task = _AlwaysFails( + task_id="retry_exhausted_long_reason", + retries=2, + retry_policy=ExceptionRetryPolicy( + rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY, reason=long_reason)] + ), + ) + ti = create_runtime_ti(task=task, try_number=3) + + state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock()) + + assert state == TaskInstanceState.FAILED + assert isinstance(msg, TaskState) + assert msg.retry_reason is not None + assert len(msg.retry_reason) == 500 + assert msg.retry_reason.endswith("; 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.""" @@ -1245,8 +1275,8 @@ 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) + task = _AlwaysFails(task_id="plain_exhausted", retries=2) + ti = create_runtime_ti(task=task, try_number=3) state, msg, error = run(ti, ti.get_template_context(), mock.MagicMock()) From d9d3a00ac151923311f5a7296d7464d9df01e696 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 21 Sep 2026 12:46:33 +0530 Subject: [PATCH 4/8] after rebase comments --- providers/common/ai/docs/retry_policies.rst | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 8d1d5f14156b0..6d003e764e4c0 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -140,10 +140,12 @@ four fields: ``category``, ``should_retry``, ``suggested_delay_seconds``, and ``reasoning``. Only ``should_retry`` and ``suggested_delay_seconds`` affect the run. -``category`` and ``reasoning`` are only recorded on a RETRY. They are written +``category`` and ``reasoning`` are recorded on both outcomes. They are written to the task instance's ``retry_reason`` (truncated to 500 characters, see -below), then cleared once the next attempt starts running. On a FAIL they are -not written anywhere -- they only show up in the task log. +below). On a RETRY the value is cleared once the next attempt starts running. +A FAIL is terminal, so there is no next attempt to clear it and the reason +stays on the row. When the model asked to retry but no attempts were left, the +stored reason ends with a ``; retries exhausted (N of M)`` note. Two limits are worth knowing about: @@ -208,7 +210,7 @@ When writing custom instructions: fields explicitly so the model fills them. - Be concrete with examples (``"'Warehouse suspended' -> transient"``) rather than vague rules ("treat warehouse issues as recoverable"). -- ``retry_reason`` is truncated to 500 chars in the audit log -- keep +- ``retry_reason`` is truncated to 500 chars on the task instance -- keep ``reasoning`` outputs concise. Parameters From 524f00dce2bd4c7b89ab4ff78e1ebf475772a289 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Tue, 22 Sep 2026 15:35:13 +0530 Subject: [PATCH 5/8] review comments from other PR --- providers/common/ai/docs/retry_policies.rst | 9 +++-- .../airflow/sdk/execution_time/task_runner.py | 10 +++--- .../execution_time/schema/test_migrator.py | 35 +++++++++---------- .../execution_time/test_task_runner.py | 32 ++++++++++------- 4 files changed, 46 insertions(+), 40 deletions(-) diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 6d003e764e4c0..58c5cd93efd58 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -144,8 +144,13 @@ the run. to the task instance's ``retry_reason`` (truncated to 500 characters, see below). On a RETRY the value is cleared once the next attempt starts running. A FAIL is terminal, so there is no next attempt to clear it and the reason -stays on the row. When the model asked to retry but no attempts were left, the -stored reason ends with a ``; retries exhausted (N of M)`` note. +stays on the row. Only the model's own words are stored -- attempt counts are +left to whatever displays the reason. + +The recorded value is exposed by the REST API as ``state_reason`` on a task +instance and on each try, and the Task Instance page in the UI shows it under +**Reason for state**. Recording it on a FAIL requires Airflow 3.4.0; on earlier +versions only the RETRY outcome is recorded. Two limits are worth knowing about: 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 77ca8edd467c2..0284e73f952cd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1935,12 +1935,10 @@ 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 = ti._ti_context_from_server.max_tries - if max_tries > 0: - # max_tries is the retry count, not the attempt count -- total attempts is max_tries + 1. - suffix = f"; retries exhausted ({ti.try_number} of {max_tries + 1})" - retry_reason = f"{retry_reason[: 500 - len(suffix)]}{suffix}" + if retry_reason is not None: + # Policy's own words only: attempt counts belong to whoever renders this, which has + # try_number and max_tries alongside and need not guess when retries was never set. + retry_reason = retry_reason[:500] log.info("Retry policy decision", action="fail", reason=retry_reason) return ( TaskState( diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 19790f81cb5ac..5fcca148b2a7b 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -48,11 +48,13 @@ _SupervisorResponse, ) +from airflow.sdk.execution_time.comms import TaskState from airflow.sdk.execution_time.schema import ( SchemaVersionMigrator, get_schema_version_migrator, resolve_body_class, ) +from airflow.utils.state import TaskInstanceState class _MockBody(BaseModel): @@ -472,35 +474,30 @@ def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): assert defaulted.value_schema.root == {"type": "integer", "format": "int64"} -class TestRealBundleRetryReasonUpgrade: +class TestRealBundleRetryReason: """ Drive the *real* supervisor bundle through the ``retry_reason`` migration. - ``TaskState`` flows foreign-runtime -> supervisor, the opposite direction from - ``arg_bindings`` above, so a runtime pinned to an older schema is exercised - through ``upgrade`` rather than ``downgrade``. + ``TaskState`` flows foreign-runtime -> supervisor, so ``upgrade`` is the direction a + pinned runtime travels. Only ``downgrade`` re-validates against the versioned class, + so that is the direction that fails if ``AddRetryReasonToTaskState`` is dropped. """ @pytest.fixture def real_migrator(self) -> SchemaVersionMigrator: return get_schema_version_migrator() - def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): - from airflow.sdk.execution_time.comms import TaskState + def test_downgrade_strips_retry_reason_for_previous_version(self, real_migrator): + msg = TaskState(state=TaskInstanceState.FAILED, retry_reason="auth error, do not retry") + out = real_migrator.downgrade(msg, "2026-06-16").model_dump() + assert "retry_reason" not in out + + def test_downgrade_keeps_retry_reason_at_head(self, real_migrator): + msg = TaskState(state=TaskInstanceState.FAILED, retry_reason="auth error, do not retry") + out = real_migrator.downgrade(msg, "2026-10-30").model_dump() + assert out["retry_reason"] == "auth error, do not retry" + def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): body = {"type": "TaskState", "state": "failed", "end_date": None, "rendered_map_index": None} out = real_migrator.upgrade(body, TaskState, "2026-06-16") assert out["retry_reason"] is None - - def test_upgrade_keeps_retry_reason_at_head(self, real_migrator): - from airflow.sdk.execution_time.comms import TaskState - - body = { - "type": "TaskState", - "state": "failed", - "end_date": None, - "rendered_map_index": None, - "retry_reason": "auth error, do not retry", - } - out = real_migrator.upgrade(body, TaskState, "2026-10-30") - assert out["retry_reason"] == "auth error, do not retry" 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 bc96aeaf3cdd2..1612ed8cc2114 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 @@ -1217,33 +1217,41 @@ def execute(self, context): 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.""" +@pytest.mark.parametrize( + ("task_id", "retries", "try_number"), + [ + pytest.param("retry_exhausted", 2, 3, id="budget-exhausted"), + # `retries` defaults to 0, so this branch is reached on the very first attempt. + pytest.param("retry_no_budget", 0, 1, id="no-budget-configured"), + ], +) +def test_retry_policy_retry_without_budget_persists_policy_reason( + create_runtime_ti, mock_supervisor_comms, task_id, retries, try_number +): + """A policy-chosen RETRY that cannot run fails, recording the reason with no counts appended.""" class _AlwaysFails(BaseOperator): def execute(self, context): raise RuntimeError("boom") task = _AlwaysFails( - task_id="retry_exhausted", - retries=2, + task_id=task_id, + retries=retries, retry_policy=ExceptionRetryPolicy( rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY, reason="rate limit")] ), ) - ti = create_runtime_ti(task=task, try_number=3) + ti = create_runtime_ti(task=task, try_number=try_number) 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)" + assert msg.retry_reason == "rate limit" -def test_retry_policy_retry_exhausted_reason_is_truncated_with_suffix_kept( - create_runtime_ti, mock_supervisor_comms -): - """A long reason is truncated to 500 chars total, with the exhausted-suffix always kept.""" +def test_retry_policy_retry_exhausted_reason_is_truncated(create_runtime_ti, mock_supervisor_comms): + """A long reason is truncated to the column width.""" class _AlwaysFails(BaseOperator): def execute(self, context): @@ -1263,9 +1271,7 @@ def execute(self, context): assert state == TaskInstanceState.FAILED assert isinstance(msg, TaskState) - assert msg.retry_reason is not None - assert len(msg.retry_reason) == 500 - assert msg.retry_reason.endswith("; retries exhausted (3 of 3)") + assert msg.retry_reason == "z" * 500 def test_plain_retries_exhausted_has_no_reason(create_runtime_ti, mock_supervisor_comms): From d2a25d6dcf7092aa445d0447f1c4c0b6064d2711 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Wed, 23 Sep 2026 15:50:02 +0530 Subject: [PATCH 6/8] easier comments first --- providers/common/ai/docs/retry_policies.rst | 6 ++--- .../airflow/sdk/execution_time/task_runner.py | 7 ++++- .../execution_time/schema/test_migrator.py | 2 +- .../execution_time/test_task_runner.py | 27 +++++++++++++++++++ 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 58c5cd93efd58..6b0c3699d9c21 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -147,10 +147,8 @@ A FAIL is terminal, so there is no next attempt to clear it and the reason stays on the row. Only the model's own words are stored -- attempt counts are left to whatever displays the reason. -The recorded value is exposed by the REST API as ``state_reason`` on a task -instance and on each try, and the Task Instance page in the UI shows it under -**Reason for state**. Recording it on a FAIL requires Airflow 3.4.0; on earlier -versions only the RETRY outcome is recorded. +Recording it on a FAIL requires Airflow 3.4.0; on earlier versions only the +RETRY outcome is recorded. Two limits are worth knowing about: 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 0284e73f952cd..b8578ba67ba60 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1939,7 +1939,12 @@ def _finalize_task_failure( # Policy's own words only: attempt counts belong to whoever renders this, which has # try_number and max_tries alongside and need not guess when retries was never set. retry_reason = retry_reason[:500] - log.info("Retry policy decision", action="fail", reason=retry_reason) + log.info( + "Retry policy requested a retry but no attempts remain", + reason=retry_reason, + try_number=ti.try_number, + max_tries=ti._ti_context_from_server.max_tries if ti._ti_context_from_server else None, + ) return ( TaskState( state=TaskInstanceState.FAILED, diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 5fcca148b2a7b..3befcf14b421b 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -48,13 +48,13 @@ _SupervisorResponse, ) +from airflow.sdk import TaskInstanceState from airflow.sdk.execution_time.comms import TaskState from airflow.sdk.execution_time.schema import ( SchemaVersionMigrator, get_schema_version_migrator, resolve_body_class, ) -from airflow.utils.state import TaskInstanceState class _MockBody(BaseModel): 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 1612ed8cc2114..b49ecebd32452 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 @@ -1402,6 +1402,33 @@ def tracking_info(msg, *args, **kwargs): ] +def test_exhausted_logs_about_retry_policy_decision(create_runtime_ti, mock_supervisor_comms): + class _AlwaysFails(BaseOperator): + def execute(self, context): + raise RuntimeError("boom") + + task = _AlwaysFails( + task_id="retry_exhausted_logging", + retries=2, + retry_policy=ExceptionRetryPolicy( + rules=[RetryRule(exception=RuntimeError, action=RetryAction.RETRY, reason="rate limit")] + ), + ) + ti = create_runtime_ti(task=task, try_number=3) + log = mock.MagicMock(spec=["info", "debug", "warning", "error", "exception", "bind"]) + + run(ti, context=ti.get_template_context(), log=log) + + events = [call.args[0] for call in log.info.call_args_list if call.args] + assert events.count("Retry policy decision") == 1 + assert log.info.call_args_list[-1] == mock.call( + "Retry policy requested a retry but no attempts remain", + reason="rate limit", + try_number=3, + max_tries=2, + ) + + 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") From c38ed53000f8fa53bcac9f55cb6dffb341fd0b0c Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 24 Sep 2026 11:21:07 +0530 Subject: [PATCH 7/8] final set of comments --- providers/common/ai/docs/retry_policies.rst | 9 ++------- task-sdk/src/airflow/sdk/definitions/retry_policy.py | 6 ++++-- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index c0e4673b88f87..cd88962473f03 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -405,8 +405,8 @@ Under ``LLMRetryPolicy`` it answers four fields: ``category``, ``should_retry``, upper limit; zero or negative means no override, so the task's own ``retry_delay`` and backoff apply. -``category`` and ``reasoning`` become the ``retry_reason``, recorded on both -outcomes. On a RETRY the value is cleared once the next attempt starts running; +``category`` and ``reasoning`` become the ``retry_reason`` (truncated to 500 +characters), recorded on both outcomes. On a RETRY the value is cleared once the next attempt starts running; a FAIL is terminal, so there is no next attempt to clear it and the reason stays on the row. Only the model's own words are stored -- attempt counts are left to whatever displays the reason. Recording on a FAIL requires Airflow 3.4.0; on @@ -420,11 +420,6 @@ the confidence, the bar, the action). A model cannot return a category the policy does not recognize, and it cannot return a category paired with an action that contradicts it. -The ``retry_reason`` is only recorded on a RETRY. It is written to the task -instance (truncated to 500 characters), then cleared once the next attempt -starts running. On a FAIL it is not written anywhere -- it only shows up in the -task log. - RETRY cannot give a task more attempts than ``retries`` allows. FAIL ends the task straight away even when attempts were left, so a wrong classification into a failing category costs the task the retries it would otherwise have had; diff --git a/task-sdk/src/airflow/sdk/definitions/retry_policy.py b/task-sdk/src/airflow/sdk/definitions/retry_policy.py index f4952023b3a23..63894103bb6f1 100644 --- a/task-sdk/src/airflow/sdk/definitions/retry_policy.py +++ b/task-sdk/src/airflow/sdk/definitions/retry_policy.py @@ -379,7 +379,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 and logs it otherwise. + worker stores it as ``retry_reason`` on a RETRY or a FAIL, and logs it when no policy + decided. :param policies: The policies to consult, in order. At least one. """ @@ -428,5 +429,6 @@ def evaluate( 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. + # 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)})") From 151fb2ee51784a541464560e969bf1ab7fd33efc Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Thu, 24 Sep 2026 12:28:42 +0530 Subject: [PATCH 8/8] final set of comments --- .../src/airflow/providers/common/compat/_retry_policy.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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 3f2c63dc8a319..3ec6f776c3073 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 @@ -74,7 +74,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 and logs it otherwise. + worker stores it as ``retry_reason`` on a RETRY, and on a FAIL from Airflow 3.4; when no + policy decided it is only logged. :param policies: The policies to consult, in order. At least one. """ @@ -123,5 +124,6 @@ def evaluate( 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. + # 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)})")