diff --git a/Dockerfile b/Dockerfile index 331cf00593c..099c0f82efe 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,8 +5,7 @@ # For reference: # - https://docs.docker.com/build/dockerfile/frontend/#stable-channel -ARG BASE_IMAGE=ubuntu:22.04 -ARG PYTHON_VERSION=3.11 +ARG BASE_IMAGE=ubuntu:24.04 FROM ${BASE_IMAGE} as dev-base RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ @@ -17,53 +16,40 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins curl \ git \ libjpeg-dev \ - libpng-dev && \ + libpng-dev \ + python3 \ + python3-pip \ + python3-venv \ + python3-dev && \ rm -rf /var/lib/apt/lists/* RUN /usr/sbin/update-ccache-symlinks RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache -ENV PATH /opt/conda/bin:$PATH +ENV PATH /opt/pytorch-venv/bin:$PATH -FROM dev-base as conda -ARG PYTHON_VERSION=3.11 -# Automatically set by buildx -ARG TARGETPLATFORM -# translating Docker's TARGETPLATFORM into miniconda arches -RUN case ${TARGETPLATFORM} in \ - "linux/arm64") MINICONDA_ARCH=aarch64 ;; \ - *) MINICONDA_ARCH=x86_64 ;; \ - esac && \ - curl -fsSL -v -o ~/miniconda.sh -O "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-${MINICONDA_ARCH}.sh" +FROM dev-base as python-venv COPY requirements.txt requirements-build.txt . -# Manually invoke bash on miniconda script per https://github.com/conda/conda/issues/10431 -RUN chmod +x ~/miniconda.sh && \ - bash ~/miniconda.sh -b -p /opt/conda && \ - rm ~/miniconda.sh && \ - /opt/conda/bin/conda install -y python=${PYTHON_VERSION} cmake conda-build pyyaml numpy ipython && \ - /opt/conda/bin/python -mpip install -r requirements.txt && \ - /opt/conda/bin/conda clean -ya +# Create virtual environment and install packages +RUN python3 -m venv /opt/pytorch-venv && \ + /opt/pytorch-venv/bin/pip install --upgrade pip setuptools wheel && \ + /opt/pytorch-venv/bin/pip install cmake pyyaml numpy ipython -r requirements.txt FROM dev-base as submodule-update WORKDIR /opt/pytorch COPY . . RUN git submodule update --init --recursive -FROM conda as conda-installs -ARG PYTHON_VERSION=3.11 +FROM python-venv as pytorch-installs ARG CUDA_PATH=cu121 ARG INSTALL_CHANNEL=whl/nightly # Automatically set by buildx -# pinning version of conda here see: https://github.com/pytorch/pytorch/issues/164574 -RUN /opt/conda/bin/conda install -y python=${PYTHON_VERSION} conda=25.7.0 - ARG TARGETPLATFORM # INSTALL_CHANNEL whl - release, whl/nightly - nightly, whl/test - test channels RUN case ${TARGETPLATFORM} in \ - "linux/arm64") pip install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchaudio ;; \ - *) pip install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH#.}/ torch torchvision torchaudio ;; \ - esac && \ - /opt/conda/bin/conda clean -ya -RUN /opt/conda/bin/pip install torchelastic + "linux/arm64") /opt/pytorch-venv/bin/pip install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchaudio ;; \ + *) /opt/pytorch-venv/bin/pip install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH#.}/ torch torchvision torchaudio ;; \ + esac +RUN /opt/pytorch-venv/bin/pip install torchelastic RUN IS_CUDA=$(python -c 'import torch ; print(torch.cuda._is_compiled())'); \ echo "Is torch compiled with cuda: ${IS_CUDA}"; \ if test "${IS_CUDA}" != "True" -a ! -z "${CUDA_VERSION}"; then \ @@ -81,12 +67,12 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins libjpeg-dev \ libpng-dev \ && rm -rf /var/lib/apt/lists/* -COPY --from=conda-installs /opt/conda /opt/conda +COPY --from=pytorch-installs /opt/pytorch-venv /opt/pytorch-venv RUN if test -n "${TRITON_VERSION}" -a "${TARGETPLATFORM}" != "linux/arm64"; then \ DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends gcc; \ rm -rf /var/lib/apt/lists/*; \ fi -ENV PATH /opt/conda/bin:$PATH +ENV PATH /opt/pytorch-venv/bin:$PATH ENV NVIDIA_VISIBLE_DEVICES all ENV NVIDIA_DRIVER_CAPABILITIES compute,utility ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64 @@ -96,5 +82,5 @@ WORKDIR /workspace FROM official as dev # Should override the already installed version from the official-image stage -COPY --from=conda /opt/conda /opt/conda +COPY --from=python-venv /opt/pytorch-venv /opt/pytorch-venv COPY --from=submodule-update /opt/pytorch /opt/pytorch diff --git a/docker.Makefile b/docker.Makefile index f510d1618a2..0d925ffbf6b 100644 --- a/docker.Makefile +++ b/docker.Makefile @@ -11,8 +11,8 @@ endif CUDA_VERSION_SHORT ?= 12.1 CUDA_VERSION ?= 12.1.1 CUDNN_VERSION ?= 9 -BASE_RUNTIME = ubuntu:22.04 -BASE_DEVEL = nvidia/cuda:$(CUDA_VERSION)-devel-ubuntu22.04 +BASE_RUNTIME = ubuntu:24.04 +BASE_DEVEL = nvidia/cuda:$(CUDA_VERSION)-devel-ubuntu24.04 CMAKE_VARS ?= # The conda channel to use to install cudatoolkit