Merge pull request #31 from implement basic timer plumbing

implement basic timer plumbing

Author: pabloem

ray_beam_runner/portability/execution.py

@@ -116,8 +116,8 @@ def ray_execute_bundle(
     ):
         if isinstance(output, beam_fn_api_pb2.Elements.Timers) and not dry_run:
             output_buffers[
-                expected_outputs[(output.transform_id, output.timer_family_id)]
-            ].append(output.data)
+                stage_timers[(output.transform_id, output.timer_family_id)]
+            ].append(output.timers)
         if isinstance(output, beam_fn_api_pb2.Elements.Data) and not dry_run:
             output_buffers[expected_outputs[output.transform_id]].append(output.data)
 
@@ -342,6 +342,9 @@ def put(self, pcoll, data_refs: List[ray.ObjectRef]):
     def get(self, pcoll) -> List[ray.ObjectRef]:
         return self.buffers[pcoll]
 
+    def clear(self, pcoll):
+        self.buffers[pcoll].clear()
+
 
 @ray.remote
 class RayWatermarkManager(watermark_manager.WatermarkManager):
@@ -450,6 +453,7 @@ def _build_timer_coders_id_map(self):
                     timer_coder_ids[
                         (transform_id, id)
                     ] = timer_family_spec.timer_family_coder_id
+        return timer_coder_ids
 
     def __reduce__(self):
         # We need to implement custom serialization for this particular class

ray_beam_runner/portability/ray_fn_runner.py

@@ -29,6 +29,7 @@
 from typing import Tuple
 from typing import Union
 
+from apache_beam.coders.coder_impl import create_OutputStream
 from apache_beam.options import pipeline_options
 from apache_beam.options.value_provider import RuntimeValueProvider
 from apache_beam.pipeline import Pipeline
@@ -42,7 +43,7 @@
 from apache_beam.runners.portability.fn_api_runner import translations
 from apache_beam.runners.portability.fn_api_runner.execution import ListBuffer
 from apache_beam.transforms import environments
-from apache_beam.utils import proto_utils
+from apache_beam.utils import proto_utils, timestamp
 
 import ray
 from ray_beam_runner.portability.context_management import RayBundleContextManager
@@ -334,6 +335,11 @@ def _run_bundle(
         )
         result = beam_fn_api_pb2.InstructionResponse.FromString(result_str)
 
+        (
+            watermarks_by_transform_and_timer_family,
+            newly_set_timers,
+        ) = self._collect_written_timers(bundle_context_manager)
+
         # TODO(pabloem): Add support for splitting of results.
 
         # After collecting deferred inputs, we 'pad' the structure with empty
@@ -346,9 +352,84 @@ def _run_bundle(
         #           coder_impl=bundle_context_manager.get_input_coder_impl(
         #               other_input))
 
-        newly_set_timers = {}
         return result, newly_set_timers, delayed_applications, output
 
+    @staticmethod
+    def _collect_written_timers(
+        bundle_context_manager: RayBundleContextManager,
+    ) -> Tuple[
+        Dict[translations.TimerFamilyId, timestamp.Timestamp],
+        Mapping[translations.TimerFamilyId, execution.PartitionableBuffer],
+    ]:
+        """Review output buffers, and collect written timers.
+        This function reviews a stage that has just been run. The stage will have
+        written timers to its output buffers. The function then takes the timers,
+        and adds them to the `newly_set_timers` dictionary, and the
+        timer_watermark_data dictionary.
+        The function then returns the following two elements in a tuple:
+        - timer_watermark_data: A dictionary mapping timer family to upcoming
+            timestamp to fire.
+        - newly_set_timers: A dictionary mapping timer family to timer buffers
+            to be passed to the SDK upon firing.
+        """
+        timer_watermark_data = {}
+        newly_set_timers = {}
+
+        execution_context = bundle_context_manager.execution_context
+        buffer_manager = execution_context.pcollection_buffers
+
+        for (
+            transform_id,
+            timer_family_id,
+        ), buffer_id in bundle_context_manager.stage_timers.items():
+            timer_buffer = ray.get(buffer_manager.get.remote(buffer_id))
+
+            coder_id = bundle_context_manager._timer_coder_ids[
+                (transform_id, timer_family_id)
+            ]
+
+            coder = execution_context.pipeline_context.coders[coder_id]
+            timer_coder_impl = coder.get_impl()
+
+            timers_by_key_tag_and_window = {}
+            if len(timer_buffer) >= 1:
+                written_timers = ray.get(timer_buffer[0])
+                # clear the timer buffer
+                buffer_manager.clear.remote(buffer_id)
+
+                # deduplicate updates to the same timer
+                for elements_timers in written_timers:
+                    for decoded_timer in timer_coder_impl.decode_all(elements_timers):
+                        key_tag_win = (
+                            decoded_timer.user_key,
+                            decoded_timer.dynamic_timer_tag,
+                            decoded_timer.windows[0],
+                        )
+                        if not decoded_timer.clear_bit:
+                            timers_by_key_tag_and_window[key_tag_win] = decoded_timer
+                        elif (
+                            decoded_timer.clear_bit
+                            and key_tag_win in timers_by_key_tag_and_window
+                        ):
+                            del timers_by_key_tag_and_window[key_tag_win]
+            if not timers_by_key_tag_and_window:
+                continue
+
+            out = create_OutputStream()
+            for decoded_timer in timers_by_key_tag_and_window.values():
+                timer_coder_impl.encode_to_stream(decoded_timer, out, True)
+                timer_watermark_data[(transform_id, timer_family_id)] = min(
+                    timer_watermark_data.get(
+                        (transform_id, timer_family_id), timestamp.MAX_TIMESTAMP
+                    ),
+                    decoded_timer.hold_timestamp,
+                )
+
+            buf = ListBuffer(coder_impl=timer_coder_impl)
+            buf.append(out.get())
+            newly_set_timers[(transform_id, timer_family_id)] = buf
+        return timer_watermark_data, newly_set_timers
+
 
 class RayRunnerResult(runner.PipelineResult):
     def __init__(self, state):

ray_beam_runner/portability/ray_runner_test.py

@@ -433,7 +433,6 @@ def process_timer(self):
             # expected = [('fired', ts) for ts in (20, 200)]
             # assert_that(actual, equal_to(expected))
 
-    @unittest.skip("Timers not yet supported")
     def test_pardo_timers(self):
         timer_spec = userstate.TimerSpec("timer", userstate.TimeDomain.WATERMARK)
         state_spec = userstate.CombiningValueStateSpec("num_called", sum)
@@ -467,7 +466,6 @@ def process_timer(
             expected = [("fired", ts) for ts in (20, 200, 40, 400)]
             assert_that(actual, equal_to(expected))
 
-    @unittest.skip("Timers not yet supported")
     def test_pardo_timers_clear(self):
         timer_spec = userstate.TimerSpec("timer", userstate.TimeDomain.WATERMARK)
         clear_timer_spec = userstate.TimerSpec(
@@ -506,15 +504,12 @@ def process_clear_timer(self):
             expected = [("fired", ts) for ts in (20, 200)]
             assert_that(actual, equal_to(expected))
 
-    @unittest.skip("Timers not yet supported")
     def test_pardo_state_timers(self):
         self._run_pardo_state_timers(windowed=False)
 
-    @unittest.skip("Timers not yet supported")
     def test_pardo_state_timers_non_standard_coder(self):
         self._run_pardo_state_timers(windowed=False, key_type=Any)
 
-    @unittest.skip("Timers not yet supported")
     def test_windowed_pardo_state_timers(self):
         self._run_pardo_state_timers(windowed=True)
 
@@ -587,7 +582,6 @@ def is_buffered_correctly(actual):
 
             assert_that(actual, is_buffered_correctly)
 
-    @unittest.skip("Timers not yet supported")
     def test_pardo_dynamic_timer(self):
         class DynamicTimerDoFn(beam.DoFn):
             dynamic_timer_spec = userstate.TimerSpec(