4. 容器化环境搭建：Docker安装与配置、NVIDIA Container Toolkit安装、拉取并运行PyTorch/TensorFlow镜像、容器资源限制

说实话，搞GPU集群运维这些年，我踩过最大的坑就是环境不一致。你想想看，开发环境跑得好好的模型，一到生产环境就报CUDA版本不对、驱动不兼容——这种问题我至少遇到过十几次。后来我学乖了，容器化成了我的标配。

这一章，咱们就聊聊怎么把GPU集群的容器化环境搭起来。说白了，就是让Docker能认你的GPU，让PyTorch和TensorFlow能在容器里跑得飞起。

4.1 Docker安装与配置

Docker的安装其实不复杂，但有几个细节我得提醒你。我个人习惯用官方脚本安装，干净利落。

# 卸载旧版本（如果有）
sudo apt-get remove docker docker-engine docker.io containerd runc

# 安装依赖
sudo apt-get update
sudo apt-get install -y \
    ca-certificates \
    curl \
    gnupg \
    lsb-release

# 添加Docker官方GPG密钥
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg

# 设置稳定版仓库
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

# 安装Docker引擎
sudo apt-get update
sudo apt-get install -y docker-ce docker-ce-cli containerd.io

# 启动并设置开机自启
sudo systemctl start docker
sudo systemctl enable docker

# 将当前用户加入docker组
sudo usermod -aG docker $USER
# 重新登录或执行 newgrp docker 使生效
newgrp docker

验证一下安装是否成功：

docker --version
docker run hello-world

看到"Hello from Docker!"就说明装好了。嗯，这里要注意，如果遇到权限问题，记得检查用户组是否生效。

4.2 NVIDIA Container Toolkit安装

光有Docker还不够，你得让容器能访问GPU。这就是NVIDIA Container Toolkit的活儿了。

我曾经在部署一个推理服务时，容器里死活调不到GPU，查了半天发现是toolkit版本和驱动不匹配。所以安装前，先确认驱动版本：

nvidia-smi

确认驱动没问题后，开始安装：

# 设置NVIDIA容器工具包仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 安装nvidia-docker2
sudo apt-get update
sudo apt-get install -y nvidia-docker2

# 重启Docker服务
sudo systemctl restart docker

验证GPU是否能在容器中使用：

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

如果能看到GPU信息，恭喜你，环境通了。

4.3 拉取并运行PyTorch/TensorFlow镜像

环境搭好了，接下来就是拉镜像。我建议直接用官方镜像，省心。但要注意版本匹配——你的驱动、CUDA版本和框架版本得能对上。

先看看官方有哪些镜像：

# PyTorch官方镜像
docker pull pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

# TensorFlow官方镜像
docker pull tensorflow/tensorflow:2.13.0-gpu

跑一个PyTorch容器试试：

docker run --gpus all -it --rm \
    -v /data:/workspace/data \
    pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime \
    python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count())"

输出应该是：

True
1

TensorFlow也一样：

docker run --gpus all -it --rm \
    -v /data:/workspace/data \
    tensorflow/tensorflow:2.13.0-gpu \
    python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

4.4 容器资源限制

在集群环境里，资源限制是必须的。不然一个容器把GPU显存吃光，其他任务全挂。我见过最惨的一次，有人跑了个训练任务没限制显存，直接把整张A100的80G显存占满，其他三个用户的推理服务全崩了。

Docker支持对GPU资源做精细控制：

# 限制使用特定GPU（比如只使用GPU 0和GPU 1）
docker run --gpus '"device=0,1"' -it --rm pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime nvidia-smi

# 限制显存使用（通过设置环境变量）
docker run --gpus all \
    -e NVIDIA_VISIBLE_DEVICES=0 \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 \
    -it --rm pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime \
    python -c "import torch; torch.cuda.set_per_process_memory_fraction(0.5, 0); print('显存限制已设置')"

除了GPU，CPU和内存也要限制：

# 限制CPU核心数和内存
docker run --gpus all \
    --cpus="4" \
    --memory="16g" \
    --memory-swap="16g" \
    -it --rm pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime \
    python -c "import os; print('CPU限制:', os.cpu_count())"

最后，给一个完整的生产级启动命令：

docker run --gpus '"device=0,1"' \
    --cpus="8" \
    --memory="32g" \
    --memory-swap="32g" \
    --shm-size="8g" \
    -v /data:/workspace/data \
    -v /models:/workspace/models \
    -e CUDA_VISIBLE_DEVICES=0,1 \
    -e NCCL_DEBUG=INFO \
    --name my_training_job \
    -d pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime \
    python train.py

好了，容器化环境这块就聊到这儿。记住一句话：环境搭得好，运维下班早。把Docker、NVIDIA Container Toolkit、镜像和资源限制这四步走扎实了，后面跑模型就顺了。