Configure MLflow for Azure Machine Learning

This article explains how you can configure MLflow to connect to an Azure Machine Learning workspace for tracking, registries, and deployment.

Azure Machine Learning workspaces are MLflow-compatible, which means they can act as an MLflow server without any extra configuration. Each workspace has an MLflow tracking URI that MLflow can use to connect to the workspace. Azure Machine Learning workspaces are already configured to work with MLflow so no extra configuration is required.

However, if you work outside of Azure Machine Learning (like your local machine, Azure Synapse Analytics, or Azure Databricks), you need to configure MLflow to point to the workspace.

Prerequisites

You need the following prerequisites to follow this tutorial:

• Install the MLflow SDK mlflow package and the Azure Machine Learning azureml-mlflow plugin for MLflow as follows:

  pip install mlflow azureml-mlflow
  

  Tip

  You can use the mlflow-skinny package, which is a lightweight MLflow package without SQL storage, server, UI, or data science dependencies. This package is recommended for users who primarily need the MLflow tracking and logging capabilities without importing the full suite of features, including deployments.

• Create an Azure Machine Learning workspace. To create a workspace, see Create resources you need to get started. Review the access permissions you need to perform your MLflow operations in your workspace.

• To do remote tracking, or track experiments running outside Azure Machine Learning, configure MLflow to point to the tracking URI of your Azure Machine Learning workspace. For more information on how to connect MLflow to your workspace, see Configure MLflow for Azure Machine Learning.

Configure MLflow tracking URI

To connect MLflow to an Azure Machine Learning workspace, you need the tracking URI for the workspace. Each workspace has its own tracking URI and it has the protocol azureml://.

1. Get the tracking URI for your workspace:

   • Azure CLI
   • Python
   • Studio
   • Manually

   APPLIES TO: Azure CLI ml extension v2 (current)

   1. Login and configure your workspace:

      az account set --subscription <subscription>
      az configure --defaults workspace=<workspace> group=<resource-group> location=<location> 
      

   2. You can get the tracking URI using the az ml workspace command:

      az ml workspace show --query mlflow_tracking_uri
      

2. Configuring the tracking URI:

   • Using MLflow SDK
   • Using environment variables

   Then the method set_tracking_uri() points the MLflow tracking URI to that URI.

   import mlflow
   
   mlflow.set_tracking_uri(mlflow_tracking_uri)
   

   Tip

   When working on shared environments, like an Azure Databricks cluster, Azure Synapse Analytics cluster, or similar, it is useful to set the environment variable MLFLOW_TRACKING_URI at the cluster level to automatically configure the MLflow tracking URI to point to Azure Machine Learning for all the sessions running in the cluster rather than to do it on a per-session basis.

Configure authentication

Once the tracking is set, you also need to configure the authentication method for the associated workspace. By default, the Azure Machine Learning plugin for MLflow performs interactive authentication by opening the default browser to prompt for credentials.

The Azure Machine Learning plugin for MLflow supports several authentication mechanisms through the package azure-identity, which is installed as a dependency for the plugin azureml-mlflow. The following authentication methods are tried one by one until one of them succeeds:

1. Environment: Reads account information specified via environment variables and uses it to authenticate.
2. Managed Identity: If the application is deployed to an Azure host with Managed Identity enabled, it authenticates with it.
3. Azure CLI: If a user signs in via the Azure CLI az login command, it authenticates as that user.
4. Azure PowerShell: If a user signs in via Azure PowerShell's Connect-AzAccount -Environment AzureChinaCloud command, it authenticates as that user.
5. Interactive browser: Interactively authenticates a user via the default browser.

For interactive jobs where there's a user connected to the session, you can rely on Interactive Authentication and hence no further action is required.

For those scenarios where unattended execution is required, you have to configure a service principal to communicate with Azure Machine Learning.

import os

os.environ["AZURE_TENANT_ID"] = "<AZURE_TENANT_ID>"
os.environ["AZURE_CLIENT_ID"] = "<AZURE_CLIENT_ID>"
os.environ["AZURE_CLIENT_SECRET"] = "<AZURE_CLIENT_SECRET>"

export AZURE_TENANT_ID="<AZURE_TENANT_ID>"
export AZURE_CLIENT_ID="<AZURE_CLIENT_ID>"
export AZURE_CLIENT_SECRET="<AZURE_CLIENT_SECRET>"

If you'd rather use a certificate instead of a secret, you can configure the environment variables AZURE_CLIENT_CERTIFICATE_PATH to the path to a PEM or PKCS12 certificate file (including private key) and AZURE_CLIENT_CERTIFICATE_PASSWORD with the password of the certificate file, if any.

Configure authorization and permission levels

Some default roles like AzureML Data Scientist or Contributor are already configured to perform MLflow operations in an Azure Machine Learning workspace. If using a custom role, you need the following permissions:

• To use MLflow tracking:

  • Microsoft.MachineLearningServices/workspaces/experiments/*
  • Microsoft.MachineLearningServices/workspaces/jobs/*

• To use MLflow model registry:

  • Microsoft.MachineLearningServices/workspaces/models/*/*

To learn how to grant access for the service principal you created or user account to your workspace, see Grant access.

Troubleshooting authentication

MLflow tries to authenticate to Azure Machine Learning on the first operation that interacts with the service, like mlflow.set_experiment() or mlflow.start_run(). If you find issues or unexpected authentication prompts during the process, you can increase the logging level to get more details about the error:

import logging

logging.getLogger("azure").setLevel(logging.DEBUG)

Set experiment name (optional)

All MLflow runs are logged to the active experiment. By default, runs are logged to an experiment named Default that is automatically created for you. You can configure the experiment where tracking is happening.

experiment_name = 'experiment_with_mlflow'
mlflow.set_experiment(experiment_name)

export MLFLOW_EXPERIMENT_NAME="experiment_with_mlflow"

Next steps

Now that your environment is connected to your workspace in Azure Machine Learning, you can start to work with it.

• Track ML experiments and models with MLflow
• Manage models registries in Azure Machine Learning with MLflow
• Train with MLflow Projects (Preview)
• Guidelines for deploying MLflow models