社区首页 >专栏 >pytorch容器化环境

pytorch容器化环境

原创

努力的gopher

发布于 2025-06-10 06:31:31

2050

文章被收录于专栏：AIOPSAIOPS

背景介绍：简化 PyTorch 训练环境的初衷

在传统的 PyTorch 开发流程中，环境配置始终是令人头疼的一环。开发者通常需要手动完成以下步骤：

安装 Miniforge 以便于 Conda 环境管理和依赖隔离；
创建虚拟环境并安装符合自身 CUDA 驱动版本的 pytorch、torchvision、torchaudio 等依赖；
避免 numpy 等库与 PyTorch 编译版本不兼容（如 numpy 2.x 导致的初始化失败）；
搭配 VS Code 的 Remote-SSH / Remote-Containers 等插件远程连接进行训练调试；
配置 GPU 驱动、CUDA 工具链、TensorBoard 和 MLflow 等训练辅助组件；
最后还要保证整个环境能够在 Kubernetes、Docker 等容器化平台中无缝运行。

上述过程不仅冗长繁琐，而且在多台机器间复现性差、兼容性容易出错。

项目目标：构建一个真正“开箱即用”的 PyTorch GPU 训练环境

为了解决上述痛点，我们基于如下需求构建了一个可高度复用的容器化训练环境：

✅ 基于官方 CUDA 镜像构建，确保 GPU 加速能力；
✅ 预装 Miniforge，支持 Conda 虚拟环境隔离；
✅ 默认集成 PyTorch（支持 GPU）、TorchVision、TorchAudio 等主流框架；
✅ 内置 TensorBoard 与 MLflow，便于训练日志与模型管理；
✅ 集成 code-server，实现 Web IDE，免去 VS Code 本地配置；
✅ 支持 Kubernetes 部署，自动挂载共享训练目录与实验数据。

通过本镜像，用户无需再手动配置复杂依赖，即可在容器启动后立即开始 GPU 加速训练，同时享受完整的 IDE 开发体验。

dockerfile如下

 FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04
 
 # ---------- 环境变量 ----------
 ENV DEBIAN_FRONTEND=noninteractive
 ENV TZ=Asia/Shanghai
 ENV LANG=C.UTF-8
 ENV LC_ALL=C.UTF-8
 ENV PATH=/opt/conda/bin:$PATH
 
 # ---------- 安装系统依赖 + Miniforge + Git ----------
 RUN apt-get update && apt-get install -y \
     wget git curl bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 \
     python3 python3-pip gnupg lsb-release \
     && wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O /tmp/miniforge.sh \
     && bash /tmp/miniforge.sh -b -p /opt/conda \
     && rm /tmp/miniforge.sh && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 # ---------- 配置 Git (可选) ----------
 #ARG GIT_USERNAME=xxx
 #ARG GIT_PASSWORD=xxx
 #ARG GIT_REPO_DEV=gitlab.xxx.com:8864
 #ARG GIT_REPO_PRO=gitlab.xxx.com:30016
 
 #RUN git config --global user.email "pytorch@pytorch.com" && \
 #    git config --global user.name "pytorch" && \
 #    git config --global credential.helper store && \
 #    git config --global http.sslVerify false && \
 #    git config --global url."https://${GIT_USERNAME}:${GIT_PASSWORD}@".insteadOf "https://" && \
 #    git config --global url."http://${GIT_USERNAME}:${GIT_PASSWORD}@".insteadOf "http://" && \
 #    echo "http://${GIT_USERNAME}:${GIT_PASSWORD}@${GIT_REPO_DEV}" >> /root/.git-credentials && \
 #    echo "http://${GIT_USERNAME}:${GIT_PASSWORD}@${GIT_REPO_PRO}" >> /root/.git-credentials
 
 # ---------- 安装 Code-Server ----------
 RUN curl -fsSL https://code-server.dev/install.sh | bash -s -- --version 4.96.4
 
 # ---------- Conda 环境安装 PyTorch 等依赖 ----------
 RUN conda create -n ml python=3.10 -y && \
     conda run -n ml conda install -y \
         pytorch=2.2.0 torchvision=0.17.0 torchaudio=2.2.0 \
         pytorch-cuda=12.1 -c pytorch -c nvidia && \
     conda run -n ml pip install tensorboard==2.19.0 mlflow==2.22.0 && \
     conda run -n ml pip install "numpy<2.0"
 
 # ---------- 自动激活环境 ----------
 SHELL ["/bin/bash", "-c"]
 RUN echo "source /opt/conda/etc/profile.d/conda.sh && conda activate ml" >> ~/.bashrc
 
 # ---------- 默认工作目录 ----------
 RUN mkdir -p /workspace /shared/logs /shared/mlruns
 WORKDIR /workspace
 
 # ---------- 暴露 code-server 端口 ----------
 EXPOSE 8080
 
 # ---------- 容器启动时默认运行 code-server ----------
 CMD ["bash", "-c", "source /opt/conda/etc/profile.d/conda.sh && conda activate ml && /usr/bin/code-server --bind-addr 0.0.0.0:8080 --auth none /workspace"]

docker 编译和启动命令

 docker build -t code-server:ml-cuda  .
 
 # 宿主机需要安装nvidia驱动和NVIDIA Container Toolkit
 docker run --rm \
   --gpus all -p 8070:8080 \
   -v /data/localpv/cuda:/workspace \
   --name train-env \
   code-server:ml-cuda

安装的必要的宿主机组件（可选）

宿主机环境配置，如果没有配置好环境，无法在容器内直接使用gpu,该步骤可选，已经配置好了则无需配置

1. NVIDIA 驱动（NVIDIA Driver）

✔ 这是最关键的依赖，必须安装并与显卡匹配。
不需要安装 CUDA toolkit 到宿主机，驱动即可。

# 安装方式（例如 Ubuntu）
sudo apt install nvidia-driver-535
reboot#

# 验证是否安装成功：
nvidia-smi

2. NVIDIA Container Toolkit（`nvidia-docker` 插件）

这个工具用于在容器中挂载 GPU、传递 CUDA 能力。

安装步骤（Ubuntu 示例）：

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
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker