Azure ML MLflow Integration

1.62.0.post2 · active · verified Sat Apr 11

The `azureml-mlflow` package provides integration code for connecting MLflow with Azure Machine Learning. It enables users to leverage MLflow's open-source capabilities for tracking machine learning experiments, logging metrics, parameters, and artifacts, and managing models directly within an Azure Machine Learning workspace. This facilitates centralized, secure, and scalable storage for MLflow-tracked assets, whether experiments run locally or on Azure compute. The library is actively maintained with regular updates.

Warnings

breaking Azure Machine Learning's MLflow integration is currently compatible with `mlflow<=2.16.x`. Versions `mlflow>=2.17` introduce breaking changes to artifact repository and `LoggedModels` API that are not yet supported by `azureml-mlflow`.
Fix: Pin your MLflow version to `<=2.16.2` (e.g., `pip install 'mlflow<=2.16.2' azureml-mlflow`).
gotcha When combining training and inference in a single environment, `azureml-defaults` (often pulled for inference) requires `Flask>=3`, while `mlflow<2.8` requires `Flask<3`. This creates a dependency deadlock, making a single 'training + inference' environment unsupported with certain MLflow versions.
Fix: Use separate environments for training (with compatible MLflow) and inference (with `azureml-defaults`). Do not force-downgrade Flask or remove `azureml-inference-server-http`.
breaking As of `azureml-mlflow` version `1.62.0`, the `azure-common` package was removed as a dependency. Code relying on `azure-common` being transitively installed by `azureml-mlflow` will now fail with `ModuleNotFoundError`.
Fix: Explicitly install `azure-common` in your environment if your code requires it (e.g., `pip install azure-common`).
gotcha When using MLflow with low-priority jobs in Azure ML that can be pre-empted and restarted, logging artifacts using `mlflow.start_run()` (without arguments, relying on Azure ML's auto-assigned run ID) can lead to failures. Azure ML's MLflow implementation does not allow overwriting existing artifacts for the same run ID upon restart.
Fix: Consider implementing custom logic to handle restarts, such as checking for existing artifacts or using a new run ID on restart if artifact overwrite is critical.
gotcha Remote MLflow tracking to Azure ML requires proper authentication. If running code outside an Azure ML compute instance (e.g., local machine), you must explicitly configure MLflow to use your Azure credentials, otherwise 403 (Not Authorized) errors will occur.
Fix: Ensure `DefaultAzureCredential` is used with `MLClient` and that you are logged into Azure (e.g., `az login`) or relevant environment variables (`AZURE_CLIENT_ID`, `AZURE_CLIENT_SECRET`, `AZURE_TENANT_ID`) are set.

Install

pip install mlflow azureml-mlflow azure-ai-ml azure-identity Recommended installation with core dependencies

Imports

mlflow
```
import mlflow
```
MLClient
```
from azure.ai.ml import MLClient
```

DefaultAzureCredential

from azure.identity import DefaultAzureCredential

Quickstart

This quickstart demonstrates how to configure MLflow to track experiments in an Azure Machine Learning workspace from a Python script. It sets the MLflow tracking URI to your Azure ML workspace and then logs a simple parameter and metric. Ensure you have `azure-ai-ml` and `azure-identity` installed and are authenticated to Azure (e.g., via `az login`).

import os
import mlflow
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

# Replace with your Azure subscription, resource group, and workspace name
# For local execution, ensure you are logged into Azure CLI or have appropriate env vars set
subscription_id = os.environ.get('AZURE_SUBSCRIPTION_ID', 'YOUR_SUBSCRIPTION_ID')
resource_group = os.environ.get('AZURE_RESOURCE_GROUP', 'YOUR_RESOURCE_GROUP')
workspace_name = os.environ.get('AZURE_ML_WORKSPACE_NAME', 'YOUR_WORKSPACE_NAME')

try:
    # Connect to Azure ML Workspace
    ml_client = MLClient(
        credential=DefaultAzureCredential(),
        subscription_id=subscription_id,
        resource_group_name=resource_group,
        workspace_name=workspace_name
    )
    azureml_tracking_uri = ml_client.workspaces.get(ml_client.workspace_name).mlflow_tracking_uri
    mlflow.set_tracking_uri(azureml_tracking_uri)
    print(f"MLflow tracking URI set to: {mlflow.get_tracking_uri()}")

    # Start an MLflow run and log a metric
    with mlflow.start_run(run_name="quickstart_run") as run:
        mlflow.log_param("alpha", 0.5)
        mlflow.log_metric("accuracy", 0.95)
        print(f"Logged metric 'accuracy' in run: {run.info.run_id}")

    print("MLflow run completed. Check Azure ML Studio -> Jobs for details.")

except Exception as e:
    print(f"An error occurred: {e}")
    print("Please ensure Azure credentials are configured (e.g., via Azure CLI 'az login')")
    print("and that the subscription, resource group, and workspace names are correct.")

view raw JSON →