moomindani commented on code in PR #69998:
URL: https://github.com/apache/airflow/pull/69998#discussion_r3745948631
##########
providers/databricks/src/airflow/providers/databricks/hooks/databricks.py:
##########
@@ -563,6 +563,28 @@ def get_run_tasks(self, run_id: int) -> list[dict[str,
Any]]:
return all_tasks
+ def get_run_failed_task_keys(self, run_id: int) -> list[str]:
+ """
+ Return the ``task_key`` of every sub-task of a run that is in a
terminal failure state.
+
+ Resolved from the live Databricks run rather than from Airflow's
metadata DB, so it
+ reflects the actual per-task state Databricks ``repair_run`` will act
on. The returned
+ keys are the values to pass as ``rerun_tasks`` to :meth:`repair_run`.
+
+ :param run_id: id of the run
+ :return: a list of Databricks ``task_key`` values for failed sub-tasks
+ """
+ failed_result_states = {"FAILED", "TIMEDOUT", "CANCELED",
"MAXIMUM_CONCURRENT_RUNS_REACHED"}
+ failed_task_keys = []
+ for task in self.get_run_tasks(run_id):
+ state = task.get("state", {})
+ if (
+ state.get("result_state") in failed_result_states
+ or state.get("life_cycle_state") == "INTERNAL_ERROR"
+ ):
+ failed_task_keys.append(task["task_key"])
+ return failed_task_keys
Review Comment:
This assumes one entry per `task_key`, but `get_run_tasks` returns one entry
**per attempt**, so a task that retried (or a run that was already repaired)
yields the same key several times. Databricks rejects that:
```
Error: rerun_tasks cannot contain duplicate task keys.
```
I confirmed this against a live run — with `max_retries=1` on a failing task
this returns `['flaky', 'flaky']` and the repair fails; deduped to `['flaky']`
it succeeds. Attempts accumulate across repairs, so the duplicate count grows
(I saw four).
There is a second, quieter symptom: because every historical attempt is
inspected, a `task_key` whose **latest** attempt succeeded is still reported as
failed. After my run finished fully green (`parent: SUCCESS`, both tasks
`SUCCESS`), this still returned `['flaky', 'flaky', 'flaky', 'flaky']`.
Databricks does *not* reject repairing an already-successful task — it accepts
it and returns a `repair_id` — so that path silently re-runs a healthy task and
clears its Airflow TI.
Both symptoms have one root cause, and the fix already exists in this repo:
`_get_current_databricks_task` in `operators/databricks.py` resolves the latest
attempt by sorting on `start_time` before building the map. Same idiom here:
```python
sorted_tasks = sorted(self.get_run_tasks(run_id), key=lambda t:
t["start_time"])
latest_by_key = {task["task_key"]: task for task in sorted_tasks}
```
then evaluate the failure states over `latest_by_key.values()`. That dedupes
and fixes the stale-attempt judgement in one move.
##########
providers/databricks/tests/unit/databricks/hooks/test_databricks.py:
##########
@@ -749,6 +749,20 @@ def test_get_run_tasks_success_multiple_pages(self,
mock_requests):
assert len(tasks) == 2
assert tasks == GET_RUN_RESPONSE["tasks"] * 2
+ def test_get_run_failed_task_keys(self):
Review Comment:
This fixture has one entry per `task_key`, which is why CI is green despite
the duplicate-key bug — the real API returns one entry per *attempt*.
Worth adding a case shaped like a live run, e.g. two `flaky` entries with
`attempt_number` 0/1 both `FAILED` plus a later attempt that `SUCCESS`ed,
asserting the key appears once and is judged by its latest attempt. That would
fail on the current implementation.
--
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]
