You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2026/04/26 16:58:13 UTC
(airflow) 04/04: [v3-2-test] Fix ti.start_date showing deferral-resume time instead of original start time (#63247)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch backport-a90202d-v3-2-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 7131cb2f81b7fbc54fbb11572a9849ae7b74db4e
Author: Henry Chen <he...@gmail.com>
AuthorDate: Sun Apr 19 21:29:51 2026 +0800
[v3-2-test] Fix ti.start_date showing deferral-resume time instead of original start time (#63247)
* Fix ti.start_date showing deferral-resume time instead of original start time
* Add API migration
* Add unit test
* add unit test to verify the resume start date behavior
(cherry picked from commit a90202d)
Co-authored-by: Henry Chen <he...@gmail.com>
---
.../execution_api/datamodels/taskinstance.py | 9 ++++
.../execution_api/routes/task_instances.py | 3 +-
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_04_06.py | 13 +++++
.../versions/head/test_task_instances.py | 61 +++++++++++++++++++++-
.../src/airflow/sdk/api/datamodels/_generated.py | 1 +
.../src/airflow/sdk/execution_time/supervisor.py | 9 +++-
.../task_sdk/execution_time/test_supervisor.py | 61 +++++++++++++++++++++-
8 files changed, 153 insertions(+), 6 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 11b6b605f0b..395f8281d69 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
@@ -384,6 +384,15 @@ class TIRunContext(BaseModel):
should_retry: bool = False
"""If the ti encounters an error, whether it should enter retry or failed state."""
+ start_date: UtcDateTime | None = None
+ """
+ The original start date of the task instance.
+
+ When resuming from deferral, this is set to the task's original ``start_date`` so the
+ supervisor uses it instead of ``datetime.now()``. This ensures ``context["ti"].start_date``
+ always reflects when the task *first* started, not when it was rescheduled/resumed.
+ """
+
class PrevSuccessfulDagRunResponse(BaseModel):
"""Schema for response with previous successful DagRun information for Task Template Context."""
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 674186a7c9b..8f447cab144 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
@@ -136,6 +136,7 @@ def ti_run(
TI.map_index,
TI.try_number,
TI.max_tries,
+ TI.start_date,
TI.next_method,
TI.hostname,
TI.unixname,
@@ -287,7 +288,7 @@ def ti_run(
if ti.next_method:
context.next_method = ti.next_method
context.next_kwargs = ti.next_kwargs
-
+ context.start_date = ti.start_date
return context
except SQLAlchemyError:
log.exception("Error marking Task Instance state as running")
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 646831501a7..9276e537b5e 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
@@ -33,6 +33,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_04_06 import (
AddDagRunDetailEndpoint,
AddNoteField,
AddPartitionKeyField,
+ AddTaskInstanceStartDateField,
MakeDagRunStartDateNullable,
ModifyDeferredTaskKwargsToJsonValue,
MovePreviousRunEndpoint,
@@ -50,6 +51,7 @@ bundle = VersionBundle(
ModifyDeferredTaskKwargsToJsonValue,
RemoveUpstreamMapIndexesField,
AddNoteField,
+ AddTaskInstanceStartDateField,
AddDagEndpoint,
),
Version("2025-11-05", AddTriggeringUserNameField),
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
index 59e671f0a24..be9e167e654 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
@@ -178,6 +178,19 @@ class AddNoteField(VersionChange):
response.body["dag_run"].pop("note", None)
+class AddTaskInstanceStartDateField(VersionChange):
+ """Add `start_date` field to TIRunContext."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("start_date").didnt_exist,)
+
+ @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type]
+ def remove_start_date_field(response: ResponseInfo) -> None: # type: ignore[misc]
+ """Remove start_date field for older API versions."""
+ response.body.pop("start_date", None)
+
+
class AddDagEndpoint(VersionChange):
"""Add the `/dags/{dag_id}` endpoint."""
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 f7e50b2d8c9..baeb2cc04f0 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
@@ -599,6 +599,7 @@ class TestTIRunState:
"xcom_keys_to_clear": [],
"next_method": "execute_complete",
"next_kwargs": expected_next_kwargs,
+ "start_date": None,
}
@pytest.mark.parametrize("resume", [True, False])
@@ -660,7 +661,10 @@ class TestTIRunState:
session.commit()
assert response.status_code == 200
- assert response.json() == {
+ result = response.json()
+ assert timezone.parse(result["start_date"]) == orig_task_start_time
+ result.pop("start_date")
+ assert result == {
"dag_run": mock.ANY,
"task_reschedule_count": 0,
"max_tries": 0,
@@ -675,6 +679,61 @@ class TestTIRunState:
ti = session.get(TaskInstance, ti.id)
assert ti.start_date == expected_start_date
+ def test_ti_run_resume_returns_original_start_date_in_context(
+ self,
+ client,
+ session,
+ create_task_instance,
+ ):
+ original_start_date = timezone.parse("2024-09-30T12:00:05Z")
+ payload_start_date = timezone.parse("2024-09-30T12:00:35Z")
+
+ ti = create_task_instance(
+ task_id="test_ti_run_resume_returns_original_start_date_in_context",
+ state=State.QUEUED,
+ session=session,
+ start_date=original_start_date,
+ dag_id=str(uuid4()),
+ )
+ ti.start_date = original_start_date
+ ti.next_method = "execute_complete"
+ ti.next_kwargs = {
+ "moment": {
+ "__classname__": "pendulum.datetime.DateTime",
+ "__version__": 2,
+ "__data__": {
+ "timestamp": 1727697605.0,
+ "tz": {
+ "__classname__": "builtins.tuple",
+ "__version__": 1,
+ "__data__": ["UTC", "pendulum.tz.timezone.Timezone", 1, True],
+ },
+ },
+ }
+ }
+ session.commit()
+
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/run",
+ json={
+ "state": "running",
+ "hostname": "random-hostname",
+ "unixname": "random-unixname",
+ "pid": 100,
+ "start_date": payload_start_date.isoformat(),
+ },
+ )
+
+ assert response.status_code == 200
+ result = response.json()
+ assert result["start_date"] is not None
+ assert timezone.parse(result["start_date"]) == original_start_date
+
+ session.expunge_all()
+ ti = session.get(TaskInstance, ti.id)
+ assert ti is not None
+ assert ti.start_date == original_start_date
+
@pytest.mark.parametrize(
"initial_ti_state",
[s for s in TaskInstanceState if s not in (TaskInstanceState.QUEUED, TaskInstanceState.RESTARTING)],
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index e08f00562d3..e61c228cd20 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -662,3 +662,4 @@ class TIRunContext(BaseModel):
next_kwargs: Annotated[dict[str, Any] | str | None, Field(title="Next Kwargs")] = None
xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None
should_retry: Annotated[bool | None, Field(title="Should Retry")] = False
+ start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index a4484199f62..1ff74c521d1 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1027,12 +1027,11 @@ class ActivitySubprocess(WatchedSubprocess):
) -> None:
"""Send startup message to the subprocess."""
self.ti = ti # type: ignore[assignment]
- start_date = datetime.now(tz=timezone.utc)
try:
# We've forked, but the task won't start doing anything until we send it the StartupDetails
# message. But before we do that, we need to tell the server it's started (so it has the chance to
# tell us "no, stop!" for any reason)
- ti_context = self.client.task_instances.start(ti.id, self.pid, start_date)
+ ti_context = self.client.task_instances.start(ti.id, self.pid, datetime.now(tz=timezone.utc))
self._should_retry = ti_context.should_retry
self._last_successful_heartbeat = time.monotonic()
except Exception:
@@ -1040,6 +1039,12 @@ class ActivitySubprocess(WatchedSubprocess):
self.kill(signal.SIGKILL)
raise
+ # ti_context.start_date is only populated by the server when resuming from a deferral (to preserve the
+ # original start_date rather than using the resume time). We fall back to now() otherwise. This ensures
+ # that `context["ti"].start_date` always reflects the *first* start time. See TIRunContext.start_date
+ # for more context. Do not remove this without updating related comments and deferral handling.
+ start_date = ti_context.start_date or datetime.now(tz=timezone.utc)
+
msg = StartupDetails.model_construct(
ti=ti,
dag_rel_path=os.fspath(dag_rel_path),
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 b486ce77766..7c62fe219d5 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -516,19 +516,29 @@ class TestWatchedSubprocess:
assert rc == -9
- def test_last_chance_exception_handling(self, capfd):
+ @pytest.mark.parametrize(
+ "start_date",
+ [
+ pytest.param(None, id="start_date_is_none"),
+ pytest.param(timezone.datetime(2025, 3, 28, tzinfo=timezone.utc), id="start_date_from_context"),
+ ],
+ )
+ def test_last_chance_exception_handling(self, capfd, start_date, make_ti_context):
def subprocess_main():
# The real main() in task_runner catches exceptions! This is what would happen if we had a syntax
# or import error for instance - a very early exception
raise RuntimeError("Fake syntax error")
+ mock_client = MagicMock(spec=sdk_client.Client)
+ mock_client.task_instances.start.return_value = make_ti_context(start_date=start_date)
+
proc = ActivitySubprocess.start(
dag_rel_path=os.devnull,
bundle_info=FAKE_BUNDLE,
what=TaskInstance(
id=uuid7(), task_id="b", dag_id="c", run_id="d", try_number=1, dag_version_id=uuid7()
),
- client=MagicMock(spec=sdk_client.Client),
+ client=mock_client,
target=subprocess_main,
)
@@ -540,6 +550,53 @@ class TestWatchedSubprocess:
assert "Last chance exception handler" in captured.err
assert "RuntimeError: Fake syntax error" in captured.err
+ @pytest.mark.parametrize(
+ "start_date",
+ [
+ pytest.param(timezone.datetime(2025, 3, 28, tzinfo=timezone.utc), id="start_date_from_context"),
+ pytest.param(None, id="start_date_missing_in_context"),
+ ],
+ )
+ def test_resume_start_date_from_context(self, mocker, make_ti_context, start_date, time_machine):
+ """Verify that start_date from ti_context (e.g. for deferral resume) is used in the
+ StartupDetails message instead of computing a new datetime.now(). This test is kept
+ separate from test_last_chance_exception_handling as their purposes do not overlap.
+ """
+ fallback_now = timezone.datetime(2026, 1, 1, tzinfo=timezone.utc)
+ time_machine.move_to(fallback_now, tick=False)
+
+ mock_client = MagicMock(spec=sdk_client.Client)
+ ti_context = make_ti_context()
+ ti_context.start_date = start_date
+ mock_client.task_instances.start.return_value = ti_context
+
+ mock_send = mocker.patch.object(ActivitySubprocess, "send_msg", autospec=True)
+ proc = ActivitySubprocess.start(
+ dag_rel_path=os.devnull,
+ bundle_info=FAKE_BUNDLE,
+ what=TaskInstance(
+ id=uuid7(),
+ task_id="b",
+ dag_id="c",
+ run_id="d",
+ try_number=1,
+ dag_version_id=uuid7(),
+ ),
+ client=mock_client,
+ target=lambda: None,
+ )
+ rc = proc.wait()
+ assert rc == 0
+
+ # The startup message is sent via send_msg(StartupDetails(...), request_id=0)
+ startup_calls = [
+ c for c in mock_send.call_args_list if len(c.args) > 1 and hasattr(c.args[1], "start_date")
+ ]
+ assert len(startup_calls) >= 1
+ msg = startup_calls[0].args[1]
+ expected_start_date = start_date or fallback_now
+ assert msg.start_date == expected_start_date
+
def test_regular_heartbeat(self, spy_agency: kgb.SpyAgency, monkeypatch, mocker, make_ti_context):
"""Test that the WatchedSubprocess class regularly sends heartbeat requests, up to a certain frequency"""
import airflow.sdk.execution_time.supervisor