From 50f6ae5406ebf2780621f90bbe35d5a2b7053891 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 26 Oct 2026 18:30:27 +0000 Subject: [PATCH 1/5] Update ONNX extras --- .../ml/inference/onnx_inference_test.py | 61 ++++++++++++------- sdks/python/setup.py | 16 +++-- .../python/test-suites/tox/py311/build.gradle | 14 +++-- sdks/python/tox.ini | 32 ++++++---- 4 files changed, 78 insertions(+), 45 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/onnx_inference_test.py b/sdks/python/apache_beam/ml/inference/onnx_inference_test.py index c301f98169fd..e573a516345f 100644 --- a/sdks/python/apache_beam/ml/inference/onnx_inference_test.py +++ b/sdks/python/apache_beam/ml/inference/onnx_inference_test.py @@ -31,10 +31,6 @@ from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to -if bool(1): # lint doesn't like an unconditional `raise`. - raise unittest.SkipTest( - 'TODO: fix https://github.com/apache/beam/issues/31254') - # Protect against environments where onnx and pytorch library is not available. # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: @@ -77,6 +73,28 @@ def generate(self, x): return out +def _build_tf_linear_model(params): + # Keras 3 (tensorflow>=2.16) no longer accepts a `weights` argument when + # constructing layers, so build the model and then set the weights. + input_dim = params[0].shape[0] + linear_layer = layers.Dense(units=1) + linear_model = tf.keras.Sequential( + [tf.keras.Input(shape=(input_dim, )), linear_layer]) + linear_layer.set_weights(params) + return linear_model + + +def _export_tf_model_to_onnx(model, input_signature, path): + # tf2onnx.convert.from_keras does not support Keras 3 models + # (tensorflow>=2.16), so convert a tf.function wrapping the model instead. + @tf.function(input_signature=input_signature) + def model_fn(x): + return model(x) + + tf2onnx.convert.from_function( + model_fn, input_signature=input_signature, opset=13, output_path=path) + + class TestDataAndModel(): def get_one_feature_samples(self): return [ @@ -121,9 +139,7 @@ def get_tf_one_feature_model(self): params = [ np.array([[2.0]], dtype="float32"), np.array([0.5], dtype="float32") ] - linear_layer = layers.Dense(units=1, weights=params) - linear_model = tf.keras.Sequential([linear_layer]) - return linear_model + return _build_tf_linear_model(params) def get_sklearn_one_feature_model(self): x = [[0], [1]] @@ -140,10 +156,11 @@ def get_torch_two_feature_model(self): return model def get_tf_two_feature_model(self): - params = [np.array([[2.0], [3]]), np.array([0.5], dtype="float32")] - linear_layer = layers.Dense(units=1, weights=params) - linear_model = tf.keras.Sequential([linear_layer]) - return linear_model + params = [ + np.array([[2.0], [3]], dtype="float32"), + np.array([0.5], dtype="float32") + ] + return _build_tf_linear_model(params) def get_sklearn_two_feature_model(self): x = [[1, 5], [3, 2], [1, 0]] @@ -211,7 +228,7 @@ def test_onnx_pytorch_run_inference(self): dummy_input, # model input path, # where to save the model export_params=True, # store the trained parameter weights - opset_version=10, # the ONNX version + opset_version=17, # the ONNX version do_constant_folding=True, # whether to execute constant- # folding for optimization input_names=['input'], # model's input names @@ -266,10 +283,7 @@ def test_onnx_tensorflow_run_inference(self): path = os.path.join(self.tmpdir, 'my_onnx_tf_path') spec = (tf.TensorSpec((None, 1), tf.float32, name="input"), ) - _, _ = tf2onnx.convert.from_keras(linear_model, - input_signature=spec, - opset=13, - output_path=path) + _export_tf_model_to_onnx(linear_model, spec, path) inference_runner = TestOnnxModelHandler(path) inference_session = ort.InferenceSession( @@ -316,7 +330,7 @@ def exportModelToOnnx(self, model, path): dummy_input, # model input path, # where to save the model export_params=True, # store the trained parameter weights - opset_version=10, # the ONNX version + opset_version=17, # the ONNX version do_constant_folding=True, # whether to execute constant # folding for optimization input_names=['input'], # odel's input names @@ -352,6 +366,7 @@ def test_model_handler_sets_env_vars(self): self.exportModelToOnnx(model, path) model_handler = OnnxModelHandlerNumpy( model_uri=path, env_vars={'FOO': 'bar'}) + os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) _ = ( pipeline @@ -360,8 +375,12 @@ def test_model_handler_sets_env_vars(self): | RunInference(model_handler)) pipeline.run() self.assertTrue('FOO' in os.environ) - self.assertTrue('bar'.equals(os.environ['FOO'])) + self.assertEqual(os.environ['FOO'], 'bar') + @unittest.skip( + 'large_model=True does not work with the default ONNX inference fn: ' + 'InferenceSession.get_inputs() returns onnxruntime NodeArg objects, ' + 'which cannot be pickled back through the multi_process_shared proxy.') def test_model_handler_large_model(self): with TestPipeline() as pipeline: @@ -384,6 +403,7 @@ def onxx_numpy_inference_fn( env_vars={'FOO': 'bar'}, inference_fn=onxx_numpy_inference_fn, large_model=True) + os.environ.pop('FOO', None) self.assertFalse('FOO' in os.environ) _ = ( pipeline @@ -427,8 +447,7 @@ def test_invalid_input_type(self): class OnnxTensorflowRunInferencePipelineTest(OnnxTestBase): def exportModelToOnnx(self, model, path): spec = (tf.TensorSpec((None, 2), tf.float32, name="input"), ) - _, _ = tf2onnx.convert.from_keras(model, - input_signature=spec, opset=13, output_path=path) + _export_tf_model_to_onnx(model, spec, path) def test_pipeline_local_model_simple(self): with TestPipeline() as pipeline: @@ -518,7 +537,7 @@ def test_pipeline_gcs_model(self): equal_to(expected_predictions, equals_fn=_compare_prediction_result)) def test_invalid_input_type(self): - with self.assertRaisesRegex(Exception, "InvalidArgument"): + with self.assertRaisesRegex(Exception, "Got invalid dimensions for input"): with TestPipeline() as pipeline: examples = [np.array([1], dtype="float32")] path = os.path.join(self.tmpdir, 'my_onnx_sklearn_path') diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 77a5b835a059..bdf65392943c 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -688,12 +688,16 @@ def get_portability_package_data(): ], 'tfrecord': ['crcmod>=1.7,<2.0'], 'onnx': [ - 'onnxruntime==1.13.1', - 'torch==1.13.1', - 'tensorflow==2.11.0', - 'tf2onnx==1.17.0', - 'skl2onnx==1.13', - 'transformers==4.25.1', # tensorflow transient dep + # Keep in line with the py311-onnx-* environments in tox.ini. + 'onnx>=1.16.0,<2', + 'onnxruntime>=1.19.0,<2', + # Converters for producing ONNX models from scikit-learn, + # TensorFlow and PyTorch models. + 'skl2onnx>=1.17.0,<2', + 'tf2onnx>=1.17.0,<1.18', + 'tensorflow>=2.18.0,<2.22', + 'torch>=2.3.0,<2.12.0', + 'onnxscript>=0.5.0,<1', 'absl-py>=0.12.0' ], 'xgboost': xgboost_dependency, diff --git a/sdks/python/test-suites/tox/py311/build.gradle b/sdks/python/test-suites/tox/py311/build.gradle index a05e9972bef4..6e6a303f0af1 100644 --- a/sdks/python/test-suites/tox/py311/build.gradle +++ b/sdks/python/test-suites/tox/py311/build.gradle @@ -110,11 +110,15 @@ toxTask "testPy311tft-113", "py311-tft-113", "${posargs}" test.dependsOn "testPy311tft-113" postCommitPyDep.dependsOn "testPy311tft-113" -// 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 the lower and upper bounds of supported onnx. +// Keep in sync with [testenv:py311-onnx-...] in tox.ini. +toxTask "testPy311onnx-119", "py311-onnx-119", "${posargs}" +test.dependsOn "testPy311onnx-119" +postCommitPyDep.dependsOn "testPy311onnx-119" + +toxTask "testPy311onnx-latest", "py311-onnx-latest", "${posargs}" +test.dependsOn "testPy311onnx-latest" +postCommitPyDep.dependsOn "testPy311onnx-latest" // Create a test task for each minor version of tensorflow toxTask "testPy311tensorflow-212", "py311-tensorflow-212", "${posargs}" diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index c69fc2e45ef7..cbb0bc62efd6 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -418,25 +418,31 @@ commands = # 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_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' -# TODO(https://github.com/apache/beam/issues/25796) - uncomment onnx tox task in tox/py311/build.gradle once onnx supports protobuf 4.x.x -[testenv:py311-onnx-113] -# TODO(https://github.com/apache/beam/issues/25443) -# apparently tox has problem when substitution key has single value. Change back to -onnx-{113,...} -# when multiple onnx versions are tested. +[testenv:py311-onnx-{119,latest}] +# Keep in line with the onnx extra in setup.py. tensorflow is pinned to 2.21 in +# both environments because tensorflow<2.20 cannot be installed alongside the +# gcp extra (protobuf<6 vs protobuf>=6.33.5). deps = - onnxruntime==1.13.1 - pandas==1.5.2 - torch==1.13.1 - tensorflow==2.11.0 - tf2onnx==1.13.0 - skl2onnx==1.13 - transformers==4.25.1 + # Lower bound of the supported range. + 119: onnxruntime>=1.19.0,<1.20 + 119: onnx>=1.16.0,<1.17 + 119: skl2onnx>=1.17.0,<1.18 + 119: torch>=2.3.0,<2.4.0 + latest: onnxruntime>=1.31.0,<2 + latest: onnx>=1.23.0,<2 + latest: skl2onnx>=1.20.0,<2 + latest: torch>=2.11.0,<2.12.0 + tf2onnx>=1.17.0,<1.18 + tensorflow>=2.21.0,<2.22 + # torch>=2.9 exports to ONNX through torch.export, which requires onnxscript. + onnxscript>=0.5.0,<1 extras = test,gcp commands = # Log onnx version for debugging /bin/sh -c "pip freeze | grep -E onnx" # Run all ONNX unit tests - pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_onnx {posargs} + # 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_onnx {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' [testenv:py311-tensorflow-212] deps = From 7eab7054fd841c7bf6590b028d928d736763445d Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 26 Oct 2026 18:32:39 +0000 Subject: [PATCH 2/5] 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..8f085d878840 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": 11 } From bfd02c0a88311db6427021a69a688d4000775d03 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 26 Oct 2026 18:54:44 +0000 Subject: [PATCH 3/5] Fix py311-ml onnx bounds --- sdks/python/tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index d09e48e0b008..91833ec61a42 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -111,7 +111,7 @@ commands = deps = pip==25.0.1 accelerate>=1.6.0 - onnx<1.19.0 + onnx<2.0.0 setenv = extras = test,gcp,dataframe,p311_ml_test commands = From 9b0cb8f09b14abbd6a0b5480721015686526c2da Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Mon, 26 Oct 2026 19:44:19 +0000 Subject: [PATCH 4/5] add onnxscript to tox environment --- sdks/python/tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 91833ec61a42..ac583f3b6c79 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -112,6 +112,7 @@ deps = pip==25.0.1 accelerate>=1.6.0 onnx<2.0.0 + onnxscript>=0.5.0,<1 setenv = extras = test,gcp,dataframe,p311_ml_test commands = From f7b5dcab1c8e84a011f42c072785f2e5dd9c04ca Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 27 Oct 2026 13:35:09 +0000 Subject: [PATCH 5/5] Add onnxscript to ml_base --- sdks/python/setup.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sdks/python/setup.py b/sdks/python/setup.py index d428062c1d17..72cff36c5773 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -183,10 +183,8 @@ def cythonize(*args, **kwargs): ml_base_core = [ 'embeddings>=0.0.4', # 0.0.3 crashes setuptools 'onnxruntime', - # onnx 1.12–1.13 cap protobuf in ways that trigger huge backtracking with - # Beam[gcp]+ml_test; pip can fall back to onnx 1.11 sdist which needs cmake. - # 1.14.1+ matches tf2onnx>=1.16 and ships manylinux wheels for py3.10. 'onnx>=1.14.1,<2', + 'onnxscript>=0.5.0,<1', 'langchain', 'sentence-transformers>=2.2.2', 'skl2onnx',