Add Python processing-time timers and clock injection

mariapython · mariapython · commit 6bcf84f2455b · 2017-12-14T14:13:47.000-08:00
diff --git a/sdks/python/apache_beam/runners/direct/clock.py b/sdks/python/apache_beam/runners/direct/clock.py
@@ -15,36 +15,35 @@
 # limitations under the License.
 #
 
-"""Clock implementations for real time processing and testing."""
+"""Clock implementations for real time processing and testing.
 
+For internal use only. No backwards compatibility guarantees.
+"""
 from __future__ import absolute_import
 
 import time
 
 
 class Clock(object):
-  """For internal use only; no backwards-compatibility guarantees."""
+  def current_time(self):
+    raise NotImplementedError()
 
-  def time(self):
-    """Returns the number of milliseconds since epoch."""
-    return int(time.time() * 1000)
+  def advance_time(self):
+    raise NotImplementedError()
 
 
-class MockClock(Clock):
-  """For internal use only; no backwards-compatibility guarantees.
+class RealClock(object):
+  def current_time(self):
+    return time.time()
 
-  Mock clock implementation for testing."""
 
-  def __init__(self, now_in_ms):
-    self._now_in_ms = now_in_ms
+class TestClock(object):
+  """Clock used for Testing"""
+  def __init__(self, current=0):
+    self._current = current
 
-  def time(self):
-    return self._now_in_ms
+  def current_time(self):
+    return self._current
 
-  def set_time(self, value_in_ms):
-    assert value_in_ms >= self._now_in_ms
-    self._now_in_ms = value_in_ms
-
-  def advance(self, duration_in_ms):
-    assert duration_in_ms >= 0
-    self._now_in_ms += duration_in_ms
+  def advance_time(self, advance_by):
+    self._current += advance_by
diff --git a/sdks/python/apache_beam/runners/direct/direct_runner.py b/sdks/python/apache_beam/runners/direct/direct_runner.py
@@ -36,6 +36,8 @@
 from apache_beam.options.value_provider import RuntimeValueProvider
 from apache_beam.pvalue import PCollection
 from apache_beam.runners.direct.bundle_factory import BundleFactory
+from apache_beam.runners.direct.clock import RealClock
+from apache_beam.runners.direct.clock import TestClock
 from apache_beam.runners.runner import PipelineResult
 from apache_beam.runners.runner import PipelineRunner
 from apache_beam.runners.runner import PipelineState
@@ -90,15 +92,14 @@ class DirectRunner(PipelineRunner):
   """Executes a single pipeline on the local machine."""
 
   # A list of PTransformOverride objects to be applied before running a pipeline
-  # using DirectRunner.
-  # Currently this only works for overrides where the input and output types do
-  # not change.
-  # For internal SDK use only. This should not be updated by Beam pipeline
-  # authors.
+  # using DirectRunner. Currently, this only works for overrides where the input
+  # and output types do not change.
+  # For internal use only; no backwards-compatibility guarantees.
   _PTRANSFORM_OVERRIDES = []
 
   def __init__(self):
     self._cache = None
+    self._use_test_clock = False  # use RealClock() in production
 
   def apply_CombinePerKey(self, transform, pcoll):
     # TODO: Move imports to top. Pipeline <-> Runner dependency cause problems
@@ -111,6 +112,10 @@ def apply_CombinePerKey(self, transform, pcoll):
     except NotImplementedError:
       return transform.expand(pcoll)
 
+  def apply_TestStream(self, transform, pcoll):
+    self._use_test_clock = True  # use TestClock() for testing
+    return transform.expand(pcoll)
+
   def apply__GroupByKeyOnly(self, transform, pcoll):
     if (transform.__class__ == _GroupByKeyOnly and
         pcoll.pipeline._options.view_as(StandardOptions).streaming):
@@ -204,14 +209,16 @@ def run(self, pipeline):
     self.consumer_tracking_visitor = ConsumerTrackingPipelineVisitor()
     pipeline.visit(self.consumer_tracking_visitor)
 
+    clock = TestClock() if self._use_test_clock else RealClock()
     evaluation_context = EvaluationContext(
         pipeline._options,
         BundleFactory(stacked=pipeline._options.view_as(DirectOptions)
                       .direct_runner_use_stacked_bundle),
         self.consumer_tracking_visitor.root_transforms,
         self.consumer_tracking_visitor.value_to_consumers,
         self.consumer_tracking_visitor.step_names,
-        self.consumer_tracking_visitor.views)
+        self.consumer_tracking_visitor.views,
+        clock)
 
     evaluation_context.use_pvalue_cache(self._cache)
 
diff --git a/sdks/python/apache_beam/runners/direct/evaluation_context.py b/sdks/python/apache_beam/runners/direct/evaluation_context.py
@@ -22,7 +22,6 @@
 import collections
 import threading
 
-from apache_beam.runners.direct.clock import Clock
 from apache_beam.runners.direct.direct_metrics import DirectMetrics
 from apache_beam.runners.direct.executor import TransformExecutor
 from apache_beam.runners.direct.watermark_manager import WatermarkManager
@@ -138,7 +137,7 @@ class EvaluationContext(object):
   """
 
   def __init__(self, pipeline_options, bundle_factory, root_transforms,
-               value_to_consumers, step_names, views):
+               value_to_consumers, step_names, views, clock):
     self.pipeline_options = pipeline_options
     self._bundle_factory = bundle_factory
     self._root_transforms = root_transforms
@@ -151,7 +150,7 @@ def __init__(self, pipeline_options, bundle_factory, root_transforms,
     self._transform_keyed_states = self._initialize_keyed_states(
         root_transforms, value_to_consumers)
     self._watermark_manager = WatermarkManager(
-        Clock(), root_transforms, value_to_consumers,
+        clock, root_transforms, value_to_consumers,
         self._transform_keyed_states)
     self._side_inputs_container = _SideInputsContainer(views)
     self._pending_unblocked_tasks = []
@@ -286,8 +285,8 @@ def create_empty_committed_bundle(self, output_pcollection):
     return self._bundle_factory.create_empty_committed_bundle(
         output_pcollection)
 
-  def extract_fired_timers(self):
-    return self._watermark_manager.extract_fired_timers()
+  def extract_all_timers(self):
+    return self._watermark_manager.extract_all_timers()
 
   def is_done(self, transform=None):
     """Checks completion of a step or the pipeline.
diff --git a/sdks/python/apache_beam/runners/direct/executor.py b/sdks/python/apache_beam/runners/direct/executor.py
@@ -545,22 +545,19 @@ def call(self):
           self._executor.executor_service.submit(self)
 
     def _should_shutdown(self):
-      """_should_shutdown checks whether pipeline is completed or not.
+      """Checks whether the pipeline is completed and should be shut down.
 
-      It will check for successful completion by checking the watermarks of all
-      transforms. If they all reached the maximum watermark it means that
-      pipeline successfully reached to completion.
+      If there is anything in the queue of tasks to do, do not shut down.
 
-      If the above is not true, it will check that at least one executor is
-      making progress. Otherwise pipeline will be declared stalled.
-
-      If the pipeline reached to a terminal state as explained above
-      _should_shutdown will request executor to gracefully shutdown.
+      Otherwise, check if all the transforms' watermarks are complete.
+      If they are not, the pipeline is not progressing (stall detected).
+      Whether the pipeline has stalled or not, the executor should shut
+      down the pipeline.
 
       Returns:
-        True if pipeline reached a terminal state and monitor task could finish.
-        Otherwise monitor task should schedule itself again for future
-        execution.
+        True only if the pipeline has reached a terminal state and should
+        be shut down.
+
       """
       if self._is_executing():
         # There are some bundles still in progress.
@@ -585,8 +582,8 @@ def _fire_timers(self):
       Returns:
         True if timers fired.
       """
-      transform_fired_timers = (
-          self._executor.evaluation_context.extract_fired_timers())
+      transform_fired_timers, _ = (
+          self._executor.evaluation_context.extract_all_timers())
       for applied_ptransform, fired_timers in transform_fired_timers:
         # Use an empty committed bundle. just to trigger.
         empty_bundle = (
@@ -602,7 +599,17 @@ def _fire_timers(self):
       return bool(transform_fired_timers)
 
     def _is_executing(self):
-      """Returns True if there is at least one non-blocked TransformExecutor."""
+      """Checks whether the job is still executing.
+
+      Returns:
+        True if there are any timers set or if there is at least
+        one non-blocked TransformExecutor active."""
+
+      watermark_manager = self._executor.evaluation_context._watermark_manager
+      _, any_unfired_timers = watermark_manager.extract_all_timers()
+      if any_unfired_timers:
+        return True
+
       executors = self._executor.transform_executor_services.executors
       if not executors:
         # Nothing is executing.
diff --git a/sdks/python/apache_beam/runners/direct/watermark_manager.py b/sdks/python/apache_beam/runners/direct/watermark_manager.py
@@ -39,7 +39,7 @@ class WatermarkManager(object):
 
   def __init__(self, clock, root_transforms, value_to_consumers,
                transform_keyed_states):
-    self._clock = clock  # processing time clock
+    self._clock = clock
     self._root_transforms = root_transforms
     self._value_to_consumers = value_to_consumers
     self._transform_keyed_states = transform_keyed_states
@@ -143,13 +143,17 @@ def _refresh_watermarks(self, applied_ptransform):
             for consumer in consumers:
               self._refresh_watermarks(consumer)
 
-  def extract_fired_timers(self):
+  def extract_all_timers(self):
+    """Extracts fired timers and reports of any timers set for all transforms."""
     all_timers = []
+    has_realtime_timer = False
     for applied_ptransform, tw in self._transform_to_watermarks.iteritems():
-      fired_timers = tw.extract_fired_timers()
+      fired_timers, had_realtime_timer = tw.extract_transform_timers()
       if fired_timers:
         all_timers.append((applied_ptransform, fired_timers))
-    return all_timers
+      if had_realtime_timer:
+        has_realtime_timer = True
+    return all_timers, has_realtime_timer
 
 
 class _TransformWatermarks(object):
@@ -244,19 +248,22 @@ def refresh(self):
 
   @property
   def synchronized_processing_output_time(self):
-    return self._clock.time()
+    return self._clock.current_time()
 
-  def extract_fired_timers(self):
+  def extract_transform_timers(self):
+    """Extracts fired timers and reports of any timers set per transform."""
     with self._lock:
-      if self._fired_timers:
-        return False
-
       fired_timers = []
+      has_realtime_timer = False
       for encoded_key, state in self._keyed_states.iteritems():
-        timers = state.get_timers(watermark=self._input_watermark)
+        timers, had_realtime_timer = state.get_timers(
+            watermark=self._input_watermark,
+            current_time=self._clock.current_time())
+        if had_realtime_timer:
+          has_realtime_timer = True
         for expired in timers:
           window, (name, time_domain, timestamp) = expired
           fired_timers.append(
               TimerFiring(encoded_key, window, name, time_domain, timestamp))
       self._fired_timers.update(fired_timers)
-      return fired_timers
+      return fired_timers, has_realtime_timer
diff --git a/sdks/python/apache_beam/transforms/trigger.py b/sdks/python/apache_beam/transforms/trigger.py
@@ -1138,20 +1138,36 @@ def clear_state(self, window, tag):
     if not self.state[window]:
       self.state.pop(window, None)
 
-  def get_timers(self, clear=False, watermark=MAX_TIMESTAMP):
+  def get_timers(self, clear=False, watermark=MAX_TIMESTAMP, current_time=None):
+    """Gets expired timers and reports if there
+    are any realtime timers set per state.
+
+    Expiration is measured against the watermark for event-time timers,
+    and against a wall clock for processing-time timers.
+    """
     expired = []
+    has_realtime_timer = False
     for window, timers in list(self.timers.items()):
       for (name, time_domain), timestamp in list(timers.items()):
-        if timestamp <= watermark:
+        if time_domain == 'REAL_TIME':
+          time_marker = current_time
+          has_realtime_timer = True
+        elif time_domain == 'WATERMARK':
+          time_marker = watermark
+        else:
+          logging.error(
+              'TimeDomain error: No timers defined for time domain %s.',
+              time_domain)
+        if timestamp <= time_marker:
           expired.append((window, (name, time_domain, timestamp)))
           if clear:
             del timers[(name, time_domain)]
       if not timers and clear:
         del self.timers[window]
-    return expired
+    return expired, has_realtime_timer
 
   def get_and_clear_timers(self, watermark=MAX_TIMESTAMP):
-    return self.get_timers(clear=True, watermark=watermark)
+    return self.get_timers(clear=True, watermark=watermark)[0]
 
   def get_earliest_hold(self):
     earliest_hold = MAX_TIMESTAMP
