This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 66b803d5efb Archive worker-reported end date and rendered map index on
task retry (#69248)
66b803d5efb is described below
commit 66b803d5efbbb743c357f2b4696f49a22aa662b8
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Mon Jul 6 19:50:54 2026 +0900
Archive worker-reported end date and rendered map index on task retry
(#69248)
* Archive worker-reported end date and rendered map index on task retry
Follow-up to #69235. The task_instance_history row for a retried try
stamped archive-time utcnow() as end_date instead of the end date the
worker reported, and missed the final rendered map index when the
mid-run update was suppressed (e.g. template errors during failure
handling). Snapshot both onto the TI before archiving, and let
record_ti() respect a pre-set end_date so the audit trail reflects
when the try actually ended.
* Add test for record_ti fallback end_date stamping
Cover the conditional branch where record_ti() archives a non-finished
TI with end_date=None, verifying it gets stamped with utcnow() and
duration is computed correctly.
* Clarify record_ti comment covers pre-set duration too
* Snapshot rendered_map_index in TIH when a retry explicitly clears it
---
.../execution_api/routes/task_instances.py | 23 ++++-----
.../src/airflow/models/taskinstancehistory.py | 7 ++-
.../versions/head/test_task_instances.py | 57 ++++++++++++++++++++--
.../tests/unit/models/test_taskinstance.py | 24 +++++++++
4 files changed, 93 insertions(+), 18 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 3b142a4e6b1..b813c4c8f7e 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
@@ -643,21 +643,22 @@ def _create_ti_state_update_query_and_update_state(
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):
+ retry_delay_override = ti_patch_payload.retry_delay_seconds
+ retry_reason = ti_patch_payload.retry_reason[:500] if
ti_patch_payload.retry_reason else None
if ti is not None:
- # Set the overrides on the TI *before* archiving so record_ti()
- # snapshots them into task_instance_history (it copies attrs
off
- # the ti object). Otherwise the per-try audit trail is always
NULL.
- ti.retry_delay_override = ti_patch_payload.retry_delay_seconds
- ti.retry_reason = (
- ti_patch_payload.retry_reason[:500] if
ti_patch_payload.retry_reason else None
- )
+ # Snapshot the finished try onto the TI *before* archiving so
record_ti()
+ # copies the values into task_instance_history (it reads attrs
off the
+ # ti object and cannot see the live-row UPDATE built below).
+ ti.retry_delay_override = retry_delay_override
+ ti.retry_reason = retry_reason
+ ti.end_date = ti_patch_payload.end_date
+ ti.set_duration()
+ if "rendered_map_index" in ti_patch_payload.model_fields_set:
+ ti._rendered_map_index =
ti_patch_payload.rendered_map_index
ti.prepare_db_for_next_try(session=session)
# Store retry policy overrides so next_retry_datetime() can read
them.
# These are cleared when the task enters RUNNING (ti_run).
- query = query.values(
- retry_delay_override=ti_patch_payload.retry_delay_seconds,
- retry_reason=(ti_patch_payload.retry_reason[:500] if
ti_patch_payload.retry_reason else None),
- )
+ query = query.values(retry_delay_override=retry_delay_override,
retry_reason=retry_reason)
elif isinstance(ti_patch_payload, TISuccessStatePayload):
if ti is not None:
TI.register_asset_changes_in_db(
diff --git a/airflow-core/src/airflow/models/taskinstancehistory.py
b/airflow-core/src/airflow/models/taskinstancehistory.py
index b0d55114bb6..47cfdb68ae5 100644
--- a/airflow-core/src/airflow/models/taskinstancehistory.py
+++ b/airflow-core/src/airflow/models/taskinstancehistory.py
@@ -211,8 +211,11 @@ class TaskInstanceHistory(Base):
ti_history_state = ti.state
if ti.state not in State.finished:
ti_history_state = TaskInstanceState.FAILED
- ti.end_date = timezone.utcnow()
- ti.set_duration()
+ # Callers that know when the try actually ended (e.g. the
Execution API
+ # retry path) pre-set end_date and duration; only stamp archive
time when unset.
+ if ti.end_date is None:
+ ti.end_date = timezone.utcnow()
+ ti.set_duration()
ti_history = TaskInstanceHistory(ti, state=ti_history_state)
session.add(ti_history)
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 4920fc001f9..da6ea74337d 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
@@ -1891,17 +1891,17 @@ class TestTIUpdateState:
def test_ti_update_state_retry_policy_overrides_persisted_in_history(
self, client, session, create_task_instance
):
- """Retry policy override + reason must be archived to
task_instance_history.
+ """The finished try's values must be archived to task_instance_history.
- record_ti() snapshots columns off the TI object, so the overrides must
be set on
- the TI before prepare_db_for_next_try() archives it. When they were
written only
- to the live-row UPDATE, the per-try audit trail in
task_instance_history was
- always NULL even though the live row was correct.
+ record_ti() snapshots columns off the TI object, so the overrides,
end_date,
+ and rendered_map_index must be set on the TI before
prepare_db_for_next_try()
+ archives it; the live-row UPDATE is not visible to it.
"""
ti = create_task_instance(
task_id="test_retry_policy_override_history",
state=State.RUNNING,
)
+ ti.start_date = DEFAULT_START_DATE
session.commit()
response = client.patch(
@@ -1909,6 +1909,7 @@ class TestTIUpdateState:
json={
"state": State.UP_FOR_RETRY,
"end_date": DEFAULT_END_DATE.isoformat(),
+ "rendered_map_index": DEFAULT_RENDERED_MAP_INDEX,
"retry_delay_seconds": 42.5,
"retry_reason": "Rate limit: backing off",
},
@@ -1925,6 +1926,52 @@ class TestTIUpdateState:
).one()
assert tih.retry_delay_override == 42.5
assert tih.retry_reason == "Rate limit: backing off"
+ assert tih.end_date == DEFAULT_END_DATE
+ assert tih.duration == (DEFAULT_END_DATE -
DEFAULT_START_DATE).total_seconds()
+ assert tih.rendered_map_index == DEFAULT_RENDERED_MAP_INDEX
+
+ def test_ti_update_state_retry_clears_rendered_map_index_in_history(
+ self, client, session, create_task_instance
+ ):
+ """An explicit ``rendered_map_index: null`` must clear the value
archived to history too.
+
+ The worker sends ``rendered_map_index`` on every retry, even when this
try never
+ (re-)computed it (e.g. it failed before rendering). That must null out
the archived
+ row along with the live one, not leave the history row snapshotting a
stale value
+ left over from an earlier try.
+ """
+ ti = create_task_instance(
+ task_id="test_retry_clears_rendered_map_index_history",
+ state=State.RUNNING,
+ )
+ ti.start_date = DEFAULT_START_DATE
+ ti._rendered_map_index = DEFAULT_RENDERED_MAP_INDEX
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": State.UP_FOR_RETRY,
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "rendered_map_index": None,
+ },
+ )
+
+ assert response.status_code == 204
+
+ ti = session.scalars(
+ select(TaskInstance).filter_by(task_id=ti.task_id,
run_id=ti.run_id, dag_id=ti.dag_id)
+ ).one()
+ assert ti.rendered_map_index is None
+
+ tih = session.scalars(
+ select(TaskInstanceHistory).where(
+ TaskInstanceHistory.dag_id == ti.dag_id,
+ TaskInstanceHistory.task_id == ti.task_id,
+ TaskInstanceHistory.run_id == ti.run_id,
+ )
+ ).one()
+ assert tih.rendered_map_index is None
def test_ti_update_state_retry_without_policy_overrides(self, client,
session, create_task_instance):
"""Without retry policy fields, the columns remain NULL."""
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py
b/airflow-core/tests/unit/models/test_taskinstance.py
index d45e367afb7..bb9f09421bd 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -2715,6 +2715,30 @@ class TestTaskInstance:
# the new try_id should be different from what's recorded in tih
assert tih[0].task_instance_id == try_id
+ def test_record_ti_stamps_end_date_when_unset_for_non_finished_state(self,
dag_maker, session):
+ """record_ti() must fill in end_date/duration when archiving a
non-finished TI with end_date=None."""
+ archive_time = pendulum.datetime(2024, 6, 15, 12, 0, 0, tz="UTC")
+ start = pendulum.datetime(2024, 6, 15, 11, 50, 0, tz="UTC")
+
+ with dag_maker(serialized=True):
+ EmptyOperator(task_id="test_record_ti_fallback")
+
+ dr = dag_maker.create_dagrun()
+ ti = dr.task_instances[0]
+ ti.state = TaskInstanceState.RUNNING
+ ti.start_date = start
+ ti.end_date = None
+ session.flush()
+
+ with time_machine.travel(archive_time, tick=False):
+ TaskInstanceHistory.record_ti(ti, session=session)
+ session.flush()
+
+ tih = session.scalars(select(TaskInstanceHistory)).one()
+ assert tih.state == str(TaskInstanceState.FAILED)
+ assert tih.end_date == archive_time
+ assert tih.duration == (archive_time - start).total_seconds()
+
@pytest.mark.parametrize(
("first_ti", "second_ti"),
[
