build PyTorch from source

Dockerfile

ARG ROOT_CONTAINER=nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
FROM $ROOT_CONTAINER

# https://gist.github.com/Birch-san/211f31f8d901dadd1025398fa1a603b8
# https://github.com/pytorch/pytorch
# https://github.com/pytorch/pytorch/issues/77939

LABEL maintainer="WhiteCat6142 <whitecat6142+git@gmail.com>"

ENV DEBIAN_FRONTEND noninteractive

RUN --mount=type=cache,target=/var/lib/apt/lists --mount=type=cache,target=/var/cache/apt \
    apt update -y && \
    apt install -y ninja-build cmake gcc-11 g++-11 git curl build-essential python3-pip

WORKDIR /root
RUN  curl -OL https://github.com/pytorch/pytorch/releases/download/v2.3.1/pytorch-v2.3.1.tar.gz

ENV _GLIBCXX_USE_CXX11_ABI 1
ENV CC "gcc-11"
ENV CXX "g++-11"
ENV CUDA_DIR "/usr/local/cuda-12"
ENV CXXFLAGS "-Wno-maybe-uninitialized -Wno-uninitialized -Wno-free-nonheap-object -Wno-nonnull"
ENV CFLAGS "-Wno-maybe-uninitialized -Wno-uninitialized -Wno-free-nonheap-object -Wno-nonnull"
ENV USE_ROCM 0
ENV TORCH_CUDA_ARCH_LIST "8.9"
ENV PATH "${CUDA_DIR}/bin:${PATH}"
ENV LD_LIBRARY_PATH "${CUDA_DIR}/lib64"

RUN --mount=type=cache,mode=0755,target=/root/.cache/pip \
    tar -zxvf pytorch-v2.3.1.tar.gz && \
    cd pytorch-v2.3.1 && \
    pip3 install -r requirements.txt && \
    python3 setup.py develop

output.log

[+] Building 5154.6s (10/10) FINISHED                            docker:default
 => [internal] load .dockerignore                                          0.1s
 => => transferring context: 2B                                            0.0s
 => [internal] load build definition from Dockerfile                       0.1s
 => => transferring dockerfile: 1.30kB                                     0.0s
 => [internal] load metadata for docker.io/nvidia/cuda:12.4.1-cudnn-devel  0.7s
 => CACHED [stage-0 1/5] FROM docker.io/nvidia/cuda:12.4.1-cudnn-devel-ub  0.0s
 => CACHED [internal] settings cache mount permissions                     0.0s
 => [stage-0 2/5] RUN --mount=type=cache,target=/var/lib/apt/lists --mou  41.7s
 => [stage-0 3/5] WORKDIR /root                                            0.2s
 => [stage-0 4/5] RUN  curl -OL https://github.com/pytorch/pytorch/releas  4.6s
 => [stage-0 5/5] RUN --mount=type=cache,mode=0755,target=/root/.cache  5021.5s
 => exporting to image                                                    85.5s
 => => exporting layers                                                   85.5s
 => => writing image sha256:aa0d898709215e5634efa18020e3190659a92d7a818e5  0.0s
 => => naming to docker.io/library/docker-torch                            0.0s