[Experimental] Update ONNX extras

.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": 8
+    "modification": 11
   }

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')

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',
@@ -702,12 +700,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,

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.

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.