Clear cached args only when finishing bundle

Author: scwhittle

sdks/python/apache_beam/runners/common.pxd

@@ -102,9 +102,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

@@ -812,16 +812,15 @@ def __init__(
       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 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.
@@ -984,9 +983,9 @@ def _invoke_process_per_window(
       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
@@ -997,10 +996,9 @@ def _invoke_process_per_window(
       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
@@ -1088,9 +1086,9 @@ 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)
@@ -1107,10 +1105,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 = (
+      if self.should_cache_args:
+        self.cached_args_for_process_batch, self.cached_kwargs_for_process_batch = (
             args_for_process_batch, kwargs_for_process_batch)
-        self.has_cached_window_batch_args = True
 
     for i, p in self.placeholders_for_process_batch:
       if core.DoFn.ElementParam == p:
@@ -1150,6 +1147,15 @@ 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, self.cached_kwargs_for_process = (None, None)
+    self.cached_args_for_process_batch, self.cached_kwargs_for_process_batch = (
+        None, None)
+
+    super(PerWindowInvoker, self).invoke_finish_bundle()
+
   @staticmethod
   def _try_split(
       fraction,

sdks/python/apache_beam/tools/map_fn_microbenchmark.py

@@ -96,11 +96,6 @@ def run_benchmark(
     starting_point=1, num_runs=10, num_elements_step=100, verbose=True, 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
       ),
@@ -110,7 +105,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,
       ),