Made changes to build and test multi-arch Beam SDK container images

.test-infra/jenkins/job_PostCommit_Java_Examples_Dataflow_V2.groovy

@@ -18,6 +18,7 @@
 
 import PostcommitJobBuilder
 import CommonJobProperties as commonJobProperties
+import java.time.LocalDateTime
 
 PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2',
     'Run Java Examples on Dataflow Runner V2', 'Google Cloud Dataflow Runner V2 Examples', this) {
@@ -30,17 +31,23 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2',
         archiveJunit('**/build/test-results/**/*.xml')
       }
 
+      // Generates a unique tag for the container as the current time.
+      def now = LocalDateTime.now()
+      def unique_tag = "${now.date}${now.hour}${now.minute}${now.second}"
       steps {
         gradle {
           rootBuildScriptDir(commonJobProperties.checkoutDir)
           tasks(':runners:google-cloud-dataflow-java:examplesJavaRunnerV2IntegrationTest')
-
           // Increase parallel worker threads above processor limit since most time is
           // spent waiting on Dataflow jobs. ValidatesRunner tests on Dataflow are slow
           // because each one launches a Dataflow job with about 3 mins of overhead.
           // 3 x num_cores strikes a good balance between maxing out parallelism without
           // overloading the machines.
           commonJobProperties.setGradleSwitches(delegate, 3 * Runtime.runtime.availableProcessors())
+          switches "-Pcontainer-architecture-list=arm64,amd64"
+          switches '-Ppush-containers'
+          switches "-Pdocker-repository-root=us.gcr.io/apache-beam-testing/java-postcommit-it"
+          switches "-Pdocker-tag=${unique_tag}"
         }
       }
     }

.test-infra/jenkins/job_PostCommit_Java_Examples_Dataflow_V2_Java11.groovy

@@ -18,6 +18,7 @@
 
 import PostcommitJobBuilder
 import CommonJobProperties as commonJobProperties
+import java.time.LocalDateTime
 
 PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_java11',
     'Run Java 11 Examples on Dataflow Runner V2', 'Google Cloud Dataflow Runner V2 Examples Java 11', this) {
@@ -30,11 +31,13 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_ja
         archiveJunit('**/build/test-results/**/*.xml')
       }
 
+      // Generates a unique tag for the container as the current time.
+      def now = LocalDateTime.now()
+      def unique_tag = "${now.date}${now.hour}${now.minute}${now.second}"
       steps {
         gradle {
           rootBuildScriptDir(commonJobProperties.checkoutDir)
           tasks(':runners:google-cloud-dataflow-java:examplesJavaRunnerV2IntegrationTest')
-
           // Increase parallel worker threads above processor limit since most time is
           // spent waiting on Dataflow jobs. ValidatesRunner tests on Dataflow are slow
           // because each one launches a Dataflow job with about 3 mins of overhead.
@@ -46,6 +49,10 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_ja
           switches '-PcompileAndRunTestsWithJava11'
           switches '-PskipCheckerFramework'
           switches "-Pjava11Home=${commonJobProperties.JAVA_11_HOME}"
+          switches '-Ppush-containers'
+          switches "-Pcontainer-architecture-list=arm64,amd64"
+          switches "-Pdocker-repository-root=us.gcr.io/apache-beam-testing/java-postcommit-it"
+          switches "-Pdocker-tag=${unique_tag}"
         }
       }
     }

.test-infra/jenkins/job_PostCommit_Java_Examples_Dataflow_V2_Java17.groovy

@@ -18,6 +18,7 @@
 
 import PostcommitJobBuilder
 import CommonJobProperties as commonJobProperties
+import java.time.LocalDateTime
 
 PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_java17',
     'Run Java 17 Examples on Dataflow Runner V2', 'Google Cloud Dataflow Runner V2 Examples Java 17', this) {
@@ -30,11 +31,13 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_ja
         archiveJunit('**/build/test-results/**/*.xml')
       }
 
+      // Generates a unique tag for the container as the current time.
+      def now = LocalDateTime.now()
+      def unique_tag = "${now.date}${now.hour}${now.minute}${now.second}"
       steps {
         gradle {
           rootBuildScriptDir(commonJobProperties.checkoutDir)
           tasks(':runners:google-cloud-dataflow-java:examplesJavaRunnerV2IntegrationTest')
-
           // Increase parallel worker threads above processor limit since most time is
           // spent waiting on Dataflow jobs. ValidatesRunner tests on Dataflow are slow
           // because each one launches a Dataflow job with about 3 mins of overhead.
@@ -46,6 +49,10 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Java_Examples_Dataflow_V2_ja
           switches '-PcompileAndRunTestsWithJava17'
           switches '-PskipCheckerFramework'
           switches "-Pjava17Home=${commonJobProperties.JAVA_17_HOME}"
+          switches "-Pcontainer-architecture-list=arm64,amd64"
+          switches '-Ppush-containers'
+          switches "-Pdocker-repository-root=us.gcr.io/apache-beam-testing/java-postcommit-it"
+          switches "-Pdocker-tag=${unique_tag}"
         }
       }
     }

.test-infra/jenkins/job_PostCommit_Python_ValidatesContainer_Dataflow.groovy

@@ -20,6 +20,7 @@ import CommonJobProperties as commonJobProperties
 import PostcommitJobBuilder
 
 import static PythonTestProperties.VALIDATES_CONTAINER_DATAFLOW_PYTHON_VERSIONS
+import java.time.LocalDateTime
 
 // This job runs the suite of Python ValidatesContainer tests against the
 // Dataflow runner.
@@ -34,11 +35,19 @@ PostcommitJobBuilder.postCommitJob('beam_PostCommit_Py_ValCont',
         archiveJunit('**/pytest*.xml')
       }
 
+      // Generates a unique tag for the container as the current time.
+      def now = LocalDateTime.now()
+      def unique_tag = "${now.date}${now.hour}${now.minute}${now.second}"
       // Execute shell command to test Python SDK.
       steps {
         gradle {
           rootBuildScriptDir(commonJobProperties.checkoutDir)
           tasks(':sdks:python:test-suites:dataflow:validatesContainerTests')
+          switches('-Pcontainer-architecture-list=arm64,amd64')
+          switches('-Ppush-containers')
+          // Push multi-arch containers to the repository set in run_validatescontainer.sh
+          switches('-Pdocker-repository-root=us.gcr.io/apache-beam-testing/jenkins')
+          switches("-Pdocker-tag=${unique_tag}")
           commonJobProperties.setGradleSwitches(delegate)
         }
       }

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -490,7 +490,7 @@ class BeamModulePlugin implements Plugin<Project> {
     project.ext.containerArchitectures = {
       if (isRelease(project)) {
         // Ensure we always publish the expected containers.
-        return ["amd64"];
+        return ["amd64", "arm64"];
       } else if (project.rootProject.findProperty("container-architecture-list") != null) {
         def containerArchitectures = project.rootProject.findProperty("container-architecture-list").split(',')
         if (containerArchitectures.size() > 1 && !project.rootProject.hasProperty("push-containers")) {

runners/google-cloud-dataflow-java/build.gradle

@@ -144,10 +144,16 @@ def dataflowLegacyWorkerJar = project.findProperty('dataflowWorkerJar') ?: proje
 def dataflowKmsKey = project.findProperty('dataflowKmsKey') ?: "projects/apache-beam-testing/locations/global/keyRings/beam-it/cryptoKeys/test"
 def firestoreDb = project.findProperty('firestoreDb') ?: 'firestoredb'
 
-def dockerImageRoot = project.findProperty('dockerImageRoot') ?: "us.gcr.io/${dataflowProject}/java-postcommit-it"
-def dockerJavaImageContainer = "${dockerImageRoot}/java"
+def javaVer = "java8"
+  if(project.hasProperty('compileAndRunTestsWithJava17')) {
+    javaVer = "java17"
+  } else if(project.hasProperty('compileAndRunTestsWithJava11')) {
+    javaVer = "java11"
+  }
+def dockerImageRoot = project.findProperty('docker-repository-root') ?: "us.gcr.io/${dataflowProject}/java-postcommit-it"
+def dockerJavaImageContainer = "${dockerImageRoot}/${project.docker_image_default_repo_prefix}java${javaVer}_sdk"
 def dockerPythonImageContainer = "${dockerImageRoot}/python"
-def dockerTag = new Date().format('yyyyMMddHHmmss')
+def dockerTag = project.findProperty('docker-tag') ?: new Date().format('yyyyMMddHHmmss')
 ext.dockerJavaImageName = "${dockerJavaImageContainer}:${dockerTag}"
 ext.dockerPythonImageName = "${dockerPythonImageContainer}:${dockerTag}"
 
@@ -272,33 +278,13 @@ def createRunnerV2ValidatesRunnerTest = { Map args ->
 // task directly ('dependsOn buildAndPushDockerJavaContainer'). This ensures the correct
 // task ordering such that the registry doesn't get cleaned up prior to task completion.
 def buildAndPushDockerJavaContainer = tasks.register("buildAndPushDockerJavaContainer") {
-  def javaVer = "java8"
-  if(project.hasProperty('compileAndRunTestsWithJava17')) {
-    javaVer = "java17"
-  } else if(project.hasProperty('compileAndRunTestsWithJava11')) {
-    javaVer = "java11"
-  }
+  // The multi-arch Java container has been pushed during build.
   dependsOn ":sdks:java:container:${javaVer}:docker"
-  def defaultDockerImageName = containerImageName(
-          name: "${project.docker_image_default_repo_prefix}${javaVer}_sdk",
-          root: "apache",
-          tag: project.sdk_version)
-  doLast {
-    exec {
-      commandLine "docker", "tag", "${defaultDockerImageName}", "${dockerJavaImageName}"
-    }
-    exec {
-      commandLine "gcloud", "docker", "--", "push", "${dockerJavaImageName}"
-    }
-  }
 }
 
 // Clean up built Java images
 def cleanUpDockerJavaImages = tasks.register("cleanUpDockerJavaImages") {
   doLast {
-    exec {
-      commandLine "docker", "rmi", "--force", "${dockerJavaImageName}"
-    }
     exec {
       commandLine "gcloud", "--quiet", "container", "images", "untag", "${dockerJavaImageName}"
     }
@@ -368,7 +354,7 @@ task printRunnerV2PipelineOptions {
     println "To run a Dataflow job with runner V2, add the following pipeline options to your command-line:"
     println runnerV2PipelineOptions.join(' ')
     println "Please delete your image upon completion with the following command:"
-    println "docker rmi ${dockerJavaImageName}; gcloud container images delete --force-delete-tags ${dockerJavaImageName}"
+    println "gcloud container images delete --force-delete-tags ${dockerJavaImageName}"
   }
 }
 

sdks/go/container/build.gradle

@@ -68,5 +68,5 @@ if (project.rootProject.hasProperty(["docker-pull-licenses"])) {
 }
 
 task pushAll {
-  dependsOn ":sdks:go:container:dockerPush"
+  dependsOn ":sdks:go:container:docker"
 }

sdks/go/test/run_validatesrunner_tests.sh

@@ -430,7 +430,7 @@ if [[ "$RUNNER" == "dataflow" ]]; then
 
   if [[ -n "$TEST_EXPANSION_ADDR" || -n "$IO_EXPANSION_ADDR" || -n "$SCHEMAIO_EXPANSION_ADDR" || -n "$DEBEZIUMIO_EXPANSION_ADDR" ]]; then
     # Delete the java cross-language container locally and remotely
-    docker rmi $JAVA_CONTAINER:$JAVA_TAG || echo "Failed to remove container"
+    docker rmi $JAVA_CONTAINER:$JAVA_TAG || echo "Built container image was not removed. Possibly, it was not not saved locally."
     gcloud --quiet container images delete $JAVA_CONTAINER:$JAVA_TAG || echo "Failed to delete container"
   fi
 

sdks/java/container/build.gradle

@@ -80,7 +80,7 @@ artifacts {
 }
 
 task pushAll {
-  dependsOn ":sdks:java:container:java8:dockerPush"
-  dependsOn ":sdks:java:container:java11:dockerPush"
-  dependsOn ":sdks:java:container:java17:dockerPush"
+  dependsOn ":sdks:java:container:java8:docker"
+  dependsOn ":sdks:java:container:java11:docker"
+  dependsOn ":sdks:java:container:java17:docker"
 }

sdks/java/container/common.gradle

@@ -104,6 +104,9 @@ task validateJavaHome {
     }
 }
 
+def useBuildx = project.containerPlatforms() != [project.nativeArchitecture()]
+def pushContainers = project.rootProject.hasProperty(["isRelease"]) || project.rootProject.hasProperty("push-containers")
+
 docker {
     name containerImageName(
             name: "${project.docker_image_default_repo_prefix}java${imageJavaVersion}_sdk",
@@ -121,8 +124,10 @@ docker {
                     project.rootProject.hasProperty(["isRelease"]),
             'java_version': imageJavaVersion,
     ])
-    buildx project.containerPlatforms() != [project.nativeArchitecture()]
+    buildx useBuildx
     platform(*project.containerPlatforms())
+    load useBuildx && !pushContainers
+    push pushContainers
 }
 
 if (project.rootProject.hasProperty(["docker-pull-licenses"]) ||

sdks/python/container/build.gradle

@@ -42,33 +42,8 @@ tasks.register("buildAll") {
   dependsOn ':sdks:python:container:py311:docker'
 }
 
-for(int i=8; i<=11; ++i) {
-  String cur = "3" + i
-  String prev = "3" + (i-1)
-  tasks.register("push" + cur) {
-    if (cur != "38") {
-      // Enforce ordering to allow the prune step to happen between runs.
-      // This will ensure we don't use up too much space (especially in CI environments)
-      mustRunAfter(":sdks:python:container:push" + prev)
-    }
-    dependsOn ':sdks:python:container:py' + cur + ':dockerPush'
-
-    doLast {
-      if (project.hasProperty("prune-images")) {
-        exec {
-          executable("docker")
-          args("system", "prune", "-a", "--force")
-        }
-      }
-    }
-  }
-}
-
 tasks.register("pushAll") {
-  dependsOn ':sdks:python:container:push38'
-  dependsOn ':sdks:python:container:push39'
-  dependsOn ':sdks:python:container:push310'
-  dependsOn ':sdks:python:container:push311'
+  dependsOn ':sdks:python:container:buildAll'
 }
 
 tasks.register("generatePythonRequirementsAll") {