Use prebuilt Snapshots SDK images for PostCommit Python Arm (#39498)

Author: aIbrahiim

.github/trigger_files/beam_PostCommit_Python_Arm.json

@@ -0,0 +1,4 @@
+{
+    "comment": "Modify this file in a trivial way to cause this test suite to run",
+    "modification": 1
+}

.github/workflows/beam_PostCommit_Python_Arm.yml

@@ -54,8 +54,9 @@ jobs:
   beam_PostCommit_Python_Arm:
     name: ${{ matrix.job_name }} ${{ matrix.python_version }}
     runs-on: [self-hosted, ubuntu-24.04, main]
-    # Python 3.14 multiarch container builds take much longer than other versions.
-    timeout-minutes: ${{ matrix.python_version == '3.14' && 600 || 480 }}
+    # Python ARM PostCommit uses prebuilt Snapshots SDK images, so it no longer
+    # spends hours on multiarch container builds. Keep headroom for the IT suite.
+    timeout-minutes: 180
     strategy:
       fail-fast: false
       matrix:
@@ -93,33 +94,23 @@ jobs:
           credentials_json: ${{ secrets.GCP_SA_KEY }}
       - name: Set up Cloud SDK
         uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c
-      - name: GCloud Docker credential helper
-        run: |
-          gcloud auth configure-docker us.gcr.io
       - name: Set PY_VER_CLEAN
         id: set_py_ver_clean
         run: |
           PY_VER=${{ matrix.python_version }}
           PY_VER_CLEAN=${PY_VER//.}
           echo "py_ver_clean=$PY_VER_CLEAN" >> $GITHUB_OUTPUT
-      - name: Generate TAG unique variable based on timestamp
-        id: set_tag
-        run: echo "TAG=$(date +'%Y%m%d-%H%M%S%N')" >> $GITHUB_OUTPUT
       - name: run PostCommit Python ${{ matrix.python_version }} script
         uses: ./.github/actions/gradle-command-self-hosted-action
         with:
           gradle-command: :sdks:python:test-suites:dataflow:py${{steps.set_py_ver_clean.outputs.py_ver_clean}}:postCommitArmIT
           arguments: |
             -PuseWheelDistribution \
             -PpythonVersion=${{ matrix.python_version }} \
-            -Pcontainer-architecture-list=arm64,amd64 \
-            -Pdocker-repository-root=us.gcr.io/apache-beam-testing/github-actions \
-            -Pdocker-tag=${{ steps.set_tag.outputs.TAG }} \
-            -Ppush-containers \
+            -PusePrebuiltSdkContainer \
         env:
-          MULTIARCH_TAG: ${{ steps.set_tag.outputs.TAG }}
+          # Use multiarch images published by Publish Beam SDK Snapshots.
+          SDK_CONTAINER_IMAGE: gcr.io/apache-beam-testing/beam-sdk/beam_python${{ matrix.python_version }}_sdk:latest
           USER: github-actions
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v7

sdks/python/scripts/run_integration_test.sh

@@ -248,8 +248,15 @@ if [[ -z $PIPELINE_OPTS ]]; then
   if [[ "$ARCH" == "ARM" ]]; then
     opts+=("--machine_type=t2a-standard-1")
 
-    IMAGE_NAME="beam_python${PY_VERSION}_sdk"
-    opts+=("--sdk_container_image=us.gcr.io/$PROJECT/$USER/$IMAGE_NAME:$MULTIARCH_TAG")
+    # Prefer an explicit image (e.g. Snapshots latest) when provided so CI can
+    # skip rebuilding multiarch SDK containers. Otherwise use the image built
+    # and pushed by the caller under MULTIARCH_TAG.
+    if [[ -n "${SDK_CONTAINER_IMAGE:-}" ]]; then
+      opts+=("--sdk_container_image=$SDK_CONTAINER_IMAGE")
+    else
+      IMAGE_NAME="beam_python${PY_VERSION}_sdk"
+      opts+=("--sdk_container_image=us.gcr.io/$PROJECT/$USER/$IMAGE_NAME:$MULTIARCH_TAG")
+    fi
   fi
 
   if [[ ! -z "$KMS_KEY_NAME" ]]; then

sdks/python/test-suites/dataflow/common.gradle

@@ -146,7 +146,13 @@ task postCommitIT {
 task postCommitArmIT {
   def pyversion = "${project.ext.pythonVersion.replace('.', '')}"
   dependsOn 'initializeForDataflowJob'
-  dependsOn ":sdks:python:container:py${pyversion}:docker"
+  // When -PusePrebuiltSdkContainer is set, skip building/pushing a multiarch
+  // SDK image and use SDK_CONTAINER_IMAGE from the environment instead (e.g.
+  // gcr.io/apache-beam-testing/beam-sdk/beam_python3.x_sdk:latest from Snapshots).
+  def usePrebuiltSdkContainer = project.rootProject.hasProperty(['usePrebuiltSdkContainer'])
+  if (!usePrebuiltSdkContainer) {
+    dependsOn ":sdks:python:container:py${pyversion}:docker"
+  }
 
   doLast {
     def testOpts = basicPytestOpts + ["--numprocesses=8", "--dist=loadfile"]