Break up OperationCounters.update() into before and after pieces

Split method update() into update_from(), called before iteration,
and update_collect(), called after iteration.  This will allow us to
attach an observer and estimate the size during the iteration.

Correspondingly, ReceiverSet.update_counters() gets split into
update_counters_start() and update_counters_finish().
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=119995822

Author: gildea

google/cloud/dataflow/worker/executor.py

@@ -66,26 +66,36 @@ def start(self, step_name):
         self.counter_factory, step_name, self.coder, self.output_index)
 
   def output(self, windowed_value):
-    self.update_counters(windowed_value)
+    self.update_counters_start(windowed_value)
     for receiver in self.receivers:
       receiver.process(windowed_value)
+    self.update_counters_finish()
 
-  def update_counters(self, windowed_value):
+  def update_counters_start(self, windowed_value):
     if self.opcounter:
-      self.opcounter.update(windowed_value)
+      self.opcounter.update_from(windowed_value)
+
+  def update_counters_finish(self):
+    if self.opcounter:
+      self.opcounter.update_collect()
 
   def __str__(self):
     return '[%s]' % ' '.join([r.str_internal(is_recursive=True)
                               for r in self.receivers])
 
+  def __repr__(self):
+    return '<%s %d %s [%s]>' % (
+        self.__class__.__name__,
+        self.output_index,
+        self.coder,
+        ' '.join([r.str_internal(is_recursive=True) for r in self.receivers]))
+
 
 class Operation(object):
   """An operation representing the live version of a work item specification.
 
   An operation can have one or more outputs and for each output it can have
   one or more receiver operations that will take that as input.
-  TODO(gildea): Refactor "receivers[OUTPUT][RECEIVER]" as
-  "outputs[INDEX][RECEIVER]"
   """
 
   def __init__(self, spec, counter_factory):
@@ -238,11 +248,12 @@ def process(self, o):
     if self.debug_logging_enabled:
       logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.receivers[0].update_counters(o)
+    self.receivers[0].update_counters_start(o)
     if self.use_windowed_value:
       self.writer.Write(o)
     else:
       self.writer.Write(o.value)
+    self.receivers[0].update_counters_finish()
 
 
 class InMemoryWriteOperation(Operation):
@@ -256,8 +267,9 @@ def process(self, o):
     if self.debug_logging_enabled:
       logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.receivers[0].update_counters(o)
+    self.receivers[0].update_counters_start(o)
     self.spec.output_buffer.append(o.value)
+    self.receivers[0].update_counters_finish()
 
 
 class GroupedShuffleReadOperation(Operation):
@@ -352,7 +364,7 @@ def process(self, o):
     if self.debug_logging_enabled:
       logging.debug('Processing [%s] in %s', o, self)
     assert isinstance(o, WindowedValue)
-    self.receivers[0].update_counters(o)
+    self.receivers[0].update_counters_start(o)
     # We typically write into shuffle key/value pairs. This is the reason why
     # the else branch below expects the value attribute of the WindowedValue
     # argument to be a KV pair. However the service may write to shuffle in
@@ -370,6 +382,7 @@ def process(self, o):
     # TODO(silviuc): Use timestamps for the secondary key to get values in
     # times-sorted order.
     self.writer.Write(k, '', v)
+    self.receivers[0].update_counters_finish()
 
 
 class DoOperation(Operation):

google/cloud/dataflow/worker/opcounters.py

@@ -16,6 +16,7 @@
 
 from __future__ import absolute_import
 
+from google.cloud.dataflow.utils.counters import Accumulator
 from google.cloud.dataflow.utils.counters import Counter
 
 
@@ -28,16 +29,28 @@ def __init__(self, counter_factory, step_name, coder, output_index):
     self.mean_byte_counter = counter_factory.get_counter(
         '%s-out%d-MeanByteCount' % (step_name, output_index), Counter.MEAN)
     self.coder = coder
+    self._active_accumulators = []
 
-  def update(self, windowed_value):
+  def update_from(self, windowed_value):
     """Add one value to this counter."""
     self.element_counter.update(1)
+    byte_size_accumulator = Accumulator(self.mean_byte_counter.name)
+    self._active_accumulators.append(byte_size_accumulator)
     # TODO(gildea):
     # Actually compute the encoded size of this value.
     # In spirit, something like this:
-    #     size = len(self.coder.encode(windowed_value))
-    #     self.mean_byte_counter.update(size)
-    # but will need to handle streams and do sampling.
+    #     self.coder.store_estimated_size(windowed_value, byte_size_accumulator)
+    # but will need to do sampling.
+
+  def update_collect(self):
+    """Collects the accumulated size estimates.
+
+    Now that the element has been processed, we ask our accumulator
+    for the total and store the result in a counter.
+    """
+    for pending in self._active_accumulators:
+      self.mean_byte_counter.update(pending.total)
+    self._active_accumulators = []
 
   def __str__(self):
     return '<%s [%s]>' % (self.__class__.__name__,

google/cloud/dataflow/worker/opcounters_test.py

@@ -49,22 +49,25 @@ def test_update_int(self):
     opcounts = OperationCounters(CounterFactory(), 'some-name',
                                  coders.PickleCoder(), 0)
     self.verify_counters(opcounts, 0)
-    opcounts.update(GlobalWindows.WindowedValue(1))
+    opcounts.update_from(GlobalWindows.WindowedValue(1))
+    opcounts.update_collect()
     self.verify_counters(opcounts, 1)
 
   def test_update_str(self):
     opcounts = OperationCounters(CounterFactory(), 'some-name',
                                  coders.PickleCoder(), 0)
     self.verify_counters(opcounts, 0)
-    opcounts.update(GlobalWindows.WindowedValue('abcde'))
+    opcounts.update_from(GlobalWindows.WindowedValue('abcde'))
+    opcounts.update_collect()
     self.verify_counters(opcounts, 1)
 
   def test_update_old_object(self):
     opcounts = OperationCounters(CounterFactory(), 'some-name',
                                  coders.PickleCoder(), 0)
     self.verify_counters(opcounts, 0)
     obj = OldClassThatDoesNotImplementLen()
-    opcounts.update(GlobalWindows.WindowedValue(obj))
+    opcounts.update_from(GlobalWindows.WindowedValue(obj))
+    opcounts.update_collect()
     self.verify_counters(opcounts, 1)
 
   def test_update_new_object(self):
@@ -73,16 +76,21 @@ def test_update_new_object(self):
     self.verify_counters(opcounts, 0)
 
     obj = ObjectThatDoesNotImplementLen()
-    opcounts.update(GlobalWindows.WindowedValue(obj))
+    opcounts.update_from(GlobalWindows.WindowedValue(obj))
+    opcounts.update_collect()
     self.verify_counters(opcounts, 1)
 
   def test_update_multiple(self):
     opcounts = OperationCounters(CounterFactory(), 'some-name',
                                  coders.PickleCoder(), 0)
     self.verify_counters(opcounts, 0)
-    opcounts.update(GlobalWindows.WindowedValue('abcde'))
-    opcounts.update(GlobalWindows.WindowedValue('defghij'))
+    opcounts.update_from(GlobalWindows.WindowedValue('abcde'))
+    opcounts.update_from(GlobalWindows.WindowedValue('defghij'))
+    opcounts.update_collect()
     self.verify_counters(opcounts, 2)
+    opcounts.update_from(GlobalWindows.WindowedValue('klmnop'))
+    opcounts.update_collect()
+    self.verify_counters(opcounts, 3)
 
 
 if __name__ == '__main__':