Skip to content

Commit

Permalink
Add DatabricksWorkflowPlugin (apache#40153)
Browse files Browse the repository at this point in the history
The DatabricksWorkflowPlugin provides with links in the Airflow
UI for tasks that allow us to see the Databricks job run in the
Databricks workspace, additionally it also provides link to 
repair task(s) in the workflow.

Databricks does not allow repairing jobs with single tasks launched
outside the workflow, hence we just provide the link for the job run.
<img width="1342" alt="Screenshot 2024-06-24 at 4 05 53 PM" src="https://github.com/apache/airflow/assets/10206082/c1ded2b7-90fb-4a3c-980d-8043cc5a459f">

Within the workflow, for each of the task, we provide links to the
job run and repair link for the single task
<img width="1368" alt="Screenshot 2024-06-24 at 5 40 27 PM" src="https://github.com/apache/airflow/assets/10206082/d98081b8-8014-4d18-8747-e5b3382db416">

And at the workflow level, for the job launch task, we provide a
link to repair all failed tasks along with the link for job run in 
the Databricks workspace that can be used to monitor the job
in the Databricks account.
<img width="1421" alt="Screenshot 2024-06-24 at 5 40 56 PM" src="https://github.com/apache/airflow/assets/10206082/9d7ad7ae-9bbf-4fef-aa52-16ac5366edf3">


---------

Co-authored-by: Wei Lee <weilee.rx@gmail.com>
  • Loading branch information
pankajkoti and Lee-W authored Jul 9, 2024
1 parent 5323471 commit 22ec726
Show file tree
Hide file tree
Showing 14 changed files with 884 additions and 5 deletions.
32 changes: 29 additions & 3 deletions airflow/providers/databricks/operators/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,10 @@
DatabricksWorkflowTaskGroup,
WorkflowRunMetadata,
)
from airflow.providers.databricks.plugins.databricks_workflow import (
WorkflowJobRepairSingleTaskLink,
WorkflowJobRunLink,
)
from airflow.providers.databricks.triggers.databricks import DatabricksExecutionTrigger
from airflow.providers.databricks.utils.databricks import _normalise_json_content, validate_trigger_event

Expand Down Expand Up @@ -958,6 +962,15 @@ def __init__(

super().__init__(**kwargs)

if self._databricks_workflow_task_group is not None:
self.operator_extra_links = (
WorkflowJobRunLink(),
WorkflowJobRepairSingleTaskLink(),
)
else:
# Databricks does not support repair for non-workflow tasks, hence do not show the repair link.
self.operator_extra_links = (DatabricksJobRunLink(),)

@cached_property
def _hook(self) -> DatabricksHook:
return self._get_hook(caller=self.caller)
Expand Down Expand Up @@ -1016,12 +1029,17 @@ def _get_run_json(self) -> dict[str, Any]:
raise ValueError("Must specify either existing_cluster_id or new_cluster.")
return run_json

def _launch_job(self) -> int:
def _launch_job(self, context: Context | None = None) -> int:
"""Launch the job on Databricks."""
run_json = self._get_run_json()
self.databricks_run_id = self._hook.submit_run(run_json)
url = self._hook.get_run_page_url(self.databricks_run_id)
self.log.info("Check the job run in Databricks: %s", url)

if self.do_xcom_push and context is not None:
context["ti"].xcom_push(key=XCOM_RUN_ID_KEY, value=self.databricks_run_id)
context["ti"].xcom_push(key=XCOM_RUN_PAGE_URL_KEY, value=url)

return self.databricks_run_id

def _handle_terminal_run_state(self, run_state: RunState) -> None:
Expand All @@ -1040,7 +1058,15 @@ def _get_current_databricks_task(self) -> dict[str, Any]:
"""Retrieve the Databricks task corresponding to the current Airflow task."""
if self.databricks_run_id is None:
raise ValueError("Databricks job not yet launched. Please run launch_notebook_job first.")
return {task["task_key"]: task for task in self._hook.get_run(self.databricks_run_id)["tasks"]}[
tasks = self._hook.get_run(self.databricks_run_id)["tasks"]

# Because the task_key remains the same across multiple runs, and the Databricks API does not return
# tasks sorted by their attempts/start time, we sort the tasks by start time. This ensures that we
# map the latest attempt (whose status is to be monitored) of the task run to the task_key while
# building the {task_key: task} map below.
sorted_task_runs = sorted(tasks, key=lambda x: x["start_time"])

return {task["task_key"]: task for task in sorted_task_runs}[
self._get_databricks_task_id(self.task_id)
]

Expand Down Expand Up @@ -1125,7 +1151,7 @@ def execute(self, context: Context) -> None:
self.databricks_run_id = workflow_run_metadata.run_id
self.databricks_conn_id = workflow_run_metadata.conn_id
else:
self._launch_job()
self._launch_job(context=context)
if self.wait_for_termination:
self.monitor_databricks_job()

Expand Down
5 changes: 5 additions & 0 deletions airflow/providers/databricks/operators/databricks_workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,10 @@
from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.providers.databricks.hooks.databricks import DatabricksHook, RunLifeCycleState
from airflow.providers.databricks.plugins.databricks_workflow import (
WorkflowJobRepairAllFailedLink,
WorkflowJobRunLink,
)
from airflow.utils.task_group import TaskGroup

if TYPE_CHECKING:
Expand Down Expand Up @@ -88,6 +92,7 @@ class _CreateDatabricksWorkflowOperator(BaseOperator):
populated after instantiation using the `add_task` method.
"""

operator_extra_links = (WorkflowJobRunLink(), WorkflowJobRepairAllFailedLink())
template_fields = ("notebook_params",)
caller = "_CreateDatabricksWorkflowOperator"

Expand Down
16 changes: 16 additions & 0 deletions airflow/providers/databricks/plugins/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
Loading

0 comments on commit 22ec726

Please sign in to comment.