Repository: astronomer/astronomer-cosmos (Python)

Issue:
Link DatabricksWorkflowTaskgroup timeout to databricks jobs timeout
Description:
To prevent infinitely hanging jobs, it would be useful to have the ability to program timeouts for Airflow tasks that are submitted to Databricks. Currently, Databricks jobs do not have timeouts by default which means that a databricks job could run forever even if the Airflow task times out. By linking Airflow task timeouts to Databricks timeouts, we can ensure that long-running tasks do not block other jobs and resources on the platform.

Proposed solution:

Add a configuration option to DatabricksWorkflowTaskGroup to specify a timeout value. When the Databricks Job is submitted, Airflow should include the timeout value as part of the request to Databricks. This will allow Databricks to enforce the timeout and kill the task if it runs for too long.

Potential challenges:

* Ensuring that the timeout value is propagated correctly from Airflow to Databricks.
* Handling cases where the task takes longer to run than the specified timeout (e.g. due to slow network or data access).
Determining an appropriate default timeout value that works for most workflows.


Discussion: @pankajkoti and I discussed this ticket and Pankaj is going to do a spike on 

1) tying the airflow timeout to the databricks job timeout when submitting
2) sending a "kill" message to databricks if the task takes longer than the timeout

and report findings here cc: @tatiana 
**Approach 1: Tying the airflow timeout to the Databricks job timeout when submitting**
The approach would use the execution_timeout param available with the BaseOperator and pass it to the Databricks task which accepts `timeout_seconds` as a param in the payload.

_Pros:_
1. Ensures that the Databricks job will eventually and definitely terminate after the set timeout has elapsed and that the job won't run forever as the timeout is set with the Databricks task specification when triggered.

_Cons:_
1. There might be a slight delay between Airflow task and Databricks task timeout. The Airflow task will be shown as failed first as it launches first and within which we trigger the Databricks task. Ideally, this would be a relatively small number of time units.


**Approach 2: Sending a "kill" message to Databricks if the task takes longer than the timeout**
In Airflow, there is no callback on timeout implementation available, but it raises an [AirflowTaskTimeout exception](https://airflow.apache.org/docs/apache-airflow/2.2.0/concepts/tasks.html#timeouts) which means for each task we will need to catch the exception and then trigger a client API call to cancel the job run.

_Pros:_
Can't think of any compared to the first approach

_Cons:_
1. There would still be a delay as we're triggering cancel after getting the exception
2. The [cancel API](https://docs.databricks.com/api-explorer/workspace/jobs/cancelrun) works asynchronously and the docs don't mention if it is a hard termination or graceful termination of the job. Ideally, the job should be cancelled irrespective and immediately but the docs don't give a hint about it
3. The Databricks API might not be available when we trigger the cancellation but it was available while creating the job. So, we might need to handle retries and make sure that the job does not remain running forever, thus adding some more implementation challenges.
4. The [cancel API](https://docs.databricks.com/api-explorer/workspace/jobs/cancelrun) only accepts runId as a payload param. Thus we would not be able to reflect the cancellation reason as Timed Out whereas, with the first approach, the job gives a status `Timed Out` upon timeout seconds elapsing


**Considering the above, I am more inclined towards approach 1 of using the timeout parameter.**
Also, `DatabricksWorkflowTaskGroup` which corresponds to the Databricks job (that encompasses notebook tasks), extends Airflow's `TaskGroup` which does not have a parameter called `execution_timeout`. Thus, **Airflow cannot apply and represent timeouts** for a `TaskGroup`. 

In Airflow, however, `tasks` have `execution_timeout` that becomes available with the `BaseOperator` which is extended by the `DatabricksNotebookOperator`. 

I believe we should support `timeout` param for the Notebook operator tasks and not at the Workflow level Taskgroup.

WDYT @dimberman @tatiana  ?



Fix (patch):
diff --git a/cosmos/providers/databricks/notebook.py b/cosmos/providers/databricks/notebook.py
index 7ce16cad3..b674ba745 100644
--- a/cosmos/providers/databricks/notebook.py
+++ b/cosmos/providers/databricks/notebook.py
@@ -91,6 +91,23 @@ def __init__(
 
         super().__init__(**kwargs)
 
+    def _get_task_base_json(self) -> dict[str, Any]:
+        """Get task base json to be used for task group tasks and single task submissions."""
+        return {
+            # Timeout seconds value of 0 for the Databricks Jobs API means the job runs forever.
+            # That is also the default behavior of Databricks jobs to run a job forever without a default timeout value.
+            "timeout_seconds": int(self.execution_timeout.total_seconds())
+            if self.execution_timeout
+            else 0,
+            "email_notifications": {},
+            "notebook_task": {
+                "notebook_path": self.notebook_path,
+                "source": self.source,
+                "base_parameters": self.notebook_params,
+            },
+            "libraries": self.notebook_packages,
+        }
+
     def convert_to_databricks_workflow_task(
         self, relevant_upstreams: list[BaseOperator]
     ):
@@ -99,6 +116,7 @@ def convert_to_databricks_workflow_task(
         """
         if hasattr(self.task_group, "notebook_packages"):
             self.notebook_packages.extend(self.task_group.notebook_packages)
+        base_task_json = self._get_task_base_json()
         result = {
             "task_key": self._get_databricks_task_id(self.task_id),
             "depends_on": [
@@ -107,14 +125,7 @@ def convert_to_databricks_workflow_task(
                 if t in relevant_upstreams
             ],
             "job_cluster_key": self.job_cluster_key,
-            "timeout_seconds": 0,
-            "email_notifications": {},
-            "notebook_task": {
-                "notebook_path": self.notebook_path,
-                "source": self.source,
-                "base_parameters": self.notebook_params,
-            },
-            "libraries": self.notebook_packages,
+            **base_task_json,
         }
         return result
 
@@ -186,13 +197,10 @@ def _get_api_client(self):
     def launch_notebook_job(self):
         """Launch the notebook as a one-time job to Databricks."""
         api_client = self._get_api_client()
+        base_task_json = self._get_task_base_json()
         run_json = {
             "run_name": self._get_databricks_task_id(self.task_id),
-            "notebook_task": {
-                "notebook_path": self.notebook_path,
-                "base_parameters": {"source": self.source},
-            },
-            "libraries": self.notebook_packages,
+            **base_task_json,
         }
         if self.new_cluster and self.existing_cluster_id:
             raise ValueError(
diff --git a/examples/databricks/example_databricks_workflow.py b/examples/databricks/example_databricks_workflow.py
index f65fad2a6..b8aaf5f9a 100644
--- a/examples/databricks/example_databricks_workflow.py
+++ b/examples/databricks/example_databricks_workflow.py
@@ -70,6 +70,7 @@
             notebook_packages=[{"pypi": {"package": "Faker"}}],
             source="WORKSPACE",
             job_cluster_key="Shared_job_cluster",
+            execution_timeout=timedelta(seconds=600),
         )
         notebook_2 = DatabricksNotebookOperator(
             task_id="notebook_2",