kaxil commented on code in PR #76295:
URL: https://github.com/apache/airflow/pull/76295#discussion_r4443942012
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -2275,6 +2286,51 @@ def _monitor_submit_run(self) -> None:
errors = extract_failed_task_errors(self._hook, run, run_state)
self._handle_terminal_run_state(run_state, errors)
+ def submit_job(self, context: Context) -> int | None:
Review Comment:
`_launch_job` makes a second `runs/get` call for the page URL and pushes two
XComs before it returns, and the mixin only persists the run id once
`submit_job` returns. If `get_run_page_url` fails after `submit_run` succeeded,
the try fails with nothing stored and the retry submits a duplicate, which is
the case this PR sets out to close. The siblings return straight after
`submit_run` and push the XComs from the poll hook, so the restructure
suggested on `poll_until_complete` would cover this too.
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -2275,6 +2286,51 @@ def _monitor_submit_run(self) -> None:
errors = extract_failed_task_errors(self._hook, run, run_state)
self._handle_terminal_run_state(run_state, errors)
+ def submit_job(self, context: Context) -> int | None:
+ return self._launch_job(context=context)
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ self.databricks_run_id = cast("int", external_id)
+ try:
+ run = self._hook.get_run(external_id)
+ except DatabricksApiError as e:
+ # Run history expired or was deleted: neither active nor
succeeded, so the mixin
+ # resubmits fresh instead of failing the task.
+ if e.http_status_code == 404:
+ return "NOT_FOUND:"
+ raise
+ # Databricks splits run state across life_cycle_state and
result_state; the mixin's status
+ # interface is a single string, so encode both and decode them in the
predicates below.
+ state = RunState(**run["state"])
+ return f"{state.life_cycle_state}:{state.result_state}"
+
+ def is_job_active(self, status: str) -> bool:
+ return status.split(":", 1)[0] in (
+ "PENDING",
+ "QUEUED",
+ "RUNNING",
+ "TERMINATING",
+ "BLOCKED",
+ "WAITING_FOR_RETRY",
+ )
+
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == "TERMINATED:SUCCESS"
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
Review Comment:
Nothing pushes `run_id` / `run_page_url` when this reconnects. The mixin's
reconnect branch returns `poll_until_complete(...)` without calling
`get_job_result`, `_launch_job` doesn't run, and `ti_run` has already cleared
the crashed try's XComs. After a crash recovery the run link is empty and a
downstream `xcom_pull(key="run_id")` gets None, deferrable or not.
`DatabricksSubmitRunOperator.poll_until_complete` pushes them itself and sets
`_run_xcoms_pushed`. Can this do the same before `monitor_databricks_job`,
sharing the push with `get_job_result`?
##########
providers/databricks/docs/operators/notebook.rst:
##########
@@ -67,7 +67,8 @@ whole Airflow task:
)
Airflow ``retries`` behaves differently depending on where the operator runs.
For a standalone
-operator, each retry submits a new Databricks run. Inside a
+operator, a retry reconnects to the run it already submitted (see `Durable
execution`_ below).
Review Comment:
Isn't this only true when the worker died mid-run? If the Databricks run
itself failed, which is the usual reason for a retry, the next try submits a
new run, as the section below says. I'd say a retry reconnects only when the
previous try was lost while the run was still going. `task.rst` has the same
sentence.
##########
providers/databricks/tests/unit/databricks/operators/test_databricks.py:
##########
@@ -5267,3 +5267,105 @@ def
test_monitor_waits_out_waiting_for_retry_state(self, mock_databricks_hook, m
operator.monitor_databricks_job()
assert mock_databricks_hook.return_value.get_run.call_count == 3
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS, reason="task_state_store (durable execution)
requires Airflow 3.3+"
+)
+class TestDatabricksTaskBaseOperatorDurable:
Review Comment:
`DatabricksTaskBaseOperator` has its own `get_job_status` / `is_job_active`
/ `is_job_succeeded`, so the SubmitRun and RunNow durable classes don't cover
them. Dropping `"BLOCKED"` from `is_job_active`, or turning the 404 branch into
a `raise`, keeps this file green. Could this class get what the SubmitRun one
has: the state predicates parametrized over the lifecycle states, a 404 stored
run and a `TERMINATED:FAILED` stored run that both resubmit, `durable=False`
never touching the store, and xcom assertions on the fresh-submit and reconnect
tests? Nothing passes `durable=` to either subclass yet, so losing the
`durable=durable` forward would go unnoticed too. `_state` is also a copy of
`TestDatabricksSubmitRunOperatorDurable._state` and could move to module level.
##########
providers/databricks/docs/operators/notebook.rst:
##########
@@ -80,3 +81,35 @@ attempt is observed, so downstream failure handling is not
delayed by long-runni
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.
+
+
+Durable execution
+-----------------
+
+Outside a
+:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`,
+this operator submits a run and polls it to completion on the worker. By
default it runs in a
+*durable* mode: the Databricks run id is persisted to
+:doc:`task state store <apache-airflow:core-concepts/task-state-store>` before
polling begins, so if
+the worker crashes or is preempted and the task is retried, the operator
reconnects to the run that
+is already executing on Databricks instead of submitting a duplicate.
+
+On retry the operator checks the prior run's state:
+
+* if it is still running, the operator reconnects and continues polling
+* if it already succeeded, the operator returns immediately without
resubmitting
+* if it failed terminally, the operator submits a fresh run
+
+Pass ``durable=False`` to always submit a fresh run on retry.
+
+Durable execution requires Airflow 3.3 or newer, since it relies on the task
state store. On
+earlier Airflow versions the flag is a no-op and the operator always submits a
fresh run on retry.
+If the task state store is unavailable at runtime, the operator logs that
crash recovery is
+disabled and behaves the same way.
+
+Inside a ``DatabricksWorkflowTaskGroup`` nothing changes: the run is submitted
by the workflow
+launch task and this operator only monitors its own task within that shared
run, so a retry already
+reconnects rather than resubmitting.
+
+See :doc:`apache-airflow:core-concepts/resumable-tasks` for how clearing a
task interacts with the
+stored run id, and for the retention caveats around ``airflow state-store
clean``.
Review Comment:
`resumable-tasks` doesn't mention `airflow state-store clean` or retention
at all. That caveat lives in
`administration-and-deployment/task-and-asset-state-store-cleanup` and
`core-concepts/task-state-store`, and `submit_run.rst` states it inline (a
`retry_delay` longer than `default_retention_days` with cleanup in between
means the retry submits fresh). Could this inline that paragraph or point the
retention half at the cleanup page? Same in `task.rst`.
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks.py:
##########
@@ -2275,6 +2286,51 @@ def _monitor_submit_run(self) -> None:
errors = extract_failed_task_errors(self._hook, run, run_state)
self._handle_terminal_run_state(run_state, errors)
+ def submit_job(self, context: Context) -> int | None:
+ return self._launch_job(context=context)
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
Review Comment:
This is the third copy of `get_job_status` / `is_job_active` /
`is_job_succeeded` in this module, with the same bodies as
`DatabricksSubmitRunOperator` and `DatabricksRunNowOperator`. The
`LIFE_CYCLE:RESULT` string and the active-state tuple are a small protocol
between the three methods, so a state added to one copy later changes reconnect
behaviour for only some of the operators. Could they move into a module-level
helper or a small private mixin all three use?
##########
providers/databricks/docs/operators/notebook.rst:
##########
@@ -80,3 +81,35 @@ attempt is observed, so downstream failure handling is not
delayed by long-runni
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.
+
+
+Durable execution
+-----------------
+
+Outside a
+:class:`~airflow.providers.databricks.operators.databricks_workflow.DatabricksWorkflowTaskGroup`,
+this operator submits a run and polls it to completion on the worker. By
default it runs in a
Review Comment:
Unlike `DatabricksSubmitRunOperator` and `DatabricksRunNowOperator`,
`execute` here goes through `execute_resumable` with `deferrable=True` as well,
so the run id is stored before `_defer_on_run` and a retry still reconnects.
The sibling pages say `durable` has no effect in deferrable mode, and "polls it
to completion on the worker" isn't true once it defers. Is the difference
intended? If so, a sentence saying durable also covers deferrable mode here
would stop readers carrying the sibling pages' rule over. Same in `task.rst`.
--
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]
