diff --git a/providers/databricks/docs/operators/notebook.rst b/providers/databricks/docs/operators/notebook.rst index 83dd899b41279..3c8ec27b8d38b 100644 --- a/providers/databricks/docs/operators/notebook.rst +++ b/providers/databricks/docs/operators/notebook.rst @@ -42,3 +42,41 @@ Running a notebook in Databricks on an existing cluster :language: python :start-after: [START howto_operator_databricks_notebook_existing_cluster] :end-before: [END howto_operator_databricks_notebook_existing_cluster] + +Configuring Databricks-native task retries +------------------------------------------- + +Use ``max_retries``, ``min_retry_interval_millis`` and ``retry_on_timeout`` to configure +`Databricks-native task retries `_. +Databricks reruns failed task attempts within the same job run, so Airflow sees only the final result. +Set ``max_retries`` to ``-1`` to retry indefinitely, or ``0`` to disable retries. + +These settings are independent of the Airflow task-level ``retries`` parameter, which retries the +whole Airflow task: + +.. code-block:: python + + DatabricksNotebookOperator( + task_id="notebook", + notebook_path="/path/to/notebook", + source="WORKSPACE", + existing_cluster_id="existing_cluster_id", + max_retries=3, + min_retry_interval_millis=2000, + retry_on_timeout=True, + ) + +Airflow ``retries`` behaves differently depending on where the operator runs. For a standalone +operator, each retry submits a new Databricks run. Inside a +:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`, +the Airflow task monitors a sub-run that was already submitted by the workflow launch task, so a +retry only re-polls the terminal sub-run. Use ``max_retries`` to retry Databricks work inside a +workflow task group. + +Inside a +:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`, +a task that exhausts a finite ``max_retries`` is reported as failed as soon as its final failed +attempt is observed, so downstream failure handling is not delayed by long-running sibling tasks. +Only unlimited retries (``max_retries=-1``) keep the Airflow task waiting (or deferring) until the +parent workflow run reaches a terminal state, because Databricks may still launch a retry attempt +under the same ``task_key`` until then. Sibling tasks in the run continue independently. diff --git a/providers/databricks/docs/operators/task.rst b/providers/databricks/docs/operators/task.rst index 5c446593531a0..b327aa20ccf97 100644 --- a/providers/databricks/docs/operators/task.rst +++ b/providers/databricks/docs/operators/task.rst @@ -44,3 +44,30 @@ Running a SQL query in Databricks using DatabricksTaskOperator :language: python :start-after: [START howto_operator_databricks_task_sql] :end-before: [END howto_operator_databricks_task_sql] + +Configuring Databricks-native task retries +------------------------------------------- + +Use ``max_retries``, ``min_retry_interval_millis`` and ``retry_on_timeout`` to configure +`Databricks-native task retries `_. +Databricks reruns failed task attempts within the same job run, so Airflow sees only the final result. +Set ``max_retries`` to ``-1`` to retry indefinitely, or ``0`` to disable retries. + +These settings are independent of the Airflow task-level ``retries`` parameter, which retries the +whole Airflow task. You can set the same fields directly in ``task_config``. When both are set, the +operator parameter takes precedence. If a field is unset, Databricks uses its default. + +Airflow ``retries`` behaves differently depending on where the operator runs. For a standalone +operator, each retry submits a new Databricks run. Inside a +:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`, +the Airflow task monitors a sub-run that was already submitted by the workflow launch task, so a +retry only re-polls the terminal sub-run. Use ``max_retries`` to retry Databricks work inside a +workflow task group. + +Inside a +:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`, +a task that exhausts a finite ``max_retries`` is reported as failed as soon as its final failed +attempt is observed, so downstream failure handling is not delayed by long-running sibling tasks. +Only unlimited retries (``max_retries=-1``) keep the Airflow task waiting (or deferring) until the +parent workflow run reaches a terminal state, because Databricks may still launch a retry attempt +under the same ``task_key`` until then. Sibling tasks in the run continue independently. diff --git a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py index a64acb7859870..d820b099c8d3a 100644 --- a/providers/databricks/src/airflow/providers/databricks/operators/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/operators/databricks.py @@ -1673,6 +1673,9 @@ class DatabricksTaskBaseOperator(BaseOperator, ABC): :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default. :param workflow_run_metadata: Metadata for the workflow run. This is used when the operator is used within a workflow. It is expected to be a dictionary containing the run_id and conn_id for the workflow. + :param max_retries: Databricks task retry count. Use ``-1`` for unlimited retries. + :param min_retry_interval_millis: Minimum interval between Databricks task retries. + :param retry_on_timeout: Whether Databricks retries timed-out tasks. """ def __init__( @@ -1690,6 +1693,9 @@ def __init__( polling_period_seconds: int = 5, wait_for_termination: bool = True, workflow_run_metadata: dict[str, Any] | None = None, + max_retries: int | str | None = None, + min_retry_interval_millis: int | str | None = None, + retry_on_timeout: bool | str | None = None, **kwargs: Any, ): self.caller = caller @@ -1705,6 +1711,9 @@ def __init__( self.polling_period_seconds = polling_period_seconds self.wait_for_termination = wait_for_termination self.workflow_run_metadata = workflow_run_metadata + self.max_retries = max_retries + self.min_retry_interval_millis = min_retry_interval_millis + self.retry_on_timeout = retry_on_timeout self.databricks_run_id: int | None = None @@ -1792,21 +1801,75 @@ 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"] = int(self.max_retries) + if self.min_retry_interval_millis is not None: + settings["min_retry_interval_millis"] = int(self.min_retry_interval_millis) + if self.retry_on_timeout is not None: + retry_on_timeout = self.retry_on_timeout + if isinstance(retry_on_timeout, str): + retry_on_timeout = retry_on_timeout.strip().lower() in ("true", "1", "yes") + settings["retry_on_timeout"] = bool(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 _resolved_max_retries(self) -> int | None: + """Return the resolved ``max_retries`` value, with operator args taking precedence.""" + 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 None + if isinstance(max_retries, str): + try: + return int(max_retries) + except ValueError: + self.log.warning( + "Ignoring unparsable max_retries value %r; treating as no retry limit configured.", + max_retries, + ) + return None + if not isinstance(max_retries, int): + return None + return max_retries + def _get_run_json(self) -> dict[str, Any]: """Get run json to be used for task submissions.""" - run_json = { - "run_name": self.databricks_task_key, - **self._get_task_base_json(), - } if self.new_cluster and self.existing_cluster_id: raise ValueError("Both new_cluster and existing_cluster_id are set. Only one should be set.") + cluster: dict[str, Any] if self.new_cluster: - run_json["new_cluster"] = self.new_cluster + cluster = {"new_cluster": self.new_cluster} elif self.existing_cluster_id: - run_json["existing_cluster_id"] = self.existing_cluster_id + cluster = {"existing_cluster_id": self.existing_cluster_id} else: raise ValueError("Must specify either existing_cluster_id or new_cluster.") - return run_json + + if not self._has_retry_settings(): + # No retry settings: keep the legacy single-task runs/submit shape unchanged. + return {"run_name": self.databricks_task_key, **self._get_task_base_json(), **cluster} + + # Retry settings are per-task SubmitTask fields, so submit the single task explicitly. + # The explicit task_key also gives monitoring a stable task to look up. + task = { + **self._get_task_base_json(), + "task_key": self.databricks_task_key, + **self._retry_settings(), + **cluster, + } + return {"run_name": self.databricks_task_key, "tasks": [task]} def _launch_job(self, context: Context | None = None) -> int | None: """Launch the job on Databricks.""" @@ -1857,13 +1920,14 @@ def _convert_to_databricks_workflow_task( base_task_json = self._get_task_base_json() result = { - "task_key": self.databricks_task_key, "depends_on": [ {"task_key": self._generate_databricks_task_key(task_id, task_dict)} for task_id in self.upstream_task_ids if task_id in relevant_upstreams ], **base_task_json, + "task_key": self.databricks_task_key, + **self._retry_settings(), } trigger_rule_value = ( @@ -1899,49 +1963,170 @@ def _convert_to_databricks_workflow_task( def monitor_databricks_job(self) -> None: """ - Monitor the Databricks job. - - Wait for the job to terminate. If deferrable, defer the task. + 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``: + + * 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 they are exhausted (or, for unlimited + 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"]) + + max_retries = self._resolved_max_retries() + if max_retries is None or (max_retries != -1 and max_retries <= 0): + 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, + max_retries: int | 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, + max_retries=max_retries, + ), + 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 + final once finite retries are exhausted; unlimited retries fall back to the workflow run's + terminal state. Each poll re-resolves the latest attempt for the task key. + """ + 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 = self._get_current_databricks_task() + current_task_run_id = current_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) + attempt_number = current_task.get("attempt_number") + + # Defer whenever the outcome is not yet conclusive: a failed attempt with retries still + # available 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, attempt_number + ): + self._defer_on_run( + current_task_run_id, + workflow_run_id=workflow_run_id, + databricks_task_key=self.databricks_task_key, + max_retries=self._resolved_max_retries(), ) - # Extract errors from the run response using utility function + while not self._workflow_task_is_conclusive(run_state, workflow_run_id, attempt_number): + time.sleep(self.polling_period_seconds) + current_task = self._get_current_databricks_task() + current_task_run_id = current_task["run_id"] + run = self._hook.get_run(current_task_run_id) + run_state = RunState(**run["state"]) + self._log_task_state(run_state) + attempt_number = current_task.get("attempt_number") + 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, attempt_number: int | None + ) -> bool: + """Whether the attempt is final: succeeded, retries exhausted, or the run is terminal.""" + if not run_state.is_terminal: + return False + if run_state.is_successful: + return True + max_retries = self._resolved_max_retries() + if ( + max_retries is not None + and max_retries != -1 + and attempt_number is not None + and attempt_number >= max_retries + ): + return True + parent_state = RunState(**self._hook.get_run(workflow_run_id)["state"]) + return parent_state.is_terminal + + def _monitor_submit_run(self) -> None: + """ + Wait for a standalone submit run to terminate, tolerating Databricks-native retries. + + The submit run owns exactly this operator's task, so its own terminal state already + accounts for every native retry attempt — we follow the run rather than any single attempt. + """ + run_id = self.databricks_run_id + if run_id is None: + raise ValueError("Databricks job not yet launched. Please run launch_notebook_job first.") + run = self._hook.get_run(run_id) + self.log.info("Check the job run in Databricks: %s", run["run_page_url"]) + run_state = RunState(**run["state"]) + self.log.info("Current state of the databricks run %s is %s", run_id, run_state.life_cycle_state) + if self.deferrable and not run_state.is_terminal: + self._defer_on_run(run_id) + + while not run_state.is_terminal: + time.sleep(self.polling_period_seconds) + run = self._hook.get_run(run_id) + run_state = RunState(**run["state"]) + self.log.info("Current state of the databricks run %s is %s", run_id, run_state.life_cycle_state) + + errors = extract_failed_task_errors(self._hook, run, run_state) self._handle_terminal_run_state(run_state, errors) def execute(self, context: Context) -> None: @@ -2033,11 +2218,17 @@ class DatabricksNotebookOperator(DatabricksTaskBaseOperator): :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default. :param workflow_run_metadata: Metadata for the workflow run. This is used when the operator is used within a workflow. It is expected to be a dictionary containing the run_id and conn_id for the workflow. + :param max_retries: Databricks task retry count. Use ``-1`` for unlimited retries. + :param min_retry_interval_millis: Minimum interval between Databricks task retries. + :param retry_on_timeout: Whether Databricks retries timed-out tasks. """ template_fields = ( "notebook_params", "workflow_run_metadata", + "max_retries", + "min_retry_interval_millis", + "retry_on_timeout", ) CALLER = "DatabricksNotebookOperator" @@ -2058,6 +2249,9 @@ def __init__( polling_period_seconds: int = 5, wait_for_termination: bool = True, workflow_run_metadata: dict | None = None, + max_retries: int | str | None = None, + min_retry_interval_millis: int | str | None = None, + retry_on_timeout: bool | str | None = None, **kwargs: Any, ): self.notebook_path = notebook_path @@ -2078,6 +2272,9 @@ def __init__( polling_period_seconds=polling_period_seconds, wait_for_termination=wait_for_termination, workflow_run_metadata=workflow_run_metadata, + max_retries=max_retries, + min_retry_interval_millis=min_retry_interval_millis, + retry_on_timeout=retry_on_timeout, **kwargs, ) @@ -2176,6 +2373,9 @@ class DatabricksTaskOperator(DatabricksTaskBaseOperator): :param new_cluster: Specs for a new cluster on which this task will be run. :param polling_period_seconds: Controls the rate which we poll for the result of this notebook job run. :param wait_for_termination: if we should wait for termination of the job run. ``True`` by default. + :param max_retries: Databricks task retry count, overriding ``task_config`` when set. + :param min_retry_interval_millis: Minimum retry interval, overriding ``task_config`` when set. + :param retry_on_timeout: Whether Databricks retries timed-out tasks, overriding ``task_config`` when set. """ CALLER = "DatabricksTaskOperator" @@ -2183,6 +2383,9 @@ class DatabricksTaskOperator(DatabricksTaskBaseOperator): "databricks_conn_id", "task_config", "workflow_run_metadata", + "max_retries", + "min_retry_interval_millis", + "retry_on_timeout", ) def __init__( @@ -2199,6 +2402,9 @@ def __init__( polling_period_seconds: int = 5, wait_for_termination: bool = True, workflow_run_metadata: dict | None = None, + max_retries: int | str | None = None, + min_retry_interval_millis: int | str | None = None, + retry_on_timeout: bool | str | None = None, **kwargs, ): self.task_config = task_config @@ -2216,6 +2422,9 @@ def __init__( polling_period_seconds=polling_period_seconds, wait_for_termination=wait_for_termination, workflow_run_metadata=workflow_run_metadata, + max_retries=max_retries, + min_retry_interval_millis=min_retry_interval_millis, + retry_on_timeout=retry_on_timeout, **kwargs, ) diff --git a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py index 57d489550f200..9706d2e4adfba 100644 --- a/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py +++ b/providers/databricks/src/airflow/providers/databricks/triggers/databricks.py @@ -43,6 +43,9 @@ class DatabricksExecutionTrigger(BaseTrigger): :param run_page_url: The run page url. :param repair_run: Repair the databricks run in case of failure. :param caller: The name of the operator that is calling the hook. + :param workflow_run_id: Parent workflow run ID for task-level monitoring. + :param databricks_task_key: Task key to monitor within ``workflow_run_id``. + :param max_retries: Resolved Databricks-native ``max_retries`` for task-level monitoring. """ def __init__( @@ -56,6 +59,9 @@ def __init__( run_page_url: str | None = None, repair_run: bool = False, caller: str = "DatabricksExecutionTrigger", + workflow_run_id: int | None = None, + databricks_task_key: str | None = None, + max_retries: int | None = None, ) -> None: super().__init__() # Trigger kwargs cross Airflow's serialization boundary, so fail before storing invalid @@ -70,6 +76,9 @@ def __init__( self.run_page_url = run_page_url self.repair_run = repair_run self.caller = caller + self.workflow_run_id = workflow_run_id + self.databricks_task_key = databricks_task_key + self.max_retries = max_retries self.hook = DatabricksHook( databricks_conn_id, retry_limit=self.retry_limit, @@ -91,19 +100,40 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "run_page_url": self.run_page_url, "repair_run": self.repair_run, "caller": self.caller, + "workflow_run_id": self.workflow_run_id, + "databricks_task_key": self.databricks_task_key, + "max_retries": self.max_retries, }, ) async def on_kill(self) -> None: """Cancel the Databricks run when the trigger is cancelled by a user action.""" - if self.run_id: - from asgiref.sync import sync_to_async + from asgiref.sync import sync_to_async + + run_id = self.run_id + if self.workflow_run_id is not None and self.databricks_task_key is not None: + # self.run_id may be an earlier, now-terminal attempt; cancel the task's latest attempt + # so a retry/repair launched under the same task_key is not left running. + tasks = await sync_to_async(self.hook.get_run_tasks)(self.workflow_run_id) + attempt = { + task["task_key"]: task for task in sorted(tasks, key=lambda task: task["start_time"]) + }.get(self.databricks_task_key) + if attempt: + run_id = attempt["run_id"] + if run_id: + self.log.info("Cancelling Databricks run %s.", run_id) + await sync_to_async(self.hook.cancel_run)(run_id) - self.log.info("Cancelling Databricks run %s.", self.run_id) - await sync_to_async(self.hook.cancel_run)(self.run_id) + def _monitors_workflow_task(self) -> bool: + """Whether this trigger follows one task inside a shared workflow run (see ``workflow_run_id``).""" + return bool(self.workflow_run_id and self.databricks_task_key) async def run(self): async with self.hook: + if self._monitors_workflow_task(): + async for event in self._run_workflow_task(): + yield event + return while True: run_state = await self.hook.a_get_run_state(self.run_id) if not run_state.is_terminal: @@ -129,6 +159,69 @@ async def run(self): ) return + async def _run_workflow_task(self): + """Monitor one task in a workflow run, tolerating in-flight retries/repairs.""" + from asgiref.sync import sync_to_async + + while True: + tasks = await sync_to_async(self.hook.get_run_tasks)(self.workflow_run_id) + sorted_task_runs = sorted(tasks, key=lambda task: task["start_time"]) + attempt = {task["task_key"]: task for task in sorted_task_runs}.get(self.databricks_task_key) + + if attempt is not None: + attempt_run_id = attempt["run_id"] + attempt_state = await self.hook.a_get_run_state(attempt_run_id) + + if attempt_state.is_terminal: + if attempt_state.is_successful: + yield TriggerEvent( + { + "run_id": attempt_run_id, + "run_page_url": self.run_page_url, + "run_state": attempt_state.to_json(), + "repair_run": self.repair_run, + "errors": [], + } + ) + return + # A failed attempt is final once finite retries are exhausted; otherwise wait + # for the parent run because another attempt may still appear. + attempt_number = attempt.get("attempt_number") + retries_exhausted = ( + self.max_retries is not None + and self.max_retries != -1 + and attempt_number is not None + and attempt_number >= self.max_retries + ) + if ( + retries_exhausted + or (await self.hook.a_get_run_state(self.workflow_run_id)).is_terminal + ): + run_info = await self.hook.a_get_run(attempt_run_id) + failed_tasks = await extract_failed_task_errors_async( + self.hook, run_info, attempt_state + ) + yield TriggerEvent( + { + "run_id": attempt_run_id, + "run_page_url": self.run_page_url, + "run_state": attempt_state.to_json(), + "repair_run": self.repair_run, + "errors": failed_tasks, + } + ) + return + + # attempt is None when the task has not yet surfaced in the run (e.g. just after launch); + # keep polling rather than crashing on a missing task_key. + self.log.info( + "databricks task %s not yet conclusive in run %s. sleeping for %s seconds", + self.databricks_task_key, + self.workflow_run_id, + self.polling_period_seconds, + ) + await asyncio.sleep(self.polling_period_seconds) + class DatabricksSQLStatementExecutionTrigger(BaseTrigger): """ diff --git a/providers/databricks/tests/unit/databricks/operators/test_databricks.py b/providers/databricks/tests/unit/databricks/operators/test_databricks.py index 65c990a592658..0f190da3661d0 100644 --- a/providers/databricks/tests/unit/databricks/operators/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/operators/test_databricks.py @@ -3687,6 +3687,103 @@ def test_execute_with_deferrable(self, mock_get_current_task, mock_databricks_ho "Trigger is not a DatabricksExecutionTrigger" ) assert exec_info.value.method_name == "execute_complete" + # Without native retries configured, the trigger keeps its original parent-unaware behavior. + assert exec_info.value.trigger.workflow_run_id is None + assert exec_info.value.trigger.databricks_task_key is None + + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_execute_standalone_with_retries_defers_on_submit_run(self, mock_databricks_hook): + # A standalone operator with native retries follows its own submit run to a terminal state, + # so the defer targets that run directly (no per-task workflow context). + mock_databricks_hook.return_value.get_run.return_value = { + "state": {"life_cycle_state": "PENDING"}, + "run_page_url": "test_url", + } + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + wait_for_termination=True, + deferrable=True, + max_retries=2, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(TaskDeferred) as exec_info: + operator.monitor_databricks_job() + assert exec_info.value.trigger.run_id == 12345 + assert exec_info.value.trigger.workflow_run_id is None + assert exec_info.value.trigger.databricks_task_key is None + + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._databricks_workflow_task_group", + new_callable=mock.PropertyMock, + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_execute_with_deferrable_passes_workflow_context_when_retries_configured( + self, mock_get_current_task, mock_databricks_hook, mock_workflow_tg + ): + # Inside a workflow task group the run is shared, so the defer carries the workflow context + # and the trigger tracks this task's own attempt within that run. + mock_workflow_tg.return_value = MagicMock() + mock_get_current_task.return_value = {"run_id": "attempt-1"} + mock_databricks_hook.return_value.get_run.return_value = { + "state": {"life_cycle_state": "PENDING"}, + "run_page_url": "test_url", + } + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + wait_for_termination=True, + deferrable=True, + max_retries=2, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(TaskDeferred) as exec_info: + operator.monitor_databricks_job() + assert exec_info.value.trigger.workflow_run_id == 12345 + assert exec_info.value.trigger.databricks_task_key == operator.databricks_task_key + + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._databricks_workflow_task_group", + new_callable=mock.PropertyMock, + ) + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_execute_with_deferrable_and_max_retries_zero_keeps_single_attempt_monitoring( + self, mock_get_current_task, mock_databricks_hook, mock_workflow_tg + ): + mock_workflow_tg.return_value = MagicMock() + mock_get_current_task.return_value = {"run_id": "attempt-1"} + mock_databricks_hook.return_value.get_run.return_value = { + "state": {"life_cycle_state": "PENDING"}, + "run_page_url": "test_url", + } + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + wait_for_termination=True, + deferrable=True, + max_retries=0, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(TaskDeferred) as exec_info: + operator.monitor_databricks_job() + assert exec_info.value.trigger.run_id == "attempt-1" + assert exec_info.value.trigger.workflow_run_id is None + assert exec_info.value.trigger.databricks_task_key is None @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") @mock.patch( @@ -3762,6 +3859,279 @@ def test_monitor_databricks_job_failed(self, mock_get_current_task, mock_databri exception_message = "Task failed. Final state FAILED. Reason: FAILURE. Errors: []" assert exception_message == str(exc_info.value) + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._databricks_workflow_task_group", + new_callable=mock.PropertyMock, + ) + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_monitor_databricks_job_retry_in_flight_succeeds( + self, mock_get_current_task, mock_databricks_hook, mock_sleep, mock_workflow_tg + ): + # Inside a workflow task group the run is shared, so the task tracks its own attempt: the + # first attempt fails while the workflow run is still active; a retried attempt then + # succeeds and the Airflow task must not fail. + mock_workflow_tg.return_value = MagicMock() + mock_get_current_task.side_effect = [{"run_id": "attempt-1"}, {"run_id": "attempt-2"}] + runs = { + "attempt-1": { + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "first attempt failed", + }, + "run_page_url": "url-1", + }, + 12345: {"state": {"life_cycle_state": "RUNNING"}, "run_page_url": "parent"}, + "attempt-2": { + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "run_page_url": "url-2", + }, + } + mock_databricks_hook.return_value.get_run.side_effect = lambda run_id: runs[run_id] + + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + max_retries=1, + ) + operator.databricks_run_id = 12345 + + operator.monitor_databricks_job() + mock_sleep.assert_called_once() + + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_monitor_standalone_submit_run_follows_run_to_terminal_state( + self, mock_databricks_hook, mock_sleep + ): + # A standalone submit run stays active while Databricks retries the task; the operator + # follows the run (not an attempt) and reports only once the run itself terminates. + run_states = iter( + [ + {"life_cycle_state": "RUNNING"}, + {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + ] + ) + mock_databricks_hook.return_value.get_run.side_effect = lambda run_id: { + "state": next(run_states), + "run_page_url": "url", + } + + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + max_retries=1, + ) + operator.databricks_run_id = 12345 + + operator.monitor_databricks_job() + # Only the submit run is polled; the per-attempt resolver is never used. + mock_databricks_hook.return_value.get_run.assert_called_with(12345) + mock_databricks_hook.return_value.get_run_tasks.assert_not_called() + + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_monitor_databricks_job_without_retries_fails_immediately( + self, mock_get_current_task, mock_databricks_hook, mock_sleep + ): + # Without native retries configured the task fails as soon as its own attempt terminates, + # even if the parent run is still active. The parent run state must not be polled. + mock_get_current_task.return_value = {"run_id": "attempt-1"} + mock_databricks_hook.return_value.get_run.return_value = { + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "attempt failed", + }, + "run_page_url": "url", + } + + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + ) + operator.databricks_run_id = 12345 + + with pytest.raises(AirflowException): + operator.monitor_databricks_job() + mock_sleep.assert_not_called() + # Only the attempt run is polled (once, before the loop); the parent run is never fetched. + mock_databricks_hook.return_value.get_run.assert_called_once_with("attempt-1") + + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._databricks_workflow_task_group", + new_callable=mock.PropertyMock, + ) + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_monitor_workflow_task_reports_failure_once_retries_exhausted( + self, mock_get_current_task, mock_databricks_hook, mock_sleep, mock_workflow_tg + ): + mock_workflow_tg.return_value = MagicMock() + mock_get_current_task.return_value = {"run_id": "attempt-2", "attempt_number": 1} + mock_databricks_hook.return_value.get_run.return_value = { + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "final attempt failed", + }, + "run_page_url": "url", + } + + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + max_retries=1, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(AirflowException): + operator.monitor_databricks_job() + mock_sleep.assert_not_called() + mock_databricks_hook.return_value.get_run.assert_called_once_with("attempt-2") + + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._databricks_workflow_task_group", + new_callable=mock.PropertyMock, + ) + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksNotebookOperator._get_current_databricks_task" + ) + def test_monitor_workflow_task_unlimited_retries_waits_for_parent( + self, mock_get_current_task, mock_databricks_hook, mock_sleep, mock_workflow_tg + ): + mock_workflow_tg.return_value = MagicMock() + mock_get_current_task.return_value = {"run_id": "attempt-1", "attempt_number": 5} + parent_states = iter( + [ + {"life_cycle_state": "RUNNING"}, + {"life_cycle_state": "TERMINATED", "result_state": "FAILED", "state_message": "failed"}, + ] + ) + + def fake_get_run(run_id): + if run_id == 12345: + return {"state": next(parent_states), "run_page_url": "parent"} + return { + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "failed", + }, + "run_page_url": "url", + } + + mock_databricks_hook.return_value.get_run.side_effect = fake_get_run + + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + max_retries=-1, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(AirflowException): + operator.monitor_databricks_job() + mock_sleep.assert_called_once() + mock_databricks_hook.return_value.get_run.assert_any_call(12345) + + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_reshaped_submit_run_without_native_retries_resolves_task_for_monitoring( + self, mock_databricks_hook, mock_sleep + ): + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + min_retry_interval_millis=2000, + ) + run_json = operator._get_run_json() + assert "tasks" in run_json + assert run_json["tasks"][0]["task_key"] == operator.databricks_task_key + assert operator._resolved_max_retries() is None + + operator.databricks_run_id = 12345 + mock_databricks_hook.return_value.get_run_tasks.return_value = [ + {"task_key": operator.databricks_task_key, "run_id": "attempt-1", "start_time": 1} + ] + mock_databricks_hook.return_value.get_run.return_value = { + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "run_page_url": "url", + } + + operator.monitor_databricks_job() + mock_databricks_hook.return_value.get_run_tasks.assert_called_once_with(12345) + mock_databricks_hook.return_value.get_run.assert_called_once_with("attempt-1") + + @pytest.mark.parametrize( + "operator", + [ + DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="WORKSPACE", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + max_retries="{{ params.max_retries }}", + min_retry_interval_millis="{{ params.min_retry_interval_millis }}", + retry_on_timeout="{{ params.retry_on_timeout }}", + ), + DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + task_config={}, + max_retries="{{ params.max_retries }}", + min_retry_interval_millis="{{ params.min_retry_interval_millis }}", + retry_on_timeout="{{ params.retry_on_timeout }}", + ), + ], + ) + def test_retry_params_render_from_templates(self, operator): + """Retry fields are templatable: Jinja values render and coerce to typed retry settings. + + Guards the ``template_fields`` membership on both subclasses (which keep separate tuples): + an un-templated field would keep its literal ``{{ ... }}`` string and fail coercion. + """ + operator.render_template_fields( + context={ + "params": { + "max_retries": 3, + "min_retry_interval_millis": 2000, + "retry_on_timeout": "true", + } + } + ) + assert operator._retry_settings() == { + "max_retries": 3, + "min_retry_interval_millis": 2000, + "retry_on_timeout": True, + } + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") def test_launch_notebook_job(self, mock_databricks_hook): operator = DatabricksNotebookOperator( @@ -3914,6 +4284,88 @@ def test_convert_to_databricks_workflow_task(self): assert task_json == expected_json + @pytest.mark.parametrize( + ("retry_kwargs", "expected"), + [ + ( + {"max_retries": -1, "min_retry_interval_millis": 2000, "retry_on_timeout": True}, + {"max_retries": -1, "min_retry_interval_millis": 2000, "retry_on_timeout": True}, + ), + ({"max_retries": 0}, {"max_retries": 0}), + ({"retry_on_timeout": False}, {"retry_on_timeout": False}), + ], + ) + def test_get_run_json_retry_settings(self, retry_kwargs, expected): + """Retry settings are added to the submitted task only when explicitly provided. + + They must live inside ``tasks[0]`` (a Databricks ``SubmitTask``); at the top level of a + runs/submit payload Databricks silently ignores them. + """ + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + **retry_kwargs, + ) + run_json = operator._get_run_json() + assert "tasks" in run_json + task = run_json["tasks"][0] + assert task["task_key"] == operator.databricks_task_key + assert task["existing_cluster_id"] == "existing_cluster_id" + assert task["notebook_task"]["notebook_path"] == "test_path" + assert "existing_cluster_id" not in run_json + for key in ("max_retries", "min_retry_interval_millis", "retry_on_timeout"): + assert key not in run_json + if key in expected: + assert task[key] == expected[key] + else: + assert key not in task + + def test_get_run_json_without_retries_uses_legacy_top_level_shape(self): + """Without native retries the payload is the legacy single-task runs/submit shape (unchanged).""" + operator = DatabricksNotebookOperator( + task_id="test_task", + notebook_path="test_path", + source="test_source", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + ) + run_json = operator._get_run_json() + assert run_json["run_name"] == operator.databricks_task_key + assert run_json["existing_cluster_id"] == "existing_cluster_id" + assert run_json["notebook_task"]["notebook_path"] == "test_path" + # No reshape into the multi-task form when retries are not configured. + assert "tasks" not in run_json + + def test_convert_to_databricks_workflow_task_includes_retry_settings(self): + """Retry settings provided to the operator are included in the workflow task JSON.""" + dag = DAG(dag_id="example_dag", schedule=None, start_date=DEFAULT_DATE) + operator = DatabricksNotebookOperator( + notebook_path="/path/to/notebook", + source="WORKSPACE", + task_id="test_task", + max_retries=2, + min_retry_interval_millis=1000, + retry_on_timeout=True, + dag=dag, + ) + + databricks_workflow_task_group = MagicMock() + databricks_workflow_task_group.notebook_packages = [] + databricks_workflow_task_group.notebook_params = {} + + operator.task_group = databricks_workflow_task_group + relevant_upstreams = [] + task_dict = {} + + task_json = operator._convert_to_databricks_workflow_task(relevant_upstreams, task_dict) + + assert task_json["max_retries"] == 2 + assert task_json["min_retry_interval_millis"] == 1000 + assert task_json["retry_on_timeout"] is True + @pytest.mark.parametrize( ("trigger_rule", "expected_run_if"), [ @@ -4054,6 +4506,69 @@ def test_get_task_base_json(self): assert operator.task_config == task_config assert task_base_json == task_config + def test_get_run_json_operator_task_key_wins_over_task_config(self): + """A task_key in task_config must not shadow the operator-managed key used for monitoring. + + The operator only injects its own task_key in the reshaped (native-retry) payload, where + monitoring re-resolves the task by that key, so the guarantee is exercised with retries set. + """ + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + max_retries=2, + task_config={"task_key": "user_supplied_key", "notebook_task": {"notebook_path": "/p"}}, + ) + task = operator._get_run_json()["tasks"][0] + assert task["task_key"] == operator.databricks_task_key + + def test_convert_to_databricks_workflow_task_includes_task_config_retry_settings(self): + """Retry settings supplied via task_config surface in the workflow task JSON.""" + dag = DAG(dag_id="example_dag", schedule=None, start_date=DEFAULT_DATE) + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + task_config={ + "notebook_task": {"notebook_path": "/path"}, + "max_retries": 5, + "min_retry_interval_millis": 1000, + "retry_on_timeout": True, + }, + dag=dag, + ) + operator.task_group = MagicMock() + + task_json = operator._convert_to_databricks_workflow_task([], {}) + + assert task_json["max_retries"] == 5 + assert task_json["min_retry_interval_millis"] == 1000 + assert task_json["retry_on_timeout"] is True + + def test_convert_to_databricks_workflow_task_operator_retry_overrides_task_config(self): + """An operator-level retry value overrides the task_config value in the workflow task JSON.""" + dag = DAG(dag_id="example_dag", schedule=None, start_date=DEFAULT_DATE) + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + task_config={ + "notebook_task": {"notebook_path": "/path"}, + "max_retries": 5, + "min_retry_interval_millis": 1000, + "retry_on_timeout": False, + }, + max_retries=2, + min_retry_interval_millis=2000, + retry_on_timeout=True, + dag=dag, + ) + operator.task_group = MagicMock() + + task_json = operator._convert_to_databricks_workflow_task([], {}) + + assert task_json["max_retries"] == 2 + assert task_json["min_retry_interval_millis"] == 2000 + assert task_json["retry_on_timeout"] is True + def test_generate_databricks_task_key(self): task_config = {} operator = DatabricksTaskOperator( @@ -4154,3 +4669,134 @@ def test_on_kill_workflow_member_get_task_raises_does_not_cancel_parent(self, db ): operator.on_kill() db_mock.cancel_run.assert_not_called() + + @pytest.mark.parametrize( + ("field", "task_config_retries", "retry_kwargs", "expected"), + [ + ("max_retries", {"max_retries": 1}, {}, 1), + ("max_retries", {"max_retries": 1}, {"max_retries": 0}, 0), + ( + "min_retry_interval_millis", + {"min_retry_interval_millis": 1000}, + {"min_retry_interval_millis": 5000}, + 5000, + ), + ("retry_on_timeout", {"retry_on_timeout": True}, {"retry_on_timeout": False}, False), + ], + ) + def test_get_run_json_retry_settings(self, field, task_config_retries, retry_kwargs, expected): + """Operator-level retry settings are applied and take precedence over task_config values.""" + task_config = { + "notebook_task": {"notebook_path": "/path", "source": "WORKSPACE", "base_parameters": {}}, + **task_config_retries, + } + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + task_config=task_config, + **retry_kwargs, + ) + run_json = operator._get_run_json() + assert run_json["tasks"][0][field] == expected + assert field not in run_json + + @pytest.mark.parametrize("max_retries", [2, "2"]) + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_monitor_uses_retry_strategy_when_retries_set_in_task_config( + self, mock_databricks_hook, mock_sleep, max_retries + ): + # max_retries supplied only through task_config must still select retry-aware monitoring + # (follow the submit run), not the single-attempt path. + run_states = iter( + [ + {"life_cycle_state": "RUNNING"}, + {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + ] + ) + mock_databricks_hook.return_value.get_run.side_effect = lambda run_id: { + "state": next(run_states), + "run_page_url": "url", + } + + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + task_config={ + "notebook_task": {"notebook_path": "/path", "source": "WORKSPACE", "base_parameters": {}}, + "max_retries": max_retries, + }, + ) + operator.databricks_run_id = 12345 + + operator.monitor_databricks_job() + mock_databricks_hook.return_value.get_run.assert_called_with(12345) + mock_databricks_hook.return_value.get_run_tasks.assert_not_called() + + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + @mock.patch( + "airflow.providers.databricks.operators.databricks.DatabricksTaskOperator._get_current_databricks_task" + ) + def test_monitor_max_retries_zero_operator_arg_overrides_task_config_retries( + self, mock_get_current_task, mock_databricks_hook, mock_sleep + ): + mock_get_current_task.return_value = {"run_id": "attempt-1"} + mock_databricks_hook.return_value.get_run.return_value = { + "state": { + "life_cycle_state": "TERMINATED", + "result_state": "FAILED", + "state_message": "attempt failed", + }, + "run_page_url": "url", + } + + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + task_config={ + "notebook_task": {"notebook_path": "/path", "source": "WORKSPACE", "base_parameters": {}}, + "max_retries": 2, + }, + max_retries=0, + ) + operator.databricks_run_id = 12345 + + with pytest.raises(AirflowException): + operator.monitor_databricks_job() + mock_sleep.assert_not_called() + mock_databricks_hook.return_value.get_run.assert_called_once_with("attempt-1") + + @mock.patch("airflow.providers.databricks.operators.databricks.time.sleep") + @mock.patch("airflow.providers.databricks.operators.databricks.DatabricksHook") + def test_monitor_waits_out_waiting_for_retry_state(self, mock_databricks_hook, mock_sleep): + # A native retry surfaces WAITING_FOR_RETRY between attempts; the poll must keep waiting + # rather than crash on an unexpected life cycle state. + run_states = iter( + [ + {"life_cycle_state": "RUNNING"}, + {"life_cycle_state": "WAITING_FOR_RETRY"}, + {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + ] + ) + mock_databricks_hook.return_value.get_run.side_effect = lambda run_id: { + "state": next(run_states), + "run_page_url": "url", + } + + operator = DatabricksTaskOperator( + task_id="test_task", + databricks_conn_id="test_conn_id", + existing_cluster_id="existing_cluster_id", + task_config={ + "notebook_task": {"notebook_path": "/path", "source": "WORKSPACE", "base_parameters": {}}, + "max_retries": 2, + }, + ) + operator.databricks_run_id = 12345 + + operator.monitor_databricks_job() + assert mock_databricks_hook.return_value.get_run.call_count == 3 diff --git a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py index 8854eb03fb5bc..ad204297dc80e 100644 --- a/providers/databricks/tests/unit/databricks/triggers/test_databricks.py +++ b/providers/databricks/tests/unit/databricks/triggers/test_databricks.py @@ -189,6 +189,9 @@ def test_serialize(self): "run_page_url": RUN_PAGE_URL, "repair_run": False, "caller": "DatabricksExecutionTrigger", + "workflow_run_id": None, + "databricks_task_key": None, + "max_retries": None, }, ) @@ -306,12 +309,210 @@ async def test_sleep_between_retries( mock_sleep.assert_called_once() mock_sleep.assert_called_with(POLLING_INTERVAL_SECONDS) + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_output") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run") + @mock.patch("airflow.providers.databricks.triggers.databricks.asyncio.sleep") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_state") + async def test_run_waits_out_waiting_for_retry_state( + self, mock_get_run_state, mock_sleep, mock_get_run, mock_get_run_output + ): + # A native retry surfaces WAITING_FOR_RETRY between attempts; the trigger must keep polling + # rather than crash on an unexpected life cycle state. + mock_get_run_state.side_effect = [ + RunState(life_cycle_state="WAITING_FOR_RETRY", state_message="", result_state=""), + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"), + ] + mock_get_run.return_value = GET_RUN_RESPONSE_TERMINATED + mock_get_run_output.return_value = GET_RUN_OUTPUT_RESPONSE + + async for event in self.trigger.run(): + assert event == TriggerEvent( + { + "run_id": RUN_ID, + "run_state": RunState( + life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS" + ).to_json(), + "run_page_url": RUN_PAGE_URL, + "repair_run": False, + "errors": [], + } + ) + mock_sleep.assert_called_once_with(POLLING_INTERVAL_SECONDS) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.triggers.databricks.asyncio.sleep") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.get_run_tasks") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_state") + async def test_run_workflow_task_retry_in_flight_succeeds( + self, mock_get_run_state, mock_get_run_tasks, mock_sleep + ): + # First attempt fails while the parent run is still active; a later attempt succeeds and + # the trigger must emit a success event rather than failing on the first attempt. + mock_get_run_tasks.side_effect = [ + [{"run_id": TASK_RUN_ID1, "task_key": TASK_RUN_ID1_KEY, "start_time": 1}], + [ + {"run_id": TASK_RUN_ID1, "task_key": TASK_RUN_ID1_KEY, "start_time": 1}, + {"run_id": TASK_RUN_ID2, "task_key": TASK_RUN_ID1_KEY, "start_time": 2}, + ], + ] + mock_get_run_state.side_effect = [ + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED"), + RunState(life_cycle_state=LIFE_CYCLE_STATE_PENDING, state_message="", result_state=""), + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS"), + ] + + trigger = DatabricksExecutionTrigger( + run_id=TASK_RUN_ID1, + databricks_conn_id=DEFAULT_CONN_ID, + polling_period_seconds=POLLING_INTERVAL_SECONDS, + run_page_url=RUN_PAGE_URL, + workflow_run_id=RUN_ID, + databricks_task_key=TASK_RUN_ID1_KEY, + ) + + events = [event async for event in trigger.run()] + assert events == [ + TriggerEvent( + { + "run_id": TASK_RUN_ID2, + "run_page_url": RUN_PAGE_URL, + "run_state": RunState( + life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="SUCCESS" + ).to_json(), + "repair_run": False, + "errors": [], + } + ) + ] + mock_sleep.assert_called_once_with(POLLING_INTERVAL_SECONDS) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.triggers.databricks.asyncio.sleep") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_output") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.get_run_tasks") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_state") + async def test_run_workflow_task_failed_attempt_waits_for_parent( + self, mock_get_run_state, mock_get_run_tasks, mock_get_run, mock_get_run_output, mock_sleep + ): + mock_get_run_tasks.return_value = [ + {"run_id": TASK_RUN_ID1, "task_key": TASK_RUN_ID1_KEY, "start_time": 1} + ] + mock_get_run_state.side_effect = [ + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED"), + RunState(life_cycle_state="RUNNING", state_message="", result_state=""), + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED"), + RunState(life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED"), + ] + mock_get_run.return_value = GET_RUN_RESPONSE_TERMINATED_WITH_FAILED + mock_get_run_output.return_value = GET_RUN_OUTPUT_RESPONSE + + trigger = DatabricksExecutionTrigger( + run_id=TASK_RUN_ID1, + databricks_conn_id=DEFAULT_CONN_ID, + polling_period_seconds=POLLING_INTERVAL_SECONDS, + run_page_url=RUN_PAGE_URL, + workflow_run_id=RUN_ID, + databricks_task_key=TASK_RUN_ID1_KEY, + ) + + events = [event async for event in trigger.run()] + assert events == [ + TriggerEvent( + { + "run_id": TASK_RUN_ID1, + "run_page_url": RUN_PAGE_URL, + "run_state": RunState( + life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED" + ).to_json(), + "repair_run": False, + "errors": [ + {"task_key": TASK_RUN_ID1_KEY, "run_id": TASK_RUN_ID1, "error": ERROR_MESSAGE}, + {"task_key": TASK_RUN_ID3_KEY, "run_id": TASK_RUN_ID3, "error": ERROR_MESSAGE}, + ], + } + ) + ] + mock_sleep.assert_called_once_with(POLLING_INTERVAL_SECONDS) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_output") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run") + @mock.patch("airflow.providers.databricks.triggers.databricks.asyncio.sleep") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.get_run_tasks") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.a_get_run_state") + async def test_run_workflow_task_reports_failure_once_retries_exhausted( + self, mock_get_run_state, mock_get_run_tasks, mock_sleep, mock_get_run, mock_get_run_output + ): + mock_get_run_tasks.return_value = [ + {"run_id": TASK_RUN_ID1, "task_key": TASK_RUN_ID1_KEY, "start_time": 1, "attempt_number": 1} + ] + mock_get_run_state.return_value = RunState( + life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED" + ) + mock_get_run.return_value = GET_RUN_RESPONSE_TERMINATED_WITH_FAILED + mock_get_run_output.return_value = GET_RUN_OUTPUT_RESPONSE + + trigger = DatabricksExecutionTrigger( + run_id=TASK_RUN_ID1, + databricks_conn_id=DEFAULT_CONN_ID, + polling_period_seconds=POLLING_INTERVAL_SECONDS, + run_page_url=RUN_PAGE_URL, + workflow_run_id=RUN_ID, + databricks_task_key=TASK_RUN_ID1_KEY, + max_retries=1, + ) + + events = [event async for event in trigger.run()] + assert events == [ + TriggerEvent( + { + "run_id": TASK_RUN_ID1, + "run_page_url": RUN_PAGE_URL, + "run_state": RunState( + life_cycle_state=LIFE_CYCLE_STATE_TERMINATED, state_message="", result_state="FAILED" + ).to_json(), + "repair_run": False, + "errors": [ + {"task_key": TASK_RUN_ID1_KEY, "run_id": TASK_RUN_ID1, "error": ERROR_MESSAGE}, + {"task_key": TASK_RUN_ID3_KEY, "run_id": TASK_RUN_ID3, "error": ERROR_MESSAGE}, + ], + } + ) + ] + mock_sleep.assert_not_called() + mock_get_run_state.assert_called_once_with(TASK_RUN_ID1) + @pytest.mark.asyncio @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.cancel_run") async def test_on_kill_cancels_run(self, mock_cancel_run): await self.trigger.on_kill() mock_cancel_run.assert_called_once_with(RUN_ID) + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.get_run_tasks") + @mock.patch("airflow.providers.databricks.hooks.databricks.DatabricksHook.cancel_run") + async def test_on_kill_workflow_task_cancels_latest_attempt(self, mock_cancel_run, mock_get_run_tasks): + # The original attempt may be terminal after a retry; on_kill must cancel the latest attempt + # (under the same task_key) instead of the now-stale run_id the trigger was created with. + mock_get_run_tasks.return_value = [ + {"run_id": TASK_RUN_ID1, "task_key": TASK_RUN_ID1_KEY, "start_time": 1}, + {"run_id": TASK_RUN_ID2, "task_key": TASK_RUN_ID1_KEY, "start_time": 2}, + ] + trigger = DatabricksExecutionTrigger( + run_id=TASK_RUN_ID1, + databricks_conn_id=DEFAULT_CONN_ID, + polling_period_seconds=POLLING_INTERVAL_SECONDS, + run_page_url=RUN_PAGE_URL, + workflow_run_id=RUN_ID, + databricks_task_key=TASK_RUN_ID1_KEY, + ) + + await trigger.on_kill() + mock_get_run_tasks.assert_called_once_with(RUN_ID) + mock_cancel_run.assert_called_once_with(TASK_RUN_ID2) + class TestDatabricksSQLStatementExecutionTrigger: @pytest.fixture(autouse=True)