From b4c5882ae8de687eed13b19b6138e99a6c08c1c0 Mon Sep 17 00:00:00 2001
From: Jake Awe <jawe@nvidia.com>
Date: Mon, 5 Aug 2024 05:03:28 -0500
Subject: [PATCH] refactor img delete

---
 .../workflows/build-test-publish-images.yml   | 105 ++++--------------
 ci/create-multiarch-manifest.sh               |   2 +
 ci/delete-temp-images.sh                      |  70 ++++++++++++
 3 files changed, 94 insertions(+), 83 deletions(-)
 create mode 100755 ci/delete-temp-images.sh

diff --git a/.github/workflows/build-test-publish-images.yml b/.github/workflows/build-test-publish-images.yml
index f93281d9..f2db6054 100644
--- a/.github/workflows/build-test-publish-images.yml
+++ b/.github/workflows/build-test-publish-images.yml
@@ -229,91 +229,30 @@ jobs:
         run: ci/create-multiarch-manifest.sh
 
   delete-temp-images:
-    needs: [compute-matrix, build, test, build-multiarch-manifest]
+    if: '!cancelled()'
+    needs: [compute-matrix, build-multiarch-manifest]
     strategy:
       matrix: ${{ fromJSON(needs.compute-matrix.outputs.MATRIX) }}
+      fail-fast: false
     runs-on: ubuntu-latest
     steps:
       - name: Remove temporary images
-        run: |
-          HUB_TOKEN=$(
-            curl -s -H "Content-Type: application/json" \
-              -X POST \
-              -d "{\"username\": \"${{ secrets.GPUCIBOT_DOCKERHUB_USER }}\", \"password\": \"${{ secrets.GPUCIBOT_DOCKERHUB_TOKEN }}\"}" \
-              https://hub.docker.com/v2/users/login/ | jq -r .token \
-          )
-          echo "::add-mask::${HUB_TOKEN}"
-
-          org="rapidsai"
-          base_repo="${{ needs.compute-matrix.outputs.BASE_IMAGE_REPO }}"
-          base_tag_array=(
-            ${{ needs.compute-matrix.outputs.BASE_TAG_PREFIX }}
-            ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
-            ${{ needs.compute-matrix.outputs.ALPHA_TAG }}-
-            cuda${{ matrix.CUDA_TAG }}-
-            py${{ matrix.PYTHON_VER }}
-          )
-          base_tag="$(printf %s "${base_tag_array[@]}" $'\n')" # Converts array to string w/o spaces
-
-
-          notebooks_repo="${{ needs.compute-matrix.outputs.NOTEBOOKS_IMAGE_REPO }}"
-          notebooks_tag_array=(
-            ${{ needs.compute-matrix.outputs.NOTEBOOKS_TAG_PREFIX }}
-            ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
-            ${{ needs.compute-matrix.outputs.ALPHA_TAG }}-
-            cuda${{ matrix.CUDA_TAG }}-
-            py${{ matrix.PYTHON_VER }}
-          )
-          notebooks_tag="$(printf %s "${notebooks_tag_array[@]}" $'\n')" # Converts array to string w/o spaces
-
-          raft_ann_bench_repo="${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_IMAGE_REPO }}"
-          raft_ann_bench_tag_array=(
-            ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_TAG_PREFIX }}
-            ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
-            ${{ needs.compute-matrix.outputs.ALPHA_TAG }}-
-            cuda${{ matrix.CUDA_TAG }}-
-            py${{ matrix.PYTHON_VER }}
-          )
-          raft_ann_bench_tag="$(printf %s "${raft_ann_bench_tag_array[@]}" $'\n')" # Converts array to string w/o spaces
-
-          raft_ann_bench_datasets_repo="${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_DATASETS_IMAGE_REPO }}"
-          raft_ann_bench_datasets_tag_array=(
-            ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_DATASETS_TAG_PREFIX }}
-            ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
-            ${{ needs.compute-matrix.outputs.ALPHA_TAG }}-
-            cuda${{ matrix.CUDA_TAG }}-
-            py${{ matrix.PYTHON_VER }}
-          )
-          raft_ann_bench_datasets_tag="$(printf %s "${raft_ann_bench_datasets_tag_array[@]}" $'\n')" # Converts array to string w/o spaces
-
-          raft_ann_bench_cpu_repo="${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_CPU_IMAGE_REPO }}"
-          raft_ann_bench_cpu_tag_array=(
-            ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_CPU_TAG_PREFIX }}
-            ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
-            ${{ needs.compute-matrix.outputs.ALPHA_TAG }}-
-            py${{ matrix.PYTHON_VER }}
-          )
-          raft_ann_bench_cpu_tag="$(printf %s "${raft_ann_bench_cpu_tag_array[@]}" $'\n')" # Converts array to string w/o spaces
-
-          for arch in $(echo '${{ toJSON(matrix.ARCHES) }}' | jq .[] -r); do
-            curl -i -X DELETE \
-              -H "Accept: application/json" \
-              -H "Authorization: JWT $HUB_TOKEN" \
-              "https://hub.docker.com/v2/repositories/$org/$base_repo/tags/$base_tag-$arch/"
-            curl -i -X DELETE \
-              -H "Accept: application/json" \
-              -H "Authorization: JWT $HUB_TOKEN" \
-              "https://hub.docker.com/v2/repositories/$org/$notebooks_repo/tags/$notebooks_tag-$arch/"
-            curl -i -X DELETE \
-              -H "Accept: application/json" \
-              -H "Authorization: JWT $HUB_TOKEN" \
-              "https://hub.docker.com/v2/repositories/$org/$raft_ann_bench_repo/tags/$raft_ann_bench_tag-$arch/"
-            curl -i -X DELETE \
-              -H "Accept: application/json" \
-              -H "Authorization: JWT $HUB_TOKEN" \
-              "https://hub.docker.com/v2/repositories/$org/$raft_ann_bench_datasets_repo/tags/$raft_ann_bench_datasets_tag-$arch/"
-            curl -i -X DELETE \
-              -H "Accept: application/json" \
-              -H "Authorization: JWT $HUB_TOKEN" \
-              "https://hub.docker.com/v2/repositories/$org/$raft_ann_bench_cpu_repo/tags/$raft_ann_bench_cpu_tag-$arch/"
-          done
+        shell: bash
+        env:
+          BASE_IMAGE_REPO: ${{ needs.compute-matrix.outputs.BASE_IMAGE_REPO }}
+          BASE_TAG_PREFIX: ${{ needs.compute-matrix.outputs.BASE_TAG_PREFIX }}
+          RAPIDS_VER: ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
+          ALPHA_TAG: ${{ needs.compute-matrix.outputs.ALPHA_TAG }}
+          CUDA_TAG: ${{ matrix.CUDA_TAG }}
+          PYTHON_VER: ${{ matrix.PYTHON_VER }}
+          NOTEBOOKS_IMAGE_REPO: ${{ needs.compute-matrix.outputs.NOTEBOOKS_IMAGE_REPO }}
+          NOTEBOOKS_TAG_PREFIX: ${{ needs.compute-matrix.outputs.NOTEBOOKS_TAG_PREFIX }}
+          RAFT_ANN_BENCH_IMAGE_REPO: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_IMAGE_REPO }}
+          RAFT_ANN_BENCH_TAG_PREFIX: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_TAG_PREFIX }}
+          RAFT_ANN_BENCH_DATASETS_IMAGE_REPO: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_DATASETS_IMAGE_REPO }}
+          RAFT_ANN_BENCH_DATASETS_TAG_PREFIX: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_DATASETS_TAG_PREFIX }}
+          RAFT_ANN_BENCH_CPU_IMAGE_REPO: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_CPU_IMAGE_REPO }}
+          RAFT_ANN_BENCH_CPU_TAG_PREFIX: ${{ needs.compute-matrix.outputs.RAFT_ANN_BENCH_CPU_TAG_PREFIX }}
+          GPUCIBOT_DOCKERHUB_USER: ${{ secrets.GPUCIBOT_DOCKERHUB_USER }}
+          GPUCIBOT_DOCKERHUB_TOKEN: ${{ secrets.GPUCIBOT_DOCKERHUB_TOKEN }}
+        run: ci/delete-temp-images.sh
diff --git a/ci/create-multiarch-manifest.sh b/ci/create-multiarch-manifest.sh
index 94992d16..7eede009 100755
--- a/ci/create-multiarch-manifest.sh
+++ b/ci/create-multiarch-manifest.sh
@@ -1,3 +1,5 @@
+set -x
+
 # Authenticate and retrieve DockerHub token
 HUB_TOKEN=$(
 curl -s -H "Content-Type: application/json" \
diff --git a/ci/delete-temp-images.sh b/ci/delete-temp-images.sh
new file mode 100755
index 00000000..1029090d
--- /dev/null
+++ b/ci/delete-temp-images.sh
@@ -0,0 +1,70 @@
+set -x
+
+# Authenticate and retrieve DockerHub token
+HUB_TOKEN=$(
+curl -s -H "Content-Type: application/json" \
+    -X POST \
+    -d "{\"username\": \"$GPUCIBOT_DOCKERHUB_USER\", \"password\": \"$GPUCIBOT_DOCKERHUB_TOKEN\"}" \
+    https://hub.docker.com/v2/users/login/ | jq -r .token \
+)
+echo "::add-mask::${HUB_TOKEN}"
+
+org="rapidsai"
+
+# Define tag arrays for different images
+base_tag="${BASE_TAG_PREFIX}${RAPIDS_VER}${ALPHA_TAG}-cuda${CUDA_TAG}-py${PYTHON_VER}"
+notebooks_tag="${NOTEBOOKS_TAG_PREFIX}${RAPIDS_VER}${ALPHA_TAG}-cuda${CUDA_TAG}-py${PYTHON_VER}"
+raft_ann_bench_tag="${RAFT_ANN_BENCH_TAG_PREFIX}${RAPIDS_VER}${ALPHA_TAG}-cuda${CUDA_TAG}-py${PYTHON_VER}"
+raft_ann_bench_datasets_tag="${RAFT_ANN_BENCH_DATASETS_TAG_PREFIX}${RAPIDS_VER}${ALPHA_TAG}-cuda${CUDA_TAG}-py${PYTHON_VER}"
+raft_ann_bench_cpu_tag="${RAFT_ANN_BENCH_CPU_TAG_PREFIX}${RAPIDS_VER}${ALPHA_TAG}-py${PYTHON_VER}"
+
+check_tag_exists() {
+    local repo="$1"
+    local tag="$2"
+    local exists
+    exists=$(curl -s -o /dev/null -w "%{http_code}" -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/${org}/${repo}/tags/${tag}/")
+
+    if [ "$exists" -ne 200 ]; then
+        echo "Error: Required image tag ${repo}:${tag} does not exist. This implies that the image was not built successfully in the build job."
+        exit 1
+    fi
+}
+
+for arch in $(echo '${{ toJSON(matrix.ARCHES) }}' | jq .[] -r); do
+    full_base_tag="${base_tag}-${arch}"
+    full_notebooks_tag="${notebooks_tag}-${arch}"
+    full_raft_ann_bench_tag="${raft_ann_bench_tag}-${arch}"
+    full_raft_ann_bench_datasets_tag="${raft_ann_bench_datasets_tag}-${arch}"
+    full_raft_ann_bench_cpu_tag="${raft_ann_bench_cpu_tag}-${arch}"
+    
+    check_tag_exists "$BASE_IMAGE_REPO" "$full_base_tag"
+    curl -i -X DELETE \
+        -H "Accept: application/json" \
+        -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/$org/$BASE_IMAGE_REPO/tags/$base_tag-$arch/"
+
+    check_tag_exists "$NOTEBOOKS_IMAGE_REPO" "$full_notebooks_tag"
+    curl -i -X DELETE \
+        -H "Accept: application/json" \
+        -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/$org/$NOTEBOOKS_IMAGE_REPO/tags/$notebooks_tag-$arch/"
+
+    check_tag_exists "$RAFT_ANN_BENCH_IMAGE_REPO" "$full_raft_ann_bench_tag"
+    curl -i -X DELETE \
+        -H "Accept: application/json" \
+        -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/$org/$RAFT_ANN_BENCH_IMAGE_REPO/tags/$raft_ann_bench_tag-$arch/"
+
+    check_tag_exists "$RAFT_ANN_BENCH_DATASETS_IMAGE_REPO" "$full_raft_ann_bench_datasets_tag"
+    curl -i -X DELETE \
+        -H "Accept: application/json" \
+        -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/$org/$RAFT_ANN_BENCH_DATASETS_IMAGE_REPO/tags/$raft_ann_bench_datasets_tag-$arch/"
+
+    check_tag_exists "$RAFT_ANN_BENCH_CPU_IMAGE_REPO" "$full_raft_ann_bench_cpu_tag"
+    curl -i -X DELETE \
+        -H "Accept: application/json" \
+        -H "Authorization: JWT $HUB_TOKEN" \
+        "https://hub.docker.com/v2/repositories/$org/$RAFT_ANN_BENCH_CPU_IMAGE_REPO/tags/$raft_ann_bench_cpu_tag-$arch/"
+done