Update Build Matrix

Author: jamesnulliu

.github/workflows/cd-docker-build-push.yml

@@ -9,11 +9,8 @@ jobs:
   build-and-push:
     runs-on: ubuntu-latest
     strategy:
-      # If any matrix job fails, cancel all others immediately.
       fail-fast: true
       matrix:
-        # Defines the different configurations to build.
-        # This will generate 16 parallel jobs (2 * 2 * 2 * 2).
         dockerfile: ["Dockerfile.cuda", "Dockerfile.common"]
         install_llvm: [true, false]
         install_rust: [true, false]
@@ -32,27 +29,25 @@ jobs:
           swap-storage: true
 
       - name: Checkout repository
-        uses: actions/checkout@v4 # Using the latest version
+        uses: actions/checkout@v4
 
       - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # Using the latest version
+        uses: docker/setup-buildx-action@v3
 
       - name: Login to Docker Hub
-        uses: docker/login-action@v3 # Using the latest version
+        uses: docker/login-action@v3
         with:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
       - name: Build Docker image
         env:
-          # Set environment variables for the build script from the matrix
           INSTALL_LLVM: ${{ matrix.install_llvm }}
           INSTALL_RUST: ${{ matrix.install_rust }}
           INSTALL_TORCH: ${{ matrix.install_torch }}
-        # Pass the correct Dockerfile name from the matrix to the build script
-        run: bash ./scripts/build.sh "${{ matrix.dockerfile }}"
+        run: |
+          bash ./scripts/build.sh "${{ matrix.dockerfile }}"
+          echo "IMAGE_NAME=$(bash ./scripts/get_image_name.sh)" >> $GITHUB_ENV
 
-      - name: Push Docker image(s)
-        # This assumes push.sh can discover and push the image(s) created
-        # in the previous step without explicit arguments.
-        run: bash ./scripts/push.sh
+      - name: Push Docker image
+        run: bash ./scripts/push.sh "${{ env.IMAGE_NAME }}"

Dockerfile.common

@@ -7,6 +7,9 @@ ARG DEBIAN_FRONTEND=noninteractive
 ENV LANGUAGE=en_US.UTF-8
 ENV LANG=en_US.UTF-8
 
+ARG INSTALL_LLVM
+ARG INSTALL_RUST
+ARG INSTALL_TORCH
 ARG LLVM_VERSION
 ARG IMAGE_VERSION
 ARG TORCH_VERSION
@@ -59,7 +62,7 @@ RUN if [[ "${INSTALL_LLVM}" == "ture" ]]; then \
     fi
 
 # Rust
-RUN if [[ "${INSTALL_RULST}" == "true" ]]; then \
+RUN if [[ "${INSTALL_RUST}" == "true" ]]; then \
         curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \
     fi
 
@@ -75,9 +78,9 @@ RUN wget -O /tmp/miniconda3.sh \
     conda upgrade libstdcxx-ng -c conda-forge -y && \
     pip3 install nvitop --no-cache-dir && \
     if [[ "${INSTALL_TORCH}" == "true" ]]; then \
-        TORCH_CU_VER=$(echo $CUDA_VERSION | cut -d'.' -f1,2 | tr -d '.') && \ 
+        TORCH_CU_VER=$(echo $CUDA_VERSION | cut -d'.' -f1,2 | tr -d '.') && \
         pip3 install torch==${TORCH_VERSION} torchvision torchaudio \
-            --index-url "https://download.pytorch.org/whl/cu${TORCH_CU_VER}" \
+            https://download.pytorch.org/whl/cpu \
             --no-cache-dir \
     fi
 

Dockerfile.cuda

@@ -7,6 +7,9 @@ ARG DEBIAN_FRONTEND=noninteractive
 ENV LANGUAGE=en_US.UTF-8
 ENV LANG=en_US.UTF-8
 
+ARG INSTALL_LLVM
+ARG INSTALL_RUST
+ARG INSTALL_TORCH
 ARG LLVM_VERSION
 ARG IMAGE_VERSION
 ARG TORCH_VERSION
@@ -59,7 +62,7 @@ RUN if [[ "${INSTALL_LLVM}" == "ture" ]]; then \
     fi
 
 # Rust
-RUN if [[ "${INSTALL_RULST}" == "true" ]]; then \
+RUN if [[ "${INSTALL_RUST}" == "true" ]]; then \
         curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \
     fi
 

scripts/.image-configs.sh

@@ -1,7 +1,6 @@
-IMAGE_VERSION=2.1.3
+IMAGE_VERSION=2.2.0
 TORCH_VERSION=2.7.1
 CUDA_VERSION=12.8.0
 UBUNTU_VERSION=24.04
 LLVM_VERSION=21
-IMAGE_TAG=v${IMAGE_VERSION}-torch${TORCH_VERSION}-cuda${CUDA_VERSION}-ubuntu${UBUNTU_VERSION}
-IMAGE_NAME=jamesnulliu/deeplearning:${IMAGE_TAG}
+IMAGE_TAG="v${IMAGE_VERSION}"

scripts/build.sh

@@ -1,10 +1,22 @@
 set -e 
 
+DOCKER_FILE=$1
+
 source ./scripts/.image-configs.sh
 
+if [[ "${DOCKER_FILE}" == *.cuda ]]; then IMAGE_TAG="${IMAGE_TAG}-cuda${CUDA_VERSION}"; fi
+if [[ "${INSTALL_TORCH}" == "true" ]]; then IMAGE_TAG="${IMAGE_TAG}-torch${TORCH_VERSION}"; fi
+if [[ "${INSTALL_LLVM}" == "true" ]]; then IMAGE_TAG="${IMAGE_TAG}-llvm${LLVM_VERSION}"; fi
+if [[ "${INSTALL_RUST}" == "true" ]]; then IMAGE_TAG="${IMAGE_TAG}-rust"; fi
+
+export IMAGE_NAME=jamesnulliu/deeplearning:${IMAGE_TAG}
+
 docker build \
-    -f Dockerfile \
-    --build-arg IMAGE_VERSION=$IMAGE_VERSION  \
+    -f $DOCKER_FILE \
+    --build-arg INSTALL_LLVM=$INSTALL_LLVM \
+    --build-arg INSTALL_RUST=$INSTALL_RUST \
+    --build-arg INSTALL_TORCH=$INSTALL_TORCH \
+    --build-arg IMAGE_VERSION=$IMAGE_VERSION \
     --build-arg TORCH_VERSION=$TORCH_VERSION \
     --build-arg CUDA_VERSION=$CUDA_VERSION \
     --build-arg UBUNTU_VERSION=$UBUNTU_VERSION \