This is an automated email from the ASF dual-hosted git repository.
eladkal 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 9326946b158 Report a retried Databricks task once, from its last
attempt (#72313)
9326946b158 is described below
commit 9326946b1583ea19abdfc2f97926bd4cdd38535c
Author: Noritaka Sekiyama <[email protected]>
AuthorDate: Thu Sep 10 05:46:37 2026 +0900
Report a retried Databricks task once, from its last attempt (#72313)
A run carries one entry per task attempt, so the failure list that reaches
a task's error message repeated a retried task once per attempt, and it
also listed tasks that had failed on an earlier attempt but succeeded on
the retry. Only the outcome of the last attempt describes what happened.
---
.../providers/databricks/utils/databricks.py | 24 ++++-
.../tests/unit/databricks/utils/test_databricks.py | 106 +++++++++++++++++++++
2 files changed, 128 insertions(+), 2 deletions(-)
diff --git
a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py
b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py
index 05b6b17710e..f2ef64b15f2 100644
--- a/providers/databricks/src/airflow/providers/databricks/utils/databricks.py
+++ b/providers/databricks/src/airflow/providers/databricks/utils/databricks.py
@@ -53,6 +53,26 @@ def normalise_json_content(content, json_path: str = "json")
-> str | bool | lis
raise AirflowException(msg)
+def _latest_attempts(tasks: list[dict]) -> list[dict]:
+ """
+ Keep only the most recent attempt of each task.
+
+ A run carries one entry per task *attempt*, so a retried task appears
several times: reporting
+ every entry both duplicates a single failure and counts a task that only
failed before its
+ retry succeeded.
+
+ :param tasks: Task entries of a run, as returned by the Databricks API
+ :return: One entry per ``task_key``, the one with the highest
``attempt_number``
+ """
+ latest: dict[str, dict] = {}
+ for task in tasks:
+ task_key = task["task_key"]
+ previous = latest.get(task_key)
+ if previous is None or task.get("attempt_number", 0) >=
previous.get("attempt_number", 0):
+ latest[task_key] = task
+ return list(latest.values())
+
+
def extract_failed_task_errors(
hook: DatabricksHook, run_info: dict, run_state: RunState
) -> list[dict[str, str | int]]:
@@ -66,7 +86,7 @@ def extract_failed_task_errors(
"""
failed_tasks = []
if run_state.result_state == "FAILED":
- for task in run_info.get("tasks", []):
+ for task in _latest_attempts(run_info.get("tasks", [])):
if task.get("state", {}).get("result_state", "") == "FAILED":
task_run_id = task["run_id"]
task_key = task["task_key"]
@@ -92,7 +112,7 @@ async def extract_failed_task_errors_async(
"""
failed_tasks = []
if run_state.result_state == "FAILED":
- for task in run_info.get("tasks", []):
+ for task in _latest_attempts(run_info.get("tasks", [])):
if task.get("state", {}).get("result_state", "") == "FAILED":
task_run_id = task["run_id"]
task_key = task["task_key"]
diff --git
a/providers/databricks/tests/unit/databricks/utils/test_databricks.py
b/providers/databricks/tests/unit/databricks/utils/test_databricks.py
index 33716879713..c4e39fb1b37 100644
--- a/providers/databricks/tests/unit/databricks/utils/test_databricks.py
+++ b/providers/databricks/tests/unit/databricks/utils/test_databricks.py
@@ -318,6 +318,76 @@ class TestExtractFailedTaskErrors:
assert result == expected
hook.get_run_output.assert_called_once_with(TASK_RUN_ID_2)
+ @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook")
+ def test_extract_failed_task_errors_reports_a_retried_task_once(self,
mock_hook_class):
+ """A task that was retried and failed again is reported once, from its
last attempt"""
+ hook = mock_hook_class.return_value
+ hook.get_run_output = mock_dict({"error": ERROR_MESSAGE})
+
+ run_state = RunState("TERMINATED", "FAILED", "Job failed")
+ run_info = {
+ "run_id": RUN_ID,
+ "state": {
+ "life_cycle_state": "TERMINATED",
+ "result_state": "FAILED",
+ "state_message": "Job failed",
+ },
+ "tasks": [
+ {
+ "run_id": TASK_RUN_ID_1,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 0,
+ "state": {"life_cycle_state": "INTERNAL_ERROR",
"result_state": "FAILED"},
+ },
+ {
+ "run_id": TASK_RUN_ID_2,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 1,
+ "state": {"life_cycle_state": "TERMINATED",
"result_state": "FAILED"},
+ },
+ ],
+ }
+
+ result = extract_failed_task_errors(hook, run_info, run_state)
+
+ assert result == [{"task_key": TASK_KEY_1, "run_id": TASK_RUN_ID_2,
"error": ERROR_MESSAGE}]
+ hook.get_run_output.assert_called_once_with(TASK_RUN_ID_2)
+
+ @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook")
+ def
test_extract_failed_task_errors_skips_task_that_succeeded_on_retry(self,
mock_hook_class):
+ """A task whose retry succeeded is not reported, even though its first
attempt failed"""
+ hook = mock_hook_class.return_value
+ hook.get_run_output = mock_dict({"error": ERROR_MESSAGE})
+
+ run_state = RunState("TERMINATED", "FAILED", "Job failed")
+ run_info = {
+ "run_id": RUN_ID,
+ "state": {
+ "life_cycle_state": "TERMINATED",
+ "result_state": "FAILED",
+ "state_message": "Job failed",
+ },
+ "tasks": [
+ {
+ "run_id": TASK_RUN_ID_1,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 0,
+ "state": {"life_cycle_state": "INTERNAL_ERROR",
"result_state": "FAILED"},
+ },
+ {
+ "run_id": TASK_RUN_ID_2,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 1,
+ "state": {"life_cycle_state": "TERMINATED",
"result_state": "SUCCESS"},
+ },
+ ],
+ }
+
+ result = extract_failed_task_errors(hook, run_info, run_state)
+
+ assert result == []
+ hook.get_run_output.assert_not_called()
+
class TestExtractFailedTaskErrorsAsync:
"""Test cases for the extract_failed_task_errors_async utility function
(asynchronous version)"""
@@ -501,3 +571,39 @@ class TestExtractFailedTaskErrorsAsync:
assert result == []
hook.a_get_run_output.assert_not_called()
+
+ @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook")
+ @pytest.mark.asyncio
+ async def
test_extract_failed_task_errors_async_reports_a_retried_task_once(self,
mock_hook_class):
+ """A task that was retried and failed again is reported once, from its
last attempt (async)"""
+ hook = mock_hook_class.return_value
+ hook.a_get_run_output = mock.AsyncMock(return_value={"error":
ERROR_MESSAGE})
+
+ run_state = RunState("TERMINATED", "FAILED", "Job failed")
+ run_info = {
+ "run_id": RUN_ID,
+ "state": {
+ "life_cycle_state": "TERMINATED",
+ "result_state": "FAILED",
+ "state_message": "Job failed",
+ },
+ "tasks": [
+ {
+ "run_id": TASK_RUN_ID_1,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 0,
+ "state": {"life_cycle_state": "INTERNAL_ERROR",
"result_state": "FAILED"},
+ },
+ {
+ "run_id": TASK_RUN_ID_2,
+ "task_key": TASK_KEY_1,
+ "attempt_number": 1,
+ "state": {"life_cycle_state": "TERMINATED",
"result_state": "FAILED"},
+ },
+ ],
+ }
+
+ result = await extract_failed_task_errors_async(hook, run_info,
run_state)
+
+ assert result == [{"task_key": TASK_KEY_1, "run_id": TASK_RUN_ID_2,
"error": ERROR_MESSAGE}]
+ hook.a_get_run_output.assert_called_once_with(TASK_RUN_ID_2)
