Databricks Migration Tool

NOTE: For a more extensive and maintained cross-workload migration solution, please use the Databricks Terraform Exporter, which creates Infrastructure-as-a-Code replicas for the entire manually-configured Databricks Workspaces.

This is a migration package to log all Databricks resources for backup and/or migrating to another Databricks workspace. Migration allows a Databricks organization to move resources between Databricks Workspaces, to move between different cloud providers, or to move to different regions / accounts.

This package is based on python 3.6 and DBR 6.x+ releases.
Python 3.7 or above is recommended if one is also exporting/importing MLflow objects.

Note: This tool does not support windows currently since path resolution is different from mac / linux.

This package also uses credentials from the Databricks CLI.

Table of Contents

• Pre-Requisites
• Setup
• Migration Components
• Import using pipeline
  • Pipeline parameters
  • Exporting the Workspace
  • Recommended parameters and checkpointing
  • Updating the AWS Account ID
  • Importing the Workspace
  • Validation
• Limitations

Pre-Requisites

To use this migration tool, you'll need:

• An environment running linux with python, pip, git, and the databricks CLI installed.
• Admin access to both the old and new databricks accounts in the form of a Personal Access Token.

Setup

Click to expand & collapse tasks

1. Generate Tokens

Generate Access Tokens for both the old and new databricks accounts

1. Click User Settings Icon Settings in the lower left corner of your Databricks workspace
2. Click on Access Tokens tab
3. Click on Generate New Token button.
4. Copy the generated token and store in a secure location.

Be sure to keep a file with the url for both the old and new databricks account Add the old and new token and the old and new Instance ID if applicable. You'll need easy access to all of these things when running the migration tool.

Migration Components

To use the migration tool see the details below to start running the tool in the order recommended to properly migrate files.

Support Matrix for Import and Export Operations:

Note on MLFlow Migration:
MLFlow asset migration is currently only partially supported; Feature Store and Model Registry will not be migrated, for example. See mlflow-export-import for comprehensive MLflow migrations.

Note on DBFS Data Migration:
DBFS is a protected object storage location on AWS and Azure. Please contact your Databricks support team for information about migrating DBFS resources.

Note on User Migration:
During user / group import, users will be notified of the new workspace and account by default. To disable this behavior, please contact your Databricks account team.

Note on Repos: Private repos cannot be imported. These should be added manually using the original user credentials.

Import using pipeline

The recommended method of exporting and importing is by using the Pipeline contained in migration_pipeline.py. This pipeline performs all export and import steps sequentially, and includes checkpointing parallelization features.

Pipeline parameters

python migration_pipeline.py -h
usage: migration_pipeline.py [-h] [--profile PROFILE] [--azure or gcp] [--silent] [--no-ssl-verification] [--debug] [--set-export-dir SET_EXPORT_DIR]
                             [--cluster-name CLUSTER_NAME] [--notebook-format {DBC,SOURCE,HTML}] [--overwrite-notebooks] [--archive-missing]
                             [--repair-metastore-tables] [--metastore-unicode] [--skip-failed] [--session SESSION] [--dry-run] [--export-pipeline] [--import-pipeline]
                             [--validate-pipeline] [--validate-source-session VALIDATE_SOURCE_SESSION] [--validate-destination-session VALIDATE_DESTINATION_SESSION]
                             [--use-checkpoint] [--skip-tasks SKIP_TASKS [SKIP_TASKS ...]] [--num-parallel NUM_PARALLEL] [--retry-total RETRY_TOTAL]
                             [--retry-backoff RETRY_BACKOFF] [--start-date START_DATE] [--hipaa] [--skip-large-nb]
                             [--exclude-work-item-prefixes EXCLUDE_WORK_ITEM_PREFIXES [EXCLUDE_WORK_ITEM_PREFIXES ...]]

Export user(s) workspace artifacts from Databricks

optional arguments for import/export pipeline:
  -h, --help            show this help message and exit
  --profile PROFILE     Profile to parse the credentials
  --azure or --gcp      Run on Azure or GCP (Default is AWS)
  --silent              Silent all logging of export operations.
  --no-ssl-verification
                        Set Verify=False when making http requests.
  --debug               Enable debug logging
  --no-prompt           Skip interactive prompt/confirmation for workspace import.
  --set-export-dir SET_EXPORT_DIR
                        Set the base directory to export artifacts
  --cluster-name CLUSTER_NAME
                        Cluster name to export the metastore to a specific cluster. Cluster will be started.
  --notebook-format {DBC,SOURCE,HTML}
                        Choose the file format to download the notebooks (default: DBC)
  --overwrite-notebooks
                        Flag to overwrite notebooks to forcefully overwrite during notebook imports
  --archive-missing     Import all missing users into the top level /Archive/ directory.
  --repair-metastore-tables
                        Repair legacy metastore tables
  --hipaa               Use HIPAA-compatible cluster profiles
  --metastore-unicode   log all the metastore table definitions including unicode characters
  --skip-failed         Skip retries for any failed hive metastore exports.
  --skip-missing-users  Skip failed principles during ACL import; for missing principles, this will result in open ACLs
  --skip-large-nb       Skip notebooks that exceed the size limit of 10485760 bytes instead of failing
  --session SESSION     If set, pipeline resumes from latest checkpoint of given session; Otherwise, pipeline starts from beginning and creates a new session.
  --dry-run             Dry run the pipeline i.e. will not execute tasks if true.
  --export-pipeline     Execute all export tasks.
  --import-pipeline     Execute all import tasks.
  --use-checkpoint      use checkpointing to restart from previous state
  --skip-tasks SKIP_TASK [SKIP_TASK ...]
                        Space-separated list of tasks to skip from the pipeline. Valid options are:
                         instance_profiles, users, groups, workspace_item_log, workspace_acls, notebooks, secrets,
                         clusters, instance_pools, jobs, metastore, metastore_table_acls, mlflow_experiments, mlflow_runs
  --keep-tasks KEEP_TASK [KEEP_TASK ...]
                        Space-separated list of tasks to run from the pipeline. See valid options in --skip-tasks. Overrides skip-tasks.
  --num-parallel NUM_PARALLEL
                        Number of parallel threads to use to export/import
  --retry-total RETRY_TOTAL
                        Total number or retries when making calls to Databricks API
  --retry-backoff RETRY_BACKOFF
                        Backoff factor to apply between retry attempts when making calls to Databricks API
  --start-date START_DATE
                        start-date format: YYYY-MM-DD. If not provided, defaults to past 30 days. Currently, only used for exporting ML runs objects.
  --groups-to-keep group [group ...]
                        List of groups to keep if selectively exporting assets. Only users (and their assets) belonging to these groups will be exported.

options for validation pipeline:
  --validate-pipeline   Validate exported data between source and destination.
  --validate-source-session VALIDATE_SOURCE_SESSION
                        Session used by exporting source workspace. Only used for --validate-pipeline.
  --validate-destination-session VALIDATE_DESTINATION_SESSION
                        Session used by exporting destination workspace. Only used for --validate-pipeline.

Exporting the Workspace

To export a workspace, run:

python3 migration_pipeline.py --profile $SRC_PROFILE --export-pipeline --use-checkpoint [--session $SESSION_ID]

Where $SRC_PROFILE is the Databricks profile for the source workspace, as configured during Setup, and $SESSION_ID is an optional session identifier used for subsequent checkpoint runs. All data is exported to a folder named according to the $SESSION_ID value under the logs folder - “logs/$SESSION_ID”. If $SESSION_ID is not specified, a random value will be generated.

Recommended parameters and checkpointing

As a starting point, we recommend using the following parameter values:

• retry-total=30
• num-parallel=8
• retry-backoff=1.0

These can be adjusted per your scenario if needed; in general, if API limits are being hit, you can increase retry-backoff, decrease num-parallel, or both.

If script failure occurs, you can safely rerun the same command with --use-checkpoint and --session $SESSION_ID to let the migration pick up from the previous checkpoint and rerun.

Updating the AWS Account ID

If your source and destination workspaces are in different accounts, you will need to update the Instance Profile ARN accordingly during the migration. To do this, run the following command after exporting the workspace assets:

python3 export_db.py --profile $SRC_PROFILE --use-checkpoint --old-account-id $OLD_AWS_ACCT_ID --update-account-id $NEW_AWS_ACCT_ID --set-export-dir $EXPORT_DIR/$SESSION_ID

Where EXPORT_DIR/SESSION_ID is the directory and session ID used by your export job, SRC_PROFILE is the profile used to export the source workspace, OLD_AWS_ACCT_ID is the source AWS account ID, and NEW_AWS_ACCT_ID is the destination AWS account ID. Note that this will only update the ARN in the Instance Profiles; the same instance profiles must still exist in the destination workspace.

Importing the Workspace

To import into a target workspace, run:

python3 migration_pipeline.py --profile $DST_PROFILE --import-pipeline --use-checkpoint [--session $SESSION_ID]

The same recommended parameters as above apply in the import workflow, and similarly, if a failure occurs, --use-checkpoint can be used to rerun from the last checkpoint.

Validation

Simple workspace object validation can be performed once the import is completed by first exporting the contents of the target workspace:

python3 migration_pipeline.py --profile $DST_PROFILE --export-pipeline --use-checkpoint --cluster-name

And then running the validate_pipeline.sh script:

./validate_pipeline.sh $SRC_EXPORT_SESSION_ID $DST_EXPORT_SESSION_ID

Once this completes, check the console summary, as well as the logs folder (where a new folder should be generated).

Limitations:

Note: To disable ssl verification pass the flag --no-ssl-verification. If still getting SSL Error add the following to your current bash shell:

export REQUESTS_CA_BUNDLE=""
export CURL_CA_BUNDLE=""

• Instance profiles (AWS only): Group access to instance profiles will take precedence. If a user is added to the role directly, and has access via a group, only the group access will be granted during a migration.
• Clusters: Cluster creator will be seen as the single admin user who migrated all the clusters. (Relevant for billing purposes)
  • Cluster creator tags cannot be updated. Added a custom tag named OriginalCreator with the original cluster creator for DBU tracking.
• Jobs: Job owners will be seen as the single admin user who migrate the job configurations. (Relevant for billing purposes)
  • Jobs with existing clusters that no longer exist will be reset to the default cluster type
  • Jobs with older legacy instances will fail with unsupported DBR or instance types. See release notes for the latest supported releases.