Databricks To Zenml Migration

Migrate Databricks Workflows to ZenML

This skill translates Databricks Workflows (Lakeflow Jobs) into idiomatic ZenML pipelines. It handles the full migration workflow: analyzing job definitions and notebook code, classifying each pattern, translating what maps cleanly, flagging what needs redesign, and producing a working ZenML project.

How migration works at a high level

Databricks Workflows and ZenML look similar on the surface -- both define a DAG of tasks/steps with dependencies. But the underlying execution models are fundamentally different. Databricks models orchestration as an explicit DAG of task objects in JSON (each with a task_key, a concrete task type like notebook_task or sql_task, and an explicit depends_on list), with substantial runtime configuration co-located in task settings (compute binding, retries, notifications). ZenML models orchestration as Python function calls forming typed artifact edges, with runtime behavior driven by step/pipeline decorators, stack components, and containerization settings.

This means migration involves two distinct challenges:

1. Structural translation: mapping Databricks' JSON-defined DAG + heterogeneous task types into Python-defined ZenML steps and pipelines
2. Semantic translation: handling the differences in data passing (string-substituted task values vs typed artifacts), execution environment (managed Spark clusters vs containerized steps), and platform-coupled features (Unity Catalog, DBFS, dbutils)

The three mapping types

Every Databricks concept falls into one of these categories:

| Type | Meaning | Action | |------|---------|--------| | Direct | Clean 1:1 mapping exists | Translate automatically | | Approximate | Conceptual equivalent exists but semantics differ | Translate with caveats noted in migration report | | Absent | No ZenML equivalent | Flag for human review with redesign suggestions |

See [references/concept-map.md](references/concept-map.md) for the full mapping tables.

The Migration Workflow

Phase 1: Receive and Analyze the Databricks Workflow

Ask the user for their Databricks job definition (JSON) and any associated notebook/script code. Databricks workflows come in multiple forms -- the user might provide:

• A Jobs API 2.1 JSON definition (the most complete representation)
• A Databricks Asset Bundle YAML (databricks.yml or resource YAML files)
• Notebook code (Python notebooks with dbutils calls, magics, widgets)
• A mix of job JSON + notebook source files

Read everything thoroughly before doing anything else. For each job, identify:

1. Tasks and their types -- What task types are used? (notebook_task, python_wheel_task, spark_python_task, sql_task, dbt_task, spark_jar_task, condition_task, for_each_task, run_job_task, pipeline_task)
2. Dependencies -- How are tasks wired? (depends_on with optional outcome conditions)
3. Data flow -- Where are dbutils.jobs.taskValues used? Are dynamic references ({{tasks..values.}}) used in parameters?
4. Control flow -- Any condition_task nodes? for_each_task iteration? run_if settings beyond ALL_SUCCESS?
5. Notebook analysis -- For each notebook_task: does it use dbutils.widgets, %sql, %pip, %run, display(), Spark temp views, DBFS paths? (See the notebook classification guide in [references/gaps-and-flags.md](references/gaps-and-flags.md))
6. Scheduling/triggers -- Cron schedules? Periodic triggers? File arrival triggers? Table update triggers? Continuous mode? If cron is present, record both the Quartz expression and the Databricks timezone.
7. Error handling -- max_retries, min_retry_interval_millis, timeout_seconds, retry_on_timeout?
8. Compute and libraries -- Job clusters, existing clusters, serverless? Per-task libraries[] entries (whl, pypi, maven, jar, requirements), workspace files, or DBFS-hosted wheels?
9. Parameters and config -- Job-level parameters with {{job.parameters.*}} references? Widget-based parameter passing? Which values are business parameters vs environment-specific settings?
10. Access control -- Job ACLs, Run-as identity, secret scopes?
11. Feature engineering / Unity Catalog -- Databricks Feature Engineering Client, Feature Store / Feature Engineering APIs, Unity Catalog feature tables, point-in-time joins, or feature lookup specs?

Phase 2: Classify and Plan

For each component identified in Phase 1, classify it using the mapping type (direct / approximate / absent). Use the decision logic below and the full tables in [references/concept-map.md](references/concept-map.md).

Quick classification guide

Direct translations (translate automatically):

• Multi-task job DAG structure → @pipeline with step calls matching depends_on edges
• task_key → step function name
• max_retries + min_retry_interval_millis → StepRetryConfig(max_retries=N, delay=M)
• dbutils.jobs.taskValues.set()/get() for simple data → step output/input artifacts
• on_success_callback / on_failure notifications → step hooks (on_success, on_failure)

Approximate translations (translate with caveats):

• notebook_task → @step wrapping refactored notebook logic (magics, dbutils, Spark session must be refactored)
• python_wheel_task → @step calling the wheel's entry point function directly (or containerized with Docker settings)
• sql_task → @step executing SQL via explicit client/connector (Databricks SQL connector + ZenML secrets)
• dbt_task → @step running dbt CLI in a container with explicit credentials
• condition_task → conditional pipeline logic (parameter-based for static pipelines, @pipeline(dynamic=True) + .load() for runtime values)
• for_each_task → @pipeline(dynamic=True) + .map() (concurrency is orchestrator-dependent)
• run_job_task → pipeline composition or API-triggered pipeline run
• Job parameters ({{job.parameters.*}}) → typed Python pipeline parameters populated from ZenML YAML configs
• Widget parameters (dbutils.widgets.get()) → step function parameters, with values supplied from pipeline config whenever they are business/configuration values
• Cron scheduling → Schedule(cron_expression=...) (orchestrator-dependent; Databricks orchestrator supports cron only and requires DatabricksOrchestratorSettings(schedule_timezone=""), e.g. UTC, America/New_York, or America/Los_Angeles)
• Job clusters / compute → Databricks-specific orchestrator or step-operator settings (DatabricksOrchestratorSettings / DatabricksStepOperatorSettings); generic ResourceSettings captures intent for many orchestrators but does not size Databricks step-operator clusters
• dbutils.secrets.get() → ZenML secrets store
• Per-task libraries → DockerSettings(requirements=[...]), pyproject.toml, or a private package/index strategy for workspace/DBFS wheel dependencies
• Databricks Feature Engineering Client / Unity Catalog feature lookups → Databricks-native feature access pattern or explicit redesign; do not blindly rewrite point-in-time feature lookup logic as ordinary SQL

Absent / needs redesign (flag for human review):

• run_if with ALL_DONE, AT_LEAST_ONE_FAILED, etc. (ZenML has pipeline-level execution modes but not per-step run_if)
• File arrival/table update triggers (Unity Catalog integration; no ZenML OSS equivalent, and current ZenML Pro platform-event triggers react to ZenML platform events rather than Databricks file/table events)
• Continuous jobs (always-on/restarting streaming workloads, not a run-to-completion ZenML pipeline pattern)
• Notebooks relying on %run, %pip, %sql magics, DBFS mounts, or shared Spark temp views across tasks
• Shared cluster state (cached tables, driver-local files, warm Spark context reused across tasks)
• DBFS-specific filesystem paths passed between tasks
• SQL/dbt tasks relying on Databricks-managed identity injection without portable auth

Present the migration plan

Before writing any code, present a summary to the user:

> "Here's what I found in your Databricks Workflow: > - Direct translations (will migrate cleanly): [list] > - Approximate translations (will work but with noted caveats): [list] > - Needs redesign (cannot auto-migrate): [list with brief explanation] > > Shall I proceed with the migration?"

If there are HIGH-severity flags, explain each one concretely: what the Databricks code does, why ZenML can't replicate it directly, and what the recommended redesign looks like.

Phase 3: Generate ZenML Code

Translate the Databricks Workflow into a ZenML project. Follow these conventions strictly.

Project structure

Every migrated project MUST use this layout:

migrated_pipeline/
├── steps/                    # One file per step
│   ├── extract.py
│   ├── transform.py
│   └── load.py
├── pipelines/
│   └── my_pipeline.py        # Pipeline definition
├── materializers/            # Custom materializers (if needed)
├── configs/
│   ├── dev.yaml
│   └── prod.yaml
├── run.py                    # CLI entry point (argparse, not click)
├── README.md
└── pyproject.toml

This matches the zenml-pipeline-authoring skill's conventions. Key rules:

• One step per file in steps/
• Separate pipeline definition from execution
• run.py uses argparse only for config selection and operational flags (for example --config, --no-cache, --dry-run), not as the main business-parameter surface
• pyproject.toml with zenml>=0.94.1 and requires-python = ">=3.12"
• Always generate populated configs/dev.yaml AND configs/prod.yaml (minimum two configs) containing business parameters plus step/pipeline/orchestrator settings discovered during migration
• Always generate a README.md explaining the migrated pipeline, how to run it, and what requires manual attention
• Include a brief ASCII DAG diagram in the pipeline file's module docstring showing the step dependency graph
• Run zenml init at project root

Configuration and CLI conventions

Prefer ZenML YAML configs as the migrated pipeline's main control surface. Put business parameters (dates, table names, feature table names, model hyperparameters), step settings, Docker settings, resource settings, schedules, and orchestrator-specific settings into configs/dev.yaml and configs/prod.yaml. The run.py entry point should mostly select which config to use and set operational flags; it should not recreate the Databricks job parameter system with a long argparse list.

A good mental model: Databricks job JSON held both the DAG and its knobs; in ZenML, Python should define the DAG, while YAML should hold the knobs that change between environments or runs.

Translation patterns

For each Databricks task, apply the appropriate translation. See [references/code-patterns.md](references/code-patterns.md) for detailed side-by-side examples covering all major patterns.

The core translation rule: Extract the task's logic (from notebook cells, wheel entry points, or SQL files) into a @step function. Type-hint all inputs and outputs. Wire steps by passing outputs to inputs in the pipeline function.

# Databricks: notebook_task with widget parameters
# dbutils.widgets.get("input_table") inside notebook

# ZenML: explicit typed parameters
@step
def extract(input_table: str, run_date: str) -> pd.DataFrame:
    # Replace Spark table read with portable data access
    return load_from_warehouse(input_table, run_date)

Task values → Artifact passing: Replace all dbutils.jobs.taskValues.set()/get() and {{tasks..values.}} references with direct function-call wiring:

# Databricks: task value set in producer, string-substituted in consumer
# dbutils.jobs.taskValues.set(key="count", value=42)
# Consumer: base_parameters: {"count": "{{tasks.producer.values.count}}"}

# ZenML: data flows naturally through function calls
@pipeline
def my_pipeline() -> None:
    count = producer_step()         # Returns int artifact
    consumer_step(count=count)      # Artifact passed directly

Retries: Map max_retries + min_retry_interval_millis to StepRetryConfig:

# Databricks: "max_retries": 3, "min_retry_interval_millis": 60000
# ZenML:
@step(retry=StepRetryConfig(max_retries=3, delay=60, backoff=1))
def my_step() -> None: ...

Notifications → Hooks: Map task-level email_notifications / webhook_notifications to ZenML hooks:

from zenml.hooks import alerter_failure_hook, alerter_success_hook

@step(on_failure=alerter_failure_hook, on_success=alerter_success_hook)
def my_step() -> None: ...

Scheduling: Map cron schedules to Schedule, and preserve the timezone separately when the target stack uses the Databricks orchestrator:

from zenml.config.schedule import Schedule
from zenml.integrations.databricks.flavors.databricks_orchestrator_flavor import (
    DatabricksOrchestratorSettings,
)

# Databricks: quartz_cron_expression "0 0 2 * * ?" (simple Quartz 6-field),
# timezone_id "America/Los_Angeles"
# ZenML: standard 5-field cron + Databricks orchestrator timezone setting
schedule = Schedule(cron_expression="0 2 * * *")
databricks_settings = DatabricksOrchestratorSettings(
    schedule_timezone="America/Los_Angeles",
)

my_pipeline.with_options(
    schedule=schedule,
    settings={"orchestrator": databricks_settings},
)()

Not all orchestrators support scheduling. The Databricks orchestrator supports cron schedules only: catchup and interval schedules are ignored, and a schedule without cron_expression is invalid. Convert only simple Quartz cron expressions mechanically: dropping a zero seconds field is safe for examples like 0 0 2 * * ?, but non-zero seconds, a year field, or Quartz-only operators such as ?, L, W, and # need human review instead of blind conversion. The Databricks step operator does not own pipeline scheduling; it only offloads selected steps while another orchestrator or ZenML Pro trigger/snapshot layer starts the pipeline run. Check [references/concept-map.md](references/concept-map.md) for the orchestrator support table.

Choosing Databricks orchestrator vs Databricks step operator

Use the Databricks orchestrator when Databricks should run the whole ZenML pipeline as a Databricks Job whose tasks mirror the pipeline steps. Put cluster sizing and job-level behavior in DatabricksOrchestratorSettings: spark_version, num_workers or autoscale, node_type_id, driver_node_type_id, policy_id, spark_conf, spark_env_vars, custom_tags, job_tags, max_concurrent_runs, max_retries, min_retry_interval_millis, retry_on_timeout, timeout_seconds, task_timeout_seconds, and schedule_timezone for cron schedules.

Use the Databricks step operator when the active orchestrator should still own the pipeline, but one or more selected steps should run on Databricks as one-time submitted runs. Put cluster sizing in DatabricksStepOperatorSettings: spark_version, num_workers or autoscale, node_type_id, driver_node_type_id, policy_id, spark_conf, spark_env_vars, custom_tags, timeout_seconds, and task_timeout_seconds. Do not put pipeline schedule settings here: the current Databricks step operator submits jobs.submit one-time runs, not persistent jobs.create jobs.

Concrete consequence: if a migrated Databricks task used a new_cluster with num_workers=8, do not assume ResourceSettings(cpu_count=8) will create an 8-worker Databricks cluster for a step-operator run. Tell the user to configure DatabricksStepOperatorSettings(num_workers=8, node_type_id="...") or the equivalent autoscale settings.

Handling Databricks Feature Engineering and Unity Catalog feature lookup

Treat Feature Engineering Client / Unity Catalog feature lookup code as Databricks-native unless you have a verified equi

…

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: zenml-io
• Source: zenml-io/skills
• License: MIT
• Homepage: https://docs.zenml.io

Install and usage instructions live in the source repository linked above.