From fbcc70dae62ba9a72b925f21a98cbe0c343da89f Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Fri, 23 Oct 2026 18:12:48 +0000 Subject: [PATCH 1/3] Update tensorflow dependencies --- .../embeddings/tensorflow_hub_test.py | 4 +- sdks/python/setup.py | 38 +++++++--- .../python/test-suites/tox/py311/build.gradle | 31 +++++---- .../python/test-suites/tox/py312/build.gradle | 4 -- sdks/python/tox.ini | 69 +++++++++---------- 5 files changed, 82 insertions(+), 64 deletions(-) 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 b9a14ddb8c53..b8ef72a93464 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 @@ 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 diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 79d585c1e74f..cc1c32f8fc42 100644 --- a/sdks/python/setup.py +++ b/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 @@ -613,8 +621,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 +669,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 +705,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'] diff --git a/sdks/python/test-suites/tox/py311/build.gradle b/sdks/python/test-suites/tox/py311/build.gradle index 4ea61cb794f4..f4a6a1dcf424 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. // 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}" diff --git a/sdks/python/test-suites/tox/py312/build.gradle b/sdks/python/test-suites/tox/py312/build.gradle index c5f1774ca46b..ccaf3b79431a 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 df5028d9f9b5..1a5ae2a8d2e8 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" @@ -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 = From 0f85d5bdbfdad57b7130e6d9768e9bad2ba3ec63 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Fri, 23 Oct 2026 18:16:55 +0000 Subject: [PATCH 2/3] Trigger postcommits --- .github/trigger_files/beam_PostCommit_Python_Dependency.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index 9d28ab31a100..0c0e8ca6a12b 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 } From 4f04020e418a03ef293d24bfe67efe275864f265 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 26 Oct 2026 14:58:58 +0000 Subject: [PATCH 3/3] Adjust ml_base_core bounds to avoid crash with triton --- sdks/python/setup.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sdks/python/setup.py b/sdks/python/setup.py index cc1c32f8fc42..a8e2e4440341 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -216,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.