Fix arm tests on Python PostCommit

.github/workflows/beam_PostCommit_Python_Arm.yml

@@ -18,10 +18,10 @@
 name: PostCommit Python Arm
 
 on:
-  # issue_comment:
-  #   types: [created]
-  # schedule:
-  #   - cron: '0 */6 * * *'
+  issue_comment:
+    types: [created]
+  schedule:
+    - cron: '0 */6 * * *'
   workflow_dispatch:
 
 # This allows a subsequently queued workflow run to interrupt previous runs
@@ -81,21 +81,35 @@ jobs:
         run: |
           sudo curl -L https://github.com/docker/compose/releases/download/1.22.0/docker-compose-$(uname -s)-$(uname -m) -o /usr/local/bin/docker-compose
           sudo chmod +x /usr/local/bin/docker-compose
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - 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 \
         env:
           CLOUDSDK_CONFIG: ${{ env.KUBELET_GCLOUD_CONFIG_PATH}}
+          MULTIARCH_TAG: ${{ steps.set_tag.outputs.TAG }}
+          USER: github-actions
       - name: Archive code coverage results
         uses: actions/upload-artifact@v3
         with:

sdks/python/scripts/run_integration_test.sh

@@ -79,6 +79,7 @@ SUITE=""
 COLLECT_MARKERS=
 REQUIREMENTS_FILE=""
 ARCH=""
+PY_VERSION=""
 
 # Default test (pytest) options.
 # Run WordCountIT.test_wordcount_it by default if no test options are
@@ -169,6 +170,11 @@ case $key in
       shift # past argument
       shift # past value
       ;;
+    --py_version)
+      PY_VERSION="$2"
+      shift # past argument
+      shift # past value
+      ;;
     *)    # unknown option
         echo "Unknown option: $1"
         exit 1
@@ -242,6 +248,9 @@ 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")
   fi
 
   if [[ ! -z "$KMS_KEY_NAME" ]]; then

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

@@ -144,7 +144,9 @@ task postCommitIT {
 }
 
 task postCommitArmIT {
+  def pyversion = "${project.ext.pythonVersion.replace('.', '')}"
   dependsOn 'initializeForDataflowJob'
+  dependsOn ":sdks:python:container:py${pyversion}:docker"
 
   doLast {
     def testOpts = basicPytestOpts + ["--numprocesses=8", "--dist=loadfile"]
@@ -153,6 +155,7 @@ task postCommitArmIT {
         "sdk_location": project.ext.sdkLocation,
         "suite": "postCommitIT-df${pythonVersionSuffix}",
         "collect": "it_postcommit",
+        "py_version": project.ext.pythonVersion,
         "arch": "ARM"
     ]
     def cmdArgs = mapToArgString(argMap)