Add support for cancelling job when deferrable Databricks operator is killed

[akaul]

trafficstars

This pull request updates the behaviour of deferrable Databricks operators to handle manual terminations more gracefully when a Databricks operator running in deferrable mode is terminated (generally by a user manually restarting or failing a task).

To handle this scenario, we catch asyncio.CancelledError that's thrown, check if the job is in a non-terminal state and then proceed with cancelling it if it is. This takes a similar approach as other providers ([1], [2]) to gracefully terminate these jobs as the standard on_kill method on the operator doesn't work.

Unlike other operators, an explicit cancel_on_kill input hasn't been added here, but happy to include a new parameter for the trigger based on feedback.

See https://github.com/apache/airflow/issues/36090 for the general issue around this.

[1] https://github.com/apache/airflow/pull/38912 [2] https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/triggers/dataproc.py#L125-L138

---

^ Add meaningful description above Read the Pull Request Guidelines for more information. In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed. In case of a new dependency, check compliance with the ASF 3rd Party License Policy. In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

[boring-cyborg[bot]]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst) Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits. Apache Airflow is a community-driven project and together we are making it better 🚀. In case of doubts contact the developers at: Mailing List: [email protected] Slack: https://s.apache.org/airflow-slack

[RNHTTR]

Won't there need to be an update to the operators too, since they currently just call the sync cancel_run method? (DatabricksSubmitRunOperator | DatabricksRunNowOperator)?

I think either that or the hook.cancel_run can accept a deferrable parameter which can be used to call a_cancel_run.

[akaul]

Won't there need to be an update to the operators too, since they currently just call the sync cancel_run method? (DatabricksSubmitRunOperator | DatabricksRunNowOperator)?

I think either that or the hook.cancel_run can accept a deferrable parameter which can be used to call a_cancel_run.

I don't believe so since the trigger executes in a completely different context than the operators.

The changes here are specifically to account for the following scenario:

• A DatabricksSubmitRunOperator/DatabricksRunNowOperator is launched with deferrable=True
• After successfully creating a Databricks job, execution is deferred by running the DatabricksExecutionTrigger
• A user goes in and manually marks the job as terminated in the UI
  • Trigger execution is cancelled when the operator is a deferred state
  • Based on local testing as well as the explanation in this comment, the trigger gets terminated and control flow is never transferred back to the operator. Due to this on_kill is never called.

Currently, terminating the task via the UI doesn't propagate that information back to Databricks that the job should be terminated in this situation. From my testing locally, the changes in this PR should remedy the situation.

Open to other implementation alternatives, but this is a pretty similar approach as other providers as I noted in the PR description. I did initially attempt to implement a Trigger.cleanup() method, but that specific function can also be executed within the context of a triggerer restart which isn't ideal behaviour since the Databricks job should continue running across triggerer restarts, especially for long-running Databricks jobs.

[thesuperzapper]

@akaul @RNHTTR this PR is NOT safe, and needs to be updated.

Please see my https://github.com/apache/airflow/issues/36090#issuecomment-2094972855 for more information.

[akaul]

@thesuperzapper Thanks for the information! This feels like it should be tackled at the Airflow core level as a high priority bug rather than each provider implementing their own solution for this. The workaround you've proposed makes sense, but it feels like too much overhead for each provider to implement in their triggers when the on_kill hook is present on operator defining the behaviour under these circumstances.

[RNHTTR]

@thesuperzapper Thanks for the information! This feels like it should be tackled at the Airflow core level as a high priority bug rather than each provider implementing their own solution for this. The workaround you've proposed makes sense, but it feels like too much overhead for each provider to implement in their triggers when the on_kill hook is present on operator defining the behaviour under these circumstances.

I agree, but I believe the fix in core Airflow is not straightforward. @sunank200 can you provide some additional info with regard to the level of difficulty in core Airflow?

[akaul]

PR's been updated to inspect for the task instance's state before cancelling the Databricks job.