Cythonize runners.common, worker.executor, and utils.counters

Also optimize DoFnRunner for the (common) case of the main output.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=117261672

Author: robertwb

google/cloud/dataflow/runners/common.pxd

@@ -0,0 +1,30 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+cdef type SideOutputValue, TimestampedValue, WindowedValue
+
+cdef class DoFnRunner(object):
+
+  cdef object dofn
+  cdef object window_fn
+  cdef object context
+  cdef object tagged_receivers
+  cdef object tagged_counters
+  cdef object logger
+  cdef object step_name
+
+  cdef list main_receivers
+  cdef object main_counters
+
+  cpdef _process_outputs(self, element, results)

google/cloud/dataflow/runners/common.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# cython: profile=True
+
 """Worker operations executor."""
 
 import sys
@@ -72,6 +74,10 @@ def finish_bundle(self, context):
     self.logger = logger or FakeLogger()
     self.step_name = step_name
 
+    # Optimize for the common case.
+    self.main_receivers = tagged_receivers[None]
+    self.main_counters = tagged_counters[None]
+
   def start(self):
     self.context.set_element(None)
     try:
@@ -145,10 +151,15 @@ def _process_outputs(self, element, results):
       else:
         windowed_value = WindowedValue(
             result, element.timestamp, element.windows)
-      for receiver in self.tagged_receivers[tag]:
-        # TODO(robertwb): Should the counters be on the context?
+      # TODO(robertwb): Should the counters be on the context?
+      if tag is None:
+        self.main_counters.update(windowed_value)
+        for receiver in self.main_receivers:
+          receiver.process(windowed_value)
+      else:
         self.tagged_counters[tag].update(windowed_value)
-        receiver.process(windowed_value)
+        for receiver in self.tagged_receivers[tag]:
+          receiver.process(windowed_value)
 
 class NoContext(WindowFn.AssignContext):
   """An uninspectable WindowFn.AssignContext."""
@@ -158,7 +169,7 @@ def __init__(self, value, timestamp=NO_VALUE):
     self._timestamp = timestamp
   @property
   def timestamp(self):
-    if self._timestamp is NO_VALUE:
+    if self._timestamp is self.NO_VALUE:
       raise ValueError('No timestamp in this context.')
     else:
       return self._timestamp

google/cloud/dataflow/utils/counters.pxd

@@ -0,0 +1,25 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# cython: profile=True
+# cython: overflowcheck=True
+
+cdef class Counter(object):
+    cdef readonly object name
+    cdef readonly object aggregation_kind
+    cdef long c_total
+    cdef object py_total
+    cdef readonly long elements
+
+    cdef _update_small(self, long delta)

google/cloud/dataflow/utils/counters.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# cython: profile=False
+# cython: overflowcheck=True
+
 """Counters collect the progress of the Worker for reporting to the service."""
 
 
@@ -61,13 +64,26 @@ def __init__(self, name, aggregation_kind):
     """
     self.name = name
     self.aggregation_kind = aggregation_kind
-    self.total = 0
+    assert aggregation_kind == self.SUM  # update only handles sum
+    self.c_total = 0
+    self.py_total = 0
     self.elements = 0
 
   def update(self, count):
-    self.total += count
+    try:
+      self._update_small(count)
+    except OverflowError:
+      self.py_total += count
     self.elements += 1
 
+  def _update_small(self, delta):
+    new_total = self.c_total + delta  # overflow is checked
+    self.c_total = new_total
+
+  @property
+  def total(self):
+    return self.c_total + self.py_total
+
   def __str__(self):
     return '<%s>' % self._str_internal()
 

google/cloud/dataflow/worker/executor.pxd

@@ -0,0 +1,32 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+cdef class Operation(object):
+  cdef public spec
+  cdef public receivers
+  cdef public counters
+
+  cdef public step_name  # initialized lazily
+
+  cpdef start(self)
+  cpdef process(self, windowed_value)
+  cpdef finish(self)
+
+cdef class PGBKCVOperation(Operation):
+  cdef public object combine_fn
+  cdef dict table
+  cdef long max_keys
+  cdef long key_count
+
+  cpdef output(self, tuple wkey, value)

google/cloud/dataflow/worker/executor.py

@@ -554,7 +554,7 @@ def process(self, wkv):
   def finish(self):
     for wkey, value in self.table.iteritems():
       self.output(wkey, value[0])
-    self.entries = {}
+    self.table = {}
     self.key_count = 0
 
   def output(self, wkey, value):
@@ -815,4 +815,4 @@ def execute(self, map_task, test_shuffle_source=None, test_shuffle_sink=None):
       logging.debug('Starting op %d %s', ix, op)
       op.start()
     for op in self._ops:
-      op.finish()
+      op.finish(*())

setup.py

@@ -105,8 +105,13 @@ def get_download_url():
     entry_points={
         'console_scripts': CONSOLE_SCRIPTS,
         },
-    ext_modules=cythonize(
-        ['**/*.pyx', 'google/cloud/dataflow/coders/coder_impl.py']),
+    ext_modules=cythonize([
+        '**/*.pyx',
+        'google/cloud/dataflow/coders/coder_impl.py',
+        'google/cloud/dataflow/runners/common.py',
+        'google/cloud/dataflow/worker/executor.py',
+        'google/cloud/dataflow/utils/counters.py',
+    ]),
     setup_requires=['nose>=1.0'],
     install_requires=REQUIRED_PACKAGES,
     test_suite='nose.collector',