diff --git a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py index 84087dee5607..bca6881a714d 100644 --- a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py +++ b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py @@ -35,6 +35,7 @@ from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.runners.runner import PipelineState from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.utils import retry # Protect against environments where bigtable library is not available. # pylint: disable=wrong-import-order, wrong-import-position @@ -106,13 +107,13 @@ class BigtableIOWriteTest(unittest.TestCase): """ DEFAULT_TABLE_PREFIX = "python-test" - instance_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] - cluster_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] - table_id = DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] number = 500 LOCATION_ID = "us-east1-b" def setUp(self): + self.instance_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] + self.cluster_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] + self.table_id = self.DEFAULT_TABLE_PREFIX + "-" + str(uuid.uuid4())[:8] try: from google.cloud.bigtable import enums self.STORAGE_TYPE = enums.StorageType.HDD @@ -136,8 +137,17 @@ def setUp(self): self.cluster_id, self.LOCATION_ID, default_storage_type=self.STORAGE_TYPE) - operation = self.instance.create(clusters=[cluster]) - operation.result(timeout=300) # Wait up to 5 min. + + @retry.with_exponential_backoff( + num_retries=3, + initial_delay_secs=5, + retry_filter=retry. + retry_on_server_errors_timeout_or_quota_issues_filter) + def _create_instance(): + operation = self.instance.create(clusters=[cluster]) + operation.result(timeout=300) # Wait up to 5 min. + + _create_instance() self.table = self.instance.table(self.table_id) diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py index 9ed27a500a7a..c67e8e90e5a0 100644 --- a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py +++ b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py @@ -66,12 +66,19 @@ def setUp(self): self.project, OUTPUT_SUB + self.uuid), topic=self.output_topic.name, ack_deadline_seconds=60) + import time + time.sleep(10) def _inject_numbers(self, topic, num_messages): """Inject numbers as test data to PubSub.""" logging.debug('Injecting %d numbers to topic %s', num_messages, topic.name) + futures = [] for n in range(num_messages): - self.pub_client.publish(self.input_topic.name, str(n).encode('utf-8')) + futures.append( + self.pub_client.publish( + self.input_topic.name, str(n).encode('utf-8'))) + for future in futures: + future.result() def tearDown(self): test_utils.cleanup_subscriptions( @@ -88,7 +95,7 @@ def test_streaming_wordcount_it(self): # Set extra options to the pipeline for test purpose state_verifier = PipelineStateMatcher(PipelineState.RUNNING) pubsub_msg_verifier = PubSubMessageMatcher( - self.project, self.output_sub.name, expected_msg, timeout=400) + self.project, self.output_sub.name, expected_msg, timeout=600) extra_opts = { 'input_subscription': self.input_sub.name, 'output_topic': self.output_topic.name, diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py index b8d212428060..d335fd9fc5d7 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio.py @@ -290,6 +290,11 @@ def process(self, direct_row): direct_row._get_mutation_pbs() if hasattr( direct_row, '_get_mutation_pbs') else direct_row._get_mutations()) for mutation in mutations: + if not hasattr(mutation, '__contains__'): + if hasattr(mutation, '_to_pb'): + mutation = mutation._to_pb() + elif hasattr(mutation, 'to_pb'): + mutation = mutation.to_pb() if mutation.__contains__("set_cell"): mutation_dict = { "type": b'SetCell', diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 4aba9a4062fd..16fd41e6d1d0 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -122,11 +122,18 @@ class TestBigtableDirectRowToBeamRow(unittest.TestCase): @staticmethod def _get_mutation_pbs(direct_row): # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python - # dataclass objects instead of protobuf messages; use _get_mutation_pbs() - # to retrieve Mutation protobuf messages. - if hasattr(direct_row, '_get_mutation_pbs'): - return direct_row._get_mutation_pbs() - return direct_row._get_mutations() + # dataclass objects instead of protobuf messages; convert each item via + # _to_pb() when needed. + mutations = ( + direct_row._get_mutation_pbs() if hasattr( + direct_row, '_get_mutation_pbs') else direct_row._get_mutations()) + return [ + m._to_pb() + if not hasattr(m, '__contains__') and hasattr(m, '_to_pb') else ( + m.to_pb() + if not hasattr(m, '__contains__') and hasattr(m, 'to_pb') else m) + for m in mutations + ] def test_set_cell(self): # create some set cell mutations diff --git a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py index 89fd4461beb3..d665f64168e1 100644 --- a/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py +++ b/sdks/python/apache_beam/io/gcp/pubsub_integration_test.py @@ -36,6 +36,7 @@ from apache_beam.testing import test_utils from apache_beam.testing.pipeline_verifiers import PipelineStateMatcher from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.utils import retry INPUT_TOPIC = 'psit_topic_input' OUTPUT_TOPIC = 'psit_topic_output' @@ -138,20 +139,27 @@ def setUp(self): self.project = self.test_pipeline.get_option('project') self.uuid = str(uuid.uuid4()) - # Set up PubSub environment. + # Set up PubSub environment with retries for transient 504/403 API errors. from google.cloud import pubsub self.pub_client = pubsub.PublisherClient() - self.input_topic = self.pub_client.create_topic( + self.sub_client = pubsub.SubscriberClient() + + gcp_retry = retry.with_exponential_backoff( + num_retries=3, + initial_delay_secs=2, + retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter + ) + + self.input_topic = gcp_retry(self.pub_client.create_topic)( name=self.pub_client.topic_path(self.project, INPUT_TOPIC + self.uuid)) - self.output_topic = self.pub_client.create_topic( + self.output_topic = gcp_retry(self.pub_client.create_topic)( name=self.pub_client.topic_path(self.project, OUTPUT_TOPIC + self.uuid)) - self.sub_client = pubsub.SubscriberClient() - self.input_sub = self.sub_client.create_subscription( + self.input_sub = gcp_retry(self.sub_client.create_subscription)( name=self.sub_client.subscription_path( self.project, INPUT_SUB + self.uuid), topic=self.input_topic.name) - self.output_sub = self.sub_client.create_subscription( + self.output_sub = gcp_retry(self.sub_client.create_subscription)( name=self.sub_client.subscription_path( self.project, OUTPUT_SUB + self.uuid), topic=self.output_topic.name) diff --git a/sdks/python/apache_beam/ml/anomaly/transforms_test.py b/sdks/python/apache_beam/ml/anomaly/transforms_test.py index 72c16841d5fb..4da9d6c91781 100644 --- a/sdks/python/apache_beam/ml/anomaly/transforms_test.py +++ b/sdks/python/apache_beam/ml/anomaly/transforms_test.py @@ -84,11 +84,14 @@ def _prediction_is_equal_to(a: AnomalyPrediction, b: AnomalyPrediction): return False if a.threshold != b.threshold: - return False + if not (a.threshold is not None and b.threshold is not None and + math.isclose(a.threshold, b.threshold, rel_tol=1e-6, abs_tol=1e-6)): + return False if a.score != b.score: if not (a.score is not None and b.score is not None and - math.isnan(a.score) and math.isnan(b.score)): + ((math.isnan(a.score) and math.isnan(b.score)) or + math.isclose(a.score, b.score, rel_tol=1e-6, abs_tol=1e-6))): return False if a.label != b.label: @@ -246,8 +249,11 @@ class TestData: ]) def test_one_detector(self, input, expected): detector = ZScore(features=["x1"], model_id="zscore_x1") - with TestPipeline() as p: - result = (p | beam.Create(input) | AnomalyDetection(detector)) + with TestPipeline( + additional_pipeline_args=["--experiments=prism_disable_sdf_split" + ]) as p: + result = ( + p | beam.Create(input, reshuffle=False) | AnomalyDetection(detector)) if isinstance(input[0], tuple): assert_that(result, equal_to(expected, _keyed_result_is_equal_to)) @@ -271,9 +277,11 @@ def test_multiple_detectors_without_aggregation(self, input, expected): threshold_criterion=FixedThreshold(2), model_id="zscore_x2")) - with beam.Pipeline() as p: + with TestPipeline( + additional_pipeline_args=["--experiments=prism_disable_sdf_split" + ]) as p: result = ( - p | beam.Create(input) + p | beam.Create(input, reshuffle=False) | AnomalyDetection(EnsembleAnomalyDetector(sub_detectors))) if isinstance(input[0], tuple): @@ -298,9 +306,11 @@ def test_multiple_sub_detectors_with_aggregation(self, input, expected): threshold_criterion=FixedThreshold(2), model_id="zscore_x2")) - with beam.Pipeline() as p: + with TestPipeline( + additional_pipeline_args=["--experiments=prism_disable_sdf_split" + ]) as p: result = ( - p | beam.Create(input) + p | beam.Create(input, reshuffle=False) | AnomalyDetection( EnsembleAnomalyDetector( sub_detectors, aggregation_strategy=AnyVote()))) @@ -328,9 +338,11 @@ def test_ensemble_preserves_model_state(self, keyed): for model_id in ('x1', 'x2') for label in [-2, -2] + [0] * (count - 2)] - with TestPipeline() as p: + with TestPipeline( + additional_pipeline_args=["--experiments=prism_disable_sdf_split" + ]) as p: result = ( - p | beam.Create(input) + p | beam.Create(input, reshuffle=False) | AnomalyDetection(EnsembleAnomalyDetector(detectors))) if not keyed: result = result | beam.WithKeys(0) @@ -648,10 +660,12 @@ def test_dofn_on_single_prediction(self): AnomalyResult( R, [AnomalyPrediction(score=30, label=1, threshold=20)]))), ] - with TestPipeline() as p: + with TestPipeline( + additional_pipeline_args=["--experiments=prism_disable_sdf_split" + ]) as p: result = ( p - | beam.Create(input) + | beam.Create(input, reshuffle=False) # use median just for test convenience | beam.ParDo( _StatefulThresholdDoFn( diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index becfaa793032..dd33a05e0665 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -107,6 +107,7 @@ def get_impl(self): from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.utils import python_callable +from apache_beam.utils import retry from apache_beam.yaml import yaml_provider from apache_beam.yaml import yaml_transform from apache_beam.yaml.conftest import yaml_test_files_dir @@ -411,6 +412,15 @@ def temp_mongodb_table(): mongo_container.start() mongo_uri = mongo_container.get_connection_url() + # MongoDbContainer's entrypoint restarts mongod after init scripts; wait + # until the server is stably accepting connections and responds to ping. + @retry.with_exponential_backoff(num_retries=5, initial_delay_secs=1) + def _wait_for_mongo(): + mongo_client = mongo_container.get_connection_client() + mongo_client.admin.command('ping') + + _wait_for_mongo() + db_name = f'db_{uuid.uuid4().hex}' collection_name = f'collection_{uuid.uuid4().hex}' diff --git a/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py b/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py index 4a87a3466689..f947a4720054 100644 --- a/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py +++ b/sdks/python/apache_beam/yaml/test_utils/datadog_test_utils.py @@ -116,16 +116,14 @@ def temp_fake_datadog_server(expected_records=None): logging.error( "Error interacting with temporary fake Datadog server: %s", err) raise err - finally: - if expected_records is not None: - canonicalize = lambda rec: json.dumps(rec, sort_keys=True) + if expected_records is not None: + canonicalize = lambda rec: json.dumps(rec, sort_keys=True) - actual_strs = sorted([canonicalize(r) for r in received]) - expected_strs = sorted([canonicalize(e) for e in expected_records]) + actual_strs = sorted(set(canonicalize(r) for r in received)) + expected_strs = sorted(set(canonicalize(e) for e in expected_records)) - assert actual_strs == expected_strs, ( - f"Mismatch in recorded Datadog events!\n" - f"Expected: {expected_strs}\n" - f"Actual: {actual_strs}" - ) + assert actual_strs == expected_strs, ( + f"Mismatch in recorded Datadog events!\n" + f"Expected: {expected_strs}\n" + f"Actual: {actual_strs}")