Maybe better syntax for calculate docker image with custom tag prefix (#6586)

Allow calculate docker image to work with custom tag prefix by just
putting the tag in the docker image name.
Old:
need custom-tag-prefix input
New:
can use custom-tag-prefix input, but can also do imagename:tagprefix


This makes it easier for some stuff in pytorch/pytorch, where we rely on
the convention that an image with the registry at the beginning is the
complete image name, and otherwise the foldersha should be added

Testing on
https://hud.pytorch.org/pytorch/pytorch/commit/e6a9e25ba6ef97810652d1c4af395e71cf08ad9d

Author: clee2000

.github/actions/calculate-docker-image/action.yml

@@ -20,7 +20,7 @@ inputs:
     default: ./build.sh
   custom-tag-prefix:
     description: |
-      The prefix to use for the docker image tag.
+      The prefix to use for the docker image tag.  You can also just do imagename:customtag
     default: ""
   working-directory:
     description: The working directory where the repo is checked out.
@@ -81,9 +81,14 @@ runs:
           echo "docker-tag=${DOCKER_TAG}" >> "${GITHUB_OUTPUT}"
           echo "docker-image=${DOCKER_IMAGE_NAME}" >> "${GITHUB_OUTPUT}"
         else
+          if [[ "${DOCKER_IMAGE_NAME}" == *:* ]]; then
+            CUSTOM_TAG_PREFIX=${DOCKER_IMAGE_NAME#*:}
+            DOCKER_IMAGE_NAME=${DOCKER_IMAGE_NAME%%:*}
+          fi
           DOCKER_TAG=${CUSTOM_TAG_PREFIX:+${CUSTOM_TAG_PREFIX}-}$(git rev-parse HEAD:"${DOCKER_BUILD_DIR}")
           echo "docker-tag=${DOCKER_TAG}" >> "${GITHUB_OUTPUT}"
           echo "docker-image=${DOCKER_REGISTRY}/${REPO_NAME}/${DOCKER_IMAGE_NAME}:${DOCKER_TAG}" >> "${GITHUB_OUTPUT}"
+          echo "custom-tag-prefix=${CUSTOM_TAG_PREFIX}" >> "${GITHUB_OUTPUT}"
         fi
 
     - name: Check if image should be built
@@ -196,7 +201,7 @@ runs:
       env:
         REPO_NAME: ${{ github.event.repository.name }}
         DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
-        CUSTOM_TAG_PREFIX: ${{ inputs.custom-tag-prefix }}
+        CUSTOM_TAG_PREFIX: ${{ steps.calculate-image.outputs.custom-tag-prefix }}
         DOCKER_REGISTRY: ${{ inputs.docker-registry }}
         WORKING_DIRECTORY: ${{ inputs.working-directory }}/${{ inputs.docker-build-dir }}
         DOCKER_BUILD_SCRIPT: ${{ inputs.docker-build-script }}