Clear cached args only when finishing bundle

Author: scwhittle

sdks/python/apache_beam/runners/common.pxd

@@ -100,9 +100,11 @@ cdef class PerWindowInvoker(DoFnInvoker):
   cdef dict kwargs_for_process_batch
   cdef list placeholders_for_process_batch
   cdef bint has_windowed_inputs
-  cdef bint recalculate_window_args
-  cdef bint has_cached_window_args
-  cdef bint has_cached_window_batch_args
+  cdef bint should_cache_args
+  cdef list cached_args_for_process
+  cdef dict cached_kwargs_for_process
+  cdef list cached_args_for_process_batch
+  cdef dict cached_kwargs_for_process_batch
   cdef object process_method
   cdef object process_batch_method
   cdef bint is_splittable

sdks/python/apache_beam/runners/common.py

@@ -761,16 +761,15 @@ def __init__(self,
       self.current_window_index = None
       self.stop_window_index = None
 
-    # TODO(https://github.com/apache/beam/issues/28776): Remove caching after
-    # fully rolling out.
-    # If true, always recalculate window args. If false, has_cached_window_args
-    # and has_cached_window_batch_args will be set to true if the corresponding
-    # self.args_for_process,have been updated and should be reused directly.
-    self.recalculate_window_args = (
-        self.has_windowed_inputs or 'disable_global_windowed_args_caching' in
-        RuntimeValueProvider.experiments)
-    self.has_cached_window_args = False
-    self.has_cached_window_batch_args = False
+    # If true, after the first process invocation the args for process will
+    # be cached in cached_args_for_process and cached_kwargs_for_process and
+    # reused on subsequent invocations in the same bundle..
+    self.should_cache_args = (not self.has_windowed_inputs)
+    self.cached_args_for_process = None
+    self.cached_kwargs_for_process = None
+    # See above, similar cached args for process_batch invocations.
+    self.cached_args_for_process_batch = None
+    self.cached_kwargs_for_process_batch = None
 
     # Try to prepare all the arguments that can just be filled in
     # without any additional work. in the process function.
@@ -932,9 +931,9 @@ def _invoke_process_per_window(self,
                                  additional_kwargs,
                                 ):
     # type: (...) -> Optional[SplitResultResidual]
-    if self.has_cached_window_args:
+    if self.cached_args_for_process:
       args_for_process, kwargs_for_process = (
-          self.args_for_process, self.kwargs_for_process)
+          self.cached_args_for_process, self.cached_kwargs_for_process)
     else:
       if self.has_windowed_inputs:
         assert len(windowed_value.windows) <= 1
@@ -945,10 +944,9 @@ def _invoke_process_per_window(self,
       side_inputs.extend(additional_args)
       args_for_process, kwargs_for_process = util.insert_values_in_args(
           self.args_for_process, self.kwargs_for_process, side_inputs)
-      if not self.recalculate_window_args:
-        self.args_for_process, self.kwargs_for_process = (
+      if self.should_cache_args:
+        self.cached_args_for_process, self.cached_kwargs_for_process = (
             args_for_process, kwargs_for_process)
-        self.has_cached_window_args = True
 
     # Extract key in the case of a stateful DoFn. Note that in the case of a
     # stateful DoFn, we set during __init__ self.has_windowed_inputs to be
@@ -1030,9 +1028,10 @@ def _invoke_process_batch_per_window(
   ):
     # type: (...) -> Optional[SplitResultResidual]
 
-    if self.has_cached_window_batch_args:
+    if self.cached_args_for_process_batch:
       args_for_process_batch, kwargs_for_process_batch = (
-          self.args_for_process_batch, self.kwargs_for_process_batch)
+          self.cached_args_for_process_batch,
+          self.cached_kwargs_for_process_batch)
     else:
       if self.has_windowed_inputs:
         assert isinstance(windowed_batch, HomogeneousWindowedBatch)
@@ -1049,10 +1048,9 @@ def _invoke_process_batch_per_window(
               side_inputs,
           )
       )
-      if not self.recalculate_window_args:
-        self.args_for_process_batch, self.kwargs_for_process_batch = (
-            args_for_process_batch, kwargs_for_process_batch)
-        self.has_cached_window_batch_args = True
+      if self.should_cache_args:
+        self.cached_args_for_process_batch = args_for_process_batch
+        self.cached_kwargs_for_process_batch = kwargs_for_process_batch
 
     for i, p in self.placeholders_for_process_batch:
       if core.DoFn.ElementParam == p:
@@ -1088,6 +1086,18 @@ def _invoke_process_batch_per_window(
             *args_for_process_batch, **kwargs_for_process_batch),
         self.threadsafe_watermark_estimator)
 
+  def invoke_finish_bundle(self):
+    # type: () -> None
+    # Clear the cached args to allow for refreshing of side inputs
+    # across bundles.
+    self.cached_args_for_process = None
+    self.cached_kwargs_for_process = None
+    self.cached_args_for_process_batch = None
+    self.cached_kwargs_for_process_batch = None
+    # super() doesn't appear to work with cython
+    # https://github.com/cython/cython/issues/3726
+    DoFnInvoker.invoke_finish_bundle(self)
+
   @staticmethod
   def _try_split(fraction,
       window_index, # type: Optional[int]

sdks/python/apache_beam/tools/map_fn_microbenchmark.py

@@ -88,11 +88,6 @@ def run_benchmark(
     profile_filename_base=None,
 ):
   suite = [
-      utils.BenchmarkConfig(
-          map_with_fixed_window_side_input_pipeline,
-          starting_point * 1000,
-          num_runs,
-      ),
       utils.LinearRegressionBenchmarkConfig(
           map_pipeline, starting_point, num_elements_step, num_runs),
       utils.BenchmarkConfig(
@@ -101,7 +96,7 @@ def run_benchmark(
           num_runs,
       ),
       utils.BenchmarkConfig(
-          map_with_global_side_input_pipeline_uncached,
+          map_with_fixed_window_side_input_pipeline,
           starting_point * 1000,
           num_runs,
       ),