[Official Docker build] Use uv instead of conda (#169873)

Fixes https://github.com/pytorch/pytorch/issues/164574

Test Plan: Pull built docker and run smoke tests

Pull Request resolved: https://github.com/pytorch/pytorch/pull/169873
Approved by: https://github.com/malfet

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

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