This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 5cc50c59997 [v3-4-test] Clear retry policy state when a TI is cleared
(#73889) (#75646)
5cc50c59997 is described below
commit 5cc50c59997ab0a72265920587be12c4975c3b52
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Oct 26 11:28:43 2026 +0000
[v3-4-test] Clear retry policy state when a TI is cleared (#73889) (#75646)
* Clear retry policy state when a TI is cleared
* fixing review from kaxil
* review comments from kaxil
(cherry picked from commit 0df550012b18a1476c797082894cf724b7599df3)
Co-authored-by: Amogh Desai <[email protected]>
---
.../core_api/datamodels/task_instance_history.py | 4 +-
.../core_api/datamodels/task_instances.py | 2 +-
.../api_fastapi/core_api/openapi/_private_ui.yaml | 6 +-
.../core_api/openapi/v2-rest-api-generated.yaml | 11 ++--
.../execution_api/routes/task_instances.py | 3 +-
airflow-core/src/airflow/models/dagrun.py | 5 ++
airflow-core/src/airflow/models/taskinstance.py | 16 ++++-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 4 +-
.../airflow/ui/openapi-gen/requests/types.gen.ts | 4 +-
airflow-core/tests/unit/models/test_dagrun.py | 41 ++++++++++++-
.../tests/unit/models/test_taskinstance.py | 68 ++++++++++++++++++++++
.../src/airflowctl/api/datamodels/generated.py | 7 +--
12 files changed, 144 insertions(+), 27 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instance_history.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instance_history.py
index 7c2a691e3d0..2d810885fe3 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instance_history.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instance_history.py
@@ -75,9 +75,7 @@ class TaskInstanceHistoryResponse(OmitsMissingRegion,
BaseModel):
state_reason: str | None = Field(
default=None,
validation_alias="retry_reason",
- description=(
- "The reason the task instance reached its current state, as
recorded by a retry policy. May describe a previous attempt: it is cleared only
when the task next starts running, so a task waiting to be retried or re-run
can still carry the reason its last attempt ended."
- ),
+ description=("The reason this try ended, as recorded by a retry
policy."),
)
@field_validator("state_reason", mode="after")
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
index a082a718511..a0c2d30d9c4 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
@@ -106,7 +106,7 @@ class TaskInstanceResponse(OmitsMissingRegion, BaseModel):
default=None,
validation_alias="retry_reason",
description=(
- "The reason the task instance reached its current state, as
recorded by a retry policy. May describe a previous attempt: it is cleared only
when the task next starts running, so a task waiting to be retried or re-run
can still carry the reason its last attempt ended."
+ "The reason the task instance reached its current state, as
recorded by a retry policy. May describe a previous attempt: a task waiting to
be retried still carries the reason its last attempt ended. It is cleared when
the task is cleared and when the next attempt is scheduled or deferred."
),
)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index 51320cabc1e..db96e77efb5 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -5990,9 +5990,9 @@ components:
- type: 'null'
title: State Reason
description: 'The reason the task instance reached its current
state, as
- recorded by a retry policy. May describe a previous attempt: it is
cleared
- only when the task next starts running, so a task waiting to be
retried
- or re-run can still carry the reason its last attempt ended.'
+ recorded by a retry policy. May describe a previous attempt: a
task waiting
+ to be retried still carries the reason its last attempt ended. It
is cleared
+ when the task is cleared and when the next attempt is scheduled or
deferred.'
type: object
required:
- id
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index aecd0bdeaf0..7c51d8b4b42 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -18016,10 +18016,7 @@ components:
- type: string
- type: 'null'
title: State Reason
- description: 'The reason the task instance reached its current
state, as
- recorded by a retry policy. May describe a previous attempt: it is
cleared
- only when the task next starts running, so a task waiting to be
retried
- or re-run can still carry the reason its last attempt ended.'
+ description: The reason this try ended, as recorded by a retry
policy.
type: object
required:
- id
@@ -18224,9 +18221,9 @@ components:
- type: 'null'
title: State Reason
description: 'The reason the task instance reached its current
state, as
- recorded by a retry policy. May describe a previous attempt: it is
cleared
- only when the task next starts running, so a task waiting to be
retried
- or re-run can still carry the reason its last attempt ended.'
+ recorded by a retry policy. May describe a previous attempt: a
task waiting
+ to be retried still carries the reason its last attempt ended. It
is cleared
+ when the task is cleared and when the next attempt is scheduled or
deferred.'
type: object
required:
- id
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 ff3ffbc021b..658c0e3055d 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
@@ -919,7 +919,8 @@ def _create_ti_state_update_query_and_update_state(
if "rendered_map_index" in ti_patch_payload.model_fields_set:
ti._rendered_map_index =
ti_patch_payload.rendered_map_index
# Store retry policy overrides so next_retry_datetime() can read
them.
- # These are cleared when the task enters RUNNING (ti_run).
+ # These are cleared on clear, when the next attempt is scheduled
or deferred, and
+ # when the task enters RUNNING (ti_run).
query = query.values(retry_delay_override=retry_delay_override,
retry_reason=retry_reason)
elif isinstance(ti_patch_payload, TISuccessStatePayload):
if ti is not None:
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index f28cd7690f4..206993e1b66 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -2717,6 +2717,9 @@ class DagRun(Base, LoggingMixin):
state=TaskInstanceState.SCHEDULED,
scheduled_dttm=timezone.utcnow(),
try_number=next_try_number,
+ # Already archived with the finished try; the new one
must not inherit it.
+ retry_reason=None,
+ retry_delay_override=None,
)
.execution_options(synchronize_session=False)
)
@@ -2767,6 +2770,8 @@ class DagRun(Base, LoggingMixin):
end_date=timezone.utcnow(),
duration=0,
try_number=next_try_number,
+ retry_reason=None,
+ retry_delay_override=None,
)
.execution_options(
synchronize_session=False,
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index eeb548e5881..ca06135a3f5 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -554,6 +554,10 @@ def clear_task_instances(
ti.max_tries = max(ti.max_tries or 0, previous_try_number)
ti.state = None
ti.external_executor_id = None
+ # retry_delay_override is the functional one:
next_retry_datetime() prefers it over
+ # the task's own retry_delay, so a stale value would retime the
next retry.
+ ti.retry_reason = None
+ ti.retry_delay_override = None
ti.clear_next_method_args()
# Match DagVersion to latest serialized DAG when running on the
latest version.
if use_latest_version:
@@ -1178,7 +1182,8 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
dag_version = relationship("DagVersion", back_populates="task_instances")
# Retry policy overrides: set by the task worker when a RetryPolicy is
configured.
- # Cleared on task start (ti_run). Read by next_retry_datetime().
+ # Cleared on clear, on scheduling or deferring the next attempt, and on
task start (ti_run).
+ # Read by next_retry_datetime().
retry_delay_override: Mapped[float | None] = mapped_column(Float,
nullable=True)
retry_reason: Mapped[str | None] = mapped_column(String(500),
nullable=True)
@@ -1608,8 +1613,14 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
self.log.debug("Setting task state for %s to %s", self, state)
if self not in session:
self.refresh_from_db(session=session)
+ previous_state = self.state
self.state = state
self.start_date = self.start_date or current_time
+ if previous_state == TaskInstanceState.UP_FOR_RETRY and state !=
TaskInstanceState.UP_FOR_RETRY:
+ # The pending attempt never runs, so it must not keep the previous
one's values:
+ # prepare_db_for_next_try() would otherwise archive them under
this try number.
+ self.retry_reason = None
+ self.retry_delay_override = None
if self.state in State.finished or self.state ==
TaskInstanceState.UP_FOR_RETRY:
self.end_date = self.end_date or current_time
self.duration = (self.end_date - self.start_date).total_seconds()
@@ -2660,6 +2671,9 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
# Keep an eye on the logic in
`check_and_change_state_before_execution()`
# depending on self.next_method semantics
self.state = TaskInstanceState.DEFERRED
+ # A start_from_trigger retry skips schedule_tis' UPDATE, so clear
here too.
+ self.retry_reason = None
+ self.retry_delay_override = None
self.trigger_id = trigger_row.id
self.next_method = start_trigger_args.next_method
self.next_kwargs =
stringify_encoding_keys(start_trigger_args.next_kwargs or {})
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index ccbdf81c024..3d64fafeab4 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -8199,7 +8199,7 @@ export const $TaskInstanceHistoryResponse = {
}
],
title: 'State Reason',
- description: 'The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: it is
cleared only when the task next starts running, so a task waiting to be retried
or re-run can still carry the reason its last attempt ended.'
+ description: 'The reason this try ended, as recorded by a retry
policy.'
}
},
type: 'object',
@@ -8534,7 +8534,7 @@ export const $TaskInstanceResponse = {
}
],
title: 'State Reason',
- description: 'The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: it is
cleared only when the task next starts running, so a task waiting to be retried
or re-run can still carry the reason its last attempt ended.'
+ description: 'The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: a task
waiting to be retried still carries the reason its last attempt ended. It is
cleared when the task is cleared and when the next attempt is scheduled or
deferred.'
}
},
type: 'object',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 0a90bf284f5..86a1680d2c5 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -2229,7 +2229,7 @@ export type TaskInstanceHistoryResponse = {
dag_version: DagVersionResponse | null;
note?: string | null;
/**
- * The reason the task instance reached its current state, as recorded by
a retry policy. May describe a previous attempt: it is cleared only when the
task next starts running, so a task waiting to be retried or re-run can still
carry the reason its last attempt ended.
+ * The reason this try ended, as recorded by a retry policy.
*/
state_reason?: string | null;
};
@@ -2280,7 +2280,7 @@ export type TaskInstanceResponse = {
dag_version: DagVersionResponse | null;
team_name?: string | null;
/**
- * The reason the task instance reached its current state, as recorded by
a retry policy. May describe a previous attempt: it is cleared only when the
task next starts running, so a task waiting to be retried or re-run can still
carry the reason its last attempt ended.
+ * The reason the task instance reached its current state, as recorded by
a retry policy. May describe a previous attempt: a task waiting to be retried
still carries the reason its last attempt ended. It is cleared when the task is
cleared and when the next attempt is scheduled or deferred.
*/
state_reason?: string | null;
};
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index ee63b2db379..cb39a14303b 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -3599,6 +3599,30 @@ def
test_schedule_tis_does_not_increment_try_number_if_ti_already_queued_by_othe
assert refreshed_ti.try_number == 1
+def test_schedule_tis_empty_operator_clears_retry_policy_columns(dag_maker,
session):
+ """The EmptyOperator UPDATE allocates an attempt too, so it must clear
like the other one."""
+ with dag_maker(session=session) as dag:
+ EmptyOperator(task_id="empty_task")
+
+ dr = dag_maker.create_dagrun(session=session)
+ ti = dr.get_task_instance("empty_task", session=session)
+ ti.refresh_from_task(dag.get_task("empty_task"))
+ ti.state = TaskInstanceState.UP_FOR_RETRY
+ ti.retry_reason = "rate limit"
+ ti.retry_delay_override = 300.0
+ session.commit()
+ ti_id = ti.id
+
+ dr.schedule_tis((ti,), session=session)
+ session.flush()
+ session.expire_all()
+
+ scheduled = session.get(TI, ti_id)
+ assert scheduled.state == TaskInstanceState.SUCCESS
+ assert scheduled.retry_reason is None
+ assert scheduled.retry_delay_override is None
+
+
def
test_schedule_tis_empty_operator_does_not_short_circuit_if_ti_already_queued(dag_maker,
session):
with dag_maker(session=session) as dag:
EmptyOperator(task_id="empty_task")
@@ -3653,6 +3677,8 @@ def
test_schedule_tis_preserves_allocated_attempt(dag_maker, session, state, try
ti.refresh_from_task(dag.get_task("task"))
ti.state = state
ti.try_number = try_number
+ ti.retry_reason = "rate limit"
+ ti.retry_delay_override = 300.0
session.commit()
ti_id = ti.id
@@ -3660,7 +3686,11 @@ def
test_schedule_tis_preserves_allocated_attempt(dag_maker, session, state, try
session.flush()
session.expire_all()
- assert session.get(TI, ti_id).try_number == try_number
+ scheduled = session.get(TI, ti_id)
+ assert scheduled.try_number == try_number
+ # The finished try keeps these in task_instance_history; the new one must
not inherit them.
+ assert scheduled.retry_reason is None
+ assert scheduled.retry_delay_override is None
def
test_schedule_tis_up_for_reschedule_does_not_increment_try_number(dag_maker,
session):
@@ -3844,7 +3874,8 @@ def
test_schedule_tis_only_one_scheduler_update_succeeds_when_competing(dag_make
@pytest.mark.need_serialized_dag
-def test_schedule_tis_start_trigger(dag_maker, session):
[email protected]("state", [None, TaskInstanceState.UP_FOR_RETRY])
+def test_schedule_tis_start_trigger(dag_maker, session, state):
"""
Test that an operator with start_trigger_args set can be directly deferred
during scheduling.
"""
@@ -3873,8 +3904,14 @@ def test_schedule_tis_start_trigger(dag_maker, session):
assert ti.state is None
ti.task = dr.dag.get_task("test_task")
+ ti.state = state
+ ti.retry_reason = "rate limit"
+ ti.retry_delay_override = 300.0
dr.schedule_tis((ti,), session=session)
assert ti.state == TaskInstanceState.DEFERRED
+ # This path defers before the UPDATE in schedule_tis, so defer_task() has
to clear.
+ assert ti.retry_reason is None
+ assert ti.retry_delay_override is None
@pytest.mark.need_serialized_dag
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index c02af00a0d5..286a04574db 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -5130,6 +5130,74 @@ def
test_clear_task_instances_resets_context_carrier(dag_maker, session):
assert dag_run.context_carrier["traceparent"] != original_dr_traceparent
[email protected]_test
[email protected](
+ "state", [TaskInstanceState.SUCCESS, TaskInstanceState.FAILED,
TaskInstanceState.SKIPPED]
+)
+def
test_set_state_clears_retry_policy_columns_when_leaving_up_for_retry(dag_maker,
session, state):
+ """A retry marked terminal by hand never runs, so it must not keep the
previous attempt's values."""
+ with dag_maker("test_set_state_retry_policy_columns"):
+ EmptyOperator(task_id="t1")
+ dag_run = dag_maker.create_dagrun()
+ ti = dag_run.get_task_instance("t1", session=session)
+ ti.state = TaskInstanceState.UP_FOR_RETRY
+ ti.retry_reason = "rate limit"
+ ti.retry_delay_override = 300.0
+ session.flush()
+
+ ti.set_state(state, session=session)
+
+ assert ti.retry_reason is None
+ assert ti.retry_delay_override is None
+
+
[email protected]_test
+def test_set_state_keeps_retry_policy_columns_while_up_for_retry(dag_maker,
session):
+ """next_retry_datetime() reads the override while the task waits, so a
re-set must not clear it."""
+ with dag_maker("test_set_state_keeps_retry_policy_columns"):
+ EmptyOperator(task_id="t1")
+ dag_run = dag_maker.create_dagrun()
+ ti = dag_run.get_task_instance("t1", session=session)
+ ti.state = TaskInstanceState.RUNNING
+ ti.retry_reason = "rate limit"
+ ti.retry_delay_override = 300.0
+ session.flush()
+
+ ti.set_state(TaskInstanceState.UP_FOR_RETRY, session=session)
+
+ assert ti.retry_reason == "rate limit"
+ assert ti.retry_delay_override == 300.0
+
+
[email protected]_test
+def test_clear_task_instances_clears_retry_policy_columns(dag_maker, session):
+ """Neither column may survive the clear."""
+ with dag_maker("test_clear_retry_policy_columns"):
+ EmptyOperator(task_id="t1")
+ dag_run = dag_maker.create_dagrun()
+ ti = dag_run.get_task_instance("t1", session=session)
+ ti.state = TaskInstanceState.FAILED
+ ti.retry_reason = "auth error, do not retry"
+ ti.retry_delay_override = 300.0
+ session.flush()
+ old_id = ti.id
+
+ clear_task_instances([ti], session)
+ session.flush()
+ session.expire_all()
+
+ # The archived try keeps its own values, so the reset cannot move above
the archive.
+ archived = session.get(TI, old_id)
+ assert archived.working_set is None
+ assert archived.retry_reason == "auth error, do not retry"
+ assert archived.retry_delay_override == 300.0
+
+ successor = dag_run.get_task_instance("t1", session=session)
+ assert successor.id != old_id
+ assert successor.retry_reason is None
+ assert successor.retry_delay_override is None
+
+
@pytest.mark.db_test
def test_clear_task_instances_preserves_detail_level(dag_maker, session):
"""clear_task_instances should produce a new context_carrier that keeps
the detail level from dag run conf."""
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
index a705d2eaf66..0693fe15742 100644
--- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
+++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -2732,10 +2732,7 @@ class TaskInstanceHistoryResponse(BaseModel):
note: Annotated[str | None, Field(title="Note")] = None
state_reason: Annotated[
str | None,
- Field(
- description="The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: it is
cleared only when the task next starts running, so a task waiting to be retried
or re-run can still carry the reason its last attempt ended.",
- title="State Reason",
- ),
+ Field(description="The reason this try ended, as recorded by a retry
policy.", title="State Reason"),
] = None
@@ -2786,7 +2783,7 @@ class TaskInstanceResponse(BaseModel):
state_reason: Annotated[
str | None,
Field(
- description="The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: it is
cleared only when the task next starts running, so a task waiting to be retried
or re-run can still carry the reason its last attempt ended.",
+ description="The reason the task instance reached its current
state, as recorded by a retry policy. May describe a previous attempt: a task
waiting to be retried still carries the reason its last attempt ended. It is
cleared when the task is cleared and when the next attempt is scheduled or
deferred.",
title="State Reason",
),
] = None
