Reland "Forces docker login on calculate-docker-image and pull-docker…

…-image" (#4999)" (#5001)
pytorch · Mar 12, 2024 · 88112de · 88112de
1 parent f79f642
commit 88112de
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 12 deletions.
diff --git a/.github/actions/calculate-docker-image/action.yml b/.github/actions/calculate-docker-image/action.yml
@@ -86,10 +86,21 @@ runs:
         BASE_REVISION: ${{ github.event.pull_request.base.sha || github.sha }}
         DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
         DOCKER_TAG: ${{ steps.calculate-image.outputs.docker-tag }}
+        DOCKER_REGISTRY: ${{ inputs.docker-registry }}
       run: |
         set +e
         set -x
 
+        login() {
+          aws ecr get-login-password --region us-east-1 | docker login -u AWS --password-stdin "$1"
+        }
+
+        retry () {
+          $*  || (sleep 1 && $*) || (sleep 2 && $*)
+        }
+
+        retry login "${DOCKER_REGISTRY}"
+
         # Check if image already exists, if it does then skip building it
         if docker manifest inspect "${DOCKER_IMAGE}"; then
           exit 0
@@ -137,25 +148,20 @@ runs:
         DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
         DOCKER_REGISTRY: ${{ inputs.docker-registry }}
       run: |
-        set -ex
+        set -x
+        set +e
 
         login() {
-          aws ecr get-authorization-token --region us-east-1 --output text --query 'authorizationData[].authorizationToken' |
-            base64 -d |
-            cut -d: -f2 |
-            docker login -u AWS --password-stdin "$1"
+          aws ecr get-login-password --region us-east-1 | docker login -u AWS --password-stdin "$1"
         }
 
         retry () {
           $*  || (sleep 1 && $*) || (sleep 2 && $*)
         }
 
-        # Only run these steps if not on github actions
-        if [[ -z "${GITHUB_ACTIONS}" ]]; then
-          retry login "${DOCKER_REGISTRY}"
-          # Logout on exit
-          trap "docker logout ${DOCKER_REGISTRY}" EXIT
-        fi
+        retry login "${DOCKER_REGISTRY}"
+
+        set -e
 
         IMAGE_NAME=$(echo ${DOCKER_IMAGE#"${DOCKER_REGISTRY}/${REPO_NAME}/"} | awk -F '[:,]' '{print $1}')
         # Build new image

diff --git a/.github/actions/pull-docker-image/action.yml b/.github/actions/pull-docker-image/action.yml
@@ -6,6 +6,9 @@ inputs:
   docker-image:
     description: the image to pull
     required: true
+  docker-registry:
+    description: The registry to store the image after it is built.
+    default: 308535385114.dkr.ecr.us-east-1.amazonaws.com
 
 runs:
   using: composite
@@ -14,8 +17,22 @@ runs:
       shell: bash
       env:
         DOCKER_IMAGE: ${{ inputs.docker-image }}
+        DOCKER_REGISTRY: ${{ inputs.docker-registry }}
       run: |
-        retry () { "$@"  || (sleep 1 && "$@") || (sleep 2 && "$@") }
+        set -x
+        set +e
+
+        login() {
+          aws ecr get-login-password --region us-east-1 | docker login -u AWS --password-stdin "$1"
+        }
+
+        retry () {
+          $*  || (sleep 1 && $*) || (sleep 2 && $*)
+        }
+
+        retry login "${DOCKER_REGISTRY}"
+
+        set -e
         # ignore output since only exit code is used for conditional
         # only pull docker image if it's not available locally
         if ! docker inspect --type=image "${DOCKER_IMAGE}" >/dev/null 2>/dev/null; then