mlflow

MLflow：ML 生命周期管理平台

何时使用此技能

在以下情况下使用 MLflow：

• 跟踪 ML 实验，记录参数、指标和制品
• 管理模型注册表，实现版本控制和阶段转换
• 部署模型到各种平台（本地、云端、服务）
• 使用项目配置复现实验
• 比较模型版本和性能指标
• 通过团队工作流协作开展 ML 项目
• 集成任何 ML 框架（与框架无关）

安装

# Install MLflow
pip install mlflow

# Install with extras
pip install mlflow[extras]  # Includes SQLAlchemy, boto3, etc.

# Start MLflow UI
mlflow ui

# Access at http://localhost:5000

快速入门

基本跟踪

import mlflow

# Start a run
with mlflow.start_run():
    # Log parameters
    mlflow.log_param("learning_rate", 0.001)
    mlflow.log_param("batch_size", 32)

    # Your training code
    model = train_model()

    # Log metrics
    mlflow.log_metric("train_loss", 0.15)
    mlflow.log_metric("val_accuracy", 0.92)

    # Log model
    mlflow.sklearn.log_model(model, "model")

自动记录（自动跟踪）

import mlflow
from sklearn.ensemble import RandomForestClassifier

# Enable autologging
mlflow.autolog()

# Train (automatically logged)
model = RandomForestClassifier(n_estimators=100, max_depth=5)
model.fit(X_train, y_train)

# Metrics, parameters, and model logged automatically!

核心概念

1. 实验和运行

实验：用于存放相关运行的逻辑容器 运行：ML 代码的单次执行（参数、指标、制品）

import mlflow

# Create/set experiment
mlflow.set_experiment("my-experiment")

# Start a run
with mlflow.start_run(run_name="baseline-model"):
    # Log params
    mlflow.log_param("model", "ResNet50")
    mlflow.log_param("epochs", 10)

    # Train
    model = train()

    # Log metrics
    mlflow.log_metric("accuracy", 0.95)

    # Log model
    mlflow.pytorch.log_model(model, "model")

# Run ID is automatically generated
print(f"Run ID: {mlflow.active_run().info.run_id}")

2. 记录参数

with mlflow.start_run():
    # Single parameter
    mlflow.log_param("learning_rate", 0.001)

    # Multiple parameters
    mlflow.log_params({
        "batch_size": 32,
        "epochs": 50,
        "optimizer": "Adam",
        "dropout": 0.2
    })

    # Nested parameters (as dict)
    config = {
        "model": {
            "architecture": "ResNet50",
            "pretrained": True
        },
        "training": {
            "lr": 0.001,
            "weight_decay": 1e-4
        }
    }

    # Log as JSON string or individual params
    for key, value in config.items():
        mlflow.log_param(key, str(value))

3. 记录指标

with mlflow.start_run():
    # Training loop
    for epoch in range(NUM_EPOCHS):
        train_loss = train_epoch()
        val_loss = validate()

        # Log metrics at each step
        mlflow.log_metric("train_loss", train_loss, step=epoch)
        mlflow.log_metric("val_loss", val_loss, step=epoch)

        # Log multiple metrics
        mlflow.log_metrics({
            "train_accuracy": train_acc,
            "val_accuracy": val_acc
        }, step=epoch)

    # Log final metrics (no step)
    mlflow.log_metric("final_accuracy", final_acc)

4. 记录制品

with mlflow.start_run():
    # Log file
    model.save('model.pkl')
    mlflow.log_artifact('model.pkl')

    # Log directory
    os.makedirs('plots', exist_ok=True)
    plt.savefig('plots/loss_curve.png')
    mlflow.log_artifacts('plots')

    # Log text
    with open('config.txt', 'w') as f:
        f.write(str(config))
    mlflow.log_artifact('config.txt')

    # Log dict as JSON
    mlflow.log_dict({'config': config}, 'config.json')

5. 记录模型

# PyTorch
import mlflow.pytorch

with mlflow.start_run():
    model = train_pytorch_model()
    mlflow.pytorch.log_model(model, "model")

# Scikit-learn
import mlflow.sklearn

with mlflow.start_run():
    model = train_sklearn_model()
    mlflow.sklearn.log_model(model, "model")

# Keras/TensorFlow
import mlflow.keras

with mlflow.start_run():
    model = train_keras_model()
    mlflow.keras.log_model(model, "model")

# HuggingFace Transformers
import mlflow.transformers

with mlflow.start_run():
    mlflow.transformers.log_model(
        transformers_model={
            "model": model,
            "tokenizer": tokenizer
        },
        artifact_path="model"
    )

自动记录

详情、代码示例和参数表：references/autologging.md。在此步骤适用时阅读。

模型注册表

详情、代码示例和参数表：references/model-registry-2.md。在此步骤适用时阅读。

搜索运行

以编程方式查找运行。

from mlflow.tracking import MlflowClient

client = MlflowClient()

# Search all runs in experiment
experiment_id = client.get_experiment_by_name("my-experiment").experiment_id
runs = client.search_runs(
    experiment_ids=[experiment_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}")
    print(f"Accuracy: {run.data.metrics['accuracy']}")
    print(f"Params: {run.data.params}")

# Search with complex filters
runs = client.search_runs(
    experiment_ids=[experiment_id],
    filter_string="""
        metrics.accuracy > 0.9 AND
        params.model = 'ResNet50' AND
        tags.dataset = 'ImageNet'
    """,
    order_by=["metrics.f1_score DESC"]
)

集成示例

详情、代码示例和参数表：references/integration-examples.md。在此步骤适用时阅读。

最佳实践

1. 使用实验进行组织

# ✅ Good: Separate experiments for different tasks
mlflow.set_experiment("sentiment-analysis")
mlflow.set_experiment("image-classification")
mlflow.set_experiment("recommendation-system")

# ❌ Bad: Everything in one experiment
mlflow.set_experiment("all-models")

2. 使用描述性运行名称

# ✅ Good: Descriptive names
with mlflow.start_run(run_name="resnet50-imagenet-lr0.001-bs32"):
    train()

# ❌ Bad: No name (auto-generated UUID)
with mlflow.start_run():
    train()

3. 记录全面的元数据

with mlflow.start_run():
    # Log hyperparameters
    mlflow.log_params({
        "learning_rate": 0.001,
        "batch_size": 32,
        "epochs": 50
    })

    # Log system info
    mlflow.set_tags({
        "dataset": "ImageNet",
        "framework": "PyTorch 2.0",
        "gpu": "A100",
        "git_commit": get_git_commit()
    })

    # Log data info
    mlflow.log_param("train_samples", len(train_dataset))
    mlflow.log_param("val_samples", len(val_dataset))

4. 跟踪模型沿袭

# Link runs to understand lineage
with mlflow.start_run(run_name="preprocessing"):
    data = preprocess()
    mlflow.log_artifact("data.csv")
    preprocessing_run_id = mlflow.active_run().info.run_id

with mlflow.start_run(run_name="training"):
    # Reference parent run
    mlflow.set_tag("preprocessing_run_id", preprocessing_run_id)
    model = train(data)

5. 使用模型注册表进行部署

# ✅ Good: Use registry for production
model_uri = "models:/my-classifier/Production"
model = mlflow.pyfunc.load_model(model_uri)

# ❌ Bad: Hard-code run IDs
model_uri = "runs:/abc123/model"
model = mlflow.pyfunc.load_model(model_uri)

部署

在本地提供模型服务

# Serve registered model
mlflow models serve -m "models:/my-classifier/Production" -p 5001

# Serve from run
mlflow models serve -m "runs:/<RUN_ID>/model" -p 5001

# Test endpoint
curl http://127.0.0.1:5001/invocations -H 'Content-Type: application/json' -d '{
  "inputs": [[1.0, 2.0, 3.0, 4.0]]
}'

部署到云端

# Deploy to AWS SageMaker
mlflow sagemaker deploy -m "models:/my-classifier/Production" --region-name us-west-2

# Deploy to Azure ML
mlflow azureml deploy -m "models:/my-classifier/Production"

配置

跟踪服务器

# Start tracking server with backend store
mlflow server \
  --backend-store-uri postgresql://user:password@localhost/mlflow \
  --default-artifact-root s3://my-bucket/mlflow \
  --host 0.0.0.0 \
  --port 5000

客户端配置

import mlflow

# Set tracking URI
mlflow.set_tracking_uri("http://localhost:5000")

# Or use environment variable
# export MLFLOW_TRACKING_URI=http://localhost:5000

资源

• 文档：https://mlflow.org/docs/latest
• GitHub：https://github.com/mlflow/mlflow
• 示例：https://github.com/mlflow/mlflow/tree/master/examples
• 社区：https://mlflow.org/community

另请参阅

• references/tracking.md - 全面的跟踪指南
• references/model-registry.md - 模型生命周期管理
• references/deployment.md - 生产部署模式

Agent 操作流程

1. 检查环境。 确认硬件、框架和服务版本、模型格式及预期负载。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。凡是会改变结果的值，不要猜测，应先询问。
3. 先运行小规模版本。 扩大规模前，先提供或记录单个请求，或端到端运行一次。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 对照参考值测量延迟、吞吐量和输出正确性；检查资源使用情况和成本。
6. 汇报。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生以下情况	请采取以下措施
服务无法启动或发生 OOM	检查模型大小与内存是否匹配，以及量化和并行设置。
这些说明中的函数、标志或端点在已安装的版本中不存在	查看该版本自带的文档（help()、--help、官方文档），进行调整并告知用户。切勿臆造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿伪造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请明确说明。
• 不要公开暴露服务或凭证；配置前确认云服务成本。
• 此处与版本相关的细节可能已过时：请根据已安装版本的官方文档进行确认。
• 对于涉及费用、占用共享 GPU 或云配额、处理个人或患者数据，或无法撤销的操作，请先征求同意。