Update docker image for cuda12.4 (#370)

Update docker image for cuda12.4
Image pushed to registry

---------

Co-authored-by: Changho Hwang <[email protected]>

Author: Binyang2014

docker/base-dev-x.dockerfile

@@ -2,7 +2,7 @@ ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="MSCCL++"
-LABEL org.opencontainers.image.source https://github.com/microsoft/mscclpp
+LABEL org.opencontainers.image.source=https://github.com/microsoft/mscclpp
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \

docker/base-x-rocm.dockerfile

@@ -2,7 +2,7 @@ ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="MSCCL++"
-LABEL org.opencontainers.image.source https://github.com/microsoft/mscclpp
+LABEL org.opencontainers.image.source=https://github.com/microsoft/mscclpp
 
 ENV DEBIAN_FRONTEND=noninteractive
 

docker/base-x.dockerfile

@@ -2,7 +2,7 @@ ARG BASE_IMAGE
 FROM ${BASE_IMAGE}
 
 LABEL maintainer="MSCCL++"
-LABEL org.opencontainers.image.source https://github.com/microsoft/mscclpp
+LABEL org.opencontainers.image.source=https://github.com/microsoft/mscclpp
 
 ENV DEBIAN_FRONTEND=noninteractive
 USER root
@@ -17,6 +17,7 @@ RUN apt-get update && \
         git \
         libcap2 \
         libnuma-dev \
+        lsb-release \
         openssh-client \
         openssh-server \
         python3-dev \
@@ -31,11 +32,13 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/* /tmp/*
 
 # Install OFED
-ENV OFED_VERSION=5.2-2.2.3.0
+ARG OFED_VERSION=5.2-2.2.3.0
 RUN cd /tmp && \
-    wget -q https://content.mellanox.com/ofed/MLNX_OFED-${OFED_VERSION}/MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu20.04-x86_64.tgz && \
-    tar xzf MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu20.04-x86_64.tgz && \
-    MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu20.04-x86_64/mlnxofedinstall --user-space-only --without-fw-update --force --all && \
+    OS_VERSION=$(lsb_release -rs) && \
+    OS_VERSION=ubuntu${OS_VERSION} && \
+    wget -q https://content.mellanox.com/ofed/MLNX_OFED-${OFED_VERSION}/MLNX_OFED_LINUX-${OFED_VERSION}-${OS_VERSION}-x86_64.tgz && \
+    tar xzf MLNX_OFED_LINUX-${OFED_VERSION}-${OS_VERSION}-x86_64.tgz && \
+    MLNX_OFED_LINUX-${OFED_VERSION}-${OS_VERSION}-x86_64/mlnxofedinstall --user-space-only --without-fw-update --without-ucx-cuda --force --all && \
     rm -rf /tmp/MLNX_OFED_LINUX-${OFED_VERSION}*
 
 # Install OpenMPI

docker/build.sh

@@ -8,6 +8,7 @@ baseImageTable=(
     ["cuda12.1"]="nvidia/cuda:12.1.1-devel-ubuntu20.04"
     ["cuda12.2"]="nvidia/cuda:12.2.2-devel-ubuntu20.04"
     ["cuda12.3"]="nvidia/cuda:12.3.2-devel-ubuntu20.04"
+    ["cuda12.4"]="nvidia/cuda:12.4.1-devel-ubuntu22.04"
     ["rocm6.2"]="rocm/rocm-terminal:6.2"
 )
 
@@ -20,11 +21,16 @@ extraLdPathTable=(
     ["rocm6.2"]="/opt/rocm/lib"
 )
 
+declare -A ofedVersionTable
+ofedVersionTable=(
+    ["cuda12.4"]="23.07-0.5.1.2"
+)
+
 GHCR="ghcr.io/microsoft/mscclpp/mscclpp"
 TARGET=${1}
 
 print_usage() {
-    echo "Usage: $0 [cuda11.8|cuda12.1|cuda12.2|cuda12.3|rocm6.2]"
+    echo "Usage: $0 [cuda11.8|cuda12.1|cuda12.2|cuda12.3|cuda12.4|rocm6.2]"
 }
 
 if [[ ! -v "baseImageTable[${TARGET}]" ]]; then
@@ -38,11 +44,18 @@ SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
 
 cd ${SCRIPT_DIR}/..
 
+DEFAULT_OFED_VERSION="5.2-2.2.3.0"
+OFED_VERSION=${ofedVersionTable[${TARGET}]}
+if [[ -z ${OFED_VERSION} ]]; then
+    OFED_VERSION=${DEFAULT_OFED_VERSION}
+fi
+
 docker build -t ${GHCR}-common:base-${TARGET} \
     -f docker/base-x.dockerfile \
     --build-arg BASE_IMAGE=${baseImageTable[${TARGET}]} \
     --build-arg EXTRA_LD_PATH=${extraLdPathTable[${TARGET}]} \
-    --build-arg TARGET=${TARGET} .
+    --build-arg TARGET=${TARGET} \
+    --build-arg OFED_VERSION=${OFED_VERSION} .
 
 if [[ ${TARGET} == rocm* ]]; then
     echo "Building ROCm base image..."