Revert back to using plugin

sdks/python/container/common.gradle

@@ -71,10 +71,16 @@ def copyLauncherDependencies = tasks.register("copyLauncherDependencies", Copy)
 }
 
 def pushContainers = project.rootProject.hasProperty(["isRelease"]) || project.rootProject.hasProperty("push-containers")
+def baseBuildTarget = 'base'
+def buildTarget = project.findProperty('container-build-target') ?: 'base'
+var imageName = project.docker_image_default_repo_prefix + "python${project.ext.pythonVersion}_sdk"
+if (buildTarget != baseBuildTarget) {
+  imageName += "_${buildTarget}"
+}
 
 docker {
   name containerImageName(
-          name: project.docker_image_default_repo_prefix + "python${project.ext.pythonVersion}_sdk",
+          name: imageName,
           root: project.rootProject.hasProperty(["docker-repository-root"]) ?
                   project.rootProject["docker-repository-root"] :
                   project.docker_image_default_repo_root,
@@ -90,7 +96,7 @@ docker {
   platform(*project.containerPlatforms())
   load project.useBuildx() && !pushContainers
   push pushContainers
-  target 'base'
+  target buildTarget
 }
 
 dockerPrepare.dependsOn copyLauncherDependencies

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

@@ -381,11 +381,10 @@ task validatesContainer() {
 }
 
 /**
- * Validates the distroless (https://github.com/GoogleContainerTools/distroless)variant of the Python SDK container
+ * Validates the distroless (https://github.com/GoogleContainerTools/distroless) variant of the Python SDK container
  * image (sdks/python/container/Dockerfile).
  * To test a single version of Python:
- * ./gradlew :sdks:python:test-suites:dataflow:py311:validatesDistrolessContainer \
- *    -PpythonVersion=3 -Pjava11Home=$JAVA_HOME -PtestJavaVersion=11
+ * ./gradlew :sdks:python:test-suites:dataflow:py311:validatesDistrolessContainer
  * See https://cwiki.apache.org/confluence/display/BEAM/Python+Tips#PythonTips-VirtualEnvironmentSetup
  * for more information on setting up different Python versions.
  */
@@ -395,30 +394,28 @@ task validatesDistrolessContainer() {
   def repository = "us.gcr.io/apache-beam-testing/${System.getenv('USER')}"
   def tag = java.time.Instant.now().getEpochSecond()
   def imageURL = "${repository}/beam_python${project.ext.pythonVersion}_sdk_${buildTarget}:${tag}"
+  project.rootProject.ext['docker-repository-root'] = repository
+  project.rootProject.ext['container-build-target'] = buildTarget
+  project.rootProject.ext['docker-tag'] = tag
+  if (project.rootProject.hasProperty('dry-run')) {
+    println "Running in dry run mode: imageURL: ${imageURL}, pyversion: ${pyversion}, buildTarget: ${buildTarget}, repository: ${repository}, tag: ${tag}, envdir: ${envdir}"
+    return
+  }
   dependsOn 'initializeForDataflowJob'
+  dependsOn ":sdks:python:container:py${pyversion}:docker"
+  dependsOn ":sdks:python:container:py${pyversion}:dockerPush"
   def testTarget = "apache_beam/examples/wordcount_it_test.py::WordCountIT::test_wordcount_it"
   def argMap = [
           "output": "gs://temp-storage-for-end-to-end-tests/py-it-cloud/output",
           "project": "apache-beam-testing",
           "region": "us-central1",
           "runner": "TestDataflowRunner",
-          "sdk_container_image": imageURL,
+          "sdk_container_image": "${imageURL}",
           "sdk_location": "container",
           "staging_location": "gs://temp-storage-for-end-to-end-tests/staging-it",
           "temp_location": "gs://temp-storage-for-end-to-end-tests/temp-it",
   ]
   def cmdArgs = mapToArgString(argMap)
-  doFirst {
-    exec {
-      workingDir = "${rootDir}/sdks/python/container/${pyversion}"
-      executable 'sh'
-      args '-c', "docker build -t ${imageURL} --target=${buildTarget} -f ../Dockerfile --build-arg=py_version=${project.ext.pythonVersion} ."
-    }
-    exec {
-      executable 'sh'
-      args '-c', "docker push ${imageURL}"
-    }
-  }
   doLast {
     exec {
       workingDir = "${rootDir}/sdks/python"