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 320e1007729 [Experimental] Update ONNX extras (#40821)
320e1007729 is described below
commit 320e1007729c63a67c0fc59f52dd7b04f53987f1
Author: Jack McCluskey <[email protected]>
AuthorDate: Thu Oct 29 09:58:12 2026 -0400
[Experimental] Update ONNX extras (#40821)
* Update ONNX extras
* Trigger postcommits
* Fix py311-ml onnx bounds
* add onnxscript to tox environment
* Add onnxscript to ml_base
---
.../beam_PostCommit_Python_Dependency.json | 2 +-
.../ml/inference/onnx_inference_test.py | 61 ++++++++++++++--------
sdks/python/setup.py | 20 +++----
sdks/python/test-suites/tox/py311/build.gradle | 14 +++--
sdks/python/tox.ini | 35 ++++++++-----
5 files changed, 82 insertions(+), 50 deletions(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json
b/.github/trigger_files/beam_PostCommit_Python_Dependency.json
index a5413bd5db9..8f085d87884 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": 10
+ "modification": 11
}
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 c301f98169f..e573a516345 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.test_pipeline import TestPipeline
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 @@ class PytorchLinearRegression(torch.nn.Module):
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 @@ class TestDataAndModel():
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 @@ class TestDataAndModel():
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 @@ class OnnxPytorchRunInferenceTest(OnnxTestBase):
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 @@ class OnnxTensorflowRunInferenceTest(OnnxTestBase):
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 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
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 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
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 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
| 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 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
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 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
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 @@ class OnnxSklearnRunInferencePipelineTest(OnnxTestBase):
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 98d5b631f13..ed39cd76ba7 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -183,10 +183,8 @@ xgboost_dependency = [
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',
@@ -703,12 +701,16 @@ if __name__ == '__main__':
],
'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 580b6fd7fba..27f7904896d 100644
--- a/sdks/python/test-suites/tox/py311/build.gradle
+++ b/sdks/python/test-suites/tox/py311/build.gradle
@@ -111,11 +111,15 @@ 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 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 the lower and upper bounds of supported tensorflow.
// Keep in sync with [testenv:py311-tensorflow-...] in tox.ini.
diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini
index 80897289733..ac583f3b6c7 100644
--- a/sdks/python/tox.ini
+++ b/sdks/python/tox.ini
@@ -111,7 +111,8 @@ commands =
deps =
pip==25.0.1
accelerate>=1.6.0
- onnx<1.19.0
+ onnx<2.0.0
+ onnxscript>=0.5.0,<1
setenv =
extras = test,gcp,dataframe,p311_ml_test
commands =
@@ -420,25 +421,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-{216,221}]
# Keep in line with the tensorflow extra in setup.py.
