Accomodate addtional stage reporting monitoring data

Author: 13MaxG

sdks/python/apache_beam/runners/direct/direct_runner_test.py

@@ -255,7 +255,8 @@ def test_flatten_two(self):
       pb_b = (
           pc_first
           | f"{label}/MapB" >> beam.Map(lambda x, y: ("b", 2), y=pv_a)
-          #| f"{label}/Reshuffle" >> beam.Reshuffle()  # beam 2.38 works without Reshuffle here
+          #| f"{label}/Reshuffle" >> beam.Reshuffle()
+          # # beam 2.38 works without Reshuffle here
       )
 
       pc_c = ((pc_a, pb_b) | f"{label}/Flatten" >> beam.Flatten())

sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner_test.py

@@ -1446,7 +1446,7 @@ def test_group_by_key_with_empty_pcoll_elements(self):
 # the sampling counter.
 class FnApiRunnerMetricsTest(unittest.TestCase):
   def assert_has_counter(
-      self, mon_infos, urn, labels, value=None, ge_value=None):
+      self, mon_infos, urn, labels, value=None, ge_value=None, total_stages=1):
     # TODO(ajamato): Consider adding a matcher framework
     found = 0
     matches = []
@@ -1464,7 +1464,7 @@ def assert_has_counter(
     ge_value_str = {'ge_value': ge_value} if ge_value else ''
     value_str = {'value': value} if value else ''
     self.assertEqual(
-        1,
+        total_stages,
         found,
         "Found (%s, %s) Expected only 1 monitoring_info for %s." % (
             found,
@@ -1473,7 +1473,15 @@ def assert_has_counter(
         ))
 
   def assert_has_distribution(
-      self, mon_infos, urn, labels, sum=None, count=None, min=None, max=None):
+      self,
+      mon_infos,
+      urn,
+      labels,
+      sum=None,
+      count=None,
+      min=None,
+      max=None,
+      total_stages=1):
     # TODO(ajamato): Consider adding a matcher framework
     sum = _matcher_or_equal_to(sum)
     count = _matcher_or_equal_to(count)
@@ -1508,7 +1516,7 @@ def assert_has_distribution(
             increment = 0
         found += increment
     self.assertEqual(
-        1,
+        total_stages,
         found,
         "Found (%s) Expected only 1 monitoring_info for %s." % (
             found,
@@ -1597,15 +1605,17 @@ def process(self, element):
           counters,
           monitoring_infos.ELEMENT_COUNT_URN,
           labels,
-          num_source_elems)
+          num_source_elems,
+          total_stages=2)
       self.assert_has_distribution(
           counters,
           monitoring_infos.SAMPLED_BYTE_SIZE_URN,
           labels,
           min=hamcrest.greater_than(0),
           max=hamcrest.greater_than(0),
           sum=hamcrest.greater_than(0),
-          count=hamcrest.greater_than(0))
+          count=hamcrest.greater_than(0),
+          total_stages=2)
 
       # GenerateTwoOutputs, "SecondOutput" output.
       labels = {
@@ -1615,15 +1625,17 @@ def process(self, element):
           counters,
           monitoring_infos.ELEMENT_COUNT_URN,
           labels,
-          2 * num_source_elems)
+          2 * num_source_elems,
+          total_stages=2)
       self.assert_has_distribution(
           counters,
           monitoring_infos.SAMPLED_BYTE_SIZE_URN,
           labels,
           min=hamcrest.greater_than(0),
           max=hamcrest.greater_than(0),
           sum=hamcrest.greater_than(0),
-          count=hamcrest.greater_than(0))
+          count=hamcrest.greater_than(0),
+          total_stages=2)
 
       # GenerateTwoOutputs, "ThirdOutput" output.
       labels = {
@@ -1633,15 +1645,17 @@ def process(self, element):
           counters,
           monitoring_infos.ELEMENT_COUNT_URN,
           labels,
-          num_source_elems)
+          num_source_elems,
+          total_stages=2)
       self.assert_has_distribution(
           counters,
           monitoring_infos.SAMPLED_BYTE_SIZE_URN,
           labels,
           min=hamcrest.greater_than(0),
           max=hamcrest.greater_than(0),
           sum=hamcrest.greater_than(0),
-          count=hamcrest.greater_than(0))
+          count=hamcrest.greater_than(0),
+          total_stages=2)
 
       # Skipping other pcollections due to non-deterministic naming for multiple
       # outputs.
@@ -1653,15 +1667,17 @@ def process(self, element):
           counters,
           monitoring_infos.ELEMENT_COUNT_URN,
           labels,
-          4 * num_source_elems)
+          4 * num_source_elems,
+          total_stages=2)
       self.assert_has_distribution(
           counters,
           monitoring_infos.SAMPLED_BYTE_SIZE_URN,
           labels,
           min=hamcrest.greater_than(0),
           max=hamcrest.greater_than(0),
           sum=hamcrest.greater_than(0),
-          count=hamcrest.greater_than(0))
+          count=hamcrest.greater_than(0),
+          total_stages=2)
 
       # PassThrough, main output
       labels = {