From e50b55cff91c8c679c070fab6003a0825880fd10 Mon Sep 17 00:00:00 2001
From: Jake Awe <jawe@nvidia.com>
Date: Mon, 12 Aug 2024 09:57:46 -0500
Subject: [PATCH] place raft-cpu-image condition in one location

---
 .github/workflows/build-image.yml               | 5 ++++-
 .github/workflows/build-test-publish-images.yml | 5 +++--
 ci/compute-matrix.jq                            | 4 ++++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/build-image.yml b/.github/workflows/build-image.yml
index 8ed8b8c9..6df806d4 100644
--- a/.github/workflows/build-image.yml
+++ b/.github/workflows/build-image.yml
@@ -39,6 +39,9 @@ on:
       RAFT_ANN_BENCH_CPU_TAG:
         required: true
         type: string
+      BUILD_RAFT_ANN_BENCH_CPU_IMAGE:
+        required: true
+        type: boolean
 
 jobs:
   build:
@@ -146,7 +149,7 @@ jobs:
             RAPIDS_VER=${{ inputs.RAPIDS_VER }}
           tags: ${{ inputs.RAFT_ANN_BENCH_DATASETS_TAG }}-${{ matrix.ARCH }}
       - name: Build RAFT ANN Benchmarks CPU image
-        if: inputs.CUDA_VER == '12.5.1' # we don't need to build CPU packages for different CUDA versions.
+        if: inputs.BUILD_RAFT_ANN_BENCH_CPU_IMAGE
         uses: docker/build-push-action@v6
         with:
           context: context
diff --git a/.github/workflows/build-test-publish-images.yml b/.github/workflows/build-test-publish-images.yml
index 32aa2df4..95f5fe36 100644
--- a/.github/workflows/build-test-publish-images.yml
+++ b/.github/workflows/build-test-publish-images.yml
@@ -133,6 +133,7 @@ jobs:
       LINUX_VER: ${{ matrix.LINUX_VER }}
       PYTHON_VER: ${{ matrix.PYTHON_VER }}
       RAPIDS_VER: ${{ needs.compute-matrix.outputs.RAPIDS_VER }}
+      BUILD_RAFT_ANN_BENCH_CPU_IMAGE: ${{ matrix.BUILD_RAFT_ANN_BENCH_CPU_IMAGE }}
       BASE_TAG:
         "rapidsai/${{ needs.compute-matrix.outputs.BASE_IMAGE_REPO }}:\
         ${{ needs.compute-matrix.outputs.BASE_TAG_PREFIX }}\
@@ -187,7 +188,7 @@ jobs:
       - name: Create multiarch manifest
         shell: bash
         env:
-          RAFT_ANN_BENCH_CPU_IMAGE_BUILT: ${{ matrix.CUDA_TAG == '12.5.1' }}
+          RAFT_ANN_BENCH_CPU_IMAGE_BUILT: ${{ matrix.BUILD_RAFT_ANN_BENCH_CPU_IMAGE }}
           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 }}
@@ -243,7 +244,7 @@ jobs:
       - name: Remove temporary images
         shell: bash
         env:
-          RAFT_ANN_BENCH_CPU_IMAGE_BUILT: ${{ matrix.CUDA_TAG == '12.5.1' }}
+          RAFT_ANN_BENCH_CPU_IMAGE_BUILT: ${{ matrix.BUILD_RAFT_ANN_BENCH_CPU_IMAGE }}
           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 }}
diff --git a/ci/compute-matrix.jq b/ci/compute-matrix.jq
index d7a4a9fc..e433bed7 100644
--- a/ci/compute-matrix.jq
+++ b/ci/compute-matrix.jq
@@ -14,6 +14,9 @@ def compute_ubuntu_version($x):
 def compute_cuda_tag($x):
   $x + {CUDA_TAG: $x.CUDA_VER | split(".") | [.[0], .[1]] | join(".") };
 
+def compute_build_raft_ann_bench_cpu_image($x):
+  $x + {BUILD_RAFT_ANN_BENCH_CPU_IMAGE: ($x.CUDA_VER == "12.5.1")}; # we don't need to build CPU packages for different CUDA versions.
+
 # Checks the current entry to see if it matches the given exclude
 def matches($entry; $exclude):
   all($exclude | to_entries | .[]; $entry[.key] == .value);
@@ -37,6 +40,7 @@ def compute_matrix($input):
     lists2dict($matrix_keys; .) |
     compute_ubuntu_version(.) |
     compute_cuda_tag(.) |
+    compute_build_raft_ann_bench_cpu_image(.) |
     filter_excludes(.; $excludes) |
     compute_arch(.)
   ] |