Jubin Soni - Portfolio & Blog

Airflow is powerful but it comes with a tax: a server to maintain, a DAG repo to manage, a scheduler to babysit, and a support burden every time a new team member needs to write their first DAG. For teams already living in Databricks, Workflows gives you 80% of what Airflow does with none of that overhead — and for ML pipelines in particular, the native integration with MLflow, Delta Lake, and cluster policies makes it the cleaner choice.

In this post I'll walk through building a production ML pipeline with Databricks Workflows — feature engineering, model training, validation, and deployment — orchestrated as a multi-task job with dependencies, conditional branching, and failure handling.

Architecture Overview

Task Execution Flow

Workflows vs Alternatives

Feature	Databricks Workflows	Apache Airflow	Azure Data Factory
Setup overhead	None (built-in)	High (server, DAG repo)	Medium (ARM templates)
MLflow integration	Native	Plugin / custom	None
Delta Lake awareness	Native	None	Limited
Cluster management	Managed per task	External	Managed
Conditional branching	`if_else_condition` task	BranchPythonOperator	If-Else activity
Task value passing	`taskValues` API	XCom	Pipeline parameters
Repair and re-run	Per-task repair	Full DAG re-run	Per-activity re-run
Cost	Included in Databricks	Infra + licensing	Pay per activity run

Step 1 — Define the Workflow as Code (DABs YAML)

Always define workflows as code, never click through the UI. This gives you version control, PR reviews, and repeatable deployments across environments.

yaml

# resources/jobs/ml_pipeline.yml
resources:
  jobs:
    ml_pipeline:
      name: ml-pipeline-${var.env}
      schedule:
        quartz_cron_expression: "0 0 2 * * ?"   # 2am daily
        timezone_id: "UTC"
        pause_status: UNPAUSED

      email_notifications:
        on_failure:
          - data-engineering@yourcompany.com
        on_success:
          - ml-team@yourcompany.com

      tasks:
        - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/01_feature_engineering.py
            base_parameters:
              env: ${var.env}
          job_cluster_key: standard_cluster

        - task_key: model_training
          depends_on:
            - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/02_train_and_register.py
          job_cluster_key: ml_cluster

        - task_key: model_validation
          depends_on:
            - task_key: model_training
          notebook_task:
            notebook_path: ./notebooks/03_validate_model.py
          job_cluster_key: standard_cluster

        - task_key: promote_model
          depends_on:
            - task_key: model_validation
          condition_task:
            op: "=="
            left: "{{tasks.model_validation.values.validation_passed}}"
            right: "true"

        - task_key: update_serving_endpoint
          depends_on:
            - task_key: promote_model
          notebook_task:
            notebook_path: ./notebooks/04_update_serving.py
          job_cluster_key: standard_cluster

        - task_key: alert_on_failure
          depends_on:
            - task_key: model_validation
          condition_task:
            op: "=="
            left: "{{tasks.model_validation.values.validation_passed}}"
            right: "false"

      job_clusters:
        - job_cluster_key: standard_cluster
          new_cluster:
            spark_version: 14.3.x-scala2.12
            node_type_id: Standard_DS3_v2
            num_workers: 4
            azure_attributes:
              availability: SPOT_WITH_FALLBACK_AZURE
              first_on_demand: 1

        - job_cluster_key: ml_cluster
          new_cluster:
            spark_version: 14.3.x-ml-scala2.12   # ML runtime with libraries
            node_type_id: Standard_DS4_v2
            num_workers: 4
            azure_attributes:
              availability: SPOT_WITH_FALLBACK_AZURE
              first_on_demand: 1

Step 2 — Passing Values Between Tasks with taskValues

taskValues is how tasks communicate in Databricks Workflows. It's the native alternative to Airflow's XCom — strongly typed and scoped to the job run.

python

# notebooks/01_feature_engineering.py
from delta.tables import DeltaTable

# ... feature engineering logic ...

# Capture Delta version and pass downstream
gold_table    = DeltaTable.forName(spark, 'churn.gold.features')
delta_version = gold_table.history(1).select('version').collect()[0][0]
feature_count = spark.table('churn.gold.features').count()

# Set task output values
dbutils.jobs.taskValues.set(key='delta_version',  value=str(delta_version))
dbutils.jobs.taskValues.set(key='feature_count',  value=str(feature_count))
dbutils.jobs.taskValues.set(key='feature_date',   value=str(date.today()))

print(f"Feature engineering complete. Delta version: {delta_version}, rows: {feature_count}")

python

# notebooks/02_train_and_register.py
import mlflow

# Read values from upstream task
delta_version = dbutils.jobs.taskValues.get(
    taskKey='feature_engineering', key='delta_version'
)
feature_date  = dbutils.jobs.taskValues.get(
    taskKey='feature_engineering', key='feature_date'
)

print(f"Training on features from Delta v{delta_version} ({feature_date})")

# ... training logic ...

with mlflow.start_run() as run:
    mlflow.log_param('delta_feature_version', delta_version)
    mlflow.log_param('feature_date',          feature_date)
    # ... log metrics and model ...

    # Pass run_id to validation task
    dbutils.jobs.taskValues.set(key='mlflow_run_id',    value=run.info.run_id)
    dbutils.jobs.taskValues.set(key='model_version',    value=str(model_version))

Step 3 — Validation Task with Conditional Branching Output

The validation task sets a taskValue that the workflow's condition_task reads to decide whether to promote or alert.

python

# notebooks/03_validate_model.py
import mlflow
from mlflow import MlflowClient

client = MlflowClient()

run_id        = dbutils.jobs.taskValues.get(taskKey='model_training', key='mlflow_run_id')
model_version = dbutils.jobs.taskValues.get(taskKey='model_training', key='model_version')

THRESHOLDS = {'roc_auc': 0.80, 'f1_score': 0.72, 'precision': 0.70}

run     = client.get_run(run_id)
metrics = run.data.metrics

failures = []
for metric, threshold in THRESHOLDS.items():
    actual = metrics.get(metric, 0.0)
    status = 'PASS' if actual >= threshold else 'FAIL'
    print(f"  {metric}: {actual:.4f} (>= {threshold}) -> {status}")
    if actual < threshold:
        failures.append(f"{metric}={actual:.4f}")

validation_passed = len(failures) == 0

# Set the value that condition_task reads
dbutils.jobs.taskValues.set(
    key='validation_passed',
    value='true' if validation_passed else 'false'
)
dbutils.jobs.taskValues.set(
    key='failure_reasons',
    value=', '.join(failures) if failures else ''
)

if not validation_passed:
    print(f"Validation FAILED: {', '.join(failures)}")
    # Don't raise — let condition_task handle branching
else:
    # Promote model to Production alias
    client.set_registered_model_alias(
        name='churn-prediction-gbm',
        alias='Production',
        version=model_version,
    )
    print(f"Model v{model_version} promoted to Production.")

Step 4 — Repair and Re-run Failed Tasks

One of the best Workflows features is per-task repair. If the training task fails, you fix the issue and re-run from training — not from feature engineering. This saves significant time and compute on long pipelines.

python

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import RepairRunRequest

w = WorkspaceClient()

# Find the last failed run
runs = w.jobs.list_runs(job_id=YOUR_JOB_ID, limit=5)
failed_run = next((r for r in runs if r.state.result_state.value == 'FAILED'), None)

if failed_run:
    print(f"Repairing run {failed_run.run_id}")
    print(f"Failed tasks: {[t.task_key for t in failed_run.tasks if t.state.result_state.value == 'FAILED']}")

    repair = w.jobs.repair_run(
        run_id=failed_run.run_id,
        rerun_tasks=["model_training", "model_validation", "promote_model"],
    )
    print(f"Repair run submitted: {repair.repair_id}")

Things to Watch in Production

Use MLRuntime for ML tasks, standard runtime for ETL. The ML runtime (14.3.x-ml-scala2.12) comes pre-installed with MLflow, scikit-learn, PyTorch, and XGBoost. Using it for ETL tasks wastes startup time loading libraries you don't need. Separate cluster keys per task type.

taskValues are strings only. The taskValues API stores everything as a string. Cast numeric values explicitly after reading them (int(dbutils.jobs.taskValues.get(...))), or you'll get type errors downstream.

Set max_retries on flaky tasks, not on the whole job. Network-dependent tasks (model registry calls, endpoint updates) sometimes fail transiently. Set max_retries: 2 on those specific tasks rather than wrapping the entire job in retry logic.

Alert on task duration, not just failure. A training task that normally takes 20 minutes but ran for 3 hours means something is wrong even if it eventually succeeded. Set up duration-based alerts in Workflows notification settings.

Wrapping Up

Databricks Workflows is the right orchestration choice when your pipeline lives entirely in Databricks. The native taskValues API, per-task cluster configs, conditional branching, and per-task repair make it significantly more productive than maintaining a separate Airflow deployment — and the MLflow and Delta Lake integration means you get reproducibility for free rather than wiring it up yourself.

Databricks Workflows: Orchestrating Multi-Task ML Pipelines Without the Airflow Overhead

Architecture Overview

Task Execution Flow

Workflows vs Alternatives

Step 1 — Define the Workflow as Code (DABs YAML)

Step 2 — Passing Values Between Tasks with taskValues

Step 3 — Validation Task with Conditional Branching Output

Step 4 — Repair and Re-run Failed Tasks

Things to Watch in Production

Wrapping Up

References