LabHub

Blog

MLflow Complete Guide

한국어English日本語

1. MLflow Overview and Installation

1.1 What Is MLflow

MLflow is an open-source platform for managing the entire machine learning lifecycle. Started at Databricks, the project is now run under the Linux Foundation, and it covers the core areas of an ML workflow: experiment tracking (Experiment Tracking), code packaging (Projects), model management (Models), and the model registry (Model Registry).

The core problems MLflow solves are as follows.

1.2 How to Install

MLflow installs easily from PyPI.

pip install mlflow

If you need integration with a particular ML framework, install the extras along with it.

# including sklearn integration
pip install mlflow[sklearn]

# install all extras
pip install mlflow[extras]

After installing, run the MLflow UI to confirm it works.

mlflow server --port 5000

Opening http://localhost:5000 in a browser brings up MLflow's web UI.

1.3 Basic Usage (QuickStart)

MLflow's basic workflow is to set up an Experiment, start a Run, and record parameters, metrics, and artifacts.

import mlflow

# Set up the Experiment
mlflow.set_experiment("MLflow Quickstart")

# Start the Run and log
with mlflow.start_run():
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("epochs", 100)
    mlflow.log_metric("accuracy", 0.95)
    mlflow.log_metric("loss", 0.05)
    mlflow.set_tag("model_type", "classification")

Running the code above stores the experiment data in a local mlruns directory, and you can inspect the results visually in the MLflow UI.


2. MLflow Tracking: Experiment Logging

2.1 Core Concepts

MLflow Tracking is the system that records and queries metadata about ML experiments. Its core building blocks are as follows.

2.2 The Logging API in Detail

mlflow.log_param() / mlflow.log_params()

Records hyperparameters. You can record a single value, or many parameters at once as a dictionary.

with mlflow.start_run():
    # Record a single parameter
    mlflow.log_param("learning_rate", 0.001)
    mlflow.log_param("batch_size", 32)

    # Record several parameters at once
    params = {
        "optimizer": "adam",
        "dropout": 0.3,
        "hidden_layers": 3,
        "activation": "relu"
    }
    mlflow.log_params(params)

mlflow.log_metric() / mlflow.log_metrics()

Records training metrics. Using the step parameter lets you track how a metric changes over the course of training as a time series.

with mlflow.start_run():
    for epoch in range(100):
        train_loss = train_one_epoch(model, train_loader)
        val_loss = evaluate(model, val_loader)

        # Record the metric per step (time-series tracking)
        mlflow.log_metric("train_loss", train_loss, step=epoch)
        mlflow.log_metric("val_loss", val_loss, step=epoch)

    # Record the final metrics all at once
    final_metrics = {"final_accuracy": 0.95, "final_f1": 0.93}
    mlflow.log_metrics(final_metrics)

mlflow.log_artifact() / mlflow.log_artifacts()

Saves a file or directory as an artifact. This suits managing large files such as model checkpoints, visualization images, and preprocessing pipelines.

import matplotlib.pyplot as plt

with mlflow.start_run():
    # Save the learning curve visualization
    plt.figure(figsize=(10, 6))
    plt.plot(train_losses, label="Train Loss")
    plt.plot(val_losses, label="Validation Loss")
    plt.legend()
    plt.savefig("learning_curve.png")

    # Record a single file as an artifact
    mlflow.log_artifact("learning_curve.png", artifact_path="plots")

    # Record an entire directory as an artifact
    mlflow.log_artifacts("./output_dir", artifact_path="results")

mlflow.log_input()

You can link the Dataset used for training to the Run and track data lineage.

import mlflow.data
from mlflow.data.pandas_dataset import PandasDataset

dataset = mlflow.data.from_pandas(df, source="s3://my-bucket/data.csv")

with mlflow.start_run():
    mlflow.log_input(dataset, context="training")

2.3 Searching and Comparing Experiments

mlflow.search_runs() lets you search and compare Runs with SQL-like syntax.

import mlflow

# Search for Runs by a specific condition
runs = mlflow.search_runs(
    experiment_names=["MLflow Quickstart"],
    filter_string="metrics.accuracy > 0.9 AND params.optimizer = 'adam'",
    order_by=["metrics.accuracy DESC"]
)
print(runs[["params.optimizer", "metrics.accuracy"]])

3. Building a Tracking Server

3.1 Architecture Overview

In production you build a centralized Tracking Server instead of using the local file system. A Tracking Server is made up of two core stores.

3.2 Tracking Server Configuration Scenarios

Scenario 1: Local Development (Default)

Stores all data in the mlruns directory with no extra configuration.

mlflow server --port 5000

Scenario 2: DB Backend + Local Artifact Store

Metadata goes to PostgreSQL and artifacts go to the local file system.

mlflow server \
    --backend-store-uri postgresql://user:password@localhost:5432/mlflow_db \
    --default-artifact-root ./mlartifacts \
    --port 5000

Metadata goes to PostgreSQL and artifacts go to S3, with an Artifact Proxy keeping clients from accessing S3 directly.

mlflow server \
    --backend-store-uri postgresql://user:password@db-host:5432/mlflow_db \
    --artifacts-destination s3://my-mlflow-bucket/artifacts \
    --port 5000

With the --artifacts-destination flag, the Tracking Server acts as a Proxy for artifact access. Clients can upload and download artifacts with plain HTTP requests, without S3 credentials.

If instead you want clients to access the Artifact Store directly, configure it as follows.

mlflow server \
    --backend-store-uri postgresql://user:password@db-host:5432/mlflow_db \
    --default-artifact-root s3://my-mlflow-bucket/artifacts \
    --no-serve-artifacts \
    --port 5000

3.3 Client Configuration

Once the Tracking Server is up, set the URI on the client.

import mlflow

mlflow.set_tracking_uri("http://tracking-server-host:5000")
mlflow.set_experiment("my-experiment")

You can also set it with an environment variable.

export MLFLOW_TRACKING_URI=http://tracking-server-host:5000

4. Analyzing the Autologging Feature

4.1 Overview

Autologging records parameters, metrics, and models automatically, with no explicit logging code. Adding a single line of code gets you experiment tracking.

# universal autolog (enables it for every supported library that is installed)
mlflow.autolog()

4.2 Scikit-learn Autologging

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris

mlflow.sklearn.autolog()

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

with mlflow.start_run():
    clf = RandomForestClassifier(n_estimators=100, max_depth=5)
    clf.fit(X_train, y_train)
    # Automatically recorded items:
    # - Parameters: n_estimators, max_depth, criterion, and so on
    # - Metrics: accuracy, precision, recall, f1-score
    # - Artifacts: the trained model, confusion matrix, feature importance

sklearn autologging automatically records every hyperparameter, the training metrics, and the trained model when fit() is called.

4.3 PyTorch Lightning Autologging

The integration with PyTorch Lightning offers the most complete autologging. Full autologging is supported for models that inherit from pytorch_lightning.LightningModule.

import mlflow
import pytorch_lightning as pl
from torch.utils.data import DataLoader

mlflow.pytorch.autolog()

class MyModel(pl.LightningModule):
    def __init__(self, lr=0.001):
        super().__init__()
        self.save_hyperparameters()
        self.model = torch.nn.Linear(10, 2)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = torch.nn.functional.cross_entropy(self.model(x), y)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, train_dataloader)
# Automatically recorded items:
# - Optimizer name and learning rate
# - training loss, validation loss
# - Model checkpoints and artifacts

4.4 Transformers Autologging

Integration with the Hugging Face Transformers library is supported too.

import mlflow
from transformers import Trainer, TrainingArguments

mlflow.transformers.autolog()

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    logging_steps=100,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

trainer.train()
# Automatically recorded: training metrics, model, tokenizer, training arguments, and so on

4.5 Full List of Supported Frameworks

The main frameworks MLflow autologging supports are as follows.

FrameworkHow to call it
Scikit-learnmlflow.sklearn.autolog()
XGBoostmlflow.xgboost.autolog()
LightGBMmlflow.lightgbm.autolog()
PyTorch (Lightning)mlflow.pytorch.autolog()
Keras / TensorFlowmlflow.tensorflow.autolog()
Transformersmlflow.transformers.autolog()
Spark MLlibmlflow.spark.autolog()
Statsmodelsmlflow.statsmodels.autolog()
CatBoostmlflow.catboost.autolog()
Prophetmlflow.prophet.autolog()

Calling mlflow.autolog() enables autologging automatically for every supported library that is installed.


5. MLflow Projects: Reproducible Experiments

5.1 Overview

MLflow Projects is a convention for packaging ML code in a reproducible, portable form. Any directory can be run as an MLflow Project as long as it contains an MLproject file, or contains .py or .sh files.

5.2 MLproject File Structure

name: My ML Project

python_env: python_env.yaml

entry_points:
  main:
    parameters:
      data_file: path
      regularization: { type: float, default: 0.1 }
      epochs: { type: int, default: 100 }
    command: 'python train.py -r {regularization} -e {epochs} {data_file}'

  validate:
    parameters:
      data_file: path
    command: 'python validate.py {data_file}'

Core Building Blocks

5.3 Supported Environments

MLflow Projects supports four execution environments.

EnvironmentDescription
VirtualenvPython virtual environment from python_env.yaml
CondaConda environment from conda.yaml
DockerContainer environment from a Dockerfile
SystemUses the current system environment as-is

An example python_env.yaml file:

python: '3.10'
build_dependencies:
  - pip
dependencies:
  - scikit-learn==1.3.0
  - pandas==2.0.3
  - mlflow

5.4 Running a Project

Running from the CLI

# Run a local project
mlflow run . -P regularization=0.5 -P data_file=data/train.csv

# Run directly from a Git URI
mlflow run https://github.com/mlflow/mlflow-example -P alpha=0.5

# Run a specific entry point
mlflow run . -e validate -P data_file=data/test.csv

Running from the Python API

import mlflow

mlflow.projects.run(
    uri="https://github.com/mlflow/mlflow-example",
    entry_point="main",
    parameters={"alpha": 0.5, "l1_ratio": 0.1}
)

Because you can point directly at a Git URI, a team can share experiment code through a single Git repository and reproduce it in an identical environment.


6. MLflow Models: The Flavor Concept and Model Signature

6.1 The Model Packaging Standard

MLflow Models is a convention for packaging machine learning models in a standard format. Following this standard lets you use the model from a variety of downstream tools, such as REST API serving and Apache Spark batch inference.

6.2 The Flavor Concept

Flavor is the core innovation of MLflow Models. A Flavor defines how a model is interpreted and used. Each ML framework has its own Flavor, and deployment tools recognize the standard Flavors and can therefore use the model without any library-specific custom integration.

For example, a model trained with sklearn carries both the sklearn flavor and the python_function flavor. Through the python_function flavor, a deployment tool can run inference in a framework-independent way.

Directory Structure of a Saved Model

my_model/
├── MLmodel              # Model metadata (YAML)
├── model.pkl            # Serialized model file
├── conda.yaml           # Conda environment definition
├── python_env.yaml      # Python environment definition
├── requirements.txt     # pip dependencies
└── input_example.json   # Input example (optional)

An example MLmodel file:

artifact_path: model
flavors:
  python_function:
    env: conda.yaml
    loader_module: mlflow.sklearn
    model_path: model.pkl
    python_version: 3.10.12
  sklearn:
    code: null
    pickled_model: model.pkl
    serialization_format: cloudpickle
    sklearn_version: 1.3.0
mlflow_version: 2.10.0
signature:
  inputs: '[{"name": "feature_1", "type": "double"}, {"name": "feature_2", "type": "double"}]'
  outputs: '[{"type": "long"}]'

6.3 Built-in Flavors

MLflow provides more than 20 built-in flavors.

6.4 Model Signature

A Model Signature defines the schema of the model's inputs, outputs, and additional inference parameters. It standardizes the model interface, and input data is validated automatically at serving time.

from mlflow.models import infer_signature
import mlflow.sklearn

X_train, y_train = load_data()
model = train_model(X_train, y_train)
predictions = model.predict(X_train)

# Infer the Signature automatically
signature = infer_signature(X_train, predictions)

with mlflow.start_run():
    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        signature=signature,
        input_example=X_train[:3]
    )

6.5 Input Example

An Input Example provides a concrete instance of valid input for the model. It lets a user of the model understand the input format easily and test with it. When an Input Example is provided and the Signature has not been specified explicitly, the Signature is inferred automatically.

Also, when an Input Example is provided, a serving_input_example.json file is generated automatically so you can refer to the payload format at serving time.

6.6 Logging and Loading Models

# Log the model
with mlflow.start_run():
    mlflow.sklearn.log_model(model, artifact_path="model")

# Load the model (framework-specific)
loaded_model = mlflow.sklearn.load_model("runs:/<run_id>/model")

# Load the model (generic pyfunc)
pyfunc_model = mlflow.pyfunc.load_model("runs:/<run_id>/model")
predictions = pyfunc_model.predict(X_test)

6.7 Models From Code

Introduced after MLflow 2.12.2, this feature lets you log a model directly from a Python script. It bypasses pickle serialization and so reduces the security risk.

# model_code.py
import mlflow

class MyCustomModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return model_input * 2

mlflow.models.set_model(MyCustomModel())

7. Model Registry: Version Management and Stage Transitions

7.1 Overview

The MLflow Model Registry is a centralized model store for version management, metadata management, and deployment workflows around trained models.

7.2 Core Concepts

7.3 Ways to Register a Model

Method 1: Register at Logging Time

with mlflow.start_run():
    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="model",
        registered_model_name="MyClassifier"  # log and register at the same time
    )

Method 2: Register from an Existing Run

result = mlflow.register_model(
    model_uri="runs:/<run_id>/model",
    name="MyClassifier"
)
print(f"Version: {result.version}")

Method 3: Register from the UI

In the MLflow web UI you can select a Run's artifact and click the "Register Model" button to register it.

7.4 Model Aliases

An Alias is a mutable name reference to a specific model version. Using an Alias instead of a version number in deployment code lets you switch the deployed model without a code change.

from mlflow import MlflowClient

client = MlflowClient()

# Set the Alias
client.set_registered_model_alias(
    name="MyClassifier",
    alias="champion",
    version=3
)

# Load the model by Alias
model = mlflow.sklearn.load_model("models:/MyClassifier@champion")

If you move the champion alias from version 3 to version 5, production code keeps using models:/MyClassifier@champion unchanged and automatically loads the new version of the model.

7.5 Tags and Annotations

Tags are key-value pairs used to classify and search models.

client = MlflowClient()

# Registered Model level tag
client.set_registered_model_tag("MyClassifier", "task", "question-answering")

# Model Version level tag
client.set_model_version_tag("MyClassifier", version=3, key="validation_status", value="approved")

Annotations support detailed descriptions in Markdown, so you can document the model's methodology, the datasets used, the algorithm, and so on.

7.6 Stage-Based Workflow (Legacy)

Note: recent versions of MLflow recommend an Alias-based workflow rather than Stages.

In the older Stage-based workflow, a model version moves through the following Stages.

NoneStagingProductionArchived
# Stage transition (Legacy API)
client.transition_model_version_stage(
    name="MyClassifier",
    version=3,
    stage="Production"
)

# Load a model by Stage (Legacy)
model = mlflow.sklearn.load_model("models:/MyClassifier/Production")

The currently recommended approach is to use Aliases. For example, you define aliases such as champion and challenger to manage the deployment workflow.


8. MLflow Deployments (formerly MLflow AI Gateway)

8.1 Overview

MLflow Deployments (previously named MLflow AI Gateway) is a gateway service for managing multiple LLM providers together across an organization. Through a single secure endpoint you can reach a range of providers, including OpenAI, Anthropic, Google Gemini, Amazon Bedrock, and Azure OpenAI.

8.2 Supported Providers

8.3 Gateway Configuration

Providers and endpoints are defined in a YAML configuration file.

endpoints:
  - name: chat-gpt4
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4
      config:
        openai_api_key: $OPENAI_API_KEY

  - name: chat-claude
    endpoint_type: llm/v1/chat
    model:
      provider: anthropic
      name: claude-3-opus-20240229
      config:
        anthropic_api_key: $ANTHROPIC_API_KEY

  - name: embeddings
    endpoint_type: llm/v1/embeddings
    model:
      provider: openai
      name: text-embedding-3-small
      config:
        openai_api_key: $OPENAI_API_KEY

8.4 Key Features

8.5 Model Serving (Traditional ML)

Traditional ML models are deployed with the mlflow models serve command.

# Serve the model as a local REST API server
mlflow models serve -m "models:/MyClassifier@champion" --port 8080

# Build a Docker image
mlflow models build-docker -m "models:/MyClassifier@champion" -n my-model-image

# Run the Docker container
docker run -p 8080:8080 my-model-image

How to send an inference request to the served model:

curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"inputs": [[1.0, 2.0, 3.0, 4.0]]}'

9. MLflow Recipes (formerly MLflow Pipelines)

9.1 Overview

MLflow Recipes (previously named MLflow Pipelines) is a framework that helps data scientists develop high-quality models quickly and get them into production. Predefined templates remove the repetitive boilerplate needed for data ingestion, feature engineering, model training and tuning, and model packaging.

9.2 Core Concepts

9.3 Templates Provided

Regression Template

Provides a standard pipeline for regression problems.

Classification Template

Provides a standard pipeline for classification problems.

Both templates follow the same Step order.

ingest → split → transform → train → evaluate → register

9.4 Configuration Structure

Recipe configuration is managed with a recipe.yaml file plus Profile YAML files.

# recipe.yaml
recipe: 'regression/v1'
target_col: 'price'
positive_class: null
primary_metric: 'root_mean_squared_error'

steps:
  ingest:
    using: 'custom'
    loader_method: load_data

  split:
    split_ratios: [0.75, 0.125, 0.125]

  transform:
    using: 'custom'
    transformer_method: transformer_fn

  train:
    using: 'custom'
    estimator_method: estimator_fn

  evaluate:
    validation_criteria:
      - metric: root_mean_squared_error
        threshold: 10

  register:
    model_name: 'my_regression_model'

9.5 Execution and Caching

from mlflow.recipes import Recipe

recipe = Recipe(profile="local")

# Run the whole Recipe
recipe.run()

# Run only up to a specific Step
recipe.run("train")

# Inspect the Step results
recipe.inspect("evaluate")

MLflow Recipes' intelligent execution engine caches the result of each Step and re-runs only the minimum set of Steps affected by a change, which speeds up development.


10. Integrating MLflow with Other Tools

10.1 MLflow + Apache Airflow

Airflow is a workflow orchestration tool, and combining it with MLflow's experiment tracking lets you build a powerful ML pipeline. The division of labor is that Airflow handles workflow scheduling and orchestration while MLflow handles experiment metric logging, model version management, and model lifecycle management.

# Example of using MLflow in an Airflow DAG
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def train_model():
    import mlflow
    mlflow.set_tracking_uri("http://mlflow-server:5000")

    with mlflow.start_run():
        model = train()
        mlflow.sklearn.log_model(model, "model", registered_model_name="MyModel")

def evaluate_model():
    import mlflow
    model = mlflow.sklearn.load_model("models:/MyModel/latest")
    metrics = evaluate(model)
    return metrics

dag = DAG("ml_pipeline", start_date=datetime(2026, 1, 1), schedule_interval="@daily")

train_task = PythonOperator(task_id="train", python_callable=train_model, dag=dag)
eval_task = PythonOperator(task_id="evaluate", python_callable=evaluate_model, dag=dag)

train_task >> eval_task

10.2 MLflow + Kubernetes

Deploying an MLflow model to a Kubernetes cluster lets you take advantage of Kubernetes infrastructure features such as automatic scaling, rolling updates, and health checks.

Building the Docker Image

mlflow models build-docker \
    -m "models:/MyClassifier@champion" \
    -n my-model-serving:v1 \
    --enable-mlserver

Kubernetes Deployment Manifest

apiVersion: apps/v1
kind: Deployment
metadata:
  name: ml-model-serving
spec:
  replicas: 3
  selector:
    matchLabels:
      app: ml-model
  template:
    metadata:
      labels:
        app: ml-model
    spec:
      containers:
        - name: model
          image: my-model-serving:v1
          ports:
            - containerPort: 8080
          resources:
            requests:
              memory: '512Mi'
              cpu: '500m'
            limits:
              memory: '1Gi'
              cpu: '1000m'
---
apiVersion: v1
kind: Service
metadata:
  name: ml-model-service
spec:
  selector:
    app: ml-model
  ports:
    - port: 80
      targetPort: 8080
  type: LoadBalancer

You can also integrate with KServe (formerly KFServing) to use more advanced serving features such as canary deployment and traffic mirroring.

10.3 MLflow + Docker

MLflow supports Docker for both project execution and model serving.

Project execution: specifying a Docker environment in the MLproject file guarantees reproducible execution inside a container.

# MLproject
name: My Docker Project
docker_env:
  image: my-ml-env:latest
  volumes: ['/data:/data']
  environment: [['MLFLOW_TRACKING_URI', 'http://mlflow-server:5000']]
entry_points:
  main:
    command: 'python train.py'

Model serving: build a serving image with mlflow models build-docker and run the same inference service in any environment.

10.4 Deploying the MLflow Tracking Server with Docker Compose

An example of a production Tracking Server composed with Docker Compose.

version: '3.8'
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:password@postgres:5432/mlflow
      --artifacts-destination s3://mlflow-artifacts/
      --host 0.0.0.0
      --port 5000
    ports:
      - '5000:5000'
    environment:
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    depends_on:
      - postgres

  postgres:
    image: postgres:15
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: password
      POSTGRES_DB: mlflow
    volumes:
      - postgres_data:/var/lib/postgresql/data

volumes:
  postgres_data:

11. References

A list of the official MLflow documentation and related resources referenced in this post.

Comments

No comments yet.

Sign in to leave a comment