Skip to content

Better error handling for running ti heartbeats after task is cleared #56443

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kriti-sc wants to merge 5 commits into
base: main
Choose a base branch
from
+8 −8
Open

Better error handling for running ti heartbeats after task is cleared #56443

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7e5ea0c
change 404 to 410
kriti-sc Oct 6, 2025
607574a
change 404 to 410
kriti-sc Oct 6, 2025
80542b5
Update airflow-core/src/airflow/api_fastapi/execution_api/routes/task…
kriti-sc Nov 9, 2025
32f289c
Update airflow-core/src/airflow/api_fastapi/execution_api/routes/task…
kriti-sc Nov 9, 2025
7341091
Update airflow-core/src/airflow/api_fastapi/execution_api/routes/task…
kriti-sc Nov 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -577,7 +577,7 @@ def ti_skip_downstream(
"/{task_instance_id}/heartbeat",
status_code=status.HTTP_204_NO_CONTENT,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Task Instance not found"},
status.HTTP_410_GONE: {"description": "Task Instance no longer exists, it may have moved to the Task Instance History table"},
Copy link
Member

@uranusjr uranusjr Nov 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think 404 should still exist, there’s a difference between a ti that never existed and a ti that existed but is gone.

status.HTTP_409_CONFLICT: {
"description": "The TI attempting to heartbeat should be terminated for the given reason"
},
Expand Down Expand Up @@ -605,12 +605,12 @@ def ti_heartbeat(
"Retrieved current task state", state=previous_state, current_hostname=hostname, current_pid=pid
)
except NoResultFound:
log.error("Task Instance not found")
status.HTTP_410_GONE: {"description": "Task Instance no longer exists, it may have moved to the Task Instance History table"},
Copy link
Contributor

@amoghrajesh amoghrajesh Nov 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is added here by mistake?

Copy link
Contributor

@amoghrajesh amoghrajesh Nov 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should actually validate that, check if it is present in the TIH table.

Two situations possible here:

  1. Task Instance doesnt exist and not in TIH table -> return 404
  2. Task Instance doesnt exist and present in TIH table -> return 410

raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
status_code=status.HTTP_410_GONE,
detail={
"reason": "not_found",
"message": "Task Instance not found",
status.HTTP_410_GONE: {"description": "Task Instance no longer exists, it may have moved to the Task Instance History table"},
Copy link
Contributor

@amoghrajesh amoghrajesh Nov 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Formatting looks off -- static checks will fail

},
)

Expand Down
6 changes: 3 additions & 3 deletions airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1264,7 +1264,7 @@ def test_ti_heartbeat(
assert response.json()["detail"] == expected_detail

def test_ti_heartbeat_non_existent_task(self, client, session, create_task_instance):
"""Test that a 404 error is returned when the Task Instance does not exist."""
"""Test that a 410 error is returned when the Task Instance does not exist."""
task_instance_id = "0182e924-0f1e-77e6-ab50-e977118bc139"

# Pre-condition: the Task Instance does not exist
Expand All @@ -1275,10 +1275,10 @@ def test_ti_heartbeat_non_existent_task(self, client, session, create_task_insta
json={"hostname": "random-hostname", "pid": 1547},
)

assert response.status_code == 404
assert response.status_code == 410
assert response.json()["detail"] == {
"reason": "not_found",
"message": "Task Instance not found",
"message": "Task Instance not found, might have moved to the Task Instance History table",
}

@pytest.mark.parametrize(
Expand Down
2 changes: 1 addition & 1 deletion task-sdk/src/airflow/sdk/execution_time/supervisor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1099,7 +1099,7 @@ def _send_heartbeat_if_needed(self):
# Reset the counter on success
self.failed_heartbeats = 0
except ServerResponseError as e:
if e.response.status_code in {HTTPStatus.NOT_FOUND, HTTPStatus.CONFLICT}:
if e.response.status_code in {HTTPStatus.GONE, HTTPStatus.CONFLICT}:
log.error(
"Server indicated the task shouldn't be running anymore",
detail=e.detail,
Expand Down
Loading