kaxil commented on code in PR #73554:
URL: https://github.com/apache/airflow/pull/73554#discussion_r4074565103
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -418,6 +416,8 @@ def clear_task_instances(
# set its state to RESTARTING so that
# the task is terminated and becomes eligible for retry.
else:
+ previous_try_number = ti.try_number
+ ti.prepare_db_for_next_try(session)
Review Comment:
Now that the retry path already allocates (B, N+1) when the failure is
reported, this branch also runs for an `up_for_retry` row, and `record_ti` has
no dedup hit for try N+1. The never-started attempt gets archived as a FAILED
history row (it copies attempt N's `start_date`, `hostname` and `pid` under try
N+1) and the rerun becomes N+2. On main the second archive of try N returned
early, so fail once then clear gave history `[1]` and a rerun at try 2; here it
gives `[1, 2]` with attempt 2 having no logs and a rerun at try 3. Same on a
second clear of a `None` row with `try_number > 0`, and in the restore branch
in `dagrun.py`. `fetch_handle_failure_context` already treats `UP_FOR_RETRY` as
already allocated; should clear do the same for `up_for_retry` / `None` rows
(reset state and `max_tries`, no archive, no increment)?
`test_task_instance_history_record[UP_FOR_RETRY]` currently pins the extra row,
and a fail-then-clear test asserting a single history row would catch this.
##########
airflow-core/src/airflow/ui/src/pages/TaskInstance/Logs/Logs.tsx:
##########
@@ -66,16 +66,19 @@ export const Logs = () => {
},
);
+ const defaultTryNumber =
+ taskInstance?.state === "up_for_retry" ? taskInstance.try_number - 1 :
taskInstance?.try_number;
Review Comment:
A cleared task now sits at `(try N+1, state null)` until the scheduler picks
it up, which can be a while when it is waiting on upstream tasks cleared
alongside it. With only `up_for_retry` special-cased, the Logs tab defaults to
the unrun attempt and shows nothing, where main showed the last attempt's logs.
Should `null` with `try_number > 1` fall back to `try_number - 1` too?
`Details.tsx` and `HITLResponse.tsx` still use the old formula, so during
`up_for_retry` they default to try N+1 (Details passes a `selectedTryNumber`
the try selector never lists) while Logs picks N, and the shared `?try_number`
param flips between tabs.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -892,7 +892,11 @@ def handle_request(request: httpx.Request) ->
httpx.Response:
if request.url.path == f"/task-instances/{ti_id}/run":
return httpx.Response(200, json=make_ti_context_dict())
if request.url.path == f"/task-instances/{ti_id}/state":
- pytest.fail("Should not have sent a state update request")
+ assert proc._process.wait(timeout=0) == -signal.SIGTERM
+ payload = json.loads(request.content)
+ assert payload["state"] == "server_terminated"
+ assert payload["pid"] == proc.pid
Review Comment:
Worth asserting `payload["hostname"] == get_hostname()` here and in the
409-restarting test below: the server answers 409 `running_elsewhere` whenever
`hostname` is missing, so dropping the `get_hostname()` line in `client.finish`
would leave every restart stuck in RESTARTING with no failing test.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py:
##########
@@ -28,6 +33,167 @@
TIMESTAMP_STR = "2024-09-30T12:00:00Z"
+
[email protected](
+ ("version", "expected_status"), [("2025-04-11", 404), ("2026-06-30", 404),
("2026-10-30", 410)]
+)
+def test_retired_state_report_response_by_version(
+ client, session, create_task_instance, version, expected_status
+):
+ ti = create_task_instance(task_id="retired_state_report",
state=State.RUNNING)
+ ti.try_number = 1
+ old_id = ti.id
+ session.commit()
+ client.headers["Airflow-API-Version"] = version
+ payload = {"state": "up_for_retry", "end_date": TIMESTAMP_STR}
+ assert client.patch(f"/execution/task-instances/{old_id}/state",
json=payload).status_code == 204
+
+ response = client.patch(f"/execution/task-instances/{old_id}/state",
json=payload)
+
+ assert response.status_code == expected_status
+ session.expunge_all()
+ replacement =
session.scalar(select(TaskInstance).where(TaskInstance.task_id ==
"retired_state_report"))
+ assert replacement.id != old_id
+ assert (replacement.try_number, replacement.state) == (2,
State.UP_FOR_RETRY)
+
+
[email protected]("version", ["2025-04-11", "2026-06-30"])
+def test_old_api_omits_stopped_report_from_schema(client, version):
+ response = client.get(f"/execution/openapi.json?version={version}")
+ assert response.status_code == 200
+ schemas = response.json()["components"]["schemas"]
+ assert "server_terminated" not in
schemas["TerminalStateNonSuccess"]["enum"]
+ assert "hostname" not in schemas["TITerminalStatePayload"]["properties"]
+ assert "pid" not in schemas["TITerminalStatePayload"]["properties"]
+
+
[email protected]("version", ["2025-04-11", "2026-06-30"])
+def test_old_api_rejects_stopped_report(client, session, create_task_instance,
version):
+ ti = create_task_instance(task_id="legacy_stopped_report",
state=State.RESTARTING)
+ ti.hostname = "worker"
+ ti.pid = 123
+ old_id = ti.id
+ session.commit()
+ client.headers["Airflow-API-Version"] = version
+
+ response = client.patch(
+ f"/execution/task-instances/{old_id}/state",
+ json={"state": "server_terminated", "end_date": TIMESTAMP_STR,
"hostname": "worker", "pid": 123},
+ )
+
+ assert response.status_code == 422
+ session.refresh(ti)
+ assert (ti.id, ti.state) == (old_id, State.RESTARTING)
+
+
[email protected]("version", ["2025-04-11", "2026-06-30"])
[email protected]("extra_field", [None, "hostname", "pid"])
+def test_legacy_worker_finish_payload(client, session, create_task_instance,
version, extra_field):
+ ti = create_task_instance(task_id="legacy_finish", state=State.RUNNING)
+ session.commit()
+ client.headers["Airflow-API-Version"] = version
+ payload = {"state": "failed", "end_date": TIMESTAMP_STR}
+ if extra_field is not None:
+ payload[extra_field] = {"hostname": "worker", "pid": 123}[extra_field]
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json=payload,
+ )
+ assert response.status_code == (204 if extra_field is None else 422)
+ if extra_field is not None:
+ assert any(
+ error["type"] == "extra_forbidden" and error["loc"][-1] ==
extra_field
+ for error in response.json()["detail"]
+ )
+ session.refresh(ti)
+ assert ti.state == (State.FAILED if extra_field is None else State.RUNNING)
+
+
[email protected]("version", ["2025-04-11", "2026-06-30"])
[email protected](
+ "payload",
+ [
+ pytest.param(
+ {
+ "state": "success",
+ "end_date": TIMESTAMP_STR,
+ "task_outlets": [],
+ "outlet_events": [],
+ },
+ id="success",
+ ),
+ pytest.param(
+ {
+ "state": "deferred",
+ "classpath": "my.trigger",
+ "trigger_kwargs": {"__type": "dict", "__var": {"key":
"value"}},
+ "trigger_timeout": None,
+ "next_method": "execute_complete",
+ "next_kwargs": {"__type": "dict", "__var": {"argument":
"value"}},
+ },
+ id="deferred",
+ ),
+ pytest.param(
+ {
+ "state": "up_for_reschedule",
+ "end_date": TIMESTAMP_STR,
+ "reschedule_date": "2024-09-30T12:05:00Z",
+ },
+ id="reschedule",
+ ),
+ ],
+)
+def test_legacy_worker_completion_payload_preserves_attempt(
+ client, session, create_task_instance, time_machine, version, payload
+):
+ time_machine.move_to("2024-09-30T11:59:00Z", tick=False)
+ ti = create_task_instance(task_id="legacy_completion", state=State.RUNNING)
+ ti.start_date = timezone.parse("2024-09-30T11:59:00Z")
+ old_identity = (ti.id, ti.try_number)
+ session.commit()
+ client.headers["Airflow-API-Version"] = version
+
+ response = client.patch(f"/execution/task-instances/{ti.id}/state",
json=payload)
+
+ assert response.status_code == 204
+ session.refresh(ti)
+ assert (ti.id, ti.try_number) == old_identity
+ assert ti.state == payload["state"]
+ if ti.state == State.DEFERRED:
+ trigger = session.get(Trigger, ti.trigger_id)
+ assert trigger.classpath == payload["classpath"]
+ assert trigger.kwargs == {"key": "value"}
+ assert ti.next_method == payload["next_method"]
+ assert ti.next_kwargs == payload["next_kwargs"]
+ assert ti.trigger_timeout is None
+ else:
+ assert ti.end_date == timezone.parse(TIMESTAMP_STR)
+ assert ti.duration == 60
+ if ti.state == State.UP_FOR_RESCHEDULE:
+ reschedule =
session.scalars(select(TaskReschedule).where(TaskReschedule.ti_id ==
ti.id)).one()
+ assert reschedule.start_date == ti.start_date
+ assert reschedule.end_date == ti.end_date
+ assert reschedule.reschedule_date ==
timezone.parse(payload["reschedule_date"])
+ assert reschedule.duration == 60
+
+
+def test_legacy_worker_clear_waits_for_executor(client, session,
create_task_instance):
Review Comment:
This is the MySQL CI failure (`got Future <Future pending> attached to a
different loop`). It is the first async-route call in this module and there is
no `reconfigure_async_db_engine` fixture here, unlike the class in the head
module. It also only exercises `ti_heartbeat`, which this PR does not touch, so
the RESTARTING 409 is main's behaviour. Adding `State.RESTARTING` to the
`test_ti_heartbeat_when_task_not_running` parametrisation under the existing
fixture would cover the same thing without the loop issue.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -379,9 +391,41 @@ def ti_update_state(
Not all state transitions are valid, and transitioning to some states
requires extra information to be
passed along. (Check out the datamodels for details, the rendered docs
might not reflect this accurately)
"""
+ # The version bundle imports the routes while registering its changes.
+ from airflow.api_fastapi.execution_api.versions.v2026_10_30 import
IdentifyRetiredTaskStateUpdates
Review Comment:
I don't think this needs to be in the function body. `v2025_08_10` imports
`routes.xcoms`, but that resolves as a submodule import against the partially
initialised package, and hoisting this to the top imports cleanly in both
orders and through `create_task_execution_api_app()` (`v2026_10_30` itself only
imports cadwyn and the datamodels). `services/task_instances.py:42` has the
same pattern, so if there is a cycle I'm missing it would be good to name it in
the comment.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py:
##########
@@ -1274,6 +1301,68 @@ def test_ti_run_creates_audit_log(self, client, session,
create_task_instance, t
class TestTIUpdateState:
+ @pytest.mark.parametrize("state", [State.SUCCESS, State.FAILED,
State.RUNNING])
+ def test_stopped_report_preserves_server_state(self, client, session,
create_task_instance, state):
+ ti = create_task_instance(task_id="stopped_preserves_state",
state=state)
+ ti.hostname = "worker"
+ ti.pid = 123
+ session.commit()
+ response = client.patch(
+ f"/execution/task-instances/{ti.id}/state",
+ json={
+ "state": "server_terminated",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "hostname": "worker",
+ "pid": 123,
+ },
+ )
+ assert response.status_code == (409 if state == State.RUNNING else 204)
+ session.refresh(ti)
+ assert ti.state == state
+
+ @pytest.mark.parametrize("matching_worker", [True, False])
+ def test_stopped_worker_completes_restart_once(
+ self, client, session, create_task_instance, matching_worker
+ ):
+ ti = create_task_instance(
+ task_id="stopped_restart", state=State.RESTARTING,
start_date=DEFAULT_START_DATE
+ )
+ ti.hostname = "original"
+ ti.pid = 123
+ ti.try_number = 3
+ ti.max_tries = 0
+ old_id = ti.id
+ session.commit()
+ payload = {
+ "state": "server_terminated",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "hostname": "original" if matching_worker else "duplicate",
+ "pid": 123,
+ }
+
+ response = client.patch(f"/execution/task-instances/{old_id}/state",
json=payload)
+
+ assert response.status_code == (204 if matching_worker else 409)
+ session.expunge_all()
+ current =
session.scalar(select(TaskInstance).where(TaskInstance.task_id ==
"stopped_restart"))
+ if not matching_worker:
+ assert (current.id, current.try_number, current.state) == (old_id,
3, State.RESTARTING)
+ return
+ assert current.id != old_id
+ assert current.try_number == 4
Review Comment:
The task here has no retries, so `complete_restart` sets `max_tries` to
exactly 3 and `>= 3` would still pass if the budget restore were dropped (3 is
also `try_number`). Giving the task `retries=2` and asserting `== 5` pins the
restore.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1580,19 +1583,21 @@ def wait(self) -> int:
# Now at the last possible moment, when all logs and comms with
the subprocess has finished,
# lets upload the remote logs. Run this in a `finally` so the logs
are uploaded even if the
# state update above raised — a failed state update is exactly
when the logs matter most.
+ self._wait_completed = True
self._upload_logs()
return self._exit_code
def update_task_state_if_needed(self):
- # If a direct-state API call (succeed / retry / defer / reschedule)
- # was attempted but raised, `_pending_terminal_state_msg` still holds
- # the original request. Re-issue the matching dedicated API call so
- # the server learns the terminal state we couldn't deliver earlier.
- # Without this recovery, a transient API failure during the direct
- # call would leave the TI stuck RUNNING on the server — `finish()`
- # cannot substitute because the server-side `finish` endpoint does
- # not accept SUCCESS / DEFERRED / SERVER_TERMINATED transitions.
+ if self.final_state == SERVER_TERMINATED:
+ self.client.task_instances.finish(
Review Comment:
Main never made a call for `SERVER_TERMINATED`, so a heartbeat 404 (TI or
DagRun deleted mid-run) just ended with the exit code. Now the ack itself 404s,
since the row is gone and not in history, and `ServerResponseError` propagates
out of `wait()`; the same happens if the API is unreachable for the ~30s the
client retries. `_replay_pending_terminal_state_msg` catches and logs for the
analogous case. Should this branch do the same, or should the server return 204
for the unknown-id case like it does for the archived one?
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1532,9 +1526,18 @@ def _on_child_started(
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:
+ except Exception as e:
# On any error kill that subprocess!
self.kill(signal.SIGKILL)
+ if (
+ isinstance(e, ServerResponseError)
+ and e.response.status_code == HTTPStatus.CONFLICT
+ and isinstance(e.detail, dict)
+ and e.detail.get("reason") == "invalid_state"
+ and e.detail.get("previous_state") == "restarting"
+ ):
+ self._terminal_state = SERVER_TERMINATED
Review Comment:
Nothing is logged on this path, so the operator sees only `Process exited
exit_code=-9 signal_sent=SIGKILL` and an empty task log for what was actually a
clear. A `log.info` plus a `process_log` line naming the reason, like the
heartbeat 409 handler does, would make the two distinguishable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
