kaxil commented on code in PR #73554:
URL: https://github.com/apache/airflow/pull/73554#discussion_r4073994129
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1933,14 +1950,7 @@ def fetch_handle_failure_context(
if task and fail_fast:
_stop_remaining_tasks(task_instance=ti, session=session)
else:
- if ti.state != TaskInstanceState.RESTARTING:
- # Record the current attempt and prepare the TI for its next
try.
- # Covers every path eligible for retry reaching
handle_failure():
- # - RUNNING: task raised an exception during execution (normal
failure)
- # - QUEUED/SCHEDULED: executor killed the task externally
before
- # it could start (e.g. pod OOMKilled in KubernetesExecutor)
- # RESTARTING is excluded: the task was cleared via the UI/API
while running;
- # prepare_db_for_next_try() was already called during that
clear operation.
+ if ti.state != TaskInstanceState.UP_FOR_RETRY:
ti.prepare_db_for_next_try(session)
Review Comment:
Now that this also bumps `try_number`, the `on_task_instance_failed` hook a
few lines down receives the TI after it has been rolled forward, so
`task_instance.try_number` is N+1 (and the id is the replacement's) for the
attempt that failed. OpenLineage's `_on_task_instance_manual_state_change`
builds the FAIL run id from `ti.try_number`, so on the scheduler-side path
(zombie or executor-reported kill of a retry-eligible task) it no longer
matches the START event the worker emitted for try N. Before this change the id
was already rotated here but the try number still matched. Could the listener
call move above the archive, or the failed attempt be captured before it?
##########
airflow-core/src/airflow/migrations/versions/0135_3_4_0_allocate_pending_task_attempt_numbers.py:
##########
@@ -0,0 +1,65 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Allocate attempt numbers for pending retries and cleared tasks.
Review Comment:
Does this want a `significant` newsfragment? After this lands, `try_number`
on an `up_for_retry` row identifies the pending attempt rather than the one
that failed, which is visible through the public REST API, the UI header, and
listeners, and there is a migration to run.
##########
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
+
bind_contextvars(ti_id=str(task_instance_id))
log.debug("Updating task instance state", new_state=ti_patch_payload.state)
+ if isinstance(ti_patch_payload, TITerminalStatePayload) and (
+ ti_patch_payload.state == TerminalStateNonSuccess.SERVER_TERMINATED
+ ):
+ ti = session.scalar(
+ select(TI)
+ .where(TI.id == task_instance_id)
+ .with_for_update(of=TI)
+ .execution_options(populate_existing=True)
+ )
+ if ti is None:
+ if
session.scalar(select(TIH.task_instance_id).where(TIH.task_instance_id ==
task_instance_id)):
+ return Response(status_code=status.HTTP_204_NO_CONTENT)
+ raise HTTPException(status_code=404, detail={"reason":
"not_found"})
+ if (ti.hostname, ti.pid) != (ti_patch_payload.hostname,
ti_patch_payload.pid) or (
+ ti_patch_payload.hostname is None or ti_patch_payload.pid is None
+ ):
+ raise HTTPException(status_code=409, detail={"reason":
"running_elsewhere"})
+ if ti.state == TaskInstanceState.RESTARTING:
+ dag = dag_bag.get_dag_for_run(dag_run=ti.dag_run, session=session)
+ if dag is None:
+ raise HTTPException(status_code=404, detail={"reason":
"dag_not_found"})
+ ti.task = dag.get_task(ti.task_id)
Review Comment:
`get_task` raises `TaskNotFound` if the task was dropped from the DAG
between the attempt starting and the clear, which would surface as a 500 here
and leave the row RESTARTING until the executor event arrives. The scheduler
fallback wraps the same lookup in try/except. Worth returning a 404 like the
`dag is None` branch above?
--
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]
