From 0cb985878b6078b2bd3b96d88e0cb0a61e9d651a Mon Sep 17 00:00:00 2001
From: Huy Do <huydo@meta.com>
Date: Fri, 25 Sep 2026 14:46:31 -0700
Subject: [PATCH] Build the docker image on an OSDC runner

The OSDC runner is a container and has no Docker daemon, so the build goes to
the shared BuildKit pool instead. build.sh gains a REMOTE_BUILDKIT switch for
it; without the variable set it behaves exactly as before.

That also drops the steps that only made sense on an EC2 host: Clean
workspace, setup-ssh, setup-linux and teardown-linux.
---
 .ci/docker/build.sh                | 11 +++++-
 .github/workflows/docker-build.yml | 58 ++++++++++++++++++------------
 2 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh
index f40c45fea3d..774f54541b1 100755
--- a/.ci/docker/build.sh
+++ b/.ci/docker/build.sh
@@ -16,7 +16,16 @@ export CUDA_VERSION="12.6.3"
 export BASE_IMAGE="nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}"
 echo "Building ${IMAGE_NAME} Docker image"
 
-docker build \
+# On OSDC the image is built by the out-of-cluster BuildKit pool through a
+# remote buildx builder. That builder has nowhere to load an image into, so the
+# result has to go straight to the registry.
+if [[ -n "${REMOTE_BUILDKIT:-}" ]]; then
+  BUILD_CMD=(docker buildx build --push)
+else
+  BUILD_CMD=(docker build)
+fi
+
+"${BUILD_CMD[@]}" \
   --no-cache \
   --progress=plain \
   -f Dockerfile \
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
index bb285233772..ead908e9097 100644
--- a/.github/workflows/docker-build.yml
+++ b/.github/workflows/docker-build.yml
@@ -15,8 +15,15 @@ on:
 
 jobs:
   docker-build:
-    runs-on: [self-hosted, linux.2xlarge]
+    # BuildKit runs out of cluster on the OSDC BuildKit pool, so this runner only
+    # streams the (small) .ci/docker context and waits.
+    runs-on: mt-l-x86iavx512-8-64
+    container:
+      image: ghcr.io/actions/actions-runner:latest
     timeout-minutes: 240
+    permissions:
+      id-token: write
+      contents: read
     strategy:
       fail-fast: false
       matrix:
@@ -25,34 +32,39 @@ jobs:
     env:
       DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/tutorials/${{ matrix.docker-image-name }}
     steps:
-      - name: Clean workspace
-        shell: bash
-        run: |
-          echo "${GITHUB_WORKSPACE}"
-          sudo rm -rf "${GITHUB_WORKSPACE}"
-          mkdir "${GITHUB_WORKSPACE}"
-
-      - name: Setup SSH (Click me for login details)
-        uses: pytorch/test-infra/.github/actions/setup-ssh@main
-        with:
-          github-secret: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Checkout
         uses: actions/checkout@v6
 
-      - name: Setup Linux
-        uses: pytorch/test-infra/.github/actions/setup-linux@main
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::308535385114:role/arc
+          aws-region: us-east-1
+          role-duration-seconds: 18000
 
-      - name: Build docker image
-        id: build-docker-image
-        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
+      # buildx forwards the client's registry auth to the remote builder, so the
+      # push at the end of build.sh authenticates with what this step writes.
+      - name: Login to ECR
+        uses: aws-actions/amazon-ecr-login@v2
         with:
-          docker-image-name: ${{ matrix.docker-image-name }}
-          push: true
+          registries: "308535385114"
 
-      - name: Teardown Linux
-        uses: pytorch/test-infra/.github/actions/teardown-linux@main
-        if: always()
+      - name: Compute the image tag
+        id: tag
+        shell: bash
+        run: |
+          set -eux
+          git config --global --add safe.directory "${GITHUB_WORKSPACE}"
+          DOCKER_TAG=$(git rev-parse HEAD:.ci/docker)
+          echo "docker-image=${DOCKER_IMAGE}:${DOCKER_TAG}" >> "${GITHUB_OUTPUT}"
+
+      - name: Build and push to ECR
+        uses: pytorch/test-infra/.github/actions/docker-build-remote-buildkit@main
+        with:
+          command: |
+            cd .ci/docker
+            REMOTE_BUILDKIT=1 ./build.sh ${{ matrix.docker-image-name }} \
+              -t ${{ steps.tag.outputs.docker-image }}
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}