DOCSEARCHbyAlgolia
Submit Docs

Databricks MLOps Stacks

repository·main·Indexed 20 days ago

https://github.com/databricks/mlops-stacks

A production-ready template for machine learning projects on Databricks. It integrates ML code, infrastructure-as-code via Databricks Asset Bundles, and CI/CD workflows (GitHub Actions, Azure DevOps, or GitLab) to automate the transition from development to production across Dev, Staging, and Prod environments.

Tokens
7.5K
Snippets
11
Records
37
Agent score
71%

What's inside Databricks MLOps Stacks

  1. Overview of Databricks MLOps Stacks

    main

    Databricks MLOps Stacks is a customizable template for starting new machine learning projects on Databricks that follow production best practices. It provides a modular structure designed to bridge the gap between data science iteration and operational deployment.

    The stack consists of three core components:

    1. ML Code: A project structure containing training and batch inference modules, notebooks, and unit tests.
    2. ML Resources as Code: ML pipeline resources (like training and batch inference jobs) defined using Databricks CLI bundles.
    3. CI/CD: Automated workflows for GitHub Actions or Azure DevOps to test and deploy code and resources.

    This approach allows data scientists to iterate quickly on code while operations engineers manage CI/CD and resource governance through pull requests.

  2. How the MLOps Stacks development loop works

    main

    The MLOps Stacks workflow uses three distinct environments to manage the lifecycle of data, code, and models:

    • Dev: Where data scientists iterate on ML code and file pull requests (PRs).
    • Staging: An isolated workspace where PRs trigger unit and integration tests. When a PR is merged into the main branch, staging jobs (training and batch inference) automatically update to run the latest code.
    • Prod: The production environment. After merging to main, users can cut a new release branch to promote specific ML code changes to production.

    This separation ensures that all production changes are automated, tested in staging, and governed via version control.

  3. Workspace isolation strategies for MLOps Stacks

    main

    MLOps Stacks supports different workspace configurations depending on your isolation needs:

    1. Recommended (Separate Workspaces): Use separate Databricks workspaces for dev, staging, and prod. This provides strong isolation and prevents issues like Databricks REST API rate limits in staging affecting production (e.g., when using Databricks Model Serving).
    2. Single Workspace Stack: You can use a single workspace by providing the same workspace URL for both input_databricks_staging_workspace_host and input_databricks_prod_workspace_host. If using this method, it is highly recommended to use different service principals for staging and production to prevent CI workloads in staging from interfering with production resources.

  4. Understand the Model Training and Batch Inference pipeline

    main

    The model training pipeline follows a specific lifecycle to ensure model quality before deployment:

    1. Training & Registration: The pipeline trains a model and registers it in the model registry.
    2. CD Workflow: Registering the model triggers a Continuous Deployment (CD) workflow.
    3. Final Checks: The CD workflow performs final validation checks on the model.
    4. Promotion: Once checks pass, the model is promoted for deployment and use in batch inference jobs.

  5. Create a custom MLOps Stacks repository

    main

    To tailor MLOps Stacks to your organization's specific needs (e.g., custom workspace environments, unique CI/CD requirements, or in-house ML code structures), follow these steps:

    1. Fork the repository: Create a fork of the MLOps Stacks repo. Use a private fork for organizational-specific tailoring or a public fork for generic enhancements.
    2. Set up CI (Optional): If you use GitHub Actions, the provided workflows in .github/ should work out of the box. If you use a different provider, you must translate the workflows from .github/ to your provider's format.
    3. Update parameters: Modify databricks_template_schema.json to add or remove parameters, and update the corresponding template variables in library/template_variables.tmpl.

    Tip: To prune parameters and replace them with hardcoded values, use the previewing changes workflow to generate an example project with substituted parameters, then copy those contents back into your fork.

  6. Preview changes by creating an example project

    main

    To see how changes to the MLOps Stacks repository affect a generated project, you can create an example project from your local checkout using databricks bundle init. You must set the MLOPS_STACKS_PATH environment variable to the path of your local MLOps Stacks checkout.

    # Example: Create an Azure project with Azure DevOps CI/CD
MLOPS_STACKS_PATH=~/mlops-stacks
databricks bundle init "$MLOPS_STACKS_PATH" --config-file "$MLOPS_STACKS_PATH/tests/example-project-configs/azure/azure-devops.json"

# Example: Create an AWS project with GitHub Actions CI/CD
MLOPS_STACKS_PATH=~/mlops-stacks
databricks bundle init "$MLOPS_STACKS_PATH" --config-file "$MLOPS_STACKS_PATH/tests/example-project-configs/aws/aws-github.json"

# Example: Create a GCP project with GitHub Actions CI/CD
MLOPS_STACKS_PATH=~/mlops-stacks
databricks bundle init "$MLOPS_STACKS_PATH" --config-file "$MLOPS_STACKS_PATH/tests/example-project-configs/gcp/gcp-github.json"

  7. Start a new project with `databricks bundle init`

    main

    To create a new ML project using the MLOps Stacks template, use the Databricks CLI bundle init command. The command will prompt you for several configuration parameters.

    Core Initialization Parameters

    • input_setup_cicd_and_project: Determines the scope of the setup:
      • CICD_and_Project (Default): Sets up both the project structure and CI/CD workflows.
      • Project_Only: Sets up only the ML code structure. Recommended for Data Scientists starting development.
      • CICD_Only: Sets up only the CI/CD workflows. Recommended for MLEs integrating into existing monorepos.
    • input_root_dir: The name of the root directory for your project.
    • input_cloud: Your cloud provider (AWS, Azure, or GCP).

    CI/CD Configuration Parameters

    To ensure automation works correctly, the following must be specified:

    • input_cicd_platform: The platform to use (e.g., GitHub Actions, Azure DevOps, GitLab).
    • input_databricks_staging_workspace_host: The URL of your staging Databricks workspace.
    • input_databricks_prod_workspace_host: The URL of your production Databricks workspace.
    • input_default_branch: The branch (e.g., main) where staging and production resources are deployed from.
    • input_release_branch: The branch from which production jobs pull ML code.

    Project and Governance Parameters

    • input_project_name: The name of the project.
    • input_read_user_group: The group granted READ permissions to project resources (defaults to users).
    • input_schema_name: The Unity Catalog schema name for model registration (defaults to the project name).
    • input_unity_catalog_read_user_group: The group granted EXECUTE privileges for registered models (defaults to account users).
    • input_include_feature_store: If selected, includes Databricks Feature Store components (modules, notebooks, and resource configs).
    databricks bundle init mlops-stacks

  8. Install development requirements for MLOps Stacks

    main

    To contribute to MLOps Stacks or run its tests, you need to install the following tools:

    • actionlint
    • databricks CLI
    • npm
    • act

    Then, install the Python dependencies using:

    pip install -r dev-requirements.txt

  9. Adopting individual MLOps Stacks components

    main

    If you do not want to adopt the full stack, you can adopt individual components by instantiating the stack via Databricks asset bundle templates and copying specific subdirectories:

    • ML Resource Configs: Found in {{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/resources and {{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/databricks.yml.
    • CI/CD: Found in .github (GitHub Actions) or .azure (Azure DevOps).

  10. Productionizing existing ML projects with MLOps Stacks

    main

    You can productionize an existing ML project by instantiating a new MLOps Stacks project and copying the relevant modular components into your current project.

    Key components to copy include:

    • CI/CD Workflows: The .github directory (for GitHub Actions) or .azure directory (for Azure DevOps).
    • ML Resource Configs: The directory {{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/resources and the databricks.yml file located in {{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/.

  11. Customize ML code and notebook interfaces

    main

    You can replace the default example ML code with your own skeleton or in-house structure. To ensure the CI/CD and ML resource components continue to function, your notebooks must adhere to the expected interface at these paths:

    • Model Training: template/{{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/training/notebooks/
    • Batch Inference: template/{{.input_root_dir}}/{{template project_name_alphanumeric_underscore .}}/deployment/batch_inference/notebooks/

    Model Promotion: If you want to integrate model promotion across environments (e.g., from dev to staging), consider using the mlflow.client.MlflowClient.copy_model_version API within your training notebooks. If using Unity Catalog, models can be shared across workspaces if they connect to the same metastore and have appropriate privileges.

  12. Deploy ML code to Staging and Production

    main

    MLOps Stacks manages deployments across different branches and environments:

    • Staging: ML jobs in the staging environment are configured to run against the default branch. They update immediately to run the latest code once PRs are merged into the default branch.
    • Production: ML jobs in production run against a release branch. To promote changes to production, you must cut a new release branch from the main branch as part of your regular release process after merging PRs.