(beam) branch master updated: Changed the retry order for test_big_query_write_temp_table_append_schema_update (#31407)

[tvalentyn]

This is an automated email from the ASF dual-hosted git repository.

tvalentyn 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 2719ca7f37a Changed the retry order for 
test_big_query_write_temp_table_append_schema_update (#31407)
2719ca7f37a is described below

commit 2719ca7f37abaa490d3e5861ae80c12d978c2c37
Author: liferoad <huxiangq...@gmail.com>
AuthorDate: Mon May 27 13:38:07 2024 -0400

    Changed the retry order for 
test_big_query_write_temp_table_append_schema_update (#31407)
    
    * changed the retry order
    
    * updated shape
    
    * fixed save_model
    
    * added load_model_args
---
 sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py   |  2 +-
 .../apache_beam/ml/inference/tensorflow_inference_test.py  | 14 +++++++++-----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py 
b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
index 4c8e8ebf489..b0140793cf7 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
@@ -506,7 +506,6 @@ class BigQueryWriteIntegrationTests(unittest.TestCase):
           equal_to(bq_result_errors))
 
   @pytest.mark.it_postcommit
-  @retry(reraise=True, stop=stop_after_attempt(3))
   @parameterized.expand([
       param(file_format=FileFormat.AVRO),
       param(file_format=FileFormat.JSON),
@@ -514,6 +513,7 @@ class BigQueryWriteIntegrationTests(unittest.TestCase):
   ])
   @mock.patch(
       "apache_beam.io.gcp.bigquery_file_loads._MAXIMUM_SOURCE_URIS", new=1)
+  @retry(reraise=True, stop=stop_after_attempt(3))
   def test_big_query_write_temp_table_append_schema_update(self, file_format):
     """
     Test that nested schema update options and schema relaxation
diff --git a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py 
b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
index 52c525cc0ea..52123516de1 100644
--- a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
@@ -65,7 +65,7 @@ class FakeTFTensorModel:
 
 
 def _create_mult2_model():
-  inputs = tf.keras.Input(shape=(3))
+  inputs = tf.keras.Input(shape=(3, ))
   outputs = tf.keras.layers.Lambda(lambda x: x * 2, dtype='float32')(inputs)
   return tf.keras.Model(inputs=inputs, outputs=outputs)
 
@@ -127,7 +127,7 @@ class TFRunInferenceTest(unittest.TestCase):
 
   def test_predict_tensor_with_batch_size(self):
     model = _create_mult2_model()
-    model_path = os.path.join(self.tmpdir, 'mult2')
+    model_path = os.path.join(self.tmpdir, 'mult2.keras')
     tf.keras.models.save_model(model, model_path)
     with TestPipeline() as pipeline:
 
@@ -146,6 +146,7 @@ class TFRunInferenceTest(unittest.TestCase):
       model_handler = TFModelHandlerTensor(
           model_uri=model_path,
           inference_fn=fake_batching_inference_fn,
+          load_model_args={'safe_mode': False},
           min_batch_size=2,
           max_batch_size=2)
       examples = [
@@ -172,7 +173,7 @@ class TFRunInferenceTest(unittest.TestCase):
 
   def test_predict_tensor_with_large_model(self):
     model = _create_mult2_model()
-    model_path = os.path.join(self.tmpdir, 'mult2')
+    model_path = os.path.join(self.tmpdir, 'mult2.keras')
     tf.keras.models.save_model(model, model_path)
     with TestPipeline() as pipeline:
 
@@ -193,6 +194,7 @@ class TFRunInferenceTest(unittest.TestCase):
       model_handler = TFModelHandlerTensor(
           model_uri=model_path,
           inference_fn=fake_batching_inference_fn,
+          load_model_args={'safe_mode': False},
           large_model=True)
       examples = [
           tf.convert_to_tensor(numpy.array([1.1, 2.2, 3.3], dtype='float32')),
@@ -218,7 +220,7 @@ class TFRunInferenceTest(unittest.TestCase):
 
   def test_predict_numpy_with_batch_size(self):
     model = _create_mult2_model()
-    model_path = os.path.join(self.tmpdir, 'mult2_numpy')
+    model_path = os.path.join(self.tmpdir, 'mult2_numpy.keras')
     tf.keras.models.save_model(model, model_path)
     with TestPipeline() as pipeline:
 
@@ -237,6 +239,7 @@ class TFRunInferenceTest(unittest.TestCase):
       model_handler = TFModelHandlerNumpy(
           model_uri=model_path,
           inference_fn=fake_batching_inference_fn,
+          load_model_args={'safe_mode': False},
           min_batch_size=2,
           max_batch_size=2)
       examples = [
@@ -260,7 +263,7 @@ class TFRunInferenceTest(unittest.TestCase):
 
   def test_predict_numpy_with_large_model(self):
     model = _create_mult2_model()
-    model_path = os.path.join(self.tmpdir, 'mult2_numpy')
+    model_path = os.path.join(self.tmpdir, 'mult2_numpy.keras')
     tf.keras.models.save_model(model, model_path)
     with TestPipeline() as pipeline:
 
@@ -280,6 +283,7 @@ class TFRunInferenceTest(unittest.TestCase):
 
       model_handler = TFModelHandlerNumpy(
           model_uri=model_path,
+          load_model_args={'safe_mode': False},
           inference_fn=fake_inference_fn,
           large_model=True)
       examples = [