Build CUDA 12.8 docker image for packaging

mlc-ai · Jan 24, 2025 · cdfccf8 · cdfccf8
1 parent 35fd2c9
commit cdfccf8
Show file tree

Hide file tree

Showing 7 changed files with 84 additions and 6 deletions.
diff --git a/.github/workflows/build_docker_package.yaml b/.github/workflows/build_docker_package.yaml
@@ -9,7 +9,13 @@ jobs:
     strategy:
       matrix:
         platform:
+          - cpu
+          - rocm61
+          - rocm62
+          - cu122
+          - cu123
           - cu124
+          # - cu128
 
     runs-on: ubuntu-latest
     steps:

diff --git a/docker/Dockerfile.package-cu124 b/docker/Dockerfile.package-cu124
@@ -40,6 +40,7 @@ RUN conda create -n py39 python=3.9 -y
 RUN conda create -n py310 python=3.10 -y
 RUN conda create -n py311 python=3.11 -y
 RUN conda create -n py312 python=3.12 -y
+RUN conda create -n py312 python=3.13 -y
 COPY install/almalinux_install_python_package.sh /install/almalinux_install_python_package.sh
 RUN bash /install/almalinux_install_python_package.sh
 

diff --git a/docker/Dockerfile.package-cu128 b/docker/Dockerfile.package-cu128
@@ -0,0 +1,54 @@
+# Docker image: mlc-ai/package-vulkan
+
+FROM quay.io/pypa/manylinux_2_28_x86_64
+
+# install core
+COPY install/almalinux_install_core.sh /install/almalinux_install_core.sh
+RUN bash /install/almalinux_install_core.sh
+
+# install cmake
+COPY install/almalinux_install_cmake.sh /install/almalinux_install_cmake.sh
+RUN bash /install/almalinux_install_cmake.sh
+
+# install llvm
+COPY install/almalinux_install_llvm.sh /install/almalinux_install_llvm.sh
+RUN bash /install/almalinux_install_llvm.sh
+
+# install cuda
+COPY install/almalinux_install_cuda_128.sh /install/almalinux_install_cuda.sh
+RUN bash /install/almalinux_install_cuda.sh
+ENV CUDA_HOME=/usr/local/cuda
+ENV PATH=${CUDA_HOME}/bin:${PATH}
+ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${PATH}
+
+# install vulkan sdk
+COPY install/almalinux_install_vulkan_sdk.sh /install/almalinux_install_vulkan_sdk.sh
+RUN bash /install/almalinux_install_vulkan_sdk.sh
+
+# upgrade patchelf due to the bug in patchelf 0.10
+# see details at https://stackoverflow.com/questions/61007071/auditwheel-repair-not-working-as-expected
+COPY install/almalinux_install_patchelf.sh /install/almalinux_install_patchelf.sh
+RUN bash /install/almalinux_install_patchelf.sh
+
+# Install Conda
+COPY install/almalinux_install_conda.sh /install/almalinux_install_conda.sh
+RUN bash /install/almalinux_install_conda.sh
+ENV PATH=/opt/conda/bin:${PATH}
+
+# Install Python
+RUN conda create -n py39 python=3.9 -y
+RUN conda create -n py310 python=3.10 -y
+RUN conda create -n py311 python=3.11 -y
+RUN conda create -n py312 python=3.12 -y
+RUN conda create -n py313 python=3.13 -y
+COPY install/almalinux_install_python_package.sh /install/almalinux_install_python_package.sh
+RUN bash /install/almalinux_install_python_package.sh
+
+# Install auditwheel
+COPY install/almalinux_install_auditwheel.sh /install/almalinux_install_auditwheel.sh
+RUN bash /install/almalinux_install_auditwheel.sh
+ENV AUDITWHEEL_PLAT=manylinux_2_28_x86_64
+
+# Install Rust
+COPY install/almalinux_install_rust.sh /install/almalinux_install_rust.sh
+RUN bash /install/almalinux_install_rust.sh
diff --git a/docker/install/almalinux_install_cuda_128.sh b/docker/install/almalinux_install_cuda_128.sh
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+set -e
+set -o pipefail
+
+dnf install epel-release -y
+dnf update -y
+rpm -q epel-release
+yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
+sed -i '2imodule_hotfixes=1' /etc/yum.repos.d/cuda-rhel8.repo
+dnf install kernel-devel -y
+dnf install cuda-12-8 -y
+
+NCCL_VERSION=$(dnf --showduplicates list libnccl | grep "cuda12.8" | tail -1 | awk '{print $2}')
+dnf install libnccl-$NCCL_VERSION libnccl-devel-$NCCL_VERSION libnccl-static-$NCCL_VERSION -y
+dnf install nvshmem -y
diff --git a/scripts/build_mlc_ai_wheel_manylinux.sh b/scripts/build_mlc_ai_wheel_manylinux.sh
@@ -6,7 +6,7 @@ source /opt/rh/gcc-toolset-11/enable # GCC-11 is the hightest GCC version compat
 function usage() {
 	echo "Usage: $0 [--gpu GPU-VERSION]"
 	echo
-	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4 rocm-6.1 rocm-6.2}"
+	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4 cuda-12.8 rocm-6.1 rocm-6.2}"
 	echo -e "\tSpecify the GPU version (CUDA/ROCm) in the TVM (default: none)."
 }
 
@@ -44,7 +44,7 @@ function audit_mlc_ai_wheel() {
 TVM_PYTHON_DIR="/workspace/tvm/python"
 PYTHON_VERSIONS_CPU=("3.9" "3.10" "3.11" "3.12")
 PYTHON_VERSIONS_GPU=("3.9" "3.10" "3.11" "3.12")
-GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "rocm-6.1" "rocm-6.2")
+GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "cuda-12.8" "rocm-6.1" "rocm-6.2")
 GPU="none"
 
 while [[ $# -gt 0 ]]; do
@@ -71,7 +71,7 @@ done
 if ! in_array "${GPU}" "${GPU_OPTIONS[*]}"; then
 	echo "Invalid GPU option: ${GPU}"
 	echo
-	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "rocm-6.1" "rocm-6.2"}'
+	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "cuda-12.8" "rocm-6.1" "rocm-6.2"}'
 	exit -1
 fi
 

diff --git a/scripts/build_mlc_llm_wheel_manylinux.sh b/scripts/build_mlc_llm_wheel_manylinux.sh
@@ -6,7 +6,7 @@ source /opt/rh/gcc-toolset-11/enable # GCC-11 is the hightest GCC version compat
 function usage() {
 	echo "Usage: $0 [--gpu GPU-VERSION]"
 	echo
-	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4 rocm-6.1 rocm-6.2}"
+	echo -e "--gpu {none cuda-11.7 cuda-11.8 cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4 cuda-12.8 rocm-6.1 rocm-6.2}"
 	echo -e "\tSpecify the GPU version (CUDA/ROCm) in the MLC-LLM (default: none)."
 }
 
@@ -44,7 +44,7 @@ function audit_mlc_llm_wheel() {
 MLC_LLM_PYTHON_DIR="/workspace/mlc-llm/python"
 PYTHON_VERSIONS_CPU=("3.9" "3.10" "3.11" "3.12")
 PYTHON_VERSIONS_GPU=("3.9" "3.10" "3.11" "3.12")
-GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "rocm-6.1" "rocm-6.2")
+GPU_OPTIONS=("none" "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "cuda-12.8" "rocm-6.1" "rocm-6.2")
 GPU="none"
 
 while [[ $# -gt 0 ]]; do
@@ -71,7 +71,7 @@ done
 if ! in_array "${GPU}" "${GPU_OPTIONS[*]}"; then
 	echo "Invalid GPU option: ${GPU}"
 	echo
-	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "rocm-6.1" "rocm-6.2"}'
+	echo 'GPU version can only be {"none", "cuda-11.7" "cuda-11.8" "cuda-12.1" "cuda-12.2" "cuda-12.3" "cuda-12.4" "cuda-12.8" "rocm-6.1" "rocm-6.2"}'
 	exit -1
 fi
 

diff --git a/scripts/sync_package.py b/scripts/sync_package.py
@@ -169,6 +169,7 @@ def main():
             "cuda-12.2",
             "cuda-12.3",
             "cuda-12.4",
+            "cuda-12.8",
             "rocm-6.1",
             "rocm-6.2",
         ],