Start pushing ml containers

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

@@ -3040,6 +3040,10 @@ class BeamModulePlugin implements Plugin<Project> {
       project.ext.pythonVersion = project.hasProperty('pythonVersion') ?
           project.pythonVersion : '3.9'
 
+      // Set min/max python versions used for containers and supported versions.
+      project.ext.minPythonVersion = 9
+      project.ext.maxPythonVersion = 13
+
       def setupVirtualenv = project.tasks.register('setupVirtualenv')  {
         doLast {
           def virtualenvCmd = [

sdks/python/container/build.gradle

@@ -18,11 +18,11 @@
 
 plugins { id 'org.apache.beam.module' }
 applyGoNature()
+applyPythonNature()
 
 description = "Apache Beam :: SDKs :: Python :: Container"
-// Keep these values in sync with sdks/python/container/distroless/build.gradle.
-int min_python_version=9
-int max_python_version=12
+int min_python_version=project.ext.minPythonVersion
+int max_python_version=project.ext.maxPythonVersion
 
 configurations {
   sdkSourceTarball
@@ -71,6 +71,7 @@ for(int i=min_python_version; i<=max_python_version; ++i) {
 
 tasks.register("pushAll") {
   dependsOn ':sdks:python:container:distroless:pushAll'
+  dependsOn ':sdks:python:container:ml:pushAll'
   for(int ver=min_python_version; ver<=max_python_version; ++ver) {
       if (!project.hasProperty("skip-python-3" + ver + "-images")) {
         dependsOn ':sdks:python:container:push3' + ver

sdks/python/container/common.gradle

@@ -41,6 +41,7 @@ def generatePythonRequirements = tasks.register("generatePythonRequirements") {
               "${project.ext.pythonVersion} " +
               "${files(configurations.sdkSourceTarball.files).singleFile} " +
               "base_image_requirements.txt " +
+              "container" +
               "[gcp,dataframe,test] " +
               "${pipExtraOptions}"
     }
@@ -51,6 +52,7 @@ def generatePythonRequirements = tasks.register("generatePythonRequirements") {
               "${project.ext.pythonVersion} " +
               "${files(configurations.sdkSourceTarball.files).singleFile} " +
               "ml_image_requirements.txt " +
+              "container/ml" +
               "[gcp,dataframe,test,tensorflow,torch,transformers] " +
               "${pipExtraOptions}"
     }

sdks/python/container/distroless/build.gradle

@@ -17,11 +17,11 @@
  */
 
 plugins { id 'org.apache.beam.module' }
+applyPythonNature()
 
 description = "Apache Beam :: SDKs :: Python :: Container :: Distroless"
-// Keep these values in sync with sdks/python/container/build.gradle.
-int min_python_version=9
-int max_python_version=12
+int min_python_version=project.ext.minPythonVersion
+int max_python_version=project.ext.maxPythonVersion
 
 
 tasks.register("buildAll") {

sdks/python/container/ml/build.gradle

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins { id 'org.apache.beam.module' }
+applyPythonNature()
+
+description = "Apache Beam :: SDKs :: Python :: Container :: ML"
+int min_python_version=project.ext.minPythonVersion
+int max_python_version=project.ext.maxPythonVersion
+
+
+tasks.register("buildAll") {
+    for(int ver=min_python_version; ver<=max_python_version; ++ver) {
+        dependsOn ':sdks:python:container:ml:py3' + ver + ':docker'
+    }
+}
+
+for(int i=min_python_version; i<=max_python_version; ++i) {
+    String min_version = "3" + min_python_version
+    String cur = "3" + i
+    String prev = "3" + (i-1)
+    tasks.register("push" + cur) {
+        if (cur != min_version) {
+            // 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)
+            if (!project.hasProperty("skip-python-3" + prev + "-images")) {
+                mustRunAfter(":sdks:python:container:ml:push" + prev)
+            }
+        }
+        dependsOn ':sdks:python:container:ml:py' + cur + ':docker'
+
+        doLast {
+            if (project.hasProperty("prune-images")) {
+                exec {
+                    executable("docker")
+                    args("system", "prune", "-a", "--force")
+                }
+            }
+        }
+    }
+}
+
+tasks.register("pushAll") {
+    for(int ver=min_python_version; ver<=max_python_version; ++ver) {
+        if (!project.hasProperty("skip-python-3" + ver + "-images")) {
+            dependsOn ':sdks:python:container:ml:push3' + ver
+        }
+    }
+}

sdks/python/container/ml/common.gradle

@@ -0,0 +1,126 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+def pythonVersionSuffix = project.ext.pythonVersion.replace('.', '')
+
+description = "Apache Beam :: SDKs :: Python :: Container :: ML :: Python ${pythonVersionSuffix} Container"
+
+configurations {
+  sdkSourceTarball
+  pythonHarnessLauncher
+}
+
+dependencies {
+  sdkSourceTarball project(path: ":sdks:python", configuration: "distTarBall")
+  pythonHarnessLauncher project(path: ":sdks:python:container", configuration: "pythonHarnessLauncher")
+}
+
+def generatePythonRequirements = tasks.register("generatePythonRequirements") {
+  dependsOn ':sdks:python:sdist'
+  def pipExtraOptions = project.hasProperty("testRCDependencies") ? "--pre" : ""
+  def runScriptsPath = "${rootDir}/sdks/python/container/run_generate_requirements.sh"
+  doLast {
+    exec {
+      executable 'sh'
+      args '-c', "cd ${rootDir} && ${runScriptsPath} " +
+              "${project.ext.pythonVersion} " +
+              "${files(configurations.sdkSourceTarball.files).singleFile} " +
+              "base_image_requirements.txt " +
+              "[gcp,dataframe,test] " +
+              "${pipExtraOptions}"
+    }
+    // Generate versions for ML dependencies
+    exec {
+      executable 'sh'
+      args '-c', "cd ${rootDir} && ${runScriptsPath} " +
+              "${project.ext.pythonVersion} " +
+              "${files(configurations.sdkSourceTarball.files).singleFile} " +
+              "ml_image_requirements.txt " +
+              "[gcp,dataframe,test,tensorflow,torch,transformers] " +
+              "${pipExtraOptions}"
+    }
+  }
+}
+
+def copyDockerfileDependencies = tasks.register("copyDockerfileDependencies", Copy) {
+  from configurations.sdkSourceTarball
+  from file("base_image_requirements.txt")
+  into "build/target"
+  if(configurations.sdkSourceTarball.isEmpty()) {
+      throw new StopExecutionException();
+  }
+}
+
+def copyLicenseScripts = tasks.register("copyLicenseScripts", Copy){
+  from ("../license_scripts")
+  into "build/target/license_scripts"
+}
+
+def copyLauncherDependencies = tasks.register("copyLauncherDependencies", Copy) {
+  from configurations.pythonHarnessLauncher
+  into "build/target/launcher"
+
+  // Avoid seemingly gradle bug stated in https://github.com/apache/beam/issues/29220
+  mustRunAfter "copyLicenses"
+
+  if(configurations.pythonHarnessLauncher.isEmpty()) {
+      throw new StopExecutionException();
+  }
+}
+
+def pushContainers = project.rootProject.hasProperty(["isRelease"]) || project.rootProject.hasProperty("push-containers")
+
+docker {
+  name containerImageName(
+          name: project.docker_image_default_repo_prefix + "python${project.ext.pythonVersion}_sdk_ml",
+          root: project.rootProject.hasProperty(["docker-repository-root"]) ?
+                  project.rootProject["docker-repository-root"] :
+                  project.docker_image_default_repo_root,
+          tag: project.rootProject.hasProperty(["docker-tag"]) ?
+                  project.rootProject["docker-tag"] : project.sdk_version)
+  // tags used by dockerTag task
+  tags containerImageTags()
+  files "../../Dockerfile", "./build"
+  buildArgs(['py_version': "${project.ext.pythonVersion}",
+             'pull_licenses': project.rootProject.hasProperty(["docker-pull-licenses"]) ||
+                     project.rootProject.hasProperty(["isRelease"])])
+  buildx project.useBuildx()
+  platform(*project.containerPlatforms())
+  load project.useBuildx() && !pushContainers
+  push pushContainers
+}
+
+dockerPrepare.dependsOn copyLauncherDependencies
+dockerPrepare.dependsOn copyDockerfileDependencies
+dockerPrepare.dependsOn copyLicenseScripts
+
+if (project.rootProject.hasProperty("docker-pull-licenses")) {
+  def copyGolangLicenses = tasks.register("copyGolangLicenses", Copy) {
+    from "${project(':release:go-licenses:py').buildDir}/output"
+    into "build/target/go-licenses"
+    dependsOn ':release:go-licenses:py:createLicenses'
+  }
+  dockerPrepare.dependsOn copyGolangLicenses
+} else {
+  def skipPullLicenses = tasks.register("skipPullLicenses", Exec) {
+    executable "sh"
+    // Touch a dummy file to ensure the directory exists.
+    args "-c", "mkdir -p build/target/go-licenses && touch build/target/go-licenses/skip"
+  }
+  dockerPrepare.dependsOn skipPullLicenses
+}

sdks/python/container/ml/py310/build.gradle

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+    id 'base'
+    id 'org.apache.beam.module'
+}
+applyDockerNature()
+applyPythonNature()
+
+pythonVersion = '3.10'
+
+apply from: "../common.gradle"

sdks/python/container/ml/py311/build.gradle

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+    id 'base'
+    id 'org.apache.beam.module'
+}
+applyDockerNature()
+applyPythonNature()
+
+pythonVersion = '3.11'
+
+apply from: "../common.gradle"

sdks/python/container/ml/py312/build.gradle

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+    id 'base'
+    id 'org.apache.beam.module'
+}
+applyDockerNature()
+applyPythonNature()
+
+pythonVersion = '3.12'
+
+apply from: "../common.gradle"

sdks/python/container/ml/py313/build.gradle

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * License); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+plugins {
+    id 'base'
+    id 'org.apache.beam.module'
+}
+applyDockerNature()
+applyPythonNature()
+
+pythonVersion = '3.13'
+
+apply from: "../common.gradle"