使用 Docker Compose 编排 8 节点 DeepSpeed 训练 - 问题详情 - 创脉思

解读

在国内大厂或 AI 初创公司的面试场景里，这道题表面问“Compose 怎么写”，实质考察三点：

是否理解 DeepSpeed 多机训练 = 多容器 + 共享文件系统 + 无密码 SSH + 统一 NCCL 通信；
能否用 Docker Compose 语法 把上述依赖固化成可一键复现的编排文件；
是否具备 镜像瘦身、启动顺序、故障自愈、资源绑定 等工程化意识。
面试官常追问：

为什么不用 K8s？
宿主机网卡选哪个驱动？
镜像里要不要装 NVIDIA Driver？
训练中断后如何续训？
回答必须体现“国内 GPU 裸金属 + 高速 RoCEv2 网络”的真实约束，否则会被判“纸上谈兵”。

知识点

DeepSpeed 启动流程：deepspeed --hostfile=hostfile --launcher=pdsh train.py，要求容器间 22 端口无密码 SSH。
Compose 3.x 语法：services.deploy.placement 只能做单机约束，多机必须借助 Docker Swarm 或外部调度，因此需声明 docker stack deploy -c deepspeed.yml ds。
GPU 支持：国内宿主机普遍预装 NVIDIA Driver 470+，镜像内只需 cuda:11.8-cudnn8-runtime-ubuntu20.04 + pytorch + deepspeed，禁止重复安装驱动；Compose 中通过 runtime: nvidia 暴露 GPU。
共享存储：训练数据与 checkpoint 必须落在 同一分布式文件系统；国内最常用 阿里云 CPFS、腾讯 CFS 或自建 RoCE+BeeGFS，容器内统一挂载到 /data。
网络优化：
- 选用 host 网络 可避开 Docker 桥接性能损耗，但端口冲突；
- 若坚持 bridge，需显式创建 macvlan 或 ipvlan，并把 NCCL_SOCKET_IFNAME=eth1 指向 RoCE 网卡。
镜像安全：
- 使用 非 root 用户 uid=1000 启动训练进程；
- 通过 BuildKit 多阶段构建 把最终镜像压到 < 2 GB；
- 敏感环境变量（阿里云 AK/SK、WandB key）统一用 Docker Config 或 Swarm Secret 注入，禁止写死 Dockerfile。
健康度与重试：Compose 文件里给 worker 加 healthcheck: test=["CMD-SHELL", "nc -z localhost 22"]，配合 restart_policy: condition=any 实现容器级自愈。
资源绑定：利用 deploy.resources.reservations.devices 把 8 卡 A100 按 UUID 独占绑定，避免调度漂移。

答案

deepspeed/
├── docker-compose.yml
├── ssh/
│   ├── id_rsa
│   └── id_rsa.pub
├── Dockerfile
├── hostfile
├── train.py
└── data/          # 软链到共享 /data

Dockerfile（国内加速 + 非 root）

# syntax=docker/dockerfile:1.4
FROM nvcr.io/nvidia/cuda:11.8-cudnn8-runtime-ubuntu20.04
RUN sed -i 's@archive.ubuntu.com@mirrors.aliyun.com@g' /etc/apt/sources.list && \
    apt-get update && apt-get install -y openssh-server pdsh python3-pip && \
    useradd -m -u 1000 deepspeed
USER deepspeed
RUN pip3 install -i https://pypi.tuna.tsinghua.edu.cn/simple torch deepspeed
COPY --chown=deepspeed:deepspeed ssh/ /home/deepspeed/.ssh/
RUN chmod 700 ~/.ssh && chmod 600 ~/.ssh/id_rsa && chmod 644 ~/.ssh/id_rsa.pub && \
    cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys
ENV NCCL_DEBUG=INFO NCCL_SOCKET_IFNAME=eth1
WORKDIR /workspace

docker-compose.yml（Swarm 模式）

version: "3.9"
services:
  ds-master:
    image: registry.cn-hangzhou.aliyuncs.com/yourrepo/deepspeed:11.8
    volumes:
      - type: bind
        source: /data
        target: /data
        read_only: false
    networks:
      - dsnet
    deploy:
      replicas: 1
      placement:
        constraints:
          - node.labels.role == master
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 8
              capabilities: [gpu]
    environment:
      - NCCL_SOCKET_IFNAME=eth1
    command: >
      bash -c "sudo service ssh start && sleep infinity"
    healthcheck:
      test: ["CMD-SHELL", "nc -z localhost 22"]
      interval: 10s
      retries: 3

  ds-worker:
    image: registry.cn-hangzhou.aliyuncs.com/yourrepo/deepspeed:11.8
    volumes:
      - type: bind
        source: /data
        target: /data
    networks:
      - dsnet
    deploy:
      replicas: 7
      placement:
        constraints:
          - node.labels.role == worker
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 8
              capabilities: [gpu]
    environment:
      - NCCL_SOCKET_IFNAME=eth1
    command: >
      bash -c "sudo service ssh start && sleep infinity"
    depends_on:
      - ds-master
    healthcheck:
      test: ["CMD-SHELL", "nc -z localhost 22"]
      interval: 10s

networks:
  dsnet:
    driver: overlay
    attachable: true
    driver_opts:
      com.docker.network.driver.mtu: 9000

启动脚本（一键）

#!/bin/bash
# 0. 初始化 Swarm
docker swarm init --advertise-addr=192.168.100.10
# 1. 给节点打标签
for i in {1..7}; do
  docker node update --label-add role=worker node-$i
done
# 2. 下发编排
docker stack deploy -c docker-compose.yml ds
# 3. 等待健康
until docker exec $(docker ps -q -f name=ds_master) nc -z localhost 22; do sleep 2; done
# 4. 生成 hostfile
docker exec $(docker ps -q -f name=ds_master) bash -c \
  'for h in $(dig +short tasks.ds_worker | sort); do echo $h slots=8; done' > hostfile
# 5. 启动训练
docker exec -u deepspeed $(docker ps -q -f name=ds_master) \
  deepspeed --hostfile=hostfile --launcher=pdsh train.py --deepspeed ds_config.json

关键验证

nvidia-smi 在 8 个容器内均显示 8×A100；
ssh ds_worker_1 无需密码；
NCCL debug 日志中出现 Ring 00 : 8 nodes；
checkpoint 实时写入 /data，掉节点后重启任务可续训。

拓展思考

如果集群规模 >8 节点，Compose+Swarm 的 overlay 网络性能瓶颈 会凸显，应迁移到 KubeFlow+PyTorchJob CRD，并用 RDMA device plugin 暴露 HCAs。
国内公有云 按卡计费 场景，可给 worker 增加 preemptible=true 标签，训练脚本捕获 SIGTERM 后主动 save checkpoint 并上传 OSS，实现 spot 实例优雅退出。
镜像更新频繁时，用 Harbor 的复制策略 把海外 nvcr.io 镜像同步到 国内 Harbor 实例，CI 侧只需 build --cache-from 即可 30 s 内完成迭代。
安全合规：金融客户要求 国密算法，可把 SSH 换成 stunnel + SM2 证书，容器内 只读根文件系统，并通过 AppArmor 模板 限制系统调用。