diff --git a/airflow-core/docs/core-concepts/tasks.rst b/airflow-core/docs/core-concepts/tasks.rst index 15f110901bd21..2561d46b5b84b 100644 --- a/airflow-core/docs/core-concepts/tasks.rst +++ b/airflow-core/docs/core-concepts/tasks.rst @@ -296,8 +296,8 @@ A policy that raises an ordinary exception, or returns something other than a ``RetryDecision``, is logged and treated as DEFAULT, so one broken policy does not take the rules after it down with it. The winning decision's reason names the policy that decided and then what the earlier ones said (``HTTPStatusRetryPolicy: HTTP 503 (after ExceptionRetryPolicy: -no decision)``). On a RETRY that string is the task's ``retry_reason``; on FAIL, or when no -policy decided, it appears in the task log as the ``Retry policy decision`` line. +no decision)``). That string is the task's ``retry_reason`` on a FAIL as well as a RETRY; when +no policy decided, it appears only in the task log as the ``Retry policy decision`` line. Custom retry policies ~~~~~~~~~~~~~~~~~~~~~ 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 a0f0f5e35531f..af3415dbfd608 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 3dd4bc11a0bfd..65afb3c132e73 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 @@ -672,6 +672,8 @@ def _create_ti_state_update_query_and_update_state( 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: + 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/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index 79f3f1833fc12..e14c3f33e344c 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 @@ -56,6 +56,7 @@ AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, AddMultiTeamToTIRunContext, + AddTerminalStateRetryReasonField, ) bundle = VersionBundle( @@ -64,6 +65,7 @@ "2026-10-30", AddArgBindingsToTIRunContext, AddCallbackRunEndpoint, + AddTerminalStateRetryReasonField, AddMultiTeamToTIRunContext, ), Version( 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 6d5730a71ab5f..979c8a4476736 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): @@ -54,6 +54,16 @@ class AddCallbackRunEndpoint(VersionChange): ) +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, + ) + + class AddMultiTeamToTIRunContext(VersionChange): """Add ``multi_team`` so a worker can determine multi-team (e.g. for plugin scoping) without needing to trust its own config.""" 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 8ecf7b1b5dd17..2d21e035c4053 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 @@ -2490,6 +2490,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/providers/common/ai/docs/retry_policies.rst b/providers/common/ai/docs/retry_policies.rst index 912d7acf31b43..cd88962473f03 100644 --- a/providers/common/ai/docs/retry_policies.rst +++ b/providers/common/ai/docs/retry_policies.rst @@ -129,9 +129,9 @@ When a task fails, either policy: it is the picked category's ``retry`` and ``delay``, unless the policy has a confidence bar and the answer is under it, in which case the answer is discarded (see `Confidence`_ below). -4. The decision is logged in the task logs and, on a RETRY, written to the task - instance's ``retry_reason``: ``: `` from - ``LLMRetryPolicy``, or one line such as +4. The decision is logged in the task logs and written to the task instance's + ``retry_reason``, on a FAIL as well as a RETRY: ``: `` + from ``LLMRetryPolicy``, or one line such as ``category=network confidence=0.91 threshold=0.60 action=retry delay=10s`` from ``ClassifierRetryPolicy``. @@ -403,8 +403,14 @@ Under ``LLMRetryPolicy`` it answers four fields: ``category``, ``should_retry``, ``suggested_delay_seconds`` and ``reasoning``, and the first two after ``category`` decide the run. A positive delay is used as returned, with no 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``. +``retry_delay`` and backoff apply. + +``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 +earlier versions only the RETRY outcome is recorded. Under ``ClassifierRetryPolicy`` it answers the category name and nothing else. It does not decide whether to retry, it does not choose the delay, and it does not explain @@ -414,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/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)})") 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 a435be7e16519..032bddf3b77b1 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/definitions/retry_policy.py b/task-sdk/src/airflow/sdk/definitions/retry_policy.py index 898047bb0494a..63894103bb6f1 100644 --- a/task-sdk/src/airflow/sdk/definitions/retry_policy.py +++ b/task-sdk/src/airflow/sdk/definitions/retry_policy.py @@ -56,7 +56,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" @@ -378,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. """ @@ -427,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)})") 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 7ae8b5db6e6d3..4072443e52c0c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -4213,6 +4213,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/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/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index d722f7d5a55c0..42d1ddf8fe5b9 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1760,6 +1760,7 @@ def _send_terminal_state_msg( state=msg.state, when=msg.end_date or datetime.now(tz=timezone.utc), rendered_map_index=self._rendered_map_index, + retry_reason=msg.retry_reason, ) elif isinstance(msg, SucceedTask): self.client.task_instances.succeed( 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 cacd6ae87ab31..46f2fdae41128 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1902,18 +1902,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[: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]: @@ -1946,9 +1948,22 @@ 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: + # 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 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, 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..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,6 +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"] == retry_reason return httpx.Response( status_code=204, ) @@ -439,7 +441,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=retry_reason, ) def test_task_instance_heartbeat(self): 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..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,6 +48,8 @@ _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, @@ -470,3 +472,32 @@ 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 TestRealBundleRetryReason: + """ + Drive the *real* supervisor bundle through the ``retry_reason`` migration. + + ``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_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 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 a55cef8aa5834..2767e9f6f06ed 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -3933,6 +3933,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", + ) + @pytest.mark.parametrize( "test_name", 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 faceabb5c72e9..198b6ec070b9c 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,101 @@ 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", + retries=2, + retry_policy=ExceptionRetryPolicy( + rules=[RetryRule(exception=RuntimeError, action=RetryAction.FAIL, reason="do not retry")] + ), + ) + ti = create_runtime_ti(task=task) + + 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" + + +@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=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=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" + + +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): + 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 == "z" * 500 + + +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", retries=2) + 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 None + + def test_run_downstream_skipped(mocked_parse, create_runtime_ti, mock_supervisor_comms, listener_manager): listener = TestTaskRunnerCallsListeners.CustomListener() listener_manager(listener) @@ -1307,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")