MLflow 2.x 实验追踪与模型注册表运维指南

# 生产追踪服务器启动命令
mlflow server \
  --backend-store-uri postgresql://mlflow_user:${DB_PASSWORD}@db.internal:5432/mlflow_prod \
  --default-artifact-root s3://company-mlflow-artifacts/prod/ \
  --artifacts-destination s3://company-mlflow-artifacts/prod/ \
  --host 0.0.0.0 \
  --port 5000 \
  --workers 4 \
  --gunicorn-opts "--timeout 120 --keep-alive 5"

# 客户端配置(正确方式)
import mlflow
import os

# 只设置追踪服务器 URI。Artifact 路径由服务器管理。
os.environ["MLFLOW_TRACKING_URI"] = "http://mlflow.internal:5000"

# S3 认证推荐使用 IAM Role(EC2/EKS 环境)
# 只有本地开发时才明确指定 credential
# os.environ["AWS_PROFILE"] = "mlflow-dev"

mlflow.set_experiment("/team-search/ranking-model-v3")

import mlflow

# 好的例子：分层命名
mlflow.set_experiment("/search-team/query-ranking/bert-finetune")
mlflow.set_experiment("/fraud-team/transaction-classifier/xgboost-baseline")
mlflow.set_experiment("/recommendation/item2vec/hyperopt-sweep")

# 不好的例子：扁平且含糊的命名
# mlflow.set_experiment("experiment_1")
# mlflow.set_experiment("test_model")
# mlflow.set_experiment("johns_experiment")

import mlflow
from datetime import datetime

with mlflow.start_run(run_name="bert-ranking-v3.2.1") as run:
    # 必需标签
    mlflow.set_tag("team", "search")
    mlflow.set_tag("owner", "jane.doe@company.com")
    mlflow.set_tag("git.commit", "a1b2c3d4")
    mlflow.set_tag("data.version", "v2026.03.05")
    mlflow.set_tag("environment", "gpu-cluster-a100")
    mlflow.set_tag("purpose", "hyperparameter-sweep")

    # 记录参数
    mlflow.log_params({
        "learning_rate": 2e-5,
        "batch_size": 32,
        "max_epochs": 10,
        "model_architecture": "bert-base-uncased",
        "optimizer": "AdamW",
        "warmup_steps": 500,
    })

    # 按 step 单位记录指标(epoch 或全局 step)
    for epoch in range(10):
        train_loss = train_one_epoch(model, train_loader)
        val_ndcg = evaluate(model, val_loader)

        mlflow.log_metrics({
            "train_loss": train_loss,
            "val_ndcg@10": val_ndcg,
        }, step=epoch)

    # 记录最终模型
    mlflow.pytorch.log_model(
        model,
        artifact_path="model",
        registered_model_name="search-ranking-bert",
    )

# 好的例子(以产品/功能为中心)
fraud-detector
search-ranker
recommendation-item2vec
churn-predictor

# 不好的例子(以算法/版本为中心)
xgboost-fraud-v3
bert-search-ranking-2026
lgbm_model_final_final

from mlflow import MlflowClient

client = MlflowClient()

# 注册新模型版本后设置 alias
model_name = "search-ranker"
version = client.create_model_version(
    name=model_name,
    source="runs:/abc123/model",
    run_id="abc123",
    description="BERT-base finetuned on 2026 Q1 query logs"
)

# 确认当前 champion
try:
    current_champion = client.get_model_version_by_alias(model_name, "champion")
    print(f"Current champion: v{current_champion.version}")
except mlflow.exceptions.MlflowException:
    print("No champion alias set yet")

# Canary 部署：给新版本赋予 challenger alias
client.set_registered_model_alias(model_name, "challenger", version.version)

# Canary 验证通过后晋升为 champion
client.set_registered_model_alias(model_name, "champion", version.version)

# 把旧 champion 移到 archived
client.set_registered_model_alias(model_name, "previous-champion", current_champion.version)

import mlflow

# 服务代码：基于 alias 加载模型
model = mlflow.pyfunc.load_model("models:/search-ranker@champion")
predictions = model.predict(input_data)

client.set_model_version_tag(model_name, version.version, "training_data", "s3://data/query-logs/2026-q1/")
client.set_model_version_tag(model_name, version.version, "training_commit", "a1b2c3d4e5f6")
client.set_model_version_tag(model_name, version.version, "validation_ndcg", "0.847")
client.set_model_version_tag(model_name, version.version, "approved_by", "jane.doe")
client.set_model_version_tag(model_name, version.version, "approval_date", "2026-03-05")

# .github/workflows/train-and-register.yml
name: Train and Register Model

on:
  push:
    paths:
      - 'models/search-ranker/**'
    branches: [main]
  workflow_dispatch:
    inputs:
      experiment_name:
        description: 'MLflow experiment name'
        required: true
        default: '/search-team/query-ranking/scheduled-retrain'

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 }}

jobs:
  train:
    runs-on: [self-hosted, gpu]
    steps:
      - uses: actions/checkout@v4

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

      - name: Install dependencies
        run: |
          pip install -r requirements.txt
          pip install mlflow[extras]

      - name: Train model
        run: |
          python models/search-ranker/train.py \
            --experiment-name "${{ github.event.inputs.experiment_name || '/search-team/query-ranking/ci-train' }}" \
            --run-name "ci-${{ github.sha }}" \
            --register-model search-ranker

      - name: Validate model
        run: |
          python models/search-ranker/validate.py \
            --model-uri "models:/search-ranker@challenger" \
            --threshold-ndcg 0.82

      - name: Promote to champion
        if: success()
        run: |
          python scripts/promote_model.py \
            --model-name search-ranker \
            --from-alias challenger \
            --to-alias champion

# scripts/validate_model.py
import mlflow
import sys
from mlflow import MlflowClient

def validate_model(model_name: str, alias: str, threshold: float) -> bool:
    """验证模型性能是否满足阈值。"""
    client = MlflowClient()

    # 按 alias 查询模型版本
    model_version = client.get_model_version_by_alias(model_name, alias)
    run = client.get_run(model_version.run_id)

    # 确认验证指标
    val_ndcg = run.data.metrics.get("val_ndcg@10")
    if val_ndcg is None:
        print(f"ERROR: val_ndcg@10 metric not found in run {model_version.run_id}")
        return False

    # 与当前 champion 比较
    try:
        champion = client.get_model_version_by_alias(model_name, "champion")
        champion_run = client.get_run(champion.run_id)
        champion_ndcg = champion_run.data.metrics.get("val_ndcg@10", 0)
        print(f"Champion v{champion.version} NDCG: {champion_ndcg:.4f}")
        print(f"Challenger v{model_version.version} NDCG: {val_ndcg:.4f}")

        # 检查相对 champion 的性能下降
        if val_ndcg < champion_ndcg * 0.98:  # 下降超过 2% 时判定失败
            print("FAIL: Challenger performs worse than champion by more than 2%")
            return False
    except Exception:
        print("No existing champion found. Proceeding with threshold check only.")

    # 绝对阈值检查
    if val_ndcg < threshold:
        print(f"FAIL: NDCG {val_ndcg:.4f} below threshold {threshold:.4f}")
        return False

    print(f"PASS: NDCG {val_ndcg:.4f} meets threshold {threshold:.4f}")
    return True

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", required=True)
    parser.add_argument("--alias", default="challenger")
    parser.add_argument("--threshold", type=float, default=0.80)
    args = parser.parse_args()

    if not validate_model(args.model_name, args.alias, args.threshold):
        sys.exit(1)

# 启动带认证的服务器
mlflow server \
  --backend-store-uri postgresql://mlflow:${DB_PASSWORD}@db:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/ \
  --app-name basic-auth \
  --host 0.0.0.0 \
  --port 5000

# 用团队前缀实现逻辑隔离
TEAM_PREFIX = {
    "search": "/search-team",
    "fraud": "/fraud-team",
    "recommendation": "/rec-team",
}

def get_experiment_name(team: str, project: str, experiment: str) -> str:
    """生成包含团队前缀的实验名称。"""
    prefix = TEAM_PREFIX.get(team)
    if prefix is None:
        raise ValueError(f"Unknown team: {team}. Allowed: {list(TEAM_PREFIX.keys())}")
    return f"{prefix}/{project}/{experiment}"

# 使用示例
experiment = get_experiment_name("search", "query-ranking", "bert-v4-sweep")
mlflow.set_experiment(experiment)  # "/search-team/query-ranking/bert-v4-sweep"

from mlflow import MlflowClient
from datetime import datetime, timedelta

def cleanup_old_runs(experiment_name: str, days_old: int = 90, dry_run: bool = True):
    """清理超过指定期限的失败/取消 run 的 Artifact。"""
    client = MlflowClient()
    experiment = client.get_experiment_by_name(experiment_name)

    if experiment is None:
        print(f"Experiment '{experiment_name}' not found")
        return

    cutoff_ts = int((datetime.now() - timedelta(days=days_old)).timestamp() * 1000)

    runs = client.search_runs(
        experiment_ids=[experiment.experiment_id],
        filter_string=f"attributes.end_time < {cutoff_ts} AND attributes.status != 'RUNNING'",
        order_by=["attributes.end_time ASC"],
        max_results=500,
    )

    deleted_count = 0
    for run in runs:
        # 用标签确认是否需要保留
        if run.data.tags.get("keep", "false").lower() == "true":
            continue

        # 跳过已在注册表中登记的模型对应的 run
        if run.data.tags.get("mlflow.registeredModelName"):
            continue

        if dry_run:
            print(f"[DRY RUN] Would delete run {run.info.run_id} "
                  f"(ended: {datetime.fromtimestamp(run.info.end_time / 1000)})")
        else:
            client.delete_run(run.info.run_id)
            deleted_count += 1

    print(f"{'Would delete' if dry_run else 'Deleted'} {deleted_count} runs "
          f"out of {len(runs)} found")

# 使用方式：先用 dry_run 确认，再实际删除
cleanup_old_runs("/search-team/query-ranking/bert-finetune", days_old=60, dry_run=True)

{
  "Rules": [
    {
      "ID": "MoveOldArtifactsToIA",
      "Status": "Enabled",
      "Filter": {
        "Prefix": "prod/"
      },
      "Transitions": [
        {
          "Days": 90,
          "StorageClass": "STANDARD_IA"
        },
        {
          "Days": 365,
          "StorageClass": "GLACIER"
        }
      ]
    },
    {
      "ID": "DeleteTempArtifacts",
      "Status": "Enabled",
      "Filter": {
        "Prefix": "tmp/"
      },
      "Expiration": {
        "Days": 7
      }
    }
  ]
}

# 启动服务器时调整连接池参数
mlflow server \
  --backend-store-uri "postgresql://mlflow:pass@db:5432/mlflow?pool_size=20&max_overflow=40" \
  --default-artifact-root s3://artifacts/ \
  --workers 8

# 延长上传超时时间
export MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT=600

# 调整 multipart upload 分块大小(S3)
export MLFLOW_S3_UPLOAD_EXTRA_ARGS='{"ServerSideEncryption": "aws:kms"}'

from mlflow import MlflowClient

client = MlflowClient()

# 强制终止卡住的 run
stuck_runs = client.search_runs(
    experiment_ids=["1"],
    filter_string="attributes.status = 'RUNNING'",
)

for run in stuck_runs:
    end_time = run.info.end_time
    # 没有 end_time 且开始时间已超过 24 小时的情况
    if end_time is None or end_time == 0:
        start_time = run.info.start_time
        if (datetime.now().timestamp() * 1000 - start_time) > 86400000:  # 24h
            client.set_terminated(run.info.run_id, status="FAILED")
            print(f"Force-terminated stuck run: {run.info.run_id}")

import redis
import time

def safe_promote_model(model_name: str, version: str, alias: str, redis_url: str):
    """使用分布式锁的安全模型晋升。"""
    r = redis.from_url(redis_url)
    lock_key = f"mlflow:promote:{model_name}:{alias}"

    # 获取 TTL 为 30 秒的分布式锁
    lock = r.lock(lock_key, timeout=30)
    if lock.acquire(blocking=True, blocking_timeout=10):
        try:
            client = MlflowClient()
            client.set_registered_model_alias(model_name, alias, version)
            print(f"Successfully promoted {model_name} v{version} to @{alias}")
        finally:
            lock.release()
    else:
        raise RuntimeError(f"Failed to acquire lock for {model_name}@{alias}")

from mlflow import MlflowClient

def rollback_model(model_name: str):
    """把 champion 模型回滚为 previous-champion。"""
    client = MlflowClient()

    try:
        previous = client.get_model_version_by_alias(model_name, "previous-champion")
    except Exception:
        print("ERROR: No previous-champion alias found. Manual intervention required.")
        return False

    current = client.get_model_version_by_alias(model_name, "champion")

    # 执行回滚
    client.set_registered_model_alias(model_name, "champion", previous.version)
    client.set_registered_model_alias(model_name, "rolled-back", current.version)

    # 标记回滚原因
    client.set_model_version_tag(
        model_name, current.version, "rollback_reason", "performance_degradation"
    )
    client.set_model_version_tag(
        model_name, current.version, "rolled_back_at", datetime.now().isoformat()
    )

    print(f"Rolled back {model_name}: v{current.version} -> v{previous.version}")
    return True

# Artifact 垃圾回收
mlflow gc \
  --backend-store-uri postgresql://mlflow:pass@db:5432/mlflow \
  --older-than 30d