Bài 40: Experiment Tracking với MLflow

Sau bài này bạn sẽ:

• ✅ Hiểu tại sao tracking là bắt buộc trong quá trình thực nghiệm ML
• ✅ Chạy MLflow tracking local và xem kết quả trên UI
• ✅ Log params, metrics, artifacts đúng cách
• ✅ Dùng autologging cho sklearn / PyTorch / Transformers
• ✅ So sánh nhiều runs theo hyperparameter
• ✅ Cấu hình remote tracking server cho team

Các tình huống thực tế xảy ra khi không có tracking:

• "Em chạy thử LR=0.01 và LR=0.001, không nhớ cái nào cho accuracy cao hơn."
• "Model production tốt nhất đang dùng là từ commit nào? Dataset nào? Hyperparameter nào?"
• "Sao notebook em chạy ra accuracy 89%, anh chạy lại ra 85%?"
• "Hôm qua train xong, hôm nay không biết checkpoint nào tương ứng config nào."

Experiment tracking giải quyết bằng cách ghi lại toàn bộ ngữ cảnh của một lần train:

• Code version: git commit hash
• Data: dataset version, split seed, tiền xử lý
• Hyperparameter: learning rate, batch size, epochs, ...
• Metrics: loss, accuracy, F1 theo từng epoch
• Artifacts: model file, tokenizer, confusion matrix

Khi có tracking, mỗi lần train là một bản ghi có thể tìm lại, so sánh, và reproduce.

MLflow là open-source platform cho machine learning lifecycle, do Databricks phát triển và donate cho Linux Foundation (MLflow 1.0 ra mắt năm 2018). Phiên bản hiện tại: 2.x (2.15+ tính đến 2025).

MLflow gồm 4 component chính:

• Tracking: log run (params, metrics, artifacts). Đây là component được dùng nhiều nhất và là nội dung chính của bài này.
• Projects: định nghĩa môi trường + entry point để reproduce experiment trên bất kỳ máy nào.
• Models: format chuẩn để save/load model (mlflow.sklearn, mlflow.pytorch, mlflow.transformers, ...).
• Model Registry: quản lý stage/version model. Bài 42 sẽ đào sâu phần này.

Bài 40 tập trung vào Tracking. Model Registry giữ nguyên cho bài 42.

pip install mlflow
# Kiểm tra version
mlflow --version
# mlflow, version 2.15.x

MLflow không có dependency nặng khi chỉ dùng tracking. Các extras (PostgreSQL backend, S3 artifact store, ...) cần cài riêng khi setup server.

# Nếu cần PostgreSQL backend
pip install mlflow[extras]
# Hoặc chỉ psycopg2
pip install psycopg2-binary

Không cần server. MLflow tự tạo thư mục ./mlruns/ để lưu dữ liệu.

import mlflow
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

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, random_state=42)

# Đặt tên experiment. Nếu chưa có sẽ tự tạo.
mlflow.set_experiment("iris-classification")

with mlflow.start_run():
    # Log hyperparameter (không thay đổi trong suốt run)
    mlflow.log_param("model_type", "logistic_regression")
    mlflow.log_param("C", 1.0)
    mlflow.log_param("max_iter", 1000)

    model = LogisticRegression(C=1.0, max_iter=1000)
    model.fit(X_train, y_train)

    acc = accuracy_score(y_test, model.predict(X_test))
    # Log metric (kết quả đánh giá)
    mlflow.log_metric("accuracy", acc)

    # Log model artifact theo format chuẩn MLflow
    mlflow.sklearn.log_model(model, "model")

print(f"Accuracy: {acc:.4f}")

Sau khi chạy, thư mục mlruns/ xuất hiện trong working directory:

mlruns/
  0/                         # Default experiment
  1/                         # iris-classification experiment
    <run_id>/
      artifacts/
        model/               # Saved sklearn model
      metrics/
        accuracy             # File chứa giá trị metric
      params/
        C
        max_iter
        model_type
      tags/
        mlflow.runName
        mlflow.source.name

Mỗi run có một run_id unique dạng UUID. Dữ liệu lưu dưới dạng file text — không cần database.

Khởi động UI từ cùng thư mục chứa mlruns/:

mlflow ui --host 0.0.0.0 --port 5000

Truy cập http://localhost:5000. Giao diện gồm:

• Experiments panel (trái): danh sách experiment theo tên.
• Runs table (giữa): mỗi hàng là một run, cột gồm run_id, thời gian bắt đầu, duration, params, metrics.
• Run detail: click vào run_id để xem toàn bộ params, metrics, artifacts, tags, source code.
• Compare: chọn nhiều run → nút "Compare" → side-by-side bảng và biểu đồ metric.

Nếu chạy từ Jupyter Notebook và không muốn mở terminal riêng:

import subprocess
subprocess.Popen(["mlflow", "ui", "--port", "5000"])
# Truy cập http://localhost:5000

Params — Hyperparameter Của Run

# Một key-value
mlflow.log_param("learning_rate", 0.001)

# Nhiều params cùng lúc
mlflow.log_params({
    "learning_rate": 0.001,
    "batch_size": 32,
    "optimizer": "adam",
    "dropout": 0.3,
})

Params là scalar (string / number / bool) không thay đổi trong suốt run. Nếu gọi log_param với cùng key hai lần, MLflow 2.x sẽ raise MlflowException.

Metrics — Kết Quả Đo Lường

# Metric đơn
mlflow.log_metric("accuracy", 0.923)

# Metric theo step (epoch)
for epoch in range(num_epochs):
    loss = train_one_epoch(model, dataloader)
    mlflow.log_metric("train_loss", loss, step=epoch)

# Nhiều metrics cùng lúc
mlflow.log_metrics({
    "precision": 0.91,
    "recall": 0.88,
    "f1": 0.895,
}, step=final_epoch)

Artifacts — File / Thư Mục

import matplotlib.pyplot as plt

# Log 1 file
mlflow.log_artifact("config.yaml")

# Log thư mục
mlflow.log_artifacts("./outputs/")

# Log matplotlib figure trực tiếp
fig, ax = plt.subplots()
ax.plot(train_losses, label="train")
ax.plot(val_losses, label="val")
ax.legend()
mlflow.log_figure(fig, "loss_curve.png")
plt.close(fig)

# Log dict dạng JSON
mlflow.log_dict({"best_epoch": 12, "final_val_loss": 0.32}, "summary.json")

# Log text
mlflow.log_text("Adam optimizer, cosine LR schedule", "notes.txt")

log_model Vs log_artifact

Luôn ưu tiên log_model thay vì log_artifact khi lưu model:

• mlflow.sklearn.log_model(model, "model") — lưu kèm signature, input example, dependencies. Có thể load lại qua mlflow.sklearn.load_model() hoặc deploy trực tiếp.
• mlflow.log_artifact("model.pkl") — chỉ lưu file. Mất signature, không thể dùng với MLflow serving.

from mlflow.models import infer_signature

# Infer signature từ data
signature = infer_signature(X_train, model.predict(X_train))

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

Autologging tự động log params, metrics, model artifact mà không cần gọi từng hàm thủ công. Kích hoạt trước khi train:

# sklearn
mlflow.sklearn.autolog()

# PyTorch Lightning
mlflow.pytorch.autolog()

# TensorFlow / Keras
mlflow.tensorflow.autolog()

# Transformers (Hugging Face)
mlflow.transformers.autolog()

# Spark MLlib
mlflow.spark.autolog()

Ví dụ với sklearn — code không cần thêm bất kỳ lệnh log nào:

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

mlflow.set_experiment("iris-rf-autolog")
mlflow.sklearn.autolog()   # bật autolog

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)

with mlflow.start_run():
    rf = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
    rf.fit(X_train, y_train)
    # MLflow tự log: n_estimators, max_depth, accuracy, f1, precision, recall,
    # model artifact, feature importances

Autologging log gì tùy framework. Xem tài liệu chính thức để biết chính xác với từng phiên bản: mlflow.org/docs/latest/tracking/autolog.html.

Chú ý: Nếu vừa bật autologging vừa gọi manual log_metric với cùng key, MLflow sẽ tạo bản ghi trùng. Chọn một trong hai cách.

Truyền step vào log_metric để MLflow vẽ được biểu đồ theo thời gian:

import mlflow

mlflow.set_experiment("pytorch-train")

with mlflow.start_run(run_name="baseline-v1"):
    mlflow.log_params({
        "lr": 1e-3,
        "batch_size": 32,
        "epochs": 20,
        "architecture": "resnet18",
    })

    for epoch in range(20):
        train_loss, train_acc = train_one_epoch(model, train_loader, optimizer)
        val_loss, val_acc = evaluate(model, val_loader)

        mlflow.log_metrics({
            "train_loss": train_loss,
            "train_acc": train_acc,
            "val_loss": val_loss,
            "val_acc": val_acc,
        }, step=epoch)

    # Log model cuối cùng
    mlflow.pytorch.log_model(model, "model")

Trên MLflow UI, mở run detail → tab Metrics → chọn metric → hiển thị line chart theo step. Có thể overlay nhiều runs trên cùng 1 chart qua tính năng Compare.

Tags là metadata về run, không phải kết quả đo lường. Dùng để gắn nhãn tìm kiếm sau:

import subprocess

with mlflow.start_run():
    # Ghi lại ngữ cảnh của run
    git_commit = subprocess.check_output(
        ["git", "rev-parse", "--short", "HEAD"]
    ).decode().strip()

    mlflow.set_tag("git_commit", git_commit)
    mlflow.set_tag("dataset_version", "v1.2")
    mlflow.set_tag("model_arch", "resnet18")
    mlflow.set_tag("author", "nam.nguyen")
    mlflow.set_tag("notes", "Thử tăng weight decay")

    # Hoặc set nhiều tags cùng lúc
    mlflow.set_tags({
        "environment": "gpu-a100",
        "purpose": "baseline",
    })

    # ... train ...

MLflow tự động gắn một số tag có tiền tố mlflow.: mlflow.source.name (file script), mlflow.source.git.commit (nếu chạy trong git repo), mlflow.user.

Qua UI

Trong MLflow UI: vào experiment → tick chọn nhiều runs → nút "Compare". Trang compare hiển thị:

• Bảng side-by-side cho params và metrics.
• Scatter plot: trục x/y chọn bất kỳ metric hoặc param.
• Contour plot cho hyperparameter sweep 2 biến.
• Line chart metrics theo step (overlay nhiều runs).

Qua Python API

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Tìm experiment theo tên
experiment = client.get_experiment_by_name("iris-classification")
exp_id = experiment.experiment_id

# Query runs, sắp xếp theo accuracy giảm dần
runs = client.search_runs(
    experiment_ids=[exp_id],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10,
)

for run in runs:
    print(
        f"run_id={run.info.run_id[:8]} "
        f"accuracy={run.data.metrics.get('accuracy', 'N/A'):.4f} "
        f"C={run.data.params.get('C', 'N/A')}"
    )

filter_string dùng cú pháp SQL-like: "metrics.val_loss < 0.5 AND params.optimizer = 'adam'".

Local file backend chỉ phù hợp khi làm một mình. Khi team cần share kết quả hoặc nhiều máy train song song, cần tracking server tập trung.

Kiến Trúc

MLflow tracking server tách biệt hai loại lưu trữ:

• Backend store: metadata (run_id, params, metrics, tags). Dùng SQLite (đơn giản), PostgreSQL, hoặc MySQL.
• Artifact store: file lớn (model checkpoint, hình ảnh, CSV). Dùng local path, S3, GCS, Azure Blob, hoặc NFS.

Khởi Động Server

# Ví dụ với PostgreSQL + S3
mlflow server \
  --backend-store-uri postgresql://user:pass@db-host:5432/mlflowdb \
  --artifacts-destination s3://my-bucket/mlflow-artifacts \
  --host 0.0.0.0 \
  --port 5000

# Ví dụ đơn giản hơn với SQLite + local artifacts
mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --artifacts-destination ./artifacts \
  --host 0.0.0.0 \
  --port 5000

Cấu Hình Client

import mlflow

# Cách 1: Set trong code
mlflow.set_tracking_uri("http://mlflow.internal:5000")

# Cách 2: Biến môi trường (khuyến nghị cho production)
# export MLFLOW_TRACKING_URI=http://mlflow.internal:5000

Khi dùng S3 artifact store, machine train cần có AWS credential (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) hoặc IAM role. MLflow client upload artifact trực tiếp lên S3, không qua tracking server.

Thêm MLflow tracking vào training script Transformers mà không thay đổi logic train:

import mlflow
from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments

mlflow.set_tracking_uri("http://mlflow.internal:5000")
mlflow.set_experiment("transformer-finetune-2025-q2")

config = {
    "model_name": "bert-base-uncased",
    "learning_rate": 2e-5,
    "num_epochs": 3,
    "batch_size": 16,
    "warmup_steps": 500,
    "weight_decay": 0.01,
}

with mlflow.start_run(run_name="bert-baseline") as run:
    # Log toàn bộ config một lần
    mlflow.log_params(config)
    mlflow.set_tag("dataset", "imdb-v2")

    model = AutoModelForSequenceClassification.from_pretrained(config["model_name"])
    tokenizer = AutoTokenizer.from_pretrained(config["model_name"])

    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=config["num_epochs"],
        per_device_train_batch_size=config["batch_size"],
        learning_rate=config["learning_rate"],
    )

    trainer = Trainer(model=model, args=training_args, ...)
    trainer.train()

    eval_results = trainer.evaluate()
    mlflow.log_metrics(eval_results)

    # Log model + tokenizer theo format MLflow
    mlflow.transformers.log_model(
        transformers_model={"model": model, "tokenizer": tokenizer},
        artifact_path="model",
        task="text-classification",
    )

    print(f"Run ID: {run.info.run_id}")

Mỗi combination hyperparameter là một run riêng biệt. Tất cả runs thuộc cùng experiment để so sánh dễ:

import itertools
import mlflow

mlflow.set_experiment("lr-batchsize-sweep")

learning_rates = [1e-3, 1e-4, 1e-5]
batch_sizes = [16, 32, 64]

for lr, batch_size in itertools.product(learning_rates, batch_sizes):
    with mlflow.start_run():
        mlflow.log_params({"lr": lr, "batch_size": batch_size})

        model, val_acc = train_and_evaluate(lr, batch_size)

        mlflow.log_metric("val_acc", val_acc)
        mlflow.sklearn.log_model(model, "model")

Sau khi chạy xong (9 runs trong ví dụ trên), vào MLflow UI → chọn tất cả → Compare → Scatter plot với trục x=lr, y=val_acc để thấy trend.

Nếu cần sweep phức tạp hơn (Bayesian optimization, pruning), MLflow tích hợp tốt với Optuna hoặc Ray Tune — các tool đó gọi mlflow.log_metric trong callback.

MLflow Projects chuẩn hóa cách chạy experiment trên bất kỳ máy nào. File MLproject (YAML) đặt ở root repo định nghĩa:

# MLproject
name: iris-experiment

conda_env: conda.yaml   # hoặc docker_env: ...

entry_points:
  main:
    parameters:
      lr: {type: float, default: 0.01}
      max_iter: {type: int, default: 1000}
    command: "python train.py --lr {lr} --max_iter {max_iter}"

Chạy experiment:

# Chạy với param mặc định
mlflow run .

# Chạy với param cụ thể
mlflow run . -P lr=0.001 -P max_iter=2000

# Chạy từ remote repo
mlflow run https://github.com/your-org/your-ml-repo -P lr=0.001

MLflow tự tạo môi trường conda/Docker, cài dependency, rồi chạy entry point. Người khác chỉ cần clone repo + mlflow run . là có kết quả y hệt.

1. Log Ngoài start_run

# SAI — log_param không có active run
mlflow.log_param("lr", 0.001)  # Sẽ tạo run auto, dễ gây nhầm lẫn

# ĐÚNG — luôn dùng context manager
with mlflow.start_run():
    mlflow.log_param("lr", 0.001)

Khi gọi log function mà không có active run, MLflow 2.x tự động tạo run mới. Đây thường là ngoài ý muốn — kết thúc bằng nhiều run rỗng trong experiment.

2. log_artifact Thay Vì log_model

import pickle

# TRÁNH — chỉ lưu file, mất signature và serving capability
with open("model.pkl", "wb") as f:
    pickle.dump(model, f)
mlflow.log_artifact("model.pkl")

# NÊN DÙNG — lưu theo format chuẩn MLflow
mlflow.sklearn.log_model(model, "model")

3. Tracking URI Sai

# Nếu quên set tracking URI, log vào ./mlruns/ local
# Không thấy run trên remote server
mlflow.set_tracking_uri("http://mlflow.internal:5000")  # Phải set trước log đầu tiên

# Kiểm tra URI đang dùng
print(mlflow.get_tracking_uri())

4. Artifact Store Đầy Do Lưu Quá Nhiều Checkpoint

Log checkpoint mỗi epoch với model lớn (GPT-2 ~500MB) trong 100 epochs = 50GB một run. Chỉ log checkpoint khi cần thiết hoặc chỉ log checkpoint tốt nhất:

best_val_loss = float("inf")

for epoch in range(num_epochs):
    val_loss = validate(model, val_loader)
    mlflow.log_metric("val_loss", val_loss, step=epoch)

    # Chỉ log model khi cải thiện
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        mlflow.pytorch.log_model(model, "best_model")

5. Autolog Và Manual Log Trùng Key

mlflow.sklearn.autolog()  # tự log "accuracy"

with mlflow.start_run():
    model.fit(X_train, y_train)
    # Autolog đã log accuracy

    # Gọi thêm sẽ tạo duplicate
    mlflow.log_metric("accuracy", acc)  # TRÁNH nếu autolog đã log cùng key

Nếu cần log metric với logic tùy chỉnh, tắt autolog cho metric đó: mlflow.sklearn.autolog(log_post_training_metrics=False).

Bài 41: Experiment Tracking với Weights & Biases (W&B) — so sánh cách tiếp cận của W&B với MLflow, khi nào nên dùng tool nào.