Resolve DaskRunner side input OOM by iteratively computing delayed partition

sdks/python/apache_beam/runners/dask/dask_runner.py

@@ -120,6 +120,13 @@ def _add_argparse_args(cls, parser: argparse.ArgumentParser) -> None:
         default=None,
         help='The length of each `dask.Bag` partition. When unspecified, '
         'an educated guess is made.')
+    parser.add_argument(
+        '--dask_lazy_side_inputs',
+        dest='lazy_side_inputs',
+        action='store_true',
+        help='Opt-in to lazy side inputs evaluation to reduce memory usage, '
+        'which may cause performance issues or deadlocks depending on the '
+        'scheduler. Fallbacks to `list()` when false.')
 
 
 @dataclasses.dataclass
@@ -166,7 +173,8 @@ def metrics(self):
 class DaskRunner(BundleBasedDirectRunner):
   """Executes a pipeline on a Dask distributed client."""
   @staticmethod
-  def to_dask_bag_visitor(bag_kwargs=None) -> PipelineVisitor:
+  def to_dask_bag_visitor(
+      bag_kwargs=None, lazy_side_inputs: bool = False) -> PipelineVisitor:
     from dask import bag as db
 
     if bag_kwargs is None:
@@ -210,7 +218,10 @@ def visit_transform(self, transform_node: AppliedPTransform) -> None:
                 SideInputMap(
                     type(si),
                     si._view_options(),
-                    DaskBagWindowedIterator(si_asbag, si._window_mapping_fn)))
+                    DaskBagWindowedIterator(
+                        si_asbag,
+                        si._window_mapping_fn,
+                        lazy_side_inputs=lazy_side_inputs)))
 
           op_kws["side_inputs"] = bag_side_inputs
 
@@ -242,7 +253,10 @@ def run_pipeline(self, pipeline, options):
 
     pipeline.replace_all(dask_overrides())
 
-    dask_visitor = self.to_dask_bag_visitor(bag_kwargs)
+    lazy_side_inputs = getattr(
+        options.view_as(DaskOptions), 'lazy_side_inputs', False)
+
+    dask_visitor = self.to_dask_bag_visitor(bag_kwargs, lazy_side_inputs)
     pipeline.visit(dask_visitor)
     # The dictionary in this visitor keeps a mapping of every Beam
     # PTransform to the equivalent Bag operation. This is highly

sdks/python/apache_beam/runners/dask/dask_runner_test.py

@@ -157,6 +157,19 @@ def mult_by(x, y):
           | beam.Map(mult_by, beam.pvalue.AsSingleton(side)))
       assert_that(pcoll, equal_to([3]))
 
+  def test_map_with_lazy_positional_side_input(self):
+    def mult_by(x, y):
+      return x * y
+
+    options = PipelineOptions(['--dask_lazy_side_inputs'])
+    with test_pipeline.TestPipeline(runner=DaskRunner(), options=options) as p:
+      side = p | "side" >> beam.Create([3])
+      pcoll = (
+          p
+          | "main" >> beam.Create([1])
+          | beam.Map(mult_by, beam.pvalue.AsSingleton(side)))
+      assert_that(pcoll, equal_to([3]))
+
   def test_map_with_keyword_side_input(self):
     def mult_by(x, y):
       return x * y

sdks/python/apache_beam/runners/dask/transform_evaluator.py

@@ -87,13 +87,24 @@ class DaskBagWindowedIterator:
 
   bag: db.Bag
   window_fn: WindowFn
+  lazy_side_inputs: bool = False
 
   def __iter__(self):
-    # FIXME(cisaacstern): list() is likely inefficient, since it presumably
-    # materializes the full result before iterating over it. doing this for
-    # now as a proof-of-concept. can we can generate results incrementally?
-    for result in list(self.bag):
-      yield get_windowed_value(result, self.window_fn)
+    if not self.lazy_side_inputs:
+      for result in list(self.bag):
+        yield get_windowed_value(result, self.window_fn)
+    else:
+      try:
+        from dask.distributed import worker_client
+        with worker_client() as client:
+          for partition in self.bag.to_delayed():
+            # explicitly submit and wait to avoid deadlock
+            for result in client.compute(partition).result():
+              yield get_windowed_value(result, self.window_fn)
+      except (ImportError, ValueError):
+        for partition in self.bag.to_delayed():
+          for result in partition.compute():
+            yield get_windowed_value(result, self.window_fn)
 
 
 @dataclasses.dataclass