mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
Add mingw to `pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11` docker image to support AOTI cross-compilation This PR will make docker container rebuild, and upgrade python version from 3.13.7 to 3.13.8. and it relies on https://github.com/pytorch/pytorch/pull/165667 Pull Request resolved: https://github.com/pytorch/pytorch/pull/165560 Approved by: https://github.com/malfet
195 lines
6.6 KiB
Docker
195 lines
6.6 KiB
Docker
ARG UBUNTU_VERSION
|
|
|
|
FROM ubuntu:${UBUNTU_VERSION} as base
|
|
|
|
ARG UBUNTU_VERSION
|
|
|
|
ENV DEBIAN_FRONTEND noninteractive
|
|
|
|
ARG CLANG_VERSION
|
|
|
|
# Install common dependencies (so that this step can be cached separately)
|
|
COPY ./common/install_base.sh install_base.sh
|
|
RUN bash ./install_base.sh && rm install_base.sh
|
|
|
|
# Install clang
|
|
ARG LLVMDEV
|
|
COPY ./common/install_clang.sh install_clang.sh
|
|
RUN bash ./install_clang.sh && rm install_clang.sh
|
|
|
|
# Install user
|
|
COPY ./common/install_user.sh install_user.sh
|
|
RUN bash ./install_user.sh && rm install_user.sh
|
|
|
|
# Install katex
|
|
ARG KATEX
|
|
COPY ./common/install_docs_reqs.sh install_docs_reqs.sh
|
|
RUN bash ./install_docs_reqs.sh && rm install_docs_reqs.sh
|
|
|
|
# Install conda and other packages (e.g., numpy, pytest)
|
|
ARG ANACONDA_PYTHON_VERSION
|
|
ARG DOCS
|
|
ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
|
|
ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
|
|
ENV DOCS=$DOCS
|
|
COPY requirements-ci.txt requirements-docs.txt /opt/conda/
|
|
COPY ./common/install_conda.sh install_conda.sh
|
|
COPY ./common/common_utils.sh common_utils.sh
|
|
COPY ./common/install_magma_conda.sh install_magma_conda.sh
|
|
RUN bash ./install_conda.sh && rm install_conda.sh install_magma_conda.sh common_utils.sh /opt/conda/requirements-ci.txt /opt/conda/requirements-docs.txt
|
|
RUN if [ -n "${UNINSTALL_DILL}" ]; then pip uninstall -y dill; fi
|
|
|
|
# Install gcc
|
|
ARG GCC_VERSION
|
|
COPY ./common/install_gcc.sh install_gcc.sh
|
|
RUN bash ./install_gcc.sh && rm install_gcc.sh
|
|
|
|
# Install lcov for C++ code coverage
|
|
COPY ./common/install_lcov.sh install_lcov.sh
|
|
RUN bash ./install_lcov.sh && rm install_lcov.sh
|
|
|
|
# Install cuda and cudnn
|
|
ARG CUDA_VERSION
|
|
COPY ./common/install_cuda.sh install_cuda.sh
|
|
COPY ./common/install_nccl.sh install_nccl.sh
|
|
COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
|
|
COPY ./common/install_cusparselt.sh install_cusparselt.sh
|
|
RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh install_nccl.sh /ci_commit_pins/nccl-cu* install_cusparselt.sh
|
|
ENV DESIRED_CUDA ${CUDA_VERSION}
|
|
ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
|
|
# No effect if cuda not installed
|
|
ENV USE_SYSTEM_NCCL=1
|
|
ENV NCCL_INCLUDE_DIR="/usr/local/cuda/include/"
|
|
ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/"
|
|
|
|
|
|
# (optional) Install UCC
|
|
ARG UCX_COMMIT
|
|
ARG UCC_COMMIT
|
|
ARG CUDA_VERSION
|
|
ENV UCX_COMMIT $UCX_COMMIT
|
|
ENV UCC_COMMIT $UCC_COMMIT
|
|
ENV UCX_HOME /usr
|
|
ENV UCC_HOME /usr
|
|
ADD ./common/install_ucc.sh install_ucc.sh
|
|
RUN if [ -n "${UCX_COMMIT}" ] && [ -n "${UCC_COMMIT}" ]; then bash ./install_ucc.sh; fi
|
|
RUN rm install_ucc.sh
|
|
|
|
# (optional) Install vision packages like OpenCV
|
|
ARG VISION
|
|
COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
|
|
RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
|
|
RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
|
|
ENV INSTALLED_VISION ${VISION}
|
|
|
|
# (optional) Install non-default Ninja version
|
|
ARG NINJA_VERSION
|
|
COPY ./common/install_ninja.sh install_ninja.sh
|
|
RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
|
|
RUN rm install_ninja.sh
|
|
|
|
COPY ./common/install_openssl.sh install_openssl.sh
|
|
RUN bash ./install_openssl.sh
|
|
ENV OPENSSL_ROOT_DIR /opt/openssl
|
|
ENV OPENSSL_DIR /opt/openssl
|
|
RUN rm install_openssl.sh
|
|
|
|
ARG INDUCTOR_BENCHMARKS
|
|
COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
|
|
COPY ./common/common_utils.sh common_utils.sh
|
|
COPY ci_commit_pins/huggingface-requirements.txt huggingface-requirements.txt
|
|
COPY ci_commit_pins/timm.txt timm.txt
|
|
COPY ci_commit_pins/torchbench.txt torchbench.txt
|
|
RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
|
|
RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt
|
|
|
|
ARG INSTALL_MINGW
|
|
COPY ./common/install_mingw.sh install_mingw.sh
|
|
RUN if [ -n "${INSTALL_MINGW}" ]; then bash ./install_mingw.sh; fi
|
|
RUN rm install_mingw.sh
|
|
|
|
ARG TRITON
|
|
ARG TRITON_CPU
|
|
|
|
# Create a separate stage for building Triton and Triton-CPU. install_triton
|
|
# will check for the presence of env vars
|
|
FROM base as triton-builder
|
|
COPY ./common/install_triton.sh install_triton.sh
|
|
COPY ./common/common_utils.sh common_utils.sh
|
|
COPY ci_commit_pins/triton.txt triton.txt
|
|
COPY ci_commit_pins/triton-cpu.txt triton-cpu.txt
|
|
RUN bash ./install_triton.sh
|
|
|
|
FROM base as final
|
|
COPY --from=triton-builder /opt/triton /opt/triton
|
|
RUN if [ -n "${TRITON}" ] || [ -n "${TRITON_CPU}" ]; then pip install /opt/triton/*.whl; chown -R jenkins:jenkins /opt/conda; fi
|
|
RUN rm -rf /opt/triton
|
|
|
|
ARG EXECUTORCH
|
|
# Build and install executorch
|
|
COPY ./common/install_executorch.sh install_executorch.sh
|
|
COPY ./common/common_utils.sh common_utils.sh
|
|
COPY ci_commit_pins/executorch.txt executorch.txt
|
|
RUN if [ -n "${EXECUTORCH}" ]; then bash ./install_executorch.sh; fi
|
|
RUN rm install_executorch.sh common_utils.sh executorch.txt
|
|
|
|
ARG HALIDE
|
|
# Build and install halide
|
|
COPY ./common/install_halide.sh install_halide.sh
|
|
COPY ./common/common_utils.sh common_utils.sh
|
|
COPY ci_commit_pins/halide.txt halide.txt
|
|
RUN if [ -n "${HALIDE}" ]; then bash ./install_halide.sh; fi
|
|
RUN rm install_halide.sh common_utils.sh halide.txt
|
|
|
|
ARG ONNX
|
|
# Install ONNX dependencies
|
|
COPY ./common/install_onnx.sh ./common/common_utils.sh ./
|
|
RUN if [ -n "${ONNX}" ]; then bash ./install_onnx.sh; fi
|
|
RUN rm install_onnx.sh common_utils.sh
|
|
|
|
# (optional) Build ACL
|
|
ARG ACL
|
|
COPY ./common/install_acl.sh install_acl.sh
|
|
RUN if [ -n "${ACL}" ]; then bash ./install_acl.sh; fi
|
|
RUN rm install_acl.sh
|
|
ENV INSTALLED_ACL ${ACL}
|
|
|
|
ARG OPENBLAS
|
|
COPY ./common/install_openblas.sh install_openblas.sh
|
|
RUN if [ -n "${OPENBLAS}" ]; then bash ./install_openblas.sh; fi
|
|
RUN rm install_openblas.sh
|
|
ENV INSTALLED_OPENBLAS ${OPENBLAS}
|
|
|
|
# Install ccache/sccache (do this last, so we get priority in PATH)
|
|
ARG SKIP_SCCACHE_INSTALL
|
|
COPY ./common/install_cache.sh install_cache.sh
|
|
ENV PATH /opt/cache/bin:$PATH
|
|
RUN if [ -z "${SKIP_SCCACHE_INSTALL}" ]; then bash ./install_cache.sh; fi
|
|
RUN rm install_cache.sh
|
|
|
|
# Add jni.h for java host build
|
|
COPY ./common/install_jni.sh install_jni.sh
|
|
COPY ./java/jni.h jni.h
|
|
RUN bash ./install_jni.sh && rm install_jni.sh
|
|
|
|
# Install Open MPI for CUDA
|
|
COPY ./common/install_openmpi.sh install_openmpi.sh
|
|
RUN if [ -n "${CUDA_VERSION}" ]; then bash install_openmpi.sh; fi
|
|
RUN rm install_openmpi.sh
|
|
|
|
# Include BUILD_ENVIRONMENT environment variable in image
|
|
ARG BUILD_ENVIRONMENT
|
|
ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
|
|
|
|
# Install LLVM dev version (Defined in the pytorch/builder github repository)
|
|
ARG SKIP_LLVM_SRC_BUILD_INSTALL
|
|
COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
|
|
RUN if [ -n "${SKIP_LLVM_SRC_BUILD_INSTALL}" ]; then set -eu; rm -rf /opt/llvm; fi
|
|
|
|
# AWS specific CUDA build guidance
|
|
ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all"
|
|
ENV CUDA_PATH /usr/local/cuda
|
|
|
|
USER jenkins
|
|
CMD ["bash"]
|