cuda-nvml-dev-12-1
libcufft-12-1
libcurand-12-1
libcusolver-dev-12-1
libcusparse-dev-12-1
libcublas-12-1
libcublas-dev-12-1
libnccl-dev=2.18.3-1+cuda12.1
libnccl2=2.18.3-1+cuda12.1
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-12=9.1.1.17-1

RUN /setup.python.sh python3.12 /builder.requirements.txt

COPY setup.sources.cudnn.sh /setup.sources.cudnn.sh
ARG INSTALL_NVIDIA_PACKAGES=false
ARG NVIDIA_PACKAGES_FILE=cuda12.1_cudnn9.1.packages.txt
COPY ${NVIDIA_PACKAGES_FILE} /nvidia.packages.txt

RUN if [ "${INSTALL_NVIDIA_PACKAGES}" = "true" ]; then \
      echo "Installing Nvidia packages"; \
      /setup.sources.cudnn.sh && /setup.packages.sh /nvidia.packages.txt; \
    else \

common:mkl_aarch64_threadpool --@compute_library//:openmp=false
common:mkl_aarch64_threadpool -c opt

# This is an alias for the mkl_aarch64_threadpool build.
common:mkl_aarch64 --config=mkl_aarch64_threadpool

# Default CUDA, CUDNN, NCCL and NVSHMEM versions.
common:cuda_version --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
common:cuda_version --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
common:cuda_version --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5"