DatabricksNotebookOperator¶

Use the DatabricksNotebookOperator to launch and monitor notebook job runs on Databricks as Airflow tasks.

Examples¶

Running a notebook in Databricks on a new cluster¶

tests/system/databricks/example_databricks.py[source]

    new_cluster_spec = {
        "cluster_name": "",
        "spark_version": "11.3.x-scala2.12",
        "aws_attributes": {
            "first_on_demand": 1,
            "availability": "SPOT_WITH_FALLBACK",
            "zone_id": "us-east-2b",
            "spot_bid_price_percent": 100,
            "ebs_volume_count": 0,
        },
        "node_type_id": "i3.xlarge",
        "spark_env_vars": {"PYSPARK_PYTHON": "/databricks/python3/bin/python3"},
        "enable_elastic_disk": False,
        "data_security_mode": "LEGACY_SINGLE_USER_STANDARD",
        "runtime_engine": "STANDARD",
        "num_workers": 8,
    }

    notebook_1 = DatabricksNotebookOperator(
        task_id="notebook_1",
        notebook_path="/Shared/Notebook_1",
        notebook_packages=[
            {
                "pypi": {
                    "package": "simplejson==3.18.0",
                    "repo": "https://pypi.org/simple",
                }
            },
            {"pypi": {"package": "Faker"}},
        ],
        source="WORKSPACE",
        new_cluster=new_cluster_spec,
    )

Running a notebook in Databricks on an existing cluster¶

tests/system/databricks/example_databricks.py[source]

    notebook_2 = DatabricksNotebookOperator(
        task_id="notebook_2",
        notebook_path="/Shared/Notebook_2",
        notebook_packages=[
            {
                "pypi": {
                    "package": "simplejson==3.18.0",
                    "repo": "https://pypi.org/simple",
                }
            },
        ],
        source="WORKSPACE",
        existing_cluster_id="existing_cluster_id",
    )

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:

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 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 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.