Cherrypick #37721

Author: Amar3tto

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

@@ -3165,14 +3165,16 @@ class BeamModulePlugin implements Plugin<Project> {
           def distTarBall = "${pythonRootDir}/build/apache-beam.tar.gz"
           def packages = "gcp,test,aws,azure,dataframe"
           def extra = project.findProperty('beamPythonExtra')
-          if (extra) {
-            packages += ",${extra}"
-          }
-
           project.exec {
             executable 'sh'
             args '-c', ". ${project.ext.envdir}/bin/activate && pip install --pre --retries 10 ${distTarBall}[${packages}]"
           }
+          if (extra) {
+            project.exec {
+              executable 'sh'
+              args '-c', ". ${project.ext.envdir}/bin/activate && pip install --pre --retries 10 ${distTarBall}[${extra}]"
+            }
+          }
         }
       }
 

sdks/python/setup.py

@@ -163,14 +163,15 @@ def cythonize(*args, **kwargs):
 milvus_dependency = ['pymilvus>=2.5.10,<3.0.0']
 
 ml_base = [
-    'embeddings',
+    'embeddings>=0.0.4', # 0.0.3 crashes setuptools
     'onnxruntime',
     'langchain',
     'sentence-transformers>=2.2.2',
     'skl2onnx',
-    'pillow',
-    'pyod',
+    'pyod>=0.7.6', # 0.7.5 crashes setuptools
     'tensorflow',
+    # tensorflow transient dep, lower versions not compatible with Python3.10+
+    'absl-py>=0.12.0',
     'tensorflow-hub',
     'tf2onnx',
     'torch',
@@ -391,6 +392,7 @@ def get_portability_package_data():
           'numpy>=1.14.3,<2.5.0',  # Update pyproject.toml as well.
           'objsize>=0.6.1,<0.8.0',
           'packaging>=22.0',
+          'pillow>=12.1.1,<13',
           'pymongo>=3.8.0,<5.0.0',
           'proto-plus>=1.7.1,<2',
           # 1. Use a tighter upper bound in protobuf dependency to make sure
@@ -467,7 +469,7 @@ def get_portability_package_data():
               'pg8000>=1.31.5',
               "PyMySQL>=1.1.0",
               'oracledb>=3.1.1'
-          ] + milvus_dependency,
+          ],
           'gcp': [
               'cachetools>=3.1.0,<7',
               'google-api-core>=2.0.0,<3',
@@ -551,7 +553,9 @@ def get_portability_package_data():
           'p312_ml_test': [
               'datatable',
           ] + ml_base,
-          'p313_ml_test': ml_base,
+          # maintainer: milvus tests only run with this extension. Make sure it
+          # is covered by docker-in-docker test when changing py version
+          'p313_ml_test': ml_base + milvus_dependency,
           'aws': ['boto3>=1.9,<2'],
           'azure': [
               'azure-storage-blob>=12.3.2,<13',
@@ -584,7 +588,11 @@ def get_portability_package_data():
           # For more info, see
           # https://docs.google.com/document/d/1c84Gc-cZRCfrU8f7kWGsNR2o8oSRjCM-dGHO9KvPWPw/edit?usp=sharing
           'torch': ['torch>=1.9.0,<2.8.0'],
-          'tensorflow': ['tensorflow>=2.12rc1,<2.21'],
+          'tensorflow': [
+              'tensorflow>=2.12rc1,<2.21',
+              # tensorflow transitive dep
+              'absl-py>=0.12.0'
+          ],
           'transformers': [
               'transformers>=4.28.0,<4.56.0',
               'tensorflow>=2.12.0',
@@ -593,7 +601,9 @@ def get_portability_package_data():
           'ml_cpu': [
               'tensorflow>=2.12.0',
               'torch==2.8.0+cpu',
-              'transformers>=4.28.0,<4.56.0'
+              'transformers>=4.28.0,<4.56.0',
+              # tensorflow transient dep
+              'absl-py>=0.12.0'
           ],
           'redis': ['redis>=5.0.0,<6'],
           'tft': [
@@ -610,7 +620,9 @@ def get_portability_package_data():
               'tensorflow==2.11.0',
               'tf2onnx==1.13.0',
               'skl2onnx==1.13',
-              'transformers==4.25.1'
+              'transformers==4.25.1',
+               # tensorflow transient dep
+              'absl-py>=0.12.0'
           ],
           'xgboost': ['xgboost>=1.6.0,<2.1.3', 'datatable==1.0.0'],
           'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'],