返回文章列表
K8sDocker深度学习DevOps

Docker + Kubernetes 部署深度学习训练平台的实践

谢工

为什么需要平台化

随着深度学习项目的增多,每个团队各自搭建环境、管理依赖的方式变得不可持续。平台化带来的好处:

  • 环境一致性:消除"在我机器上能跑"的问题
  • 资源共享:GPU 资源统一调度,提高利用率
  • 可复现性:实验环境和代码版本可追溯
  • 自动化:训练、评估、部署流程自动化

容器化训练环境

Dockerfile 最佳实践

FROM nvidia/cuda:12.1-runtime-ubuntu22.04

RUN apt-get update && apt-get install -y \
    python3.10 python3-pip git curl && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /workspace
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 使用非 root 用户运行
RUN useradd -m -u 1000 trainuser
USER trainuser

COPY --chown=trainuser:trainuser . .

ENTRYPOINT ["python", "train.py"]

关键点

  • 使用 NVIDIA CUDA 基础镜像
  • 最小化镜像层数
  • 使用非 root 用户

多阶段构建

# 构建阶段
FROM nvidia/cuda:12.1-devel-ubuntu22.04 AS builder
RUN pip install torch torchvision --no-cache-dir

# 运行阶段
FROM nvidia/cuda:12.1-runtime-ubuntu22.04
COPY --from=builder /usr/local/lib/python3.10 /usr/local/lib/python3.10

Kubernetes 集群配置

GPU 节点配置

apiVersion: v1
kind: Node
metadata:
  labels:
    nvidia.com/gpu: "true"
spec:
  taints:
  - key: nvidia.com/gpu
    effect: NoSchedule

训练任务 CRD

apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: distributed-train
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      template:
        spec:
          containers:
          - name: pytorch
            image: train-image:latest
            resources:
              limits:
                nvidia.com/gpu: 8
            env:
            - name: NCCL_IB_DISABLE
              value: "1"
            - name: NCCL_SOCKET_IFNAME
              value: "eth0"

实验管理

MLflow 集成

import mlflow

mlflow.set_tracking_uri("http://mlflow-service:5000")

with mlflow.start_run():
    mlflow.log_params({
        "learning_rate": 0.001,
        "batch_size": 64,
        "model": "resnet50"
    })
    mlflow.log_metric("accuracy", 0.95)
    mlflow.pytorch.log_model(model, "model")

监控与告警

GPU 监控

# Prometheus GPU exporter
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-dcgm-exporter
spec:
  template:
    spec:
      containers:
      - name: dcgm-exporter
        image: nvidia/dcgm-exporter:latest
        ports:
        - containerPort: 9400

关键监控指标

指标 说明 告警阈值
GPU 利用率 核心使用率 > 90% (持续1h)
GPU 内存 显存使用率 > 95%
GPU 温度 核心温度 > 85°C
PCIe 带宽 数据传输 < 10%

总结

构建深度学习训练平台不是一蹴而就的,需要从环境容器化、集群调度、实验管理和监控告警等多个维度逐步建设。建议按照 容器化 → 单机调度 → 集群调度 → 平台化 的路径逐步演进。


欢迎在评论区分享你的训练平台搭建经验!

谢工

AI Agent / HPC 工程师

更多文章