FROM quay.io/pypa/manylinux2014_x86_64:2024-09-09-f386546
LABEL maintainer "Joakim Andén"

ENV CUDA_MAJOR 12
ENV CUDA_MINOR 0

ENV CUDART_VERSION 12.0.146
ENV CUFFT_VERSION 11.0.1.95
ENV CURAND_VERSION 10.3.1.124
ENV PROFILER_VERSION 12.0.140
ENV NVTX_VERSION 12.0.140
ENV NVCC_VERSION 12.0.140

ENV CUDA_DASH_VERSION ${CUDA_MAJOR}-${CUDA_MINOR}
ENV CUDA_DOT_VERSION ${CUDA_MAJOR}.${CUDA_MINOR}

# ---- The following block adds layers for CUDA --- #
# base
RUN NVIDIA_GPGKEY_SUM=d0664fbbdb8c32356d45de36c5984617217b2d0bef41b93ccecd326ba3b80c87 && \
    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/D42D0685.pub | sed '/^Version/d' > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
    echo "$NVIDIA_GPGKEY_SUM  /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA" | sha256sum -c --strict -

COPY tools/cufinufft/docker/cuda${CUDA_DOT_VERSION}/cuda.repo /etc/yum.repos.d/cuda.repo

# For libraries in the cuda-compat-* package: https://docs.nvidia.com/cuda/eula/index.html#attachment-a
RUN yum install -y \
        cuda-cudart-${CUDA_DASH_VERSION}-${CUDART_VERSION}-1 \
        cuda-compat-${CUDA_DASH_VERSION} && \
    ln -s cuda-${CUDA_DOT_VERSION} /usr/local/cuda && \
    rm -rf /var/cache/yum/*

# nvidia-docker 1.0
RUN echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf && \
    echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf

ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
ENV LD_LIBRARY_PATH ${LD_LIBRARY_PATH}:/usr/local/nvidia/lib:/usr/local/nvidia/lib64

# nvidia-container-runtime
ENV NVIDIA_VISIBLE_DEVICES all
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
ENV NVIDIA_REQUIRE_CUDA "cuda>=${CUDA_DOT_VERSION} brand=tesla,driver>=418,driver<419 brand=tesla,driver>=440,driver<441"

# runtime
RUN yum install -y \
        libcufft-${CUDA_DASH_VERSION}-${CUFFT_VERSION}-1 \
        libcurand-${CUDA_DASH_VERSION}-${CURAND_VERSION}-1 \
        cuda-nvtx-${CUDA_DASH_VERSION}-${NVTX_VERSION}-1 && \
    rm -rf /var/cache/yum/*

# devel
RUN yum install -y \
        cuda-cudart-devel-${CUDA_DASH_VERSION}-${CUDART_VERSION}-1 \
        libcufft-devel-${CUDA_DASH_VERSION}-${CUFFT_VERSION}-1 \
        libcurand-devel-${CUDA_DASH_VERSION}-${CURAND_VERSION}-1 \
        cuda-profiler-api-${CUDA_DASH_VERSION}-${PROFILER_VERSION}-1 \
        cuda-nvcc-${CUDA_DASH_VERSION}-${NVCC_VERSION}-1 && \
    rm -rf /var/cache/yum/*

ENV LIBRARY_PATH /usr/local/cuda/lib64/stubs

# /CUDA #

# finufft reqs
RUN yum install -y \
        cmake && \
    rm -rf /var/cache/yum/*

# pick py312 as default
RUN ln -s /opt/python/cp312-cp312/bin/python3 /usr/bin/python3

# upgrade pip, install auditwheel
RUN python3 -m pip install --no-cache-dir --upgrade pip && \
    python3 -m pip install --no-cache-dir --upgrade auditwheel

CMD ["/bin/bash"]
