Docker + GPU 部署 AI 模型

nvidia-docker / CUDA 镜像选择、多模型容器编排、GPU 资源监控

Docker GPU 部署概述

为什么用 Docker 部署 AI

在 AI 模型的开发与部署过程中，环境一致性是最令人头疼的问题之一。训练时用 CUDA 11.8 + PyTorch 2.0，生产环境却只有 CUDA 11.0，或者缺少某个底层系统库——这些差异往往导致模型推理失败或性能下降。Docker 通过容器化技术将整个运行时环境（操作系统依赖、CUDA 驱动库、Python 包、模型文件）打包在一起，确保开发、测试和生产环境完全一致。

Docker 部署 AI 的核心优势包括：

• 环境一致性：容器内包含完整的 CUDA、cuDNN、Python 和依赖库，消除环境差异带来的问题。
• 依赖隔离：不同模型可以使用不同的 CUDA 版本和 Python 环境，互不干扰。可以在同一台机器上同时运行 PyTorch 1.x 和 PyTorch 2.x 的推理服务。
• 快速扩缩：结合容器编排工具，可以在数秒内启动或销毁推理服务实例，应对流量波动。

nvidia-container-toolkit 工作原理

Docker 本身无法直接访问 GPU 设备。要让容器内识别 NVIDIA GPU，核心组件是 nvidia-container-toolkit。它的工作原理如下：

1. 在容器启动时，nvidia-container-toolkit 通过 LD_PRELOAD 机制拦截容器内的 CUDA Runtime API 调用。
2. 将调用转发到宿主机上的 NVIDIA 驱动层，实现 GPU 设备透传。
3. 将宿主机的 GPU 设备文件（/dev/nvidia0、/dev/nvidiactl 等）以及必要的驱动库挂载到容器中。

与早期的 nvidia-docker 1.0 方案（需要安装 nvidia-docker-plugin）不同，nvidia-container-toolkit 是 Docker 的原生运行时扩展，通过 --gpus 参数即可使用，无需额外启动守护进程。

nvidia-docker 环境搭建

nvidia-container-toolkit 安装流程

以下以 Ubuntu 22.04 为例展示安装步骤：

# 添加 NVIDIA 容器工具包仓库
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
    sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# 安装
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

# 配置 Docker 运行时
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

对于 CentOS / RHEL 系统，只需将包管理器替换为 yum，仓库源替换为对应的 RHEL 版本即可。

Docker 运行时配置

安装完成后，通过 --gpus 参数指定 GPU 资源。常用方式如下：

# 使用所有 GPU
docker run --gpus all nvidia/cuda:12.2-base nvidia-smi

# 使用指定 GPU（通过 UUID 或索引）
docker run --gpus '"device=0,1"' nvidia/cuda:12.2-base nvidia-smi

# 使用指定数量的 GPU
docker run --gpus '"count=2"' nvidia/cuda:12.2-base nvidia-smi

验证 GPU 可用性

运行以下命令验证容器内能否识别 GPU：

docker run --rm --gpus all nvidia/cuda:12.2-base nvidia-smi

正常输出应显示 GPU 型号、驱动版本和 CUDA 版本信息，类似宿主机上执行 nvidia-smi 的效果。

检查项	预期结果	常见问题
nvidia-smi 输出	显示 GPU 型号、显存、驱动版本	驱动未安装或版本过低
nvcc --version	显示 CUDA 编译版本	nvidia-container-toolkit 未正确配置
容器内运行 PyTorch	torch.cuda.is_available() 返回 True	PyTorch CUDA 版本与驱动不匹配

CUDA 镜像选择

基础镜像层级

NVIDIA 官方提供了多个层级的 CUDA 基础镜像，选择合适的镜像对容器体积和启动速度有显著影响。

镜像标签	内容	体积	适用场景
nvidia/cuda:12.2-base	CUDA 运行时动态库	~800 MB	使用 CUDA 编译好的二进制程序
nvidia/cuda:12.2-runtime	base + CUDA 工具库（cuBLAS、cuFFT 等）	~1.2 GB	多数 AI 推理场景
nvidia/cuda:12.2-devel	runtime + 头文件、编译器（nvcc）	~2.5 GB	需要编译 CUDA 代码的训练场景

对于推理服务，推荐使用 runtime 镜像——它包含了运行已编译 CUDA 程序所需的全部库文件，又不会像 devel 镜像那样包含大量编译工具导致镜像臃肿。

PyTorch 官方镜像 vs 自定义 Dockerfile

PyTorch 官方提供了预装好 PyTorch 的 Docker 镜像（pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime），适合快速原型验证。但对于生产环境，建议使用自定义 Dockerfile 以获得更精细的控制——包括选择最小基础镜像、仅安装生产所需的 Python 包、配置非 root 用户等。

Python 依赖管理与多阶段构建减镜像体积

多阶段构建是减少最终镜像体积的关键手段。在构建阶段安装编译依赖和 Python 包，在运行阶段只保留运行时所需的文件和库。

# ============ 构建阶段 ============
FROM nvidia/cuda:12.2-runtime AS builder

RUN apt-get update && apt-get install -y python3 python3-pip && \
    pip install --user torch torchvision --index-url https://download.pytorch.org/whl/cu121

# ============ 运行阶段 ============
FROM nvidia/cuda:12.2-runtime

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

# 复制构建阶段安装的 Python 包
COPY --from=builder /root/.local /root/.local
ENV PATH=/root/.local/bin:$PATH

WORKDIR /app
COPY requirements.txt .
RUN pip install --user -r requirements.txt

COPY model/ ./model/
COPY app.py .

# 使用非 root 用户运行
RUN useradd -m -u 1000 appuser && chown -R appuser:appuser /app
USER appuser

EXPOSE 8000
CMD ["python3", "app.py"]

AI 模型推理服务的 Dockerfile

以下是一个完整的多阶段构建示例，用于部署基于 FastAPI 的 BERT 模型推理服务：

# Stage 1: 构建依赖
FROM nvidia/cuda:12.2-runtime AS builder

RUN apt-get update && apt-get install -y python3 python3-pip
COPY requirements.txt .
RUN pip install --user -r requirements.txt

# Stage 2: 运行镜像
FROM nvidia/cuda:12.2-runtime

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

COPY --from=builder /root/.local /root/.local
ENV PATH=/root/.local/bin:$PATH

WORKDIR /app
COPY server.py .
COPY models/ ./models/

EXPOSE 8000
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

对应的 requirements.txt：

torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121
transformers==4.36.0
fastapi==0.108.0
uvicorn[standard]==0.25.0

构建命令：

docker build -t ai-bert-service:latest .
docker run --gpus all -p 8000:8000 ai-bert-service:latest

多模型容器编排

当需要同时部署多个 AI 模型服务（如文本分类模型 A、图像生成模型 B），再加上前端界面和缓存中间件时，使用 docker-compose 进行编排是最直接有效的方案。

docker-compose 编排多个服务

version: "3.8"

services:
  model-a:
    image: ai-text-classifier:latest
    build:
      context: ./model-a
    container_name: text-classifier
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0"]
              capabilities: [gpu]
    shm_size: 4g
    environment:
      - CUDA_VISIBLE_DEVICES=0
    ports:
      - "8001:8000"
    networks:
      - ai-net
    restart: unless-stopped

  model-b:
    image: ai-image-generator:latest
    build:
      context: ./model-b
    container_name: image-generator
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["1"]
              capabilities: [gpu]
    shm_size: 8g
    environment:
      - CUDA_VISIBLE_DEVICES=1
    volumes:
      - ./model-b/cache:/app/.cache
    ports:
      - "8002:8000"
    networks:
      - ai-net
    restart: unless-stopped

  redis:
    image: redis:7-alpine
    container_name: ai-cache
    ports:
      - "6379:6379"
    networks:
      - ai-net
    restart: unless-stopped

  frontend:
    image: ai-web-console:latest
    build:
      context: ./frontend
    container_name: web-console
    ports:
      - "8080:80"
    depends_on:
      - model-a
      - model-b
    networks:
      - ai-net
    restart: unless-stopped

networks:
  ai-net:
    driver: bridge

GPU 资源分配要点

上述 docker-compose 配置中体现了几个关键参数：

• device_ids：通过 device_ids: ["0"] 将 GPU 0 分配给 model-a，GPU 1 分配给 model-b，实现物理隔离，避免显存争抢。
• shm_size：PyTorch 的 DataLoader 依赖共享内存进行进程间通信，默认的 64 MB 往往不够。对于 Transformer 模型，建议设置为 4 GB 以上。
• CUDA_VISIBLE_DEVICES：配合环境变量确保容器内的 CUDA 程序只识别指定的 GPU。

网络配置与服务发现

所有服务接入同一个 ai-net 桥接网络后，容器之间通过服务名称进行通信。例如，前端服务可以通过 http://model-a:8000 访问模型 A 的推理接口，无需关心 IP 地址。这样不仅简化了配置，还实现了服务间调用的解耦。

GPU 资源监控

常用监控工具

工具	用途	特点
nvidia-smi	实时查看 GPU 利用率、显存、温度	Docker 自带支持，无需额外安装
nvtop	类 htop 的 GPU 监控界面	交互式，支持多 GPU 同时查看
dcgm-exporter	NVIDIA 官方指标导出器	输出 Prometheus 格式指标，适合集成

容器内使用 nvidia-smi：

docker run --rm --gpus all nvidia/cuda:12.2-base nvidia-smi \
    --query-gpu=index,name,utilization.gpu,memory.used,temperature.gpu \
    --format=csv

Prometheus + Grafana 监控方案

将 DCGM Exporter 与 Prometheus + Grafana 集成，可以搭建完整的 GPU 监控看板：

# 启动 DCGM Exporter
docker run -d --gpus all --rm \
    -p 9400:9400 \
    nvidia/dcgm-exporter:3.3.0

# 验证指标输出
curl http://localhost:9400/metrics | grep -E "DCGM_FI_DEV_GPU_UTIL|DCGM_FI_DEV_FB_USED"

Prometheus 配置中增加以下 job：

scrape_configs:
  - job_name: "dcgm"
    static_configs:
      - targets: ["localhost:9400"]

DCGM Exporter 可导出的关键指标包括：

指标名	含义	PromQL 查询示例
DCGM_FI_DEV_GPU_UTIL	GPU 计算核心利用率	avg(DCGM_FI_DEV_GPU_UTIL)
DCGM_FI_DEV_FB_USED	已用显存（字节）	DCGM_FI_DEV_FB_USED / 1024^3 转换为 GB
DCGM_FI_DEV_POWER_USAGE	GPU 功耗（瓦）	DCGM_FI_DEV_POWER_USAGE
DCGM_FI_DEV_TEMPERATURE	GPU 温度（摄氏度）	DCGM_FI_DEV_TEMPERATURE

Docker 容器 GPU 资源限制

精细控制容器可以使用的 GPU 资源，避免单个容器耗尽所有显存：

# 指定特定 GPU 设备
docker run --gpus '"device=0,2"' ai-bert-service:latest

# 限制 GPU 数量（自动选择空闲设备）
docker run --gpus '"count=1"' ai-bert-service:latest

# 结合 nvidia-smi 查询进行高级调度
available_gpu=$(nvidia-smi --query-gpu=index,memory.free --format=csv,noheader,nounits \
    | sort -t, -k2 -rn \
    | head -1 \
    | awk -F, '{print $1}')

docker run --gpus "\"device=$available_gpu\"" ai-bert-service:latest

需要注意的是，--gpus 参数虽然可以控制容器能访问哪些 GPU 设备，但无法直接限制单个容器在 GPU 上的显存上限。如果需要严格限制显存使用，需要在模型层面（如 PyTorch 的 torch.cuda.set_per_process_memory_fraction）或使用 MPS（Multi-Process Service）来实现。

总结

Docker + GPU 部署方案已经成为 AI 模型生产化部署的事实标准。通过 nvidia-container-toolkit 实现容器 GPU 访问，合理选择 CUDA 镜像层级并利用多阶段构建控制镜像体积，借助 docker-compose 编排多模型服务，再配合 Prometheus + Grafana 构建 GPU 监控体系，可以搭建一套高效、稳定、可观测的 AI 推理基础设施。这套方案既能满足中小规模团队的快速迭代需求，也为未来迁移到 K8s 集群奠定了容器化基础。