MLflow 实验管理完全指南：实验追踪·模型注册表·部署流水线搭建

# tracking_server_config.py
"""
MLflow Tracking Server 生产环境配置
Backend Store: PostgreSQL
Artifact Store: S3
"""

import os

TRACKING_CONFIG = {
    "backend_store_uri": "postgresql://mlflow:password@db-host:5432/mlflow",
    "default_artifact_root": "s3://mlflow-artifacts/experiments",
    "host": "0.0.0.0",
    "port": 5000,
    "workers": 4,
}

# 运行 MLflow Tracking Server
mlflow server \
  --backend-store-uri postgresql://mlflow:password@db-host:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/experiments \
  --host 0.0.0.0 \
  --port 5000 \
  --workers 4

# 用 Docker Compose 运行
docker compose up -d mlflow-server

# docker-compose.yaml
version: '3.8'
services:
  mlflow-db:
    image: postgres:15
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: mlflow_password
    volumes:
      - pgdata:/var/lib/postgresql/data
    ports:
      - '5432:5432'

  mlflow-server:
    build: ./mlflow
    depends_on:
      - mlflow-db
    environment:
      MLFLOW_BACKEND_STORE_URI: postgresql://mlflow:mlflow_password@mlflow-db:5432/mlflow
      MLFLOW_DEFAULT_ARTIFACT_ROOT: s3://mlflow-artifacts/experiments
      AWS_ACCESS_KEY_ID: your-access-key
      AWS_SECRET_ACCESS_KEY: your-secret-key
    ports:
      - '5000:5000'
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:mlflow_password@mlflow-db:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/experiments
      --host 0.0.0.0
      --port 5000
      --workers 4

volumes:
  pgdata:

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.datasets import load_iris

# 连接 Tracking Server
mlflow.set_tracking_uri("http://mlflow-server:5000")

# 创建实验或选择已有实验
mlflow.set_experiment("iris-classification")

# 准备数据
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)

# 运行实验
with mlflow.start_run(run_name="rf-baseline-v1") as run:
    # 记录超参数
    params = {
        "n_estimators": 100,
        "max_depth": 5,
        "min_samples_split": 2,
        "random_state": 42,
    }
    mlflow.log_params(params)

    # 训练模型
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 预测并计算指标
    y_pred = model.predict(X_test)
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1_macro": f1_score(y_test, y_pred, average="macro"),
        "precision_macro": precision_score(y_test, y_pred, average="macro"),
        "recall_macro": recall_score(y_test, y_pred, average="macro"),
    }
    mlflow.log_metrics(metrics)

    # 记录模型产物
    mlflow.sklearn.log_model(
        model,
        artifact_path="model",
        registered_model_name="iris-classifier",
    )

    # 记录额外产物（例如混淆矩阵图像）
    import matplotlib.pyplot as plt
    from sklearn.metrics import ConfusionMatrixDisplay

    fig, ax = plt.subplots(figsize=(8, 6))
    ConfusionMatrixDisplay.from_predictions(y_test, y_pred, ax=ax)
    fig.savefig("/tmp/confusion_matrix.png")
    mlflow.log_artifact("/tmp/confusion_matrix.png", "plots")

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

import mlflow
import mlflow.sklearn
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score

# 启用自动日志记录
mlflow.sklearn.autolog(
    log_input_examples=True,      # 保存输入数据示例
    log_model_signatures=True,     # 自动检测模型签名
    log_models=True,               # 自动保存模型产物
    log_datasets=True,             # 保存训练数据集信息
    silent=False,                  # 显示日志消息
)

mlflow.set_experiment("iris-autolog-experiment")

with mlflow.start_run(run_name="gbc-autolog"):
    model = GradientBoostingClassifier(
        n_estimators=200,
        max_depth=3,
        learning_rate=0.1,
        random_state=42,
    )
    # autolog 会在调用 fit 时自动记录参数/指标/模型
    model.fit(X_train, y_train)

    # cross-validation 分数也会被自动记录
    cv_scores = cross_val_score(model, X_train, y_train, cv=5)
    mlflow.log_metric("cv_mean_accuracy", cv_scores.mean())
    mlflow.log_metric("cv_std_accuracy", cv_scores.std())

import mlflow
import mlflow.pytorch
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

mlflow.set_experiment("pytorch-classification")

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.3)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 训练配置
config = {
    "input_dim": 4,
    "hidden_dim": 64,
    "output_dim": 3,
    "learning_rate": 0.001,
    "epochs": 50,
    "batch_size": 16,
}

with mlflow.start_run(run_name="pytorch-simplenet"):
    mlflow.log_params(config)

    model = SimpleNet(config["input_dim"], config["hidden_dim"], config["output_dim"])
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config["learning_rate"])

    X_tensor = torch.FloatTensor(X_train)
    y_tensor = torch.LongTensor(y_train)
    dataset = TensorDataset(X_tensor, y_tensor)
    dataloader = DataLoader(dataset, batch_size=config["batch_size"], shuffle=True)

    for epoch in range(config["epochs"]):
        model.train()
        total_loss = 0
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        # 按 epoch 记录指标
        mlflow.log_metric("train_loss", avg_loss, step=epoch)

        # 验证
        model.eval()
        with torch.no_grad():
            X_test_tensor = torch.FloatTensor(X_test)
            test_outputs = model(X_test_tensor)
            _, predicted = torch.max(test_outputs, 1)
            val_acc = (predicted.numpy() == y_test).mean()
            mlflow.log_metric("val_accuracy", val_acc, step=epoch)

    # 保存模型
    mlflow.pytorch.log_model(model, "pytorch-model")

import mlflow
from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="http://mlflow-server:5000")

# 查询特定实验的全部 Run
experiment = client.get_experiment_by_name("iris-classification")
runs = client.search_runs(
    experiment_ids=[experiment.experiment_id],
    filter_string="metrics.accuracy > 0.9 AND params.n_estimators = '100'",
    order_by=["metrics.f1_macro DESC"],
    max_results=10,
)

# 输出结果
for run in runs:
    print(f"Run ID: {run.info.run_id}")
    print(f"  Accuracy: {run.data.metrics.get('accuracy', 'N/A')}")
    print(f"  F1 Score: {run.data.metrics.get('f1_macro', 'N/A')}")
    print(f"  Params: {run.data.params}")
    print("---")

# 比较两个 Run
run1 = runs[0]
run2 = runs[1] if len(runs) > 1 else None

if run2:
    print("=== Run Comparison ===")
    for metric_key in run1.data.metrics:
        v1 = run1.data.metrics[metric_key]
        v2 = run2.data.metrics.get(metric_key, "N/A")
        print(f"  {metric_key}: {v1} vs {v2}")

from mlflow.tracking import MlflowClient

client = MlflowClient()

# 注册模型（直接从训练 Run 注册）
model_name = "iris-classifier"
result = mlflow.register_model(
    model_uri=f"runs:/{run.info.run_id}/model",
    name=model_name,
)
print(f"Model Version: {result.version}")

# 为模型版本添加描述
client.update_model_version(
    name=model_name,
    version=result.version,
    description="RandomForest baseline model with 100 trees, accuracy 0.95",
)

# 为模型版本添加标签
client.set_model_version_tag(
    name=model_name,
    version=result.version,
    key="validation_status",
    value="approved",
)

from mlflow.tracking import MlflowClient

client = MlflowClient()
model_name = "iris-classifier"

# Alias 方式（MLflow 2.x 推荐）
# 设置 champion alias
client.set_registered_model_alias(
    name=model_name,
    alias="champion",
    version=3,
)

# 设置 challenger alias
client.set_registered_model_alias(
    name=model_name,
    alias="challenger",
    version=4,
)

# 通过 Alias 加载模型
champion_model = mlflow.pyfunc.load_model(f"models:/{model_name}@champion")
challenger_model = mlflow.pyfunc.load_model(f"models:/{model_name}@challenger")

# 比较预测结果
champion_pred = champion_model.predict(X_test)
challenger_pred = challenger_model.predict(X_test)

print(f"Champion Accuracy: {accuracy_score(y_test, champion_pred)}")
print(f"Challenger Accuracy: {accuracy_score(y_test, challenger_pred)}")

# 若 Challenger 更优，则晋升为 Champion
if accuracy_score(y_test, challenger_pred) > accuracy_score(y_test, champion_pred):
    client.set_registered_model_alias(
        name=model_name,
        alias="champion",
        version=4,
    )
    print("Challenger promoted to Champion!")

def model_approval_workflow(model_name, version):
    """模型审批工作流"""
    client = MlflowClient()

    # 第 1 步：确认模型验证指标
    model_version = client.get_model_version(model_name, version)
    run = client.get_run(model_version.run_id)
    accuracy = run.data.metrics.get("accuracy", 0)
    f1 = run.data.metrics.get("f1_macro", 0)

    # 第 2 步：确认质量门槛
    quality_gates = {
        "accuracy >= 0.90": accuracy >= 0.90,
        "f1_macro >= 0.85": f1 >= 0.85,
    }

    all_passed = all(quality_gates.values())
    print("=== Quality Gate Results ===")
    for gate, passed in quality_gates.items():
        status = "PASS" if passed else "FAIL"
        print(f"  {gate}: {status}")

    # 第 3 步：根据审批结果设置 Alias
    if all_passed:
        client.set_model_version_tag(
            name=model_name, version=version,
            key="approval_status", value="approved"
        )
        # 赋予 Staging Alias
        client.set_registered_model_alias(
            name=model_name, alias="staging", version=version
        )
        print(f"Model v{version} approved and moved to staging")
        return True
    else:
        client.set_model_version_tag(
            name=model_name, version=version,
            key="approval_status", value="rejected"
        )
        print(f"Model v{version} rejected - quality gates not met")
        return False

# 执行工作流
model_approval_workflow("iris-classifier", 5)

# Dockerfile.mlflow-serve
FROM python:3.11-slim

RUN pip install mlflow[extras] boto3 psycopg2-binary

ENV MLFLOW_TRACKING_URI=http://mlflow-server:5000
ENV MODEL_NAME=iris-classifier
ENV MODEL_ALIAS=champion

EXPOSE 8080

CMD mlflow models serve \
    --model-uri "models:/${MODEL_NAME}@${MODEL_ALIAS}" \
    --host 0.0.0.0 \
    --port 8080 \
    --workers 2 \
    --no-conda

# 构建并运行 Docker 镜像
docker build -t mlflow-model-serve -f Dockerfile.mlflow-serve .
docker run -p 8080:8080 \
  -e AWS_ACCESS_KEY_ID=your-key \
  -e AWS_SECRET_ACCESS_KEY=your-secret \
  mlflow-model-serve

# 测试预测请求
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

# k8s/mlflow-model-deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: iris-classifier-serving
  labels:
    app: iris-classifier
spec:
  replicas: 3
  selector:
    matchLabels:
      app: iris-classifier
  template:
    metadata:
      labels:
        app: iris-classifier
    spec:
      containers:
        - name: model-server
          image: mlflow-model-serve:latest
          ports:
            - containerPort: 8080
          env:
            - name: MLFLOW_TRACKING_URI
              value: 'http://mlflow-server.mlflow.svc.cluster.local:5000'
            - name: MODEL_NAME
              value: 'iris-classifier'
            - name: MODEL_ALIAS
              value: 'champion'
          resources:
            requests:
              cpu: '500m'
              memory: '512Mi'
            limits:
              cpu: '1000m'
              memory: '1Gi'
          readinessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 30
            periodSeconds: 10
          livenessProbe:
            httpGet:
              path: /health
              port: 8080
            initialDelaySeconds: 60
            periodSeconds: 30
---
apiVersion: v1
kind: Service
metadata:
  name: iris-classifier-service
spec:
  selector:
    app: iris-classifier
  ports:
    - protocol: TCP
      port: 80
      targetPort: 8080
  type: ClusterIP
---
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
  name: iris-classifier-ingress
  annotations:
    nginx.ingress.kubernetes.io/rewrite-target: /
spec:
  rules:
    - host: model.example.com
      http:
        paths:
          - path: /
            pathType: Prefix
            backend:
              service:
                name: iris-classifier-service
                port:
                  number: 80

# .github/workflows/model-deploy.yaml
name: Model Deployment Pipeline

on:
  workflow_dispatch:
    inputs:
      model_name:
        description: 'Model name in registry'
        required: true
        default: 'iris-classifier'
      model_version:
        description: 'Model version to deploy'
        required: true

jobs:
  validate:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: pip install mlflow boto3 scikit-learn

      - name: Validate model
        env:
          MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_TRACKING_URI }}
          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
        run: |
          python scripts/validate_model.py \
            --model-name ${{ github.event.inputs.model_name }} \
            --model-version ${{ github.event.inputs.model_version }}

  deploy-staging:
    needs: validate
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Deploy to staging
        run: |
          kubectl apply -f k8s/staging/
          kubectl set image deployment/model-serving \
            model-server=registry.example.com/model:v${{ github.event.inputs.model_version }}

  deploy-production:
    needs: deploy-staging
    runs-on: ubuntu-latest
    environment: production
    steps:
      - uses: actions/checkout@v4

      - name: Deploy to production
        run: |
          kubectl apply -f k8s/production/
          kubectl set image deployment/model-serving \
            model-server=registry.example.com/model:v${{ github.event.inputs.model_version }}

      - name: Update MLflow alias
        env:
          MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_TRACKING_URI }}
        run: |
          python -c "
          from mlflow.tracking import MlflowClient
          client = MlflowClient()
          client.set_registered_model_alias(
              name='${{ github.event.inputs.model_name }}',
              alias='champion',
              version=${{ github.event.inputs.model_version }}
          )
          "

import mlflow
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
)
from datasets import load_dataset

mlflow.set_experiment("sentiment-analysis")

# 准备数据集
dataset = load_dataset("imdb", split="train[:1000]")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=256)

tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset = tokenized_dataset.train_test_split(test_size=0.2)

# 启用 MLflow 自动日志记录
mlflow.transformers.autolog(log_models=True)

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    warmup_steps=100,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
)

# 开始训练（自动记录到 MLflow）
with mlflow.start_run(run_name="distilbert-sentiment"):
    trainer.train()

    # 记录额外指标
    eval_results = trainer.evaluate()
    mlflow.log_metrics(eval_results)

import mlflow
import os

def setup_mlflow_distributed():
    """分布式训练环境下的 MLflow 配置"""

    rank = int(os.environ.get("RANK", 0))
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))

    # 只有 Rank 0 进程记录到 MLflow
    if rank == 0:
        mlflow.set_tracking_uri("http://mlflow-server:5000")
        mlflow.set_experiment("distributed-training")
        run = mlflow.start_run(run_name=f"dist-train-{world_size}gpu")
        mlflow.log_param("world_size", world_size)
        return run
    else:
        # 其他进程禁用日志记录
        os.environ["MLFLOW_TRACKING_URI"] = ""
        return None


def log_distributed_metrics(metrics, step, rank=0):
    """仅在 Rank 0 记录指标"""
    if rank == 0:
        mlflow.log_metrics(metrics, step=step)

from mlflow.tracking import MlflowClient
from mlflow.exceptions import MlflowException
import time

def safe_transition_model(model_name, version, target_alias, max_retries=3):
    """安全的模型阶段切换（含重试逻辑）"""
    client = MlflowClient()

    for attempt in range(max_retries):
        try:
            # 确认当前 champion
            try:
                current_champion = client.get_model_version_by_alias(
                    model_name, target_alias
                )
                print(f"Current {target_alias}: v{current_champion.version}")
            except MlflowException:
                print(f"No current {target_alias} found")

            # 切换 Alias
            client.set_registered_model_alias(
                name=model_name,
                alias=target_alias,
                version=version,
            )
            print(f"Successfully set v{version} as {target_alias}")
            return True

        except MlflowException as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            if attempt < max_retries - 1:
                time.sleep(2 ** attempt)  # 指数退避

    print(f"Failed to transition model after {max_retries} attempts")
    return False

import boto3
from botocore.exceptions import ClientError

def diagnose_artifact_access(bucket_name, prefix="experiments/"):
    """诊断 S3 Artifact Store 访问情况"""
    s3 = boto3.client("s3")

    checks = {}

    # 1. 确认 bucket 访问权限
    try:
        s3.head_bucket(Bucket=bucket_name)
        checks["bucket_access"] = "OK"
    except ClientError as e:
        checks["bucket_access"] = f"FAIL: {e.response['Error']['Code']}"

    # 2. 确认对象列表
    try:
        response = s3.list_objects_v2(
            Bucket=bucket_name, Prefix=prefix, MaxKeys=5
        )
        count = response.get("KeyCount", 0)
        checks["list_objects"] = f"OK ({count} objects found)"
    except ClientError as e:
        checks["list_objects"] = f"FAIL: {e.response['Error']['Code']}"

    # 3. 确认写入权限
    try:
        test_key = f"{prefix}_health_check"
        s3.put_object(Bucket=bucket_name, Key=test_key, Body=b"test")
        s3.delete_object(Bucket=bucket_name, Key=test_key)
        checks["write_access"] = "OK"
    except ClientError as e:
        checks["write_access"] = f"FAIL: {e.response['Error']['Code']}"

    print("=== S3 Artifact Store Diagnosis ===")
    for check, result in checks.items():
        print(f"  {check}: {result}")

    return checks