This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 49a6a2ecfa4 Update tensorflow dependencies (#40752)
49a6a2ecfa4 is described below
commit 49a6a2ecfa4679c025e434deb90293dc9cfee042
Author: Jack McCluskey <[email protected]>
AuthorDate: Mon Oct 26 13:22:11 2026 -0400
Update tensorflow dependencies (#40752)
* Update tensorflow dependencies
* Trigger postcommits
* Adjust ml_base_core bounds to avoid crash with triton
---
.../beam_PostCommit_Python_Dependency.json | 2 +-
.../transforms/embeddings/tensorflow_hub_test.py | 4 +-
sdks/python/setup.py | 44 ++++++++++----
sdks/python/test-suites/tox/py311/build.gradle | 31 +++++-----
sdks/python/test-suites/tox/py312/build.gradle | 4 --
sdks/python/tox.ini | 69 +++++++++++-----------
6 files changed, 87 insertions(+), 67 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json
b/.github/trigger_files/beam_PostCommit_Python_Dependency.json
index 9d28ab31a10..0c0e8ca6a12 100644
--- a/.github/trigger_files/beam_PostCommit_Python_Dependency.json
+++ b/.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
}
diff --git
a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
index b9a14ddb8c5..b8ef72a9346 100644
--- a/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
+++ b/sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py
@@ -211,7 +211,9 @@ class TFHubImageEmbeddingsTest(unittest.TestCase):
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
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 77a5b835a05..8236944aea1 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -177,9 +177,9 @@ xgboost_dependency = [
'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 @@ ml_base_core = [
'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 @@ ml_base_core = [
# 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 @@ if __name__ == '__main__':
# 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 @@ if __name__ == '__main__':
#
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': [
@@ -697,7 +711,15 @@ if __name__ == '__main__':
'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']
diff --git a/sdks/python/test-suites/tox/py311/build.gradle
b/sdks/python/test-suites/tox/py311/build.gradle
index a05e9972bef..580b6fd7fba 100644
--- a/sdks/python/test-suites/tox/py311/build.gradle
+++ b/sdks/python/test-suites/tox/py311/build.gradle
@@ -106,9 +106,10 @@ 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
@@ -116,10 +117,15 @@ postCommitPyDep.dependsOn "testPy311tft-113"
// 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.
// Keep in sync with [testenv:py311-transformers-...] in tox.ini.
@@ -150,13 +156,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}"
diff --git a/sdks/python/test-suites/tox/py312/build.gradle
b/sdks/python/test-suites/tox/py312/build.gradle
index c5f1774ca46..ccaf3b79431 100644
--- a/sdks/python/test-suites/tox/py312/build.gradle
+++ b/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"
diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini
index c69fc2e45ef..80897289733 100644
--- a/sdks/python/tox.ini
+++ b/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"
@@ -543,15 +541,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 =
