Retry docker build step to handle network error (#5759)

I'm frequently seeing network errors when building Docker image which
were subsequently fixed by retrying. So, it's a good idea to support
retry in this step. For example,

* Failure
https://github.com/pytorch/pytorch/actions/runs/11329053976/job/31503695304#step:6:5836
* Retry
https://github.com/pytorch/pytorch/actions/runs/11329053976/job/31509384226

### Testing

Using PyTorch PR pytorch/pytorch#137896 to test
the change at
https://github.com/pytorch/pytorch/actions/runs/11331057887

Author: huydhn

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

@@ -137,15 +137,11 @@ runs:
 
         echo "rebuild=true" >> "${GITHUB_OUTPUT}"
 
-    - name: Build and push docker image
+    - name: Login to ECR
+      if: ${{ steps.calculate-image.outputs.skip != 'true' && (inputs.always-rebuild || steps.check-image.outputs.rebuild) }}
       shell: bash
       working-directory: ${{ inputs.working-directory }}/${{ inputs.docker-build-dir }}
-      if: ${{ steps.calculate-image.outputs.skip != 'true' && (inputs.always-rebuild || steps.check-image.outputs.rebuild) }}
       env:
-        REPO_NAME: ${{ github.event.repository.name }}
-        DOCKER_PUSH: ${{ inputs.push }}
-        DOCKER_FORCE_PUSH: ${{ inputs.force-push }}
-        DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
         DOCKER_REGISTRY: ${{ inputs.docker-registry }}
       run: |
         set -x
@@ -161,11 +157,42 @@ runs:
 
         retry login "${DOCKER_REGISTRY}"
 
-        set -e
-
-        IMAGE_NAME=$(echo ${DOCKER_IMAGE#"${DOCKER_REGISTRY}/${REPO_NAME}/"} | awk -F '[:,]' '{print $1}')
-        # Build new image
-        ./build.sh "${IMAGE_NAME}" -t "${DOCKER_IMAGE}"
+    - name: Build docker image
+      if: ${{ steps.calculate-image.outputs.skip != 'true' && (inputs.always-rebuild || steps.check-image.outputs.rebuild) }}
+      env:
+        REPO_NAME: ${{ github.event.repository.name }}
+        DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
+        DOCKER_REGISTRY: ${{ inputs.docker-registry }}
+        WORKING_DIRECTORY: ${{ inputs.working-directory }}/${{ inputs.docker-build-dir }}
+      # NB: Retry here as this step frequently fails with network error downloading various stuffs
+      uses: nick-fields/[email protected]
+      with:
+        shell: bash
+        timeout_minutes: 90
+        max_attempts: 3
+        retry_wait_seconds: 90
+        command: |
+          set -ex
+
+          # NB: Setting working directory on the step doesn't work with nick-fields/retry https://github.com/nick-fields/retry/issues/89
+          pushd "${WORKING_DIRECTORY}"
+
+          IMAGE_NAME=$(echo ${DOCKER_IMAGE#"${DOCKER_REGISTRY}/${REPO_NAME}/"} | awk -F '[:,]' '{print $1}')
+          # Build new image
+          ./build.sh "${IMAGE_NAME}" -t "${DOCKER_IMAGE}"
+
+          popd
+
+    - name: Push to ECR
+      if: ${{ steps.calculate-image.outputs.skip != 'true' && (inputs.always-rebuild || steps.check-image.outputs.rebuild) }}
+      shell: bash
+      working-directory: ${{ inputs.working-directory }}/${{ inputs.docker-build-dir }}
+      env:
+        DOCKER_PUSH: ${{ inputs.push }}
+        DOCKER_FORCE_PUSH: ${{ inputs.force-push }}
+        DOCKER_IMAGE: ${{ steps.calculate-image.outputs.docker-image }}
+      run: |
+        set -ex
 
         if [ "${DOCKER_PUSH:-false}" == "true" ]; then
           # Only push if docker image doesn't exist already