如何将你的机器学习运营管道从训练阶段到部署阶段进行容器化处理

HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1

将模型与容器分离

这是本文中最重要的设计模式之一，也是初学者最容易犯错的地方。

人们往往会忍不住在构建Docker镜像时直接将训练好的模型文件（即包含模型权重数据的.pkl、.pt或.onnx文件）复制到镜像中。但千万不要这样做。一旦将模型文件嵌入到Docker镜像中，每次模型更新都需要重新构建并推送镜像。对于一个大小为2GB的模型来说，这意味着需要重新构建容器、将2GB大小的文件上传到注册中心，然后再进行部署——而实际上只有模型本身发生了变化，代码并没有任何改变。

正确的做法是让服务容器在启动时从模型注册平台（如MLflow）或云存储服务（如S3）中下载模型文件。这样，容器镜像的大小就可以保持较小且通用性更强。模型的更新只需要通过配置文件的修改来指明新的模型版本即可，而无需重新进行部署操作。

下面是一个采用了这种设计模式的完整服务应用示例，它使用了FastAPI框架。如果您之前使用过Flask，那么FastAPI的功能类似，但运行速度更快，并且还内置了请求验证机制：

import os
from contextlib import asynccontextmanager

import mlflow
from fastapi import FastAPI

# MODEL_URI指向MLflow注册中心中的特定模型版本
# 格式为："models://"，其中stage可以是Staging或Production
MODEL URI = os.environ.get("MODEL_URI", "models:/fraud-detector/production")
model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 服务器启动时执行一次此代码
    global model
    print(f"正在从{MODEL_URI}加载模型...")
    model = mlflow.pyfunc.load_model(MODEL URI)
    print("模型加载成功。")
    yield
    # 服务器关闭时执行此代码
    print("正在关闭模型服务器。")


app = FastAPI(lifespan=lifespan)


@app.get("/health")
async def health():
    """用于供Docker的健康检查机制使用，以确认服务器是否已准备好运行。"""
    if model is None:
        return {"status": "loading"}, 503
    return {"status": "healthy"}


@app.post("/predict")
async def predict(features: dict):
    """接受以JSON格式传递的输入数据，并返回模型预测结果。"""
    import pandas as pd

    # 将输入数据转换为DataFrame格式（大多数sklearn/mlflow模型都要求这样处理）
    df = pd.DataFrame([features])
    prediction = model.predict(df)
    return {"prediction": prediction.tolist()}

当客户向/predict发送包含如下JSON数据的POST请求时：{"amount": 500, "merchant_category": "electronics", "hour": 23}，模型会返回相应的预测结果。在模型加载过程中，/health端点会返回503状态码；而一旦模型加载完成，就会返回200状态码——这正是Docker的HEALTHCHECK命令所检测的内容。

要推广新的模型版本，只需更新MODEL_URI环境变量，然后重新启动容器即可。MLflow模型注册系统支持模型阶段的转换（测试阶段、生产阶段、归档阶段），因此你可以在MLflow用户界面中推进模型的升级流程，之后再将服务容器指向新版本。

为了实现模型更新的零停机时间目标，可以开发一个无需重启容器即可切换模型的功能：

@app.post("/admin/reload")
async def reload_model():
    global model
    model = mlflow.pyfunc.load_model(MODEL_URI)
    return {"status": "reloaded"}

如何配置用于训练的GPU直通功能

默认情况下，Docker容器无法访问主机上的GPU硬件。“GPU直通”功能意味着让容器能够使用主机的GPU，这样像PyTorch和TensorFlow这样的库就可以利用这些GPU来加速计算。

在主机上需要满足以下两个条件才能启用这一功能（这里指的是运行Docker的主机，而不是容器内部）：

1. NVIDIA GPU驱动程序已安装且能够正常使用。可以通过nvidia-smi命令进行验证。如果该命令能显示你的GPU信息，说明驱动程序安装成功。

2. NVIDIA容器工具包也已安装。这个工具包是Docker与GPU驱动程序之间的桥梁。你可以从NVIDIA官方文档中下载并安装它，然后通过docker run --rm --gpus all nvidia/cuda:12.6.3-base-ubuntu22.04 nvidia-smi命令进行测试。如果命令输出中列出了你的GPU信息，说明工具包安装成功。

完成主机配置后，在Docker Compose中配置GPU直通功能的代码如下：

services:
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./data:/app/data
    environment:
      - MLFLOW TRACKING_URI=http://mlflow:5000

deploy.resources.reservationsdevices这一配置表示：“这个容器需要使用所有可用的NVIDIA GPU。”在容器内部，PyTorch和TensorFlow会自动识别这些GPU并加以使用。你可以通过在训练脚本中添加print(torch.cuda.is_available())这条代码来验证这一功能，如果输出结果为True，说明配置成功。

如果你使用的是Compose v2.30.0或更高版本，也可以使用更简洁的gpus语法进行配置：

services:
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    gpus: all
    volumes:
      - ./data:/app/data
    environment:
      - MLFLOW TRACKING_URI=http://mlflow:5000

在使用PyTorch的DistributedDataParallel等框架进行多GPU训练时，可以通过device_ids来指定使用哪些特定的GPU。当同时运行多个训练任务时，这一点尤为重要：

services:
  training-job-1:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0", "1"]
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_devices=0,1

  training-job-2:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["2", "3"]
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_devices=0,1

需要注意的是，容器内部的CUDAVISIBLE_DEVICES所指定的设备是相对于Docker分配的设备而言的，而不是主机上的GPU索引。即使这两个容器使用的是不同的物理GPU，它们也会将自己对应的GPU视为设备0和设备1。

如何利用Compose配置文件来整合这些组件

如果你还不熟悉Compose配置文件的话：默认情况下，docker compose up会启动docker-compose.yml中定义的所有服务。但有时候你并不需要所有服务都同时运行。例如，你的MLflow服务器和API应该一直处于运行状态，而训练容器则只应在你需要进行模型训练时才被启动（而且训练容器需要GPU，而你的笔记本电脑可能并没有GPU）。

配置文件就可以解决这个问题。当你为某个服务添加profiles: ["train"]时，该服务在默认情况下就不会被docker compose up启动。只有当你通过docker compose --profile train明确指定要启动该配置文件时，它才会被运行。这意味着，虽然只有一个文件定义了你的整个ML基础设施，但你完全可以控制哪些服务会在什么时候运行。

以下是整合了本文中所有内容的完整docker-compose.yml文件：

services:
  # --- 始终处于运行状态的组件 ---
  db:
    image: postgres:16-alpine
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: secret
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mlflow"]
      interval: 5s
      timeout: 2s
      retries: 5
      start_period: 10s
    volumes:
      - postgres-data:/var/lib/postgresql/data

  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.19.0
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:secret@db/mlflow
      --default-artifact-root /mlflow/artifacts
      --host 0.0.0.0
    ports:
      - "5000:5000"
    volumes:
      - mlflow-artifacts:/mlflow/artifacts
    depends_on:
      db: { condition: service_healthy }

  serving:
    build: { context: ., dockerfile: Dockerfile.serve }
    ports:
      - "8000:8000"
    environment:
      - MODEL_URI=models:/fraud-detector/production
      - MLFLOW TRACKING_URI=http://mlflow:5000
    depends_on:
      mlflow: { condition: service_started }

  # --- 根据需求启动的训练组件 ---
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    profiles: ["train"]
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./data:/app/data
      - ./configs:/app/configs
    environment:
      - MLFLOW TRACKING_URI=http://mlflow:5000
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    depends_on:
      mlflow: { condition: service_started }

volumes:
  postgres-data:
  mlflow-artifacts:

使用该文件进行日常操作的工作流程如下：

# 第一步：启动基础设施（MLflow + Postgres + API服务）
# 使用-d参数可让所有进程在后台运行
docker compose up -d

# 第二步：打开MLflow用户界面查看之前的实验结果
open http://localhost:5000    # macOS系统
# xdg-open http://localhost:5000  # Linux系统

# 第三步：检查API服务是否正常运行
curl http://localhost:8000/health
# 应返回：“{"status":"healthy"}”

# 第四步：运行训练任务（通过DVC获取数据，将日志记录到MLflow系统中）
# 由于使用了--profile参数，因此仅会启动“训练”服务
docker compose --profile train run training

# 第五步：在localhost:5000处的MLflow用户界面中查看训练进度
# 如果你的训练代码包含了日志记录功能，那么你会看到各项指标实时更新

# 第六步：训练完成后，在MLflow用户界面中将模型推送到生产环境
# 点击该模型，进入“注册模型”页面，将阶段设置为“生产”

# 第七步：重新启动API服务容器，以便使用新的模型版本
docker compose restart serving

# 第八步：测试新模型
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"amount": 500, "merchant_category": "electronics", "hour": 23}'

采用这种单文件配置方式，新的团队成员只需克隆代码仓库，运行docker compose up -d命令，就能在几分钟内让整个ML基础设施在本地正常运行。这些容器在部署到测试环境或生产环境时，只需要调整一些环境变量即可（例如数据库登录信息、模型URL地址、GPU资源配置等）。

可复现性：这才是关键所在

本文所介绍的所有方法都是为了实现一个目标：确保实验结果的可复现性。也就是说，无论使用哪个Git提交哈希值进行训练，构建出的容器、获取的数据以及最终得到的模型都应该是完全相同的。

以下是一些有助于实现这一目标的实践方法：

固定所有相关配置

请将基础镜像固定在特定的版本号上，而不仅仅是标签上；对于Python包，也要使用pip freeze > requirements.txt命令将其版本明确指定下来。在训练代码中要使用固定的随机种子，并将这些信息记录到MLflow系统中。

记录所有操作细节

每次训练运行时，都应将所使用的库版本(pip freeze)、Git提交哈希值、DVC数据版本、所有超参数以及所有的评估指标记录到MLflow系统中。你可以使用自动化脚本来完成这项工作：

import subprocess
import mlflow

with mlflow.start_run():
    # 自动记录环境配置信息
    pip_freeze = subprocess.check_output(["pip", "freeze")).decode()
    mlflow.log_text(pip_freeze, "pip_freeze.txt")

    git_hash = subprocess.check_output(["git", "rev-parse", "HEAD")).decode().strip()
    mlflow.log_param("git_commit", git_hash)

    # ...继续执行训练代码 ...

为所有内容指定版本号

使用Git来管理代码，用DVC来管理数据，通过MLflow记录实验过程，再用Docker的镜像版本号来区分不同的运行环境。这种组合方式能够形成一条完整、可追溯的链。当有利益相关者询问为什么某个模型会得出特定的预测结果时，你就可以追踪到具体的代码、数据和超参数。对于金融、医疗保健等受监管的行业来说，这种可追溯性是一项强制性的合规要求，而不仅仅是一种便利功能。

这种方法存在的局限性

这种方案对于在单台主机或小型集群上运行的中小型团队来说效果不错，但以下情况下会遇到问题：

大型数据集。不要将规模达数TB的数据集放入容器中。应使用对象存储服务（如S3、GCS），并在训练过程中以流式方式传输数据。虽然DVC可以处理版本控制问题，但数据本身仍应存储在Docker之外。

GPU驱动程序不兼容。容器的CUDA版本必须与主机上的驱动程序相匹配。请在与生产环境完全相同的硬件和驱动程序配置上进行测试，并在README文件中明确指定最低所需的驱动程序版本。

多节点训练。当需要将训练任务分布在多台机器上时，Compose就不再适用了。此时，使用Kubernetes结合Kubeflow或KServe才是实现分布式训练和自动扩展服务的标准方案。

大规模服务部署。单个容器使用uvicorn处理中等规模的流量尚可，但对于高吞吐量的推理任务来说，就需要负载均衡器、多个副本，以及专门的服务器框架（如NVIDIA Triton Inference Server或TensorFlow Serving）。虽然Compose可以通过`docker compose up --scale serving=3`命令创建多个副本，但它无法提供真正的路由管理、基于健康状况的负载均衡功能，也无法实现自动更新机制。

生产环境中的密码管理。在上述示例中，Compose文件使用明文密码进行本地开发，但在生产环境中应使用Docker Secrets、HashiCorp Vault或云服务提供商提供的秘密管理工具。切勿将敏感信息提交到代码仓库中。

总结

通过将机器学习运营流程容器化，原本依赖特定环境的脆弱模型就能被转化为可复现、可部署的应用程序。多阶段构建机制有助于保持镜像文件的体积小巧。MLflow能够帮助您跟踪实验过程并追踪模型的演变历程，而DVC则实现了代码与数据之间的紧密关联。GPU直通技术还能确保训练性能不受影响。一个包含所有配置信息的Compose文件就能将整个工作流程整合起来。

我之前提到过的那个欺诈检测模型？我们最终将其整个开发流程都容器化了。之后开发的另一个模型，从“开发完成”到“投入生产”，只用了两天时间，而不是原来的三周。其中大部分时间都被用来进行评估和审查工作，而非解决环境配置问题。

容器化本身并不会让模型变得更好，但它能消除基础设施带来的阻碍，让您能够专注于核心业务的发展。

尽管存在这些限制，但容器化的机器学习运营流程确实有效解决了机器学习项目中最常见的延迟问题——即开发环境和生产环境之间的不匹配。我们曾经为调试那个欺诈检测模型的部署花费了三周时间，而现在这种情况再也不会发生了。

如果您觉得这篇文章有帮助，可以在我的博客上找到我关于机器学习运营、容器化工作流程以及生产环境中的AI系统的更多文章。