Configure MLflow for Azure Machine Learning

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 can be used by MLflow 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 are working 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. In this article, you'll learn how you can configure MLflow to connect to an Azure Machine Learning for tracking, registries, and deployment.

Important

When running on Azure Compute (Azure Machine Learning Notebooks, Jupyter notebooks hosted on Azure Machine Learning Compute Instances, or jobs running on Azure Machine Learning compute clusters) you don't have to configure the tracking URI. It's automatically configured for you.

Prerequisites

You need the following prerequisites to follow this tutorial:

• Install Mlflow SDK package mlflow and Azure Machine Learning plug-in for MLflow azureml-mlflow.

  pip install mlflow azureml-mlflow
  

  Tip

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

• You need an Azure Machine Learning workspace. You can create one following this tutorial.

  • See which access permissions you need to perform your MLflow operations with your workspace.

• If you're doing remote tracking (tracking experiments running outside Azure Machine Learning), configure MLflow to point to your Azure Machine Learning workspace's tracking URI as explained at 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

   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
      

   Python

   APPLIES TO: Python SDK azure-ai-ml v2 (current)

   You can get the Azure ML MLflow tracking URI using the Azure Machine Learning SDK v2 for Python. Ensure you have the library azure-ai-ml installed in the compute you are using. The following sample gets the unique MLFLow tracking URI associated with your workspace.

   1. Login into your workspace using the MLClient. The easier way to do that is by using the workspace config file:

      from azure.ai.ml import MLClient
      from azure.identity import DefaultAzureCredential
      
      ml_client = MLClient.from_config(credential=DefaultAzureCredential())
      

      Tip

      You can download the workspace configuration file by:

      1. Navigate to Azure ML studio
      2. Click on the upper-right corner of the page -> Download config file.
      3. Save the file config.json in the same directory where you are working on.

   2. Alternatively, you can use the subscription ID, resource group name and workspace name to get it:

      from azure.ai.ml import MLClient
      from azure.identity import DefaultAzureCredential
      
      #Enter details of your AzureML workspace
      subscription_id = '<SUBSCRIPTION_ID>'
      resource_group = '<RESOURCE_GROUP>'
      workspace_name = '<WORKSPACE_NAME>'
      
      ml_client = MLClient(credential=DefaultAzureCredential(),
                              subscription_id=subscription_id, 
                              resource_group_name=resource_group,
                              workspace_name=workspace_name)
      

      Important

      DefaultAzureCredential will try to pull the credentials from the available context. If you want to specify credentials in a different way, for instance using the web browser in an interactive way, you can use InteractiveBrowserCredential or any other method available in azure.identity package.

   3. Get the Azure Machine Learning Tracking URI:

      mlflow_tracking_uri = ml_client.workspaces.get(ml_client.workspace_name).mlflow_tracking_uri
      

   Studio

   Use the Azure Machine Learning portal to get the tracking URI:

   1. Open the Azure Machine Learning studio portal and log in using your credentials.

   2. In the upper right corner, click on the name of your workspace to show the Directory + Subscription + Workspace blade.

   3. Click on View all properties in Azure Portal.

   4. On the Essentials section, you will find the property MLflow tracking URI.

   Manually

   The Azure Machine Learning Tracking URI can be constructed using the subscription ID, region of where the resource is deployed, resource group name and workspace name. The following code sample shows how:

   Warning

   If you are working in a private link-enabled workspace, the MLflow endpoint will also use a private link to communicate with Azure Machine Learning. As a consequence, the tracking URI will look different as proposed here. You need to get the tracking URI using the Azure ML SDK or CLI v2 on those cases.

   region = "<LOCATION>"
   subscription_id = '<SUBSCRIPTION_ID>'
   resource_group = '<RESOURCE_GROUP>'
   workspace_name = '<AML_WORKSPACE_NAME>'
   
   mlflow_tracking_uri = f"azureml://{region}.api.azureml.ms/mlflow/v1.0/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.MachineLearningServices/workspaces/{workspace_name}"
   

2. Configuring the tracking URI:

   MLflow SDK

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

   import mlflow
   
   mlflow.set_tracking_uri(mlflow_tracking_uri)
   

   Using environment variables

   You can set the MLflow environment variables MLFLOW_TRACKING_URI in your compute to make any interaction with MLflow in that compute to point by default to Azure Machine Learning.

   MLFLOW_TRACKING_URI=$(az ml workspace show --query mlflow_tracking_uri | sed 's/"//g') 
   

   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'll also need to configure how the authentication needs to happen to the associated workspace. By default, the Azure Machine Learning plugin for MLflow will perform 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: it reads account information specified via environment variables and use 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 has signed in via the Azure CLI az login command, it authenticates as that user.
4. Azure PowerShell: if a user has signed in via Azure PowerShell's Connect-AzAccount command, it authenticates as that user.
5. Interactive browser: it 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.

Warning

Interactive browser authentication will block code execution when prompting for credentials. It is not a suitable option for authentication in unattended environments like training jobs. We recommend to configure other authentication mode.

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

MLflow SDK

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>"

Using environment variables

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

Tip

When working on shared environments, it is advisable to configure these environment variables at the compute. As a best practice, manage them as secrets in an instance of Azure Key Vault whenever possible. For instance, in Azure Databricks you can use secrets in environment variables as follows in the cluster configuration: AZURE_CLIENT_SECRET={{secrets/<scope-name>/<secret-name>}}. See Reference a secret in an environment variable for how to do it in Azure Databricks or refer to similar documentation in your platform.

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 roles, 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/*/*

Grant access for the service principal you created or user account to your workspace as explained at Grant access.

Troubleshooting authentication

MLflow will try to authenticate to Azure Machine Learning on the first operation interacting 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.

Tip

When submitting jobs using Azure Machine Learning CLI v2, you can set the experiment name using the property experiment_name in the YAML definition of the job. You don't have to configure it on your training script. See YAML: display name, experiment name, description, and tags for details.

MLflow SDK

To configure the experiment you want to work on use MLflow command mlflow.set_experiment().

experiment_name = 'experiment_with_mlflow'
mlflow.set_experiment(experiment_name)

Using environment variables

You can also set one of the MLflow environment variables MLFLOW_EXPERIMENT_NAME or MLFLOW_EXPERIMENT_ID with the experiment name.

export MLFLOW_EXPERIMENT_NAME="experiment_with_mlflow"

Non-public Azure Clouds support

The Azure Machine Learning plugin for MLflow is configured by default to work with the global Azure cloud. However, you can configure the Azure cloud you are using by setting the environment variable AZUREML_CURRENT_CLOUD.

MLflow SDK

import os

os.environ["AZUREML_CURRENT_CLOUD"] = "AzureChinaCloud"

Using environment variables

export AZUREML_CURRENT_CLOUD="AzureChinaCloud"

You can identify the cloud you are using with the following Azure CLI command:

az cloud list

The current cloud has the value IsActive set to True.

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