From d5a0fd6a9c00122e24c77afa49ac994d2efa94b9 Mon Sep 17 00:00:00 2001
From: ayodeawe <jawe@nvidia.com>
Date: Tue, 19 Sep 2023 11:15:47 -0500
Subject: [PATCH] generate additional tag

---
 .github/workflows/build-and-publish-images.yml |  1 +
 .github/workflows/build-image.yml              | 12 +++++++++++-
 ci/generate-tags.sh                            | 18 ++++++++++++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100755 ci/generate-tags.sh

diff --git a/.github/workflows/build-and-publish-images.yml b/.github/workflows/build-and-publish-images.yml
index 854a9b9..82ec399 100644
--- a/.github/workflows/build-and-publish-images.yml
+++ b/.github/workflows/build-and-publish-images.yml
@@ -35,6 +35,7 @@ jobs:
     uses: ./.github/workflows/build-image.yml
     with:
       ARCHES: ${{ toJSON(matrix.ARCHES) }}
+      BUILD_TYPE: ${{ inputs.build_type }}
       CUDA_VER: ${{ matrix.CUDA_VER }}
       LINUX_VER: ${{ matrix.LINUX_VER }}
       PYTHON_VER: ${{ matrix.PYTHON_VER }}
diff --git a/.github/workflows/build-image.yml b/.github/workflows/build-image.yml
index 9825626..0ef7e65 100644
--- a/.github/workflows/build-image.yml
+++ b/.github/workflows/build-image.yml
@@ -21,6 +21,9 @@ on:
       IMAGE_NAME:
         required: true
         type: string
+      build_type:
+        required: true
+        type: string
 
 jobs:
   run:
@@ -57,6 +60,12 @@ jobs:
           LINUX_VER: ${{ inputs.LINUX_VER }}
           PYTHON_VER: ${{ inputs.PYTHON_VER }}
           ARCH: ${{ matrix.ARCH }}
+      - name: Generate Additional Tag
+        id: generate-additional-tag
+        run: ci/generate-tags.sh
+        env:
+          CURRENT_TAG: ${{ inputs.IMAGE_NAME }}-${{ matrix.ARCH }}
+          build_type: ${{ inputs.build_type}}
       - name: Build image
         uses: docker/build-push-action@v4
         with:
@@ -66,4 +75,5 @@ jobs:
           pull: true
           build-args: |
             ${{ steps.generate-build-args.outputs.ARGS }}
-          tags: ${{ inputs.IMAGE_NAME }}-${{ matrix.ARCH }}
+          tags: |
+            ${{ steps.generate-additional-tag.outputs.TAGS }}
diff --git a/ci/generate-tags.sh b/ci/generate-tags.sh
new file mode 100755
index 0000000..bdbd51a
--- /dev/null
+++ b/ci/generate-tags.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+prefix=$(echo "$CURRENT_TAG" | awk -F':' '{print $1}')
+suffix=$(echo "$CURRENT_TAG" | awk -F':' '{print $2}')
+
+if [ "$build_type" == "branch" ]; then
+  cat <<EOF > "${GITHUB_OUTPUT:-/dev/stdout}"
+TAGS<<EOT
+$(printf "%s\n" "${prefix}:${suffix}" "${prefix}-conda:${suffix}")
+EOT
+EOF
+else
+  cat <<EOF > "${GITHUB_OUTPUT:-/dev/stdout}"
+TAGS<<EOT
+$(printf "%s\n" "$CURRENT_TAG")
+EOT
+EOF
+fi