Add per_element_output_counter to python SDK

Boyuan Zhang · Boyuan Zhang · commit 4b2f31fa7902 · 2018-05-09T10:13:40.000-07:00
diff --git a/sdks/python/apache_beam/runners/common.pxd b/sdks/python/apache_beam/runners/common.pxd
@@ -19,6 +19,9 @@ cimport cython
 
 from apache_beam.utils.windowed_value cimport WindowedValue
 from apache_beam.metrics.execution cimport ScopedMetricsContainer
+from apache_beam.transforms.cy_dataflow_distribution_counter cimport DataflowDistributionCounter
+
+from libc.stdint cimport int64_t
 
 
 cdef type TaggedOutput, TimestampedValue
@@ -87,15 +90,16 @@ cdef class DoFnRunner(Receiver):
 
 
 cdef class OutputProcessor(object):
-  @cython.locals(windowed_value=WindowedValue)
+  @cython.locals(windowed_value=WindowedValue,
+                 output_element_count=int64_t)
   cpdef process_outputs(self, WindowedValue element, results)
 
 
 cdef class _OutputProcessor(OutputProcessor):
   cdef object window_fn
   cdef Receiver main_receivers
   cdef object tagged_receivers
-
+  cdef DataflowDistributionCounter per_element_output_counter
 
 cdef class DoFnContext(object):
   cdef object label
diff --git a/sdks/python/apache_beam/runners/common.py b/sdks/python/apache_beam/runners/common.py
@@ -28,13 +28,16 @@
 import six
 
 from apache_beam.internal import util
+from apache_beam.options.value_provider import RuntimeValueProvider
 from apache_beam.pvalue import TaggedOutput
 from apache_beam.transforms import DoFn
 from apache_beam.transforms import core
 from apache_beam.transforms.core import RestrictionProvider
 from apache_beam.transforms.window import GlobalWindow
 from apache_beam.transforms.window import TimestampedValue
 from apache_beam.transforms.window import WindowFn
+from apache_beam.utils.counters import Counter
+from apache_beam.utils.counters import CounterName
 from apache_beam.utils.windowed_value import WindowedValue
 
 
@@ -518,7 +521,8 @@ def __init__(self,
                step_name=None,
                logging_context=None,
                state=None,
-               scoped_metrics_container=None):
+               scoped_metrics_container=None,
+               operation_name=None):
     """Initializes a DoFnRunner.
 
     Args:
@@ -532,6 +536,7 @@ def __init__(self,
       logging_context: a LoggingContext object
       state: handle for accessing DoFn state
       scoped_metrics_container: Context switcher for metrics container
+      operation_name: The system name assigned by the runner for this operation.
     """
     # Need to support multiple iterations.
     side_inputs = list(side_inputs)
@@ -548,8 +553,22 @@ def __init__(self,
 
     # Optimize for the common case.
     main_receivers = tagged_receivers[None]
+
+    # TODO(BEAM-3937): Remove if block after output counter released.
+    experiments = RuntimeValueProvider.get_value('experiments', str, [])
+    # Experimental flag format: experimental-name_version-number.
+    if 'outputs_per_element_counter_v0' in experiments:
+      # TODO(BEAM-3955): Make step_name and operation_name less confused.
+      output_counter_name = (CounterName('per-element-output-count',
+                                         step_name=operation_name))
+      per_element_output_counter = state._counter_factory.get_counter(
+          output_counter_name, Counter.DATAFLOW_DISTRIBUTION).accumulator
+    else:
+      per_element_output_counter = None
+
     output_processor = _OutputProcessor(
-        windowing.windowfn, main_receivers, tagged_receivers)
+        windowing.windowfn, main_receivers, tagged_receivers,
+        per_element_output_counter)
 
     self.do_fn_invoker = DoFnInvoker.create_invoker(
         do_fn_signature, output_processor, self.context, side_inputs, args,
@@ -617,17 +636,24 @@ def process_outputs(self, windowed_input_element, results):
 class _OutputProcessor(OutputProcessor):
   """Processes output produced by DoFn method invocations."""
 
-  def __init__(self, window_fn, main_receivers, tagged_receivers):
+  def __init__(self,
+               window_fn,
+               main_receivers,
+               tagged_receivers,
+               per_element_output_counter):
     """Initializes ``_OutputProcessor``.
 
     Args:
       window_fn: a windowing function (WindowFn).
       main_receivers: a dict of tag name to Receiver objects.
       tagged_receivers: main receiver object.
+      per_element_output_counter: per_element_output_counter of one work_item.
+                                  could be none if experimental flag turn off
     """
     self.window_fn = window_fn
     self.main_receivers = main_receivers
     self.tagged_receivers = tagged_receivers
+    self.per_element_output_counter = per_element_output_counter
 
   def process_outputs(self, windowed_input_element, results):
     """Dispatch the result of process computation to the appropriate receivers.
@@ -636,9 +662,17 @@ def process_outputs(self, windowed_input_element, results):
     then dispatched to the appropriate indexed output.
     """
     if results is None:
+      # TODO(BEAM-3937): Remove if block after output counter released.
+      # Only enable per_element_output_counter when counter cythonized.
+      if (self.per_element_output_counter is not None and
+          self.per_element_output_counter.is_cythonized):
+        self.per_element_output_counter.add_input(0)
       return
 
+    output_element_count = 0
     for result in results:
+      # results here may be a generator, which cannot call len on it.
+      output_element_count += 1
       tag = None
       if isinstance(result, TaggedOutput):
         tag = result.tag
@@ -663,6 +697,11 @@ def process_outputs(self, windowed_input_element, results):
         self.main_receivers.receive(windowed_value)
       else:
         self.tagged_receivers[tag].receive(windowed_value)
+    # TODO(BEAM-3937): Remove if block after output counter released.
+    # Only enable per_element_output_counter when counter cythonized
+    if (self.per_element_output_counter is not None and
+        self.per_element_output_counter.is_cythonized):
+      self.per_element_output_counter.add_input(output_element_count)
 
   def start_bundle_outputs(self, results):
     """Validate that start_bundle does not output any elements"""
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
@@ -391,7 +391,8 @@ def start(self):
           logging_context=logger.PerThreadLoggingContext(
               step_name=self.name_context.logging_name()),
           state=state,
-          scoped_metrics_container=None)
+          scoped_metrics_container=None,
+          operation_name=self.name_context.metrics_name())
       self.dofn_receiver = (self.dofn_runner
                             if isinstance(self.dofn_runner, Receiver)
                             else DoFnRunnerReceiver(self.dofn_runner))
diff --git a/sdks/python/apache_beam/transforms/cy_dataflow_distribution_counter.pxd b/sdks/python/apache_beam/transforms/cy_dataflow_distribution_counter.pxd
@@ -42,3 +42,4 @@ cdef class DataflowDistributionCounter(object):
   cpdef void translate_to_histogram(self, histogram)
   cpdef bint add_inputs_for_test(self, elements) except -1
   cpdef int64_t calculate_bucket_index(self, int64_t element)
+  cpdef merge(self, accumulators)
diff --git a/sdks/python/apache_beam/transforms/cy_dataflow_distribution_counter.pyx b/sdks/python/apache_beam/transforms/cy_dataflow_distribution_counter.pyx
@@ -133,3 +133,6 @@ cdef class DataflowDistributionCounter(object):
     cdef calculate_bucket_index cannot be called directly from def.
     """
     return self._fast_calculate_bucket_index(element)
+
+  cpdef merge(self, accumulators):
+    pass
diff --git a/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py b/sdks/python/apache_beam/transforms/py_dataflow_distribution_counter.py
@@ -113,3 +113,6 @@ def translate_to_histogram(self, histogram):
     histogram.firstBucketOffset = first_bucket_offset
     histogram.bucketCounts = (
         self.buckets[first_bucket_offset:last_bucket_offset + 1])
+
+  def merge(self, accumulators):
+    pass
