use default SDK image and aarch64 wheel

.github/trigger_files/beam_PostCommit_Python_Arm.json

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

.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 ARM PostCommit uses prebuilt Snapshots SDK images, so it no longer
-    # spends hours on multiarch container builds. Keep headroom for the IT suite.
+    # Skip local multiarch SDK container builds. ARM workers use the default
+    # Dataflow image plus an aarch64 wheel from this checkout. Keep headroom
+    # for the IT suite and qemu wheel build.
     timeout-minutes: 180
     strategy:
       fail-fast: false
@@ -94,6 +95,9 @@ jobs:
           credentials_json: ${{ secrets.GCP_SA_KEY }}
       - name: Set up Cloud SDK
         uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db
+      # aarch64 wheels are built with cibuildwheel on this x86 runner.
+      - uses: docker/setup-qemu-action@96fe6ef7f33517b61c61be40b68a1882f3264fb8
+        name: Set up QEMU
       - name: Set PY_VER_CLEAN
         id: set_py_ver_clean
         run: |
@@ -106,11 +110,10 @@ jobs:
           gradle-command: :sdks:python:test-suites:dataflow:py${{steps.set_py_ver_clean.outputs.py_ver_clean}}:postCommitArmIT
           arguments: |
             -PuseWheelDistribution \
+            -PcibwArchs=aarch64 \
             -PpythonVersion=${{ matrix.python_version }} \
             -PusePrebuiltSdkContainer \
         env:
-          # 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/build.gradle

@@ -214,12 +214,13 @@ def platform_identifiers_map = [
 ]
 
 platform_identifiers_map.each { platform, idsuffix ->
-  def archs = 'auto'
+  def archs = project.rootProject.findProperty('cibwArchs') ?: 'auto'
   // note: A fix for arm64 platform in gradle environment. For some reason the
   // task fails with "Invalid archs option {<Architecture.arm64: 'arm64'>}."
   // even though os.arch is 'aarch64'
   // Running cibuildwheel command directly in shell it succeeded, however
-  if (platform == 'linux' && 'aarch64'.equalsIgnoreCase(System.getProperty("os.arch"))) {
+  if (archs == 'auto' && platform == 'linux' &&
+      'aarch64'.equalsIgnoreCase(System.getProperty("os.arch"))) {
     archs = 'aarch64'
   }
   getVersionsAsList('python_versions').each { it ->

sdks/python/scripts/run_integration_test.sh

@@ -248,12 +248,11 @@ if [[ -z $PIPELINE_OPTS ]]; then
   if [[ "$ARCH" == "ARM" ]]; then
     opts+=("--machine_type=t2a-standard-1")
 
-    # 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.
+    # Default to the Dataflow SDK container for this version. Official images
+    # are multi-arch, so ARM workers do not need a custom image.
     if [[ -n "${SDK_CONTAINER_IMAGE:-}" ]]; then
       opts+=("--sdk_container_image=$SDK_CONTAINER_IMAGE")
-    else
+    elif [[ -n "${MULTIARCH_TAG:-}" ]]; then
       IMAGE_NAME="beam_python${PY_VERSION}_sdk"
       opts+=("--sdk_container_image=us.gcr.io/$PROJECT/$USER/$IMAGE_NAME:$MULTIARCH_TAG")
     fi

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

@@ -50,8 +50,12 @@ task initializeForDataflowJob{
     dependsOn ":sdks:python:bdistPy${pythonVersionNumber}linux"
 
     doLast {
+      def cibwArchs = project.rootProject.findProperty('cibwArchs')
+      def wheelInclude = cibwArchs == 'aarch64' ?
+          "**/apache_beam-*cp${pythonVersionNumber}*aarch64*.whl" :
+          "**/apache_beam-*cp${pythonVersionNumber}*manylinux*.whl"
       def collection = project.fileTree(project.project(':sdks:python').buildDir){
-          include "**/apache_beam-*cp${pythonVersionNumber}*manylinux*.whl"
+          include wheelInclude
       }
       // sdkLocation ext is set at execution time
       String packageFilename = collection.singleFile.toString()
@@ -147,8 +151,8 @@ task postCommitArmIT {
   def pyversion = "${project.ext.pythonVersion.replace('.', '')}"
   dependsOn 'initializeForDataflowJob'
   // 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).
+  // SDK image. Dataflow then uses the default image for this SDK version
+  // unless SDK_CONTAINER_IMAGE is set.
   def usePrebuiltSdkContainer = project.rootProject.hasProperty(['usePrebuiltSdkContainer'])
   if (!usePrebuiltSdkContainer) {
     dependsOn ":sdks:python:container:py${pyversion}:docker"