Skip to main content

Quickstart: connect your development environment to MLflow

This page shows you how to create an MLflow Experiment and connect your development environment to it.

An MLflow Experiment is the container for your gen AI application. Learn more about MLflow Experiments in the Experiment data model concept guide.

Go the section relevant to your development environment:

  1. Locally in an IDE or notebook
  2. Databricks-hosted Notebook

Local development environment​

Step 1: Install MLflow​

Install MLflow with Databricks connectivity:

Bash
pip install --upgrade "mlflow[databricks]>=3.1"

Step 2: Create an MLflow Experiment​

  1. Open your Databricks workspace.
  2. In the left sidebar, under AI/ML, click Experiments.
  3. At the top of the Experiments page, click GenAI apps & agents.

create experiment

Step 3: Configure authentication​

note

These steps describes using a Databricks Personal Access Token. MLflow also works with the other Databricks-supported authentication methods.

Choose one of the following authentication methods:

  1. In your MLflow Experiment, click Generate API Key.
  2. Copy and run the generated code in your terminal:
Bash
export DATABRICKS_TOKEN=<databricks-personal-access-token>
export DATABRICKS_HOST=https://<workspace-name>.cloud.databricks.com
export MLFLOW_TRACKING_URI=databricks
export MLFLOW_EXPERIMENT_ID=<experiment-id>

Step 4: Verify your connection​

Create a test file and run this code to verify your connection and log a test trace to your MLflow Experiment:

Python
import mlflow
import os

experiment_id = os.environ.get("MLFLOW_EXPERIMENT_ID")
databricks_host = os.environ.get("DATABRICKS_HOST")
mlflow_tracking_uri = os.environ.get("MLFLOW_TRACKING_URI")

if experiment_id is None or databricks_host is None or mlflow_tracking_uri is None:
    raise Exception("Environment variables are not configured correctly.")

@mlflow.trace
def hello_mlflow(message: str):

    hello_data = {
        "experiment_url": f"{databricks_host}/mlflow/experiments/{experiment_id}",
        "experiment_name": mlflow.get_experiment(experiment_id=experiment_id).name,
        "message": message,
    }
    return hello_data

result = hello_mlflow("hello, world!")
print(result)

hello mlflow

Develop in a Databricks-hosted Notebook​

Step 1: Create a notebook​

Creating a Databricks Notebook creates an MLflow Experiment that is the container for your GenAI application. To learn more about experiments, see data model.

  1. Open your Databricks workspace.
  2. Go to New at the top of the left sidebar.
  3. Click Notebook.

Step 2: Install MLflow​

Databricks runtimes include MLflow, but for the best experience with GenAI capabilities, update to the latest version:

Python
%pip install --upgrade "mlflow[databricks]>=3.1"
dbutils.library.restartPython()

Step 3: Configure authentication​

No additional authentication configuration is needed when working within a Databricks Notebook. The notebook automatically has access to your workspace and the associated MLflow Experiment.

Step 4: Verify your connection​

Run this code in a notebook cell to verify your connection. You will see an MLflow trace appear below your notebook cell.

Python
import mlflow
import os

@mlflow.trace
def hello_mlflow(message: str):
    hello_data = {
        "message": message,
    }
    return hello_data

result = hello_mlflow("hello, world!")
print(result)

hello mlflow

Next steps​

Continue your journey with these recommended actions and tutorials.

Reference guides​

Explore detailed documentation for concepts and features mentioned in this guide.

Last updated on