moomindani commented on code in PR #69182:
URL: https://github.com/apache/airflow/pull/69182#discussion_r3511205030
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -1542,21 +1551,68 @@ def _get_task_base_json(self) -> dict[str, Any]:
"""Get the base json for the task."""
raise NotImplementedError()
+ def _retry_settings(self) -> dict[str, Any]:
+ """Databricks-native task retry settings that were explicitly
provided."""
+ settings: dict[str, Any] = {}
+ if self.max_retries is not None:
+ settings["max_retries"] = self.max_retries
+ if self.min_retry_interval_millis is not None:
+ settings["min_retry_interval_millis"] =
self.min_retry_interval_millis
+ if self.retry_on_timeout is not None:
+ settings["retry_on_timeout"] = self.retry_on_timeout
+ return settings
+
+ def _has_retry_settings(self) -> bool:
+ """Whether any Databricks-native retry field is explicitly
configured."""
+ if self._retry_settings():
+ return True
+ task_config = getattr(self, "task_config", {}) or {}
+ return any(
+ task_config.get(key) is not None
+ for key in ("max_retries", "min_retry_interval_millis",
"retry_on_timeout")
+ )
+
+ def _has_native_retry_attempts(self) -> bool:
+ """Whether Databricks can launch another native task attempt for this
operator."""
+ task_config = getattr(self, "task_config", {}) or {}
+ max_retries = self.max_retries if self.max_retries is not None else
task_config.get("max_retries")
+
+ if isinstance(max_retries, bool) or max_retries is None:
+ return False
+ if isinstance(max_retries, str):
Review Comment:
This str-coercion implies templated values were anticipated, but the new
params are not in `template_fields`, so a Jinja string would reach the API
unrendered. Either add the three params to `template_fields` (and keep this
branch) or drop the str handling — as-is it only masks a misconfiguration.
---
Drafted-by: Claude Code (Fable 5); reviewed by @moomindani before posting
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -1649,49 +1706,154 @@ def _convert_to_databricks_workflow_task(
def monitor_databricks_job(self) -> None:
"""
- Monitor the Databricks job.
+ Monitor the Databricks job until it terminates and surface its result.
+
+ Picks one of three monitoring strategies, depending on whether
Databricks-native retry
+ attempts can occur and whether the operator runs inside a
``DatabricksWorkflowTaskGroup``:
- Wait for the job to terminate. If deferrable, defer the task.
+ * standalone with native retries -> :meth:`_monitor_submit_run`,
following the submit run
+ whose own terminal state already accounts for every retry attempt.
+ * workflow task with native retries -> :meth:`_monitor_workflow_task`,
following the task's
+ latest attempt and tolerating in-flight retries until the shared run
is terminal.
+ * otherwise -> :meth:`_monitor_single_attempt`, following one attempt
and reporting as soon
+ as it terminates (the historical behaviour, unchanged).
"""
if self.databricks_run_id is None:
raise ValueError("Databricks job not yet launched. Please run
launch_notebook_job first.")
- current_task_run_id = self._get_current_databricks_task()["run_id"]
- run = self._hook.get_run(current_task_run_id)
- run_page_url = run["run_page_url"]
- self.log.info("Check the task run in Databricks: %s", run_page_url)
- run_state = RunState(**run["state"])
+
+ if not self._has_native_retry_attempts():
+ self._monitor_single_attempt()
+ elif self._databricks_workflow_task_group is None:
+ self._monitor_submit_run()
+ else:
+ self._monitor_workflow_task()
+
+ def _log_task_state(self, run_state: RunState) -> None:
self.log.info(
"Current state of the databricks task %s is %s",
self.databricks_task_key,
run_state.life_cycle_state,
)
+
+ def _defer_on_run(
+ self,
+ run_id: int,
+ *,
+ workflow_run_id: int | None = None,
+ databricks_task_key: str | None = None,
+ ) -> None:
+ """Defer monitoring of ``run_id`` to the trigger, optionally with
workflow-task context."""
+ self.defer(
+ trigger=DatabricksExecutionTrigger(
+ run_id=run_id,
+ databricks_conn_id=self.databricks_conn_id,
+ polling_period_seconds=self.polling_period_seconds,
+ retry_limit=self.databricks_retry_limit,
+ retry_delay=self.databricks_retry_delay,
+ retry_args=self.databricks_retry_args,
+ caller=self.caller,
+ workflow_run_id=workflow_run_id,
+ databricks_task_key=databricks_task_key,
+ ),
+ method_name=DEFER_METHOD_NAME,
+ )
+
+ def _monitor_single_attempt(self) -> None:
+ """Follow this task's attempt run and report as soon as it terminates
(no native retries)."""
+ current_task_run_id = self._get_current_databricks_task()["run_id"]
+ run = self._hook.get_run(current_task_run_id)
+ self.log.info("Check the task run in Databricks: %s",
run["run_page_url"])
+ run_state = RunState(**run["state"])
+ self._log_task_state(run_state)
+
if self.deferrable and not run_state.is_terminal:
- self.defer(
- trigger=DatabricksExecutionTrigger(
- run_id=current_task_run_id,
- databricks_conn_id=self.databricks_conn_id,
- polling_period_seconds=self.polling_period_seconds,
- retry_limit=self.databricks_retry_limit,
- retry_delay=self.databricks_retry_delay,
- retry_args=self.databricks_retry_args,
- caller=self.caller,
- ),
- method_name=DEFER_METHOD_NAME,
- )
+ self._defer_on_run(current_task_run_id)
+
while not run_state.is_terminal:
time.sleep(self.polling_period_seconds)
run = self._hook.get_run(current_task_run_id)
run_state = RunState(**run["state"])
+ self._log_task_state(run_state)
- self.log.info(
- "Current state of the databricks task %s is %s",
- self.databricks_task_key,
- run_state.life_cycle_state,
+ errors = extract_failed_task_errors(self._hook, run, run_state)
+ self._handle_terminal_run_state(run_state, errors)
+
+ def _monitor_workflow_task(self) -> None:
+ """
+ Follow this task's latest attempt within a shared workflow run,
tolerating native retries.
+
+ Inside a ``DatabricksWorkflowTaskGroup`` the run holds sibling tasks,
so the operator must
+ report when its own task finishes rather than wait for the whole run.
A failed attempt is
+ treated as final only once the workflow run is itself terminal,
because Databricks may still
+ launch another retry attempt under the same ``task_key``; each poll
re-resolves the latest
+ attempt for that key.
+
+ Tradeoff: a task that has exhausted its native retries keeps waiting
(or deferring) until the
+ parent run is terminal, since there is no per-task "retries exhausted"
signal before then.
+ Sibling tasks in the same run continue independently in the meantime.
+ """
+ workflow_run_id = self.databricks_run_id
+ if workflow_run_id is None:
+ raise ValueError("Databricks job not yet launched. Please run
launch_notebook_job first.")
+ current_task_run_id = self._get_current_databricks_task()["run_id"]
+ run = self._hook.get_run(current_task_run_id)
+ self.log.info("Check the task run in Databricks: %s",
run["run_page_url"])
+ run_state = RunState(**run["state"])
+ self._log_task_state(run_state)
+
+ # Defer whenever the outcome is not yet conclusive: a failed attempt
with the workflow run
+ # still active means a retry may follow, and the trigger waits for it
without blocking a worker.
+ if self.deferrable and not
self._workflow_task_is_conclusive(run_state, workflow_run_id):
+ self._defer_on_run(
+ current_task_run_id,
+ workflow_run_id=workflow_run_id,
+ databricks_task_key=self.databricks_task_key,
)
- # Extract errors from the run response using utility function
+ while not self._workflow_task_is_conclusive(run_state,
workflow_run_id):
+ time.sleep(self.polling_period_seconds)
+ current_task_run_id = self._get_current_databricks_task()["run_id"]
+ run = self._hook.get_run(current_task_run_id)
+ run_state = RunState(**run["state"])
+ self._log_task_state(run_state)
+
errors = extract_failed_task_errors(self._hook, run, run_state)
+ self._handle_terminal_run_state(run_state, errors)
+
+ def _workflow_task_is_conclusive(self, run_state: RunState,
workflow_run_id: int) -> bool:
+ """Whether the attempt is final: succeeded, or failed with the
workflow run also terminal."""
+ if not run_state.is_terminal:
+ return False
+ if run_state.is_successful:
+ return True
+ parent_state = RunState(**self._hook.get_run(workflow_run_id)["state"])
+ return parent_state.is_terminal
Review Comment:
This wait-for-parent gate is more conservative than it needs to be.
`runs/get` exposes `attempt_number` per attempt (verified on a live workspace:
attempts arrive as `attempt=0`, `attempt=1`, ... under the same `task_key`),
and the operator knows its own `max_retries` (operator arg or `task_config`).
So when the latest attempt is terminal-failed and `attempt_number >=
max_retries` (with `max_retries != -1`), retries are exhausted and the failure
is already conclusive — no need to keep the Airflow task waiting (or the
trigger deferring) until every sibling task finishes, which with long-running
siblings delays downstream failure handling by the length of the whole run. The
parent-terminal fallback would then only be needed for `max_retries=-1`. If you
take this, the "there is no per-task retries-exhausted signal" paragraph in
`notebook.rst`/`task.rst` and this docstring should be updated to match.
---
Drafted-by: Claude Code (Fable 5); reviewed by @moomindani before posting
--
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]
