面向 AI 的 Docker:把「我这能跑」变成「到处能跑」

面向 AI 的 Docker:把「我这能跑」变成「到处能跑」

本节摘要:容器(Container)让「works on my machine(在我机器上能跑)」成为历史。你笔记本上用 PyTorch 2.3 + CUDA 12.4 + Python 3.12 训练的模型,同事的机器是 PyTorch 2.1 + CUDA 11.8 + Python 3.10,模型一跑就崩。AI 项目是依赖噩梦的重灾区:Python、PyTorch、CUDA 驱动、cuDNN、系统级 C 库,还有像 flash-attn 这种需要精确编译器版本的特殊包,任何一环错位都崩。Docker 把这一切打包成单个镜像(Image),在任何地方都一致运行。本节带你吃透容器化的核心概念、AI 项目为何比一般项目更需要 Docker、如何写一份 GPU 版 Dockerfile、如何用 Volume 持久化模型与数据、如何用 NVIDIA Container Toolkit 把 GPU 透传进容器,以及如何用 Docker Compose 编排「推理服务器 + 向量数据库」这类多服务 AI 应用。

对应原课程:Phase 00 · Lesson 07 · docker-for-ai(原英文 phases/00-setup-and-tooling/07-docker-for-ai/docs/en.md)。前置:第 01、03 节。

学习目标

阅读完本节,你应当能够:

1. 用 Dockerfile 构建一个带 CUDA、PyTorch、AI 库的 GPU 镜像。
2. 用 Volume 挂载把宿主机目录映射进容器,持久化模型、数据集、代码,跨容器重建不丢。
3. 配置 NVIDIA Container Toolkit,把宿主 GPU 暴露给容器。
4. 用 Docker Compose 编排多服务 AI 应用(推理服务器 + 向量数据库)。

一、问题与直觉

你笔记本上用 PyTorch 2.3、CUDA 12.4、Python 3.12 训练了模型;同事是 PyTorch 2.1、CUDA 11.8、Python 3.10;模型在他机器上一跑就崩。而你的 Dockerfile 在两边都能跑。

AI 项目是依赖噩梦:典型栈包括 Python、PyTorch、CUDA 驱动、cuDNN、系统级 C 库,还有 flash-attn 这种要精确编译器版本的特殊包。Docker 把这些全打包进单个镜像,到处一致。

为什么 AI 项目比别人更需要 Docker

1. GPU 驱动脆弱:CUDA 12.4 的代码在 CUDA 11.8 上跑不了。Docker 把 CUDA 工具链隔离在容器内,同时通过 NVIDIA Container Toolkit 共享宿主 GPU 驱动。
2. 模型权重巨大:7B 参数模型 fp16 就 14 GB,你绝不想每次重建容器都重下一遍。Docker Volume 让你把宿主的 models 目录挂进容器。
3. 多服务架构常见:真实 AI 应用不是一个 Python 脚本,而是推理服务器 + RAG 的向量数据库 + 也许还有 Web 前端。Docker Compose 一条命令编排这些。

关键词汇

AI 中常见的容器模式

开发容器(Dev Container)
  全工具链、编辑器支持、Jupyter、调试工具。开发与实验阶段用。

训练容器(Training Container)
  最小化,只有训练脚本与依赖。跑在 GPU 集群上,无编辑器、无 Jupyter。

推理容器(Inference Container)
  为服务优化,镜像小、冷启动快。生产环境跑在负载均衡器后。
​

二、从零实现

Step 1 装 Docker

# macOS
brew install --cask docker && open /Applications/Docker.app

# Ubuntu
curl -fsSL https://get.docker.com | sh
sudo usermod -aG docker $USER     # 登出再登入让用户组生效
​

验证:docker --version 与 docker run hello-world。

Step 2 装 NVIDIA Container Toolkit(Linux + NVIDIA GPU)

这让 Docker 容器能访问你的 GPU。macOS 与 Windows(WSL2)用户可跳过——Docker Desktop 在这些平台用别的方式处理 GPU 透传。

安装后用一条命令验证容器内能看到 GPU:

docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
​

能看到你的 GPU 信息就说明 toolkit 工作正常。关键就是 --gpus all 这个标志。

Step 3 选对基础镜像

选错基础镜像会浪费几小时排查。常见选择(完整安装命令见原课程 code/Dockerfile):

nvidia/cuda:12.4.1-devel-ubuntu22.04
  完整 CUDA 工具链,含编译器。装需要 nvcc 的包(flash-attn、bitsandbytes)用。约 4 GB。

nvidia/cuda:12.4.1-runtime-ubuntu22.04
  仅 CUDA 运行时,无编译器。跑已构建好的代码用。约 1.5 GB。

pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime
  在 CUDA 之上预装好 PyTorch。想省去装 PyTorch 这步用。约 6 GB。

python:3.12-slim
  无 CUDA,CPU only。CPU 推理、轻量工具用。约 150 MB。
​

💡 选镜像心法:要编译用 devel,只运行用 runtime,想要预装 PyTorch 用官方 torch 镜像,CPU-only 用 slim。镜像越小,构建与拉取越快,生产环境越该小。

Step 4 写 AI 开发用的 Dockerfile

原课程 code/Dockerfile 给出完整版。关键结构如下(命令简化呈现,完整文件见原仓库):

FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# 1. 系统依赖 + Python 3.12
RUN apt-get update && apt-get install -y git curl build-essential ... python3.12 python3.12-venv python3.12-dev

# 2. 装 pip
RUN curl -sSL <get-pip.py> | python

# 3. 装 CUDA 版 PyTorch(关键:带 --index-url)
RUN python -m pip install torch==2.6.0+cu124 torchvision==0.21.0+cu124 \
    --index-url https://download.pytorch.org/whl/cu124

# 4. 装其他 AI 库
RUN python -m pip install numpy pandas scikit-learn matplotlib jupyter \
    transformers datasets accelerate safetensors

WORKDIR /workspace
VOLUME ["/workspace", "/models"]
EXPOSE 8888
CMD ["python"]
​

构建与运行:

docker build -t ai-dev -f .../code/Dockerfile .
docker run --rm -it --gpus all \
    -v $(pwd):/workspace \
    -v ~/models:/models \
    ai-dev python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
​

首次构建会慢(下载 CUDA 基础镜像 + PyTorch);后续构建复用缓存层会快很多。

设计要点:Dockerfile 的每一行 RUN 都是一层(layer),不变的层会被缓存。把变动少的层放前面(系统依赖、PyTorch),变动多的放后面(你的代码、新增的小包),这样改代码时只重建后面的几层,构建飞快。

Step 5 用 Volume 持久化数据与模型

Volume 挂载对 AI 工作至关重要。没有它,你那 14 GB 的模型在容器一停就没了:

-v $(pwd):/workspace      # 挂载你的代码
-v ~/models:/models       # 挂载共享模型目录
-v ~/datasets:/data       # 挂载数据集
​

训练脚本里从挂载路径加载:

from transformers import AutoModel
model = AutoModel.from_pretrained("/models/llama-7b")
​

模型躺在宿主文件系统上,容器想重建多少次都行,不用重下。

Step 6 Docker Compose 编排多服务

一个真实 RAG 应用需要推理服务器 + 向量数据库。Docker Compose 一条命令拉起全部(完整 code/docker-compose.yml):

services:
  ai-dev:
    build: { context: ., dockerfile: Dockerfile }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ../../../:/workspace
      - ~/models:/models
    ports: ["8888:8888"]
    command: jupyter notebook --ip=0.0.0.0 --port=8888 --no-browser --allow-root

  qdrant:
    image: qdrant/qdrant:v1.12.5
    ports: ["6333:6333", "6334:6334"]
    volumes:
      - qdrant_data:/qdrant/storage

volumes:
  qdrant_data:
​

docker compose up -d        # 启动全部
docker compose down         # 停止
docker compose down -v      # 停止并删除数据卷
​

启动后,AI 容器能通过服务名 http://qdrant:6333 访问向量数据库——Docker Compose 会自动创建共享网络。

三、框架对比:三种 AI 容器模式

💡 三者本质是「同一份代码,三种打包方式」。开发容器追求体验,训练容器追求稳定可复现,推理容器追求体积与延迟。学会用多阶段构建(multi-stage build)从开发镜像产出一个精简推理镜像,是工业界的标准做法。

四、可复用产物

本节产出(原课程 code/ 与 outputs/):

• Dockerfile:一份 GPU 版 AI 开发镜像的完整构建脚本,可直接 docker build 使用;改基础镜像与 PyTorch 版本即可适配你的 CUDA 版本。
• docker-compose.yml:一份「AI 开发容器 + Qdrant 向量数据库」的编排模板,任何 RAG / Agent 项目都能复制改用。
• 一组常用 Docker 命令:docker ps、docker images、docker system prune -a、docker exec -it <id> nvidia-smi、docker cp、docker logs -f,贴在桌面随时查。

五、练习

1. (Easy) 构建本节 Dockerfile,在容器内跑 python -c "import torch; print(torch.__version__, torch.cuda.is_available())",确认 GPU 被识别;再 docker images 看镜像大小。
2. (Medium) 用 docker compose up -d 启动开发容器 + Qdrant,在 AI 容器内用 qdrant_client 连 http://qdrant:6333,建一个集合、插几条向量、做一次检索;完成后 docker compose down。
3. (Hard) 给 Dockerfile 加一个 flask 依赖,重建镜像,在容器内跑一个返回模型推理结果的简单 API 服务(端口 5000),用 -p 5000:5000 把端口映射到宿主,从宿主 curl 调用它;再尝试把基础镜像从 devel 换成 runtime,对比镜像大小变化,理解「构建期 vs 运行期」分离的价值。

本节要点回顾

1. 容器终结「works on my machine」:把代码、运行时、库、系统工具打包成一个到处一致的镜像。
2. AI 比一般项目更需要 Docker:GPU 驱动脆弱、模型权重巨大、多服务架构常见。
3. GPU 透传靠 NVIDIA Container Toolkit:用 --gpus all 把宿主 GPU 暴露给容器,宿主只共享驱动,工具链隔离在容器内。
4. 基础镜像选型:devel(要编译)、runtime(只运行)、官方 torch 镜像(预装 PyTorch)、slim(CPU-only)。
5. Volume 挂载是 AI 关键:模型与数据集挂载进容器,跨重建不丢;14 GB 模型不能靠容器内重下。
6. 层缓存优化:变动少的放前面(系统依赖、PyTorch),变动多的放后面(代码、小包),改代码时只重建后几层。
7. Docker Compose 编排多服务:推理服务器 + 向量数据库 + 前端,一条 up -d 拉起,服务名互通。
8. 三种容器模式:开发(全工具链)、训练(最小化)、推理(小镜像快冷启动),用多阶段构建从开发镜像产出推理镜像。

下一节,我们将进入「编辑器配置」——把 VS Code 调成顺手的 AI 开发环境,让 Docker 容器、Python、Notebook、Git 在编辑器里无缝协作。