Update tensorflow dependencies

.github/trigger_files/beam_PostCommit_Python_Dependency.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 7
+    "modification": 8
   }

sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py

@@ -211,7 +211,9 @@ def assert_element(element):
   def test_with_str_data_types(self):
     embedding_config = TensorflowHubImageEmbeddings(
         hub_url=hub_img_url, columns=[test_query_column])
-    with self.assertRaises(TypeError):
+    # Portable runners surface the worker's TypeError wrapped in a
+    # RuntimeError, so match on the message like the other embeddings tests.
+    with self.assertRaisesRegex(Exception, "Embeddings can only be generated"):
       with beam.Pipeline() as pipeline:
         _ = (
             pipeline

sdks/python/setup.py

@@ -177,9 +177,9 @@ def cythonize(*args, **kwargs):
     'datatable>=1.1.0,<2',
 ]
 
-# google-adk / OpenTelemetry require protobuf>=5; tensorflow-transform in
-# ml_test is pinned to versions that require protobuf<5 on Python 3.10. Those
-# cannot be installed together, so ADK deps stay out of ml_test (use ml_base).
+# tensorflow-transform 1.21 (via tfx-bsl 1.21) requires
+# google-api-python-client<2, so tensorflow-transform stays out of ml_base and
+# is only added to ml_test.
 ml_base_core = [
     'embeddings>=0.0.4',  # 0.0.3 crashes setuptools
     'onnxruntime',
@@ -191,10 +191,18 @@ def cythonize(*args, **kwargs):
     'sentence-transformers>=2.2.2',
     'skl2onnx',
     'pyod>=0.7.6',  # 0.7.5 crashes setuptools
-    'tensorflow',
+    # Keep in line with the tensorflow extra and the py*-tensorflow-*
+    # environments in tox.ini.
+    'tensorflow>=2.16.0,<2.22',
     # tensorflow transitive dep, lower versions not compatible with Python3.10+
     'absl-py>=0.12.0',
-    'tensorflow-hub',
+    # tensorflow-hub<0.16 imports tf.compat.v1.estimator, which was removed in
+    # tensorflow 2.16.
+    'tensorflow-hub>=0.16.1,<0.17',
+    # tensorflow-hub imports pkg_resources, which was removed in setuptools 82.
+    # Without this cap, Python 3.12+ environments (no bundled setuptools)
+    # resolve setuptools>=82 and the TensorFlow tests are skipped.
+    'setuptools<82',
     # tokenizers 0.23.0rc0 renamed the PyO3 kwarg of
     # processors.RobertaProcessing (and BertProcessing) from `cls` to
     # `cls_token` -- the rename was a drive-by inside huggingface/tokenizers
@@ -208,8 +216,10 @@ def cythonize(*args, **kwargs):
     # Drop this cap once transformers updates the CLIP call site to
     # `cls_token=` or tokenizers reinstates `cls=` as a deprecation alias.
     'tokenizers<0.23',
-    # Keep in line with the py*-pytorch-* environments in tox.ini.
-    'torch>=2.3.0,<2.15.0',
+    # torch>=2.12 Linux wheels require triton>=3.7, which segfaults/aborts when
+    # imported after tensorflow. Keep this cap while tensorflow is part of
+    # ml_base_core.
+    'torch>=2.3.0,<2.12.0',
     # Match tested transformers range.
     'transformers>=4.28.0,<4.56.0',
     # Keep tokenizers compatible with this transformers range.
@@ -613,8 +623,8 @@ def get_portability_package_data():
               # tests due to tag check introduced since pip 24.2
               # https://github.com/apache/beam/issues/31285
               # 'xgboost<2.0',  # https://github.com/apache/beam/issues/31252
-              # tft needs protobuf<5; tf2onnx 1.17+ allows protobuf 5 on the
-              # ADK-only path.
+              # tf2onnx 1.16.x pins protobuf~=3.20; 1.17+ allows protobuf 6,
+              # which tensorflow-transform 1.21 requires.
               'tf2onnx>=1.17.0,<1.18',
           ] + ml_base_core,
           'p310_ml_test': [
@@ -661,7 +671,11 @@ def get_portability_package_data():
           # https://docs.google.com/document/d/1c84Gc-cZRCfrU8f7kWGsNR2o8oSRjCM-dGHO9KvPWPw/edit?usp=sharing
           'torch': ['torch>=2.3.0,<2.15.0'],
           'tensorflow': [
-              'tensorflow>=2.12rc1,<2.22',  # tensorflow transitive dep
+              # tensorflow 2.16 is the first release with Keras 3 and
+              # Python 3.12 wheels. tensorflow<2.20 requires protobuf<6, which
+              # cannot be installed alongside the gcp extra (protobuf>=6.33.5
+              # via google-cloud-bigtable).
+              'tensorflow>=2.16.0,<2.22',
               'absl-py>=0.12.0'
           ],
           'transformers': [
@@ -693,7 +707,15 @@ def get_portability_package_data():
               'absl-py>=0.12.0'
           ],
           'xgboost': xgboost_dependency,
-          'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'],
+          'tensorflow-hub': [
+              # tensorflow-hub<0.16 imports tf.compat.v1.estimator, which was
+              # removed in tensorflow 2.16. tensorflow-hub 0.16 depends on
+              # tf-keras.
+              'tensorflow-hub>=0.16.1,<0.17',
+              # tensorflow-hub imports pkg_resources, which was removed in
+              # setuptools 82.
+              'setuptools<82',
+          ],
           'milvus': milvus_dependency,
           'qdrant': qdrant_dependency,
           'vllm': ['openai==1.107.1', 'vllm==0.10.1.1', 'triton==3.3.1']

sdks/python/test-suites/tox/py311/build.gradle

@@ -106,20 +106,26 @@ toxTask "testPy311pytorch-214", "py311-pytorch-214", "${posargs}"
 test.dependsOn "testPy311pytorch-214"
 postCommitPyDep.dependsOn "testPy311pytorch-214"
 
-toxTask "testPy311tft-113", "py311-tft-113", "${posargs}"
-test.dependsOn "testPy311tft-113"
-postCommitPyDep.dependsOn "testPy311tft-113"
+// Keep in sync with [testenv:py311-tft-...] in tox.ini.
+toxTask "testPy311tft-121", "py311-tft-121", "${posargs}"
+test.dependsOn "testPy311tft-121"
+postCommitPyDep.dependsOn "testPy311tft-121"
 
 // TODO(https://github.com/apache/beam/issues/25796) - uncomment onnx tox task once onnx supports protobuf 4.x.x
 // Create a test task for each minor version of onnx
 // toxTask "testPy311onnx-113", "py311-onnx-113", "${posargs}"
 // test.dependsOn "testPy311onnx-113"
 // postCommitPyDep.dependsOn "testPy311onnx-113"
 
-// Create a test task for each minor version of tensorflow
-toxTask "testPy311tensorflow-212", "py311-tensorflow-212", "${posargs}"
-test.dependsOn "testPy311tensorflow-212"
-postCommitPyDep.dependsOn "testPy311tensorflow-212"
+// Create a test task for the lower and upper bounds of supported tensorflow.
+// Keep in sync with [testenv:py311-tensorflow-...] in tox.ini.
+toxTask "testPy311tensorflow-216", "py311-tensorflow-216", "${posargs}"
+test.dependsOn "testPy311tensorflow-216"
+postCommitPyDep.dependsOn "testPy311tensorflow-216"
+
+toxTask "testPy311tensorflow-221", "py311-tensorflow-221", "${posargs}"
+test.dependsOn "testPy311tensorflow-221"
+postCommitPyDep.dependsOn "testPy311tensorflow-221"
 
 // Create a test task for each minor version of transformers.
 // Dropped transformers-428: sentence-transformers==2.2.2 + torch<1.14 is
@@ -147,13 +153,10 @@ postCommitPyDep.dependsOn "testPy311embeddingsMLTransform"
 
 // Part of MLTransform embeddings test suite but requires tensorflow hub, which we need to test on
 // mutliple versions so keeping this suite separate.
-toxTask "testPy311TensorflowHubEmbeddings-014", "py311-TFHubEmbeddings-014", "${posargs}"
-test.dependsOn "testPy311TensorflowHubEmbeddings-014"
-postCommitPyDep.dependsOn "testPy311TensorflowHubEmbeddings-014"
-
-toxTask "testPy311TensorflowHubEmbeddings-015", "py311-TFHubEmbeddings-015", "${posargs}"
-test.dependsOn "testPy311TensorflowHubEmbeddings-015"
-postCommitPyDep.dependsOn "testPy311TensorflowHubEmbeddings-015"
+// Keep in sync with [testenv:py311-TFHubEmbeddings-...] in tox.ini.
+toxTask "testPy311TensorflowHubEmbeddings-016", "py311-TFHubEmbeddings-016", "${posargs}"
+test.dependsOn "testPy311TensorflowHubEmbeddings-016"
+postCommitPyDep.dependsOn "testPy311TensorflowHubEmbeddings-016"
 
 toxTask "whitespacelint", "whitespacelint", "${posargs}"
 

sdks/python/test-suites/tox/py312/build.gradle

@@ -30,9 +30,5 @@ pythonVersion = '3.12'
 
 project.tasks.register("postCommitPyDep") {}
 
-toxTask "testPy312tensorflow-216", "py312-tensorflow-216", "${posargs}"
-test.dependsOn "testPy312tensorflow-216"
-postCommitPyDep.dependsOn "testPy312tensorflow-216"
-
 apply from: "../common.gradle"
 

sdks/python/tox.ini

@@ -379,11 +379,13 @@ commands =
   # Run all DataFrame API unit tests
   bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/dataframe'
 
-[testenv:py311-tft-{113,114}]
+[testenv:py311-tft-121]
+# Keep in line with the tft extra in setup.py. tensorflow-transform<1.21
+# requires pyarrow<11, which is incompatible with Beam's pyarrow>=14.
 deps =
-  # Help pip resolve conflict with typing-extensions due to an old version of tensorflow https://github.com/apache/beam/issues/30852
-  113: pydantic<2.0
-  114: tensorflow_transform>=1.14.0,<1.15.0
+  tensorflow_transform>=1.21.0,<1.22.0
+  dill
+extras = test,gcp
 commands =
   bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py'
 
@@ -438,32 +440,28 @@ commands =
   # Run all ONNX unit tests
   pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_onnx {posargs}
 
-[testenv:py311-tensorflow-212]
-deps =
-  212:
-    googleapis-common-protos<1.70
-    tensorflow>=2.12rc1,<2.13
-    # Help pip resolve conflict with typing-extensions for old version of TF https://github.com/apache/beam/issues/30852
-    pydantic<2.7
-extras = test,gcp
-commands_pre =
-  pip install -U 'protobuf==4.25.5'
-commands =
-  # Log tensorflow version for debugging
-  /bin/sh -c "pip freeze | grep -E tensorflow"
-  # Run all Tensorflow unit tests
-  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
-  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_tf {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
-
-[testenv:py312-tensorflow-216]
+[testenv:py311-tensorflow-{216,221}]
+# Keep in line with the tensorflow extra in setup.py.
+# tensorflow<2.20 requires protobuf<6, which cannot be installed alongside the
+# gcp extra (protobuf>=6.33.5 via google-cloud-bigtable), so the lower-bound
+# environment only installs the test extra.
+# tensorflow-hub is needed by the TensorFlow inference tests; it imports
+# pkg_resources, which was removed in setuptools 82.
 deps =
-  216:
-    tensorflow>=2.16.1,<2.17
-    # Help pip resolve conflict with typing-extensions for old version of TF https://github.com/apache/beam/issues/30852
-    pydantic<2.7
-extras = test,gcp
-commands_pre =
-  pip install -U 'protobuf==4.25.5'
+  216: tensorflow>=2.16.0,<2.17
+  216: tf-keras>=2.16.0,<2.17
+  216: protobuf>=4.25.2,<5
+  221: tensorflow>=2.21.0,<2.22
+  221: tf-keras>=2.21.0,<2.22
+  tensorflow-hub>=0.16.1,<0.17
+  setuptools<82
+extras =
+  216: test
+  221: test,gcp
+# Apply the deps above as constraints when installing Beam's own dependencies,
+# so that they cannot upgrade e.g. protobuf past tensorflow 2.16's protobuf<5
+# cap.
+constrain_package_deps = true
 commands =
   # Log tensorflow version for debugging
   /bin/sh -c "pip freeze | grep -E tensorflow"
@@ -537,15 +535,14 @@ commands =
   /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
 
-[testenv:py311-TFHubEmbeddings-{014,015}]
+[testenv:py311-TFHubEmbeddings-016]
+# Keep in line with the tensorflow-hub extra in setup.py. tensorflow-hub<0.16
+# imports tf.compat.v1.estimator, which was removed in tensorflow 2.16.
 deps =
-  014: tensorflow-hub>=0.14.0,<0.15.0
-  # Help pip resolve conflict with typing-extensions due to an old version of tensorboard https://github.com/apache/beam/issues/30852
-  014: pydantic<2.7
-  015: tensorflow-hub>=0.15.0,<0.16.0
-  # Help pip resolve conflict with typing-extensions due to an old version of tensorboard https://github.com/apache/beam/issues/30852
-  015: pydantic<2.7
+  tensorflow-hub>=0.16.1,<0.17
   tensorflow-text # required to register ops for text embedding models.
+  # tensorflow-hub imports pkg_resources, which was removed in setuptools 82.
+  setuptools<82
 
 extras = test,gcp
 commands =