Vamsi-klu commented on code in PR #69442:
URL: https://github.com/apache/airflow/pull/69442#discussion_r3635909791
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -1941,6 +1940,25 @@ def execute_complete(self, context: dict | None, event:
dict) -> None:
errors = event.get("errors", [])
self._handle_terminal_run_state(run_state, errors)
+ def on_kill(self) -> None:
+ if self.databricks_run_id is None:
+ return
+ if self._databricks_workflow_task_group:
+ # Workflow member: cancel only this task's child run, not the
shared parent workflow run.
+ # Cancelling the parent would also stop all sibling tasks.
+ try:
+ run_id_to_cancel =
self._get_current_databricks_task()["run_id"]
+ except Exception:
Review Comment:
Thanks for routing workflow members through
`_get_current_databricks_task()["run_id"]` - that matches
`monitor_databricks_job`.
This fallback undoes that guarantee though. On any failure resolving the
child run (API blip, missing task_key, kill racing shutdown), we cancel
`self.databricks_run_id`, which for workflow members is the shared parent run
and stops siblings - the bug this path is meant to avoid. `on_kill` is also one
of the more failure-prone moments to make a fresh API call.
Prefer: log the exception and return without calling `cancel_run`, rather
than falling back to the parent. A unit test where
`_get_current_databricks_task` raises and `cancel_run` is not called with the
parent id would lock that in.
--
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]
