5.1 Configure Experiment Tracking with MLflow

Configure Experiment Tracking with MLflow

Quick Answer: Azure Machine Learning SDK v2 has no native logging API. Install mlflow and azureml-mlflow, then log with mlflow.log_metric, log_param, log_artifact, log_model, or mlflow.autolog(). Submitted jobs auto-start an MLflow run; notebooks need mlflow.start_run(). The workspace is the tracking server (azureml:// URI). Use async logging for high-volume metrics.

Domain 2 of Exam AI-300 asks you to configure experiment tracking with MLflow. That skill is the difference between a training job you can compare, promote, and debug six months later and a GPU bill with no lineage. Machine Learning Operations (MLOps) lives on metrics, parameters, artifacts, and models that are stored against a run identity — not on print() statements in std_log.txt.

SDK v2 has no logging API of its own

If you learned Azure Machine Learning SDK v1, you remember Run.get_context() and run.log(). That API is gone in v2. Microsoft’s current guidance is explicit: there is no logging functionality in the Azure Machine Learning SDK for Python (v2). You track experiments with MLflow Tracking plus the azureml-mlflow plugin that teaches MLflow how to talk to a workspace.

Install both packages in the environment that actually trains:

pip install mlflow azureml-mlflow

For a smaller footprint on tracking-only nodes, mlflow-skinny is enough. Asynchronous metric logging additionally needs MLflow 2.8.0 or later and azureml-mlflow 1.55 or later.

Your training script then imports MLflow, not azure.ai.ml logging helpers:

import mlflow

mlflow.log_param("num_epochs", 20)
mlflow.log_metric("val_auc", 0.91, step=3)
mlflow.log_artifact("plots/roc.png")
mlflow.sklearn.log_model(model, artifact_path="model")

Parameters can be any type (log_param / log_params with a dictionary). Metrics are numeric. Booleans are logged as 0 or 1. Curves are the same metric logged many times with a step (or a log_batch of Metric entities). Files are artifacts. An MLflow model is a folder that packages the weights plus the flavor files needed to load them — that package is what you later register and deploy with no-code scoring.

Experiments, runs, jobs, and where the URI lives

MLflow organizes work as experiments that contain runs. In Azure Machine Learning, a run is a job. The workspace is already an MLflow server. Each workspace has a tracking URI that starts with azureml://. You do not stand up a separate tracking cluster.

How you get that URI depends on where the process runs:

Get the URI with CLI v2 (az ml workspace show --query mlflow_tracking_uri), with SDK v2 (ml_client.workspaces.get(name).mlflow_tracking_uri), from the Azure portal Essentials blade, or by constructing azureml://{region}.api.azureml.ms/mlflow/v1.0/subscriptions/.... Private-link workspaces use a different URI shape — do not hand-build the public template; ask the SDK or CLI. You can also export MLFLOW_TRACKING_URI so every process on a shared cluster points at the same workspace.

Jobs submitted with CLI v2 set the experiment from YAML (experiment_name:). You do not have to call mlflow.set_experiment inside the training script. Interactive sessions should call mlflow.set_experiment("claims-fraud-2026") or export MLFLOW_EXPERIMENT_NAME. If you set nothing, runs land in an experiment named Default.

Interactive notebooks versus training jobs

On a compute-instance notebook the pattern is: set the experiment, start a run, log, end the run. A context manager is the cleanest form:

mlflow.set_experiment("iris-classifier")
with mlflow.start_run(run_name="iris-rf-depth12") as run:
    mlflow.log_params({"n_estimators": 200, "max_depth": 12})
    mlflow.log_metric("accuracy", 0.97)

run_name becomes the display name in studio, which is how humans find the trial later. Technically a logging API creates a run if none is active, and mlflow.active_run() returns it — still prefer an explicit start_run so you control the name and the lifetime.

Inside a command job, the platform already opened the run. Call mlflow.autolog() and log_metric directly. Calling start_run() there creates a nested run, which is almost never what the exam (or a production parent job) wants.

Exam scenario

A fraud team trains LightGBM on a compute cluster with CLI v2. The YAML has experiment_name: fraud-lgbm. The script calls mlflow.autolog() and mlflow.log_metric("pr_auc", score). Six child jobs later, studio Jobs shows one experiment with six jobs, each with parameters, a metrics chart, and an MLflow model folder. The MLOps engineer did not call start_run and did not set a tracking URI in the script — both were inherited from the workspace.

Common trap

Do not reach for a v1 Run.log helper, and do not assume mlflow.get_run(id).data.metrics["loss"] is a time series. That dictionary returns only the most recently logged value for each metric name. Use MlflowClient.get_metric_history() when you need the full curve. Another trap: logging a 4 GB model with the default 300-second artifact timeout (AZUREML_ARTIFACTS_DEFAULT_TIMEOUT). Raise it before log_model on large artifacts.

Autolog, explicit logs, and large models

mlflow.autolog() before training tells supported frameworks (scikit-learn, LightGBM, XGBoost, PyTorch, TensorFlow, and others) to log typical metrics, parameters, and a model. You can keep the convenience and skip the model with mlflow.autolog(log_models=False) when you want to call mlflow.sklearn.log_model (or another flavor) yourself after extra packaging. Some flavors disable automatic model logging when the trained model exceeds an internal size boundary. If autolog produced metrics but no model folder, that is expected — log the model explicitly.

Images go through log_image (numpy or PIL) or log_figure (matplotlib). Text, dictionaries, and existing files use log_text, log_dict, log_artifact, and log_artifacts. Prefer logging an MLflow model over a raw pickle: the model folder is what no-code deployment and the Responsible AI dashboard expect.

Asynchronous logging for high volume

Synchronous log_metric waits until the backend accepts the value. That is fine for a dozen scalars. It is a bottleneck when tens of nodes log hundreds of thousands of points. Enable async globally with mlflow.config.enable_async_logging() or export MLFLOW_ENABLE_ASYNC_LOGGING=True, or pass synchronous=False on a single call. Control returns as soon as the operation is accepted; the value is not immediately readable. You can wait() on the returned operation if a later line must read that metric. Azure Machine Learning still guarantees metric order, and it waits for pending async metrics when the job is about to finish — completed jobs have a complete metric store.

Do not call log_metric in a tight loop if you can batch with MlflowClient.log_batch.

Authentication and authorization

The plugin authenticates with azure-identity. The chain is environment variables, then managed identity, Azure CLI, Azure PowerShell, then an interactive browser. Interactive browser blocks and is wrong for unattended jobs. For those, set a service principal: AZURE_TENANT_ID, AZURE_CLIENT_ID, and AZURE_CLIENT_SECRET (or a certificate path). Store secrets in Azure Key Vault, not in the training script. Built-in roles such as AzureML Data Scientist and Contributor already cover MLflow. A custom role needs Microsoft.MachineLearningServices/workspaces/experiments/* and .../jobs/* for tracking, plus .../models/*/* if you also use the MLflow model registry. First contact with the service is often set_experiment or start_run — turn on logging.getLogger("azure").setLevel(logging.DEBUG) when a prompt appears at a surprising time.

Studio Jobs UI versus the MLflow client

Both views read the same workspace store.

• Studio: Jobs tab, filter by experiment, open a job, use Metrics to chart and share a layout, and Outputs and logs for files. user_logs/std_log.txt is stdout/stderr from your script (the first place you look). system_logs is platform diagnostics. Multi-node jobs add a folder per node IP.
• MLflow client: mlflow.get_run(run_id) for the latest metrics, params, and tags; MlflowClient.list_artifacts / mlflow.artifacts.download_artifacts for files; search_runs to compare trials in Python.

Use studio when a human needs a chart. Use the client when a promotion script must pick the best pr_auc without clicking.