You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "moomindani (via GitHub)" <gi...@apache.org> on 2026/07/02 07:24:48 UTC
Re: [PR] Add Databricks-native retry settings to task operators [airflow]
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org