Enable specifying whether or not to pull the docker image from ECR (#6836)

We originally tried to simplify the number of parameters users had to
pass in to linux_job_v2 by using the existence of the docker build
script as a heuristic.

That worked for a while, but that heuristic is now starting to break. If
a repo has an unrelated docker build script located in
`.ci/docker/build.sh`, we assume it wants to pull from our custom ECR
path instead of docker hub.

Our ideal end state would now look like:

1. Have an explicit parameter that needs to be set if a repo wants to
use our ECR registries instead of assuming the desire based on the
existence of a specific file. We default to docker hub
2. Migration to this end state would include fix all existing domain
repos that have that docker build file defined to pass in that
parameter.

That migration requires many repos to be edited with a fix however, so
for now we're having the following non-breaking change:

1. Have an explicit parameter on linux_job_v2 that needs to be set if a
repo wants to use our ECR registries. *Set it to TRUE by default*, with
the calculate-docker-image logic being "If param is set to True AND
special docker file exists" then we use the ECR registry. If either is
false, use docker hub

2. No migrations needed in the short term. We can later migrate domain
repos to explicitly set this setting and then change the default to
FALSE, but that's P2

Testing: Ran [this
job](https://github.com/pytorch/pytorch/actions/runs/15884049575/job/44791525110)
against it, and it only failed due to an error in the script inputted to
the workflow, which is ignorable

Author: ZainRizvi

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

@@ -3,6 +3,10 @@ name: Calculate docker image
 description: Determine docker image to pull, building a new one if necessary.
 
 inputs:
+  use-custom-docker-registry:
+    description: "Use the custom ECR registry. Applies only if build.sh exists in docker-build-dir."
+    default: true
+    type: boolean
   docker-image-name:
     description: |
       The name of a docker image, like `pytorch-linux-focal-linter`. A fullname
@@ -57,21 +61,22 @@ runs:
         DOCKER_BUILD_DIR: ${{ inputs.docker-build-dir }}
         DOCKER_BUILD_SCRIPT: ${{ inputs.docker-build-script }}
         DOCKER_REGISTRY: ${{ inputs.docker-registry }}
+        USE_CUSTOM_DOCKER_REGISTRY: ${{ inputs.use-custom-docker-registry }}
         CUSTOM_TAG_PREFIX: ${{ inputs.custom-tag-prefix }}
       run: |
         set -ex
 
         # If the docker build directory or the build script doesn't exist, the action will
         # gracefully return the docker image name as it is.  Pulling docker image in Linux
         # job could then download the pre-built image as usual
-        if [[ ! -d "${DOCKER_BUILD_DIR}" ]] || [[ ! -f "${DOCKER_BUILD_DIR}/${DOCKER_BUILD_SCRIPT}" ]]; then
+        if [[ -d "${DOCKER_BUILD_DIR}" ]] && [[ -f "${DOCKER_BUILD_DIR}/${DOCKER_BUILD_SCRIPT}" ]] && [[ "${USE_CUSTOM_DOCKER_REGISTRY}" == "true" ]]; then
+          echo "skip=false" >> "${GITHUB_OUTPUT}"
+        else
           echo "skip=true" >> "${GITHUB_OUTPUT}"
           echo "docker-image=${DOCKER_IMAGE_NAME}" >> "${GITHUB_OUTPUT}"
 
-          echo "There is no Docker build script in ${REPO_NAME} repo, skipping..."
+          echo "Not using custom ECR registry.  Either it was not requested or there is no Docker build script in the ${REPO_NAME} repo..."
           exit 0
-        else
-          echo "skip=false" >> "${GITHUB_OUTPUT}"
         fi
 
         if [[ "${DOCKER_IMAGE_NAME}" == *"${DOCKER_REGISTRY}/${REPO_NAME}"* ]]; then

.github/workflows/linux_job_v2.yml

@@ -58,6 +58,10 @@ on:
         description: "Test infra reference to use"
         default: ""
         type: string
+      use-custom-docker-registry:
+        description: "Use the custom ECR registry. Applies only if build.sh exists in docker-build-dir."
+        default: true
+        type: boolean
       docker-image:
         description: Identifies the Docker image by name.
         default: "pytorch/almalinux-builder"
@@ -189,6 +193,7 @@ jobs:
         id: calculate-docker-image
         uses: ./test-infra/.github/actions/calculate-docker-image
         with:
+          use-custom-docker-registry: ${{ inputs.use-custom-docker-registry }}
           docker-image-name: ${{ env.DOCKER_IMAGE }}
           docker-build-dir: ${{ inputs.docker-build-dir }}
           # This needs to be where the repository is checked out