Files
pytorch/Dockerfile
Nikita Shulga 6718b1caeb Add python3-dev to Dockerfile dependencies (#170456)
Fixes dev image:

Test:
```
apt-get install python3-dev
find / -name "Python.h" 2>/dev/null
/usr/include/python3.12/Python.h
```

Failing  Test:
```
python smoke_test.py --package torchonly --runtime-error-check disabled
torch: 2.10.0.dev20251213+cu129
ATen/Parallel:
	at::get_num_threads() : 4
	at::get_num_interop_threads() : 4
OpenMP 201511 (a.k.a. OpenMP 4.5)
	omp_get_max_threads() : 4
Intel(R) oneAPI Math Kernel Library Version 2024.2-Product Build 20240605 for Intel(R) 64 architecture applications
	mkl_get_max_threads() : 4
Intel(R) MKL-DNN v3.7.1 (Git Hash 8d263e693366ef8db40acc569cc7d8edf644556d)
std::thread::hardware_concurrency() : 8
Environment variables:
	OMP_NUM_THREADS : [not set]
	MKL_NUM_THREADS : [not set]
ATen parallel backend: OpenMP

Skip version check for channel None as stable version is None
Testing smoke_test_conv2d
Testing smoke_test_linalg on cpu
Testing SDPA on cpu using type torch.float16
Testing smoke_test_compile for cpu and torch.float16
W1215 21:06:22.064000 886 torch/utils/cpp_extension.py:118] [0/0] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
/usr/local/lib/python3.12/dist-packages/torch/_inductor/cpp_builder.py:1198: UserWarning: Can't find Python.h in /usr/include/python3.12
  warnings.warn(f"Can't find Python.h in {str(include_dir)}")
```

Passing Test:
```
python smoke_test.py --package torchonly --runtime-error-check disabled
torch: 2.10.0.dev20251213+cu129
ATen/Parallel:
	at::get_num_threads() : 4
	at::get_num_interop_threads() : 4
OpenMP 201511 (a.k.a. OpenMP 4.5)
	omp_get_max_threads() : 4
Intel(R) oneAPI Math Kernel Library Version 2024.2-Product Build 20240605 for Intel(R) 64 architecture applications
	mkl_get_max_threads() : 4
Intel(R) MKL-DNN v3.7.1 (Git Hash 8d263e693366ef8db40acc569cc7d8edf644556d)
std::thread::hardware_concurrency() : 8
Environment variables:
	OMP_NUM_THREADS : [not set]
	MKL_NUM_THREADS : [not set]
ATen parallel backend: OpenMP

Skip version check for channel None as stable version is None
Testing smoke_test_conv2d
Testing smoke_test_linalg on cpu
Testing SDPA on cpu using type torch.float16
Testing smoke_test_compile for cpu and torch.float16
W1215 21:04:14.211000 754 torch/utils/cpp_extension.py:118] [0/0] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
Testing smoke_test_compile for cpu and torch.float32
Testing smoke_test_compile for cpu and torch.float64
Picked CPU ISA VecAVX512 bit width 512
Testing smoke_test_compile with mode 'max-autotune' for torch.float32
cudagraph partition due to non gpu ops
cudagraph partition due to non gpu ops. Found from :
   File "/workspace/pytorch/.ci/pytorch/smoke_test/smoke_test.py", line 57, in forward
    x = self.conv1(x)

cudagraph partition due to non gpu ops
cudagraph partition due to non gpu ops. Found from :
   File "/workspace/pytorch/.ci/pytorch/smoke_test/smoke_test.py", line 58, in forward
    x = self.conv2(x)

cudagraph partition due to non gpu ops. Found from :
   File "/workspace/pytorch/.ci/pytorch/smoke_test/smoke_test.py", line 61, in forward
    output = self.fc1(x)

Windows platform or CUDA is not available, skipping NVSHMEM test

```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/170456
Approved by: https://github.com/atalman
2025-12-15 22:41:31 +00:00

93 lines
3.3 KiB
Docker

# syntax=docker/dockerfile:1
# NOTE: Building this image require's docker version >= 23.0.
#
# For reference:
# - https://docs.docker.com/build/dockerfile/frontend/#stable-channel
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 \
build-essential \
ca-certificates \
ccache \
cmake \
curl \
git \
libjpeg-dev \
libpng-dev \
python3 \
python3-pip \
python-is-python3 \
python3-dev && \
rm -rf /var/lib/apt/lists/*
# Remove PEP 668 restriction (safe in containers)
RUN rm -f /usr/lib/python*/EXTERNALLY-MANAGED
RUN /usr/sbin/update-ccache-symlinks
RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache
FROM dev-base as python-deps
COPY requirements.txt requirements-build.txt .
# Install Python packages to system Python
RUN pip3 install --upgrade --ignore-installed pip setuptools wheel && \
pip3 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 python-deps as pytorch-installs
ARG CUDA_PATH=cu121
ARG INSTALL_CHANNEL=whl/nightly
# Automatically set by buildx
ARG TARGETPLATFORM
# INSTALL_CHANNEL whl - release, whl/nightly - nightly, whl/test - test channels
RUN case ${TARGETPLATFORM} in \
"linux/arm64") pip3 install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchaudio ;; \
*) pip3 install --index-url https://download.pytorch.org/${INSTALL_CHANNEL}/${CUDA_PATH#.}/ torch torchvision torchaudio ;; \
esac
RUN pip3 install torchelastic
RUN IS_CUDA=$(python3 -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 \
exit 1; \
fi
FROM ${BASE_IMAGE} as official
ARG PYTORCH_VERSION
ARG TRITON_VERSION
ARG TARGETPLATFORM
ARG CUDA_VERSION
LABEL com.nvidia.volumes.needed="nvidia_driver"
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
ca-certificates \
libjpeg-dev \
libpng-dev \
python-is-python3 \
python3 \
python3-dev \
python3-pip \
&& rm -rf /var/lib/apt/lists/*
# Copy Python packages from pytorch-installs stage
COPY --from=pytorch-installs /usr/local/lib/python3.12 /usr/local/lib/python3.12
COPY --from=pytorch-installs /usr/local/bin /usr/local/bin
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 NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
ENV PYTORCH_VERSION ${PYTORCH_VERSION}
WORKDIR /workspace
FROM official as dev
# Should override the already installed version from the official-image stage
COPY --from=python-deps /usr/local/lib/python3.12 /usr/local/lib/python3.12
COPY --from=python-deps /usr/local/bin /usr/local/bin
COPY --from=submodule-update /opt/pytorch /opt/pytorch