| # Start from the NVIDIA official image (ubuntu-22.04 + python-3.10) | |
| # https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html | |
| FROM nvcr.io/nvidia/pytorch:24.08-py3 | |
| # Define environments | |
| ENV MAX_JOBS=32 | |
| ENV VLLM_WORKER_MULTIPROC_METHOD=spawn | |
| ENV DEBIAN_FRONTEND=noninteractive | |
| ENV NODE_OPTIONS="" | |
| # Install systemctl and tini | |
| RUN apt-get update && \ | |
| apt-get install -y -o Dpkg::Options::="--force-confdef" systemd tini && \ | |
| apt-get clean || { echo "Installation failed"; exit 1; } | |
| RUN apt-get install -y tzdata \ | |
| && ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime \ | |
| && dpkg-reconfigure -f noninteractive tzdata | |
| # Change pip source | |
| RUN python -m pip install --upgrade pip | |
| # Install torch-2.5.1 + vllm-0.7.3 | |
| RUN pip install --no-cache-dir vllm==0.7.3 torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 tensordict torchdata \ | |
| transformers>=4.49.0 accelerate datasets peft hf-transfer diffusers \ | |
| codetiming hydra-core pandas pyarrow>=15.0.0 pylatexenc qwen-vl-utils wandb ninja liger-kernel \ | |
| pytest yapf py-spy pyext pre-commit ruff packaging | |
| # Install flux | |
| RUN pip install --no-cache-dir byte-flux | |
| # Install flash-attn and triton | |
| RUN pip install --no-cache-dir flash-attn triton>=3.1.0 | |