[Interactive Beam] Fix caching deadlock, wait race conditions, and stale graph in notebooks

sdks/python/apache_beam/runners/interactive/interactive_beam_test.py

@@ -853,12 +853,12 @@ def test_compute_dependency_wait_true(self):
       spy_wait.assert_called_with({pcoll2}, async_res2)
 
       # Let pcoll1 finish
-      async_res1.result(timeout=60)
+      async_res1.wait_for_completion()
       self.assertTrue(pcoll1 in self.env.computed_pcollections)
       self.assertFalse(self.env.is_pcollection_computing(pcoll1))
 
       # pcoll2 should now run and complete
-      async_res2.result(timeout=60)
+      async_res2.wait_for_completion()
       self.assertTrue(pcoll2 in self.env.computed_pcollections)
 
   @patch.object(ie.InteractiveEnvironment, 'is_pcollection_computing')

sdks/python/apache_beam/runners/interactive/recording_manager.py

@@ -86,6 +86,7 @@ def __init__(
             bar_style='info',
         ) if IS_IPYTHON else None)
     self._cancel_requested = False
+    self._completed_event = threading.Event()
 
     if IS_IPYTHON:
       self._cancel_button.on_click(self._cancel_clicked)
@@ -151,27 +152,33 @@ def exception(self, timeout=None):
     except TimeoutError:
       return None
 
-  def _on_done(self, future: Future):
-    self._env.unmark_pcollection_computing(self._pcolls)
-    self._recording_manager._async_computations.pop(self._display_id, None)
-
-    if future.cancelled():
-      self.update_display('Computation Cancelled.', 1.0)
-      return
+  def wait_for_completion(self):
+    self._completed_event.wait()
 
-    exc = future.exception()
-    if exc:
-      self.update_display(f'Error: {exc}', 1.0)
-      _LOGGER.error('Asynchronous computation failed: %s', exc, exc_info=exc)
-    else:
-      self.update_display('Computation Finished Successfully.', 1.0)
-      res = future.result()
-      if res and res.state == PipelineState.DONE:
-        self._env.mark_pcollection_computed(self._pcolls)
+  def _on_done(self, future: Future):
+    try:
+      if future.cancelled():
+        self.update_display('Computation Cancelled.', 1.0)
+        return
+
+      exc = future.exception()
+      if exc:
+        self.update_display(f'Error: {exc}', 1.0)
+        _LOGGER.error('Asynchronous computation failed: %s', exc, exc_info=exc)
       else:
-        _LOGGER.warning(
-            'Async computation finished but state is not DONE: %s',
-            res.state if res else 'Unknown')
+        self.update_display('Computation Finished Successfully.', 1.0)
+        res = future.result()
+        if res and res.state == PipelineState.DONE:
+          self._env.mark_pcollection_computed(self._pcolls)
+        else:
+          _LOGGER.warning(
+              'Async computation finished but state is not DONE: %s',
+              res.state if res else 'Unknown')
+    finally:
+      self._env.unmark_pcollection_computing(self._pcolls)
+      with self._recording_manager._lock:
+        self._recording_manager._async_computations.pop(self._display_id, None)
+      self._completed_event.set()
 
   def cancel(self):
     if self._future.done():
@@ -438,7 +445,7 @@ def __init__(
     self._executor = ThreadPoolExecutor(max_workers=os.cpu_count())
     self._env = ie.current_env()
     self._async_computations: dict[str, AsyncComputationResult] = {}
-    self._pipeline_graph = None
+    self._lock = threading.Lock()
 
   def _execute_pipeline_fragment(
       self,
@@ -694,7 +701,8 @@ def compute_async(
       future = Future()
       async_result = AsyncComputationResult(
           future, pcolls_to_compute, self.user_pipeline, self)
-      self._async_computations[async_result._display_id] = async_result
+      with self._lock:
+        self._async_computations[async_result._display_id] = async_result
       self._env.mark_pcollection_computing(pcolls_to_compute)
 
       def task():
@@ -710,19 +718,13 @@ def task():
       return async_result
 
   def _get_pipeline_graph(self):
-    """Lazily initializes and returns the PipelineGraph."""
-    if self._pipeline_graph is None:
-      try:
-        # Try to create the graph.
-        self._pipeline_graph = PipelineGraph(self.user_pipeline)
-      except (ImportError, NameError, AttributeError):
-        # If pydot is missing, PipelineGraph() might crash.
-        _LOGGER.warning(
-            "Could not create PipelineGraph (pydot missing?). " \
-            "Async features disabled."
-        )
-        self._pipeline_graph = None
-    return self._pipeline_graph
+    """Initializes and returns the PipelineGraph."""
+    try:
+      # Try to create the graph.
+      return PipelineGraph(self.user_pipeline)
+    except (ImportError, NameError, AttributeError):
+      # If pydot is missing, PipelineGraph() might crash.
+      return None
 
   def _get_pcoll_id_map(self):
     """Creates a map from PCollection object to its ID in the proto."""
@@ -800,11 +802,19 @@ def _wait_for_dependencies(
     """Waits for any dependencies of the given
     PCollections that are currently being computed."""
     dependencies = self._get_all_dependencies(pcolls)
+    if async_result is None:
+      pcolls_to_check = dependencies.union(pcolls)
+    else:
+      pcolls_to_check = dependencies
     computing_deps: dict[beam.pvalue.PCollection, AsyncComputationResult] = {}
 
-    for dep in dependencies:
-      if self._env.is_pcollection_computing(dep):
-        for comp in self._async_computations.values():
+    with self._lock:
+      async_computations_copy = list(self._async_computations.values())
+
+    for dep in pcolls_to_check:
+      is_computing = self._env.is_pcollection_computing(dep)
+      if is_computing:
+        for comp in async_computations_copy:
           if dep in comp._pcolls:
             computing_deps[dep] = comp
             break
@@ -820,17 +830,16 @@ def _wait_for_dependencies(
         len(computing_deps),
         computing_deps.keys())
 
-    futures_to_wait = list(
-        set(comp._future for comp in computing_deps.values()))
+    results_to_wait = list(set(comp for comp in computing_deps.values()))
 
     try:
-      for i, future in enumerate(futures_to_wait):
+      for i, comp in enumerate(results_to_wait):
         if async_result:
           async_result.update_display(
-              f'Waiting for dependency {i + 1}/{len(futures_to_wait)}...',
-              progress=0.05 + 0.05 * (i / len(futures_to_wait)),
+              f'Waiting for dependency {i + 1}/{len(results_to_wait)}...',
+              progress=0.05 + 0.05 * (i / len(results_to_wait)),
           )
-        future.result()
+        comp.wait_for_completion()
       if async_result:
         async_result.update_display('Dependencies finished.', progress=0.1)
       _LOGGER.info('Dependencies finished successfully.')
@@ -891,23 +900,32 @@ def record(
             'Cannot record because a dependency failed to compute'
             ' asynchronously.')
 
-      self._clear()
-
-      merged_options = pipeline_options.PipelineOptions(
-          **{
-              **self.user_pipeline.options.get_all_options(
-                  drop_default=True, retain_unknown_options=True),
-              **options.get_all_options(
-                  drop_default=True, retain_unknown_options=True)
-          }) if options else self.user_pipeline.options
-
-      cache_path = ie.current_env().options.cache_root
-      is_remote_run = cache_path and ie.current_env(
-      ).options.cache_root.startswith('gs://')
-      pf.PipelineFragment(
-          list(uncomputed_pcolls), merged_options,
-          runner=runner).run(blocking=is_remote_run)
-      result = ie.current_env().pipeline_result(self.user_pipeline)
+      # Recalculate uncomputed PCollections because some may have finished computing during the wait
+      computed_pcolls = set(
+          pcoll for pcoll in pcolls
+          if pcoll in ie.current_env().computed_pcollections)
+      uncomputed_pcolls = set(pcolls).difference(computed_pcolls)
+
+      if uncomputed_pcolls:
+        self._clear()
+
+        merged_options = pipeline_options.PipelineOptions(
+            **{
+                **self.user_pipeline.options.get_all_options(
+                    drop_default=True, retain_unknown_options=True),
+                **options.get_all_options(
+                    drop_default=True, retain_unknown_options=True)
+            }) if options else self.user_pipeline.options
+
+        cache_path = ie.current_env().options.cache_root
+        is_remote_run = cache_path and ie.current_env(
+        ).options.cache_root.startswith('gs://')
+        pf.PipelineFragment(
+            list(uncomputed_pcolls), merged_options,
+            runner=runner).run(blocking=is_remote_run)
+        result = ie.current_env().pipeline_result(self.user_pipeline)
+      else:
+        result = ie.current_env().pipeline_result(self.user_pipeline)
     else:
       result = None
 

sdks/python/apache_beam/runners/interactive/recording_manager_test.py

@@ -1062,9 +1062,83 @@ def test_wait_for_dependencies(self, mock_async_result_cls):
     ie.current_env().mark_pcollection_computing({p1})
 
     self.assertTrue(rm._wait_for_dependencies({p2}))
-    mock_future.result.assert_called_once()
+    mock_async_res.wait_for_completion.assert_called_once()
     ie.current_env().unmark_pcollection_computing({p1})
 
+  def test_wait_for_dependencies_async_result(self):
+    p = beam.Pipeline(InteractiveRunner())
+    pcoll = p | beam.Create([1])
+    rm = RecordingManager(p)
+
+    # Mock a background computation for pcoll
+    async_res = MagicMock(spec=AsyncComputationResult)
+    async_res._pcolls = {pcoll}
+    rm._async_computations['id1'] = async_res
+    ie.current_env().mark_pcollection_computing({pcoll})
+
+    # Case 1: Called from main thread (async_result=None)
+    # It should wait for pcoll's background job completion event
+    with patch.object(async_res, 'wait_for_completion') as mock_wait:
+      with patch.object(rm, '_get_all_dependencies', return_value=set()):
+        rm._wait_for_dependencies({pcoll}, async_result=None)
+        mock_wait.assert_called_once()
+
+    # Case 2: Called from background thread (async_result=async_res)
+    # It should NOT wait (prevents self-deadlock)
+    with patch.object(async_res, 'wait_for_completion') as mock_wait:
+      with patch.object(rm, '_get_all_dependencies', return_value=set()):
+        rm._wait_for_dependencies({pcoll}, async_result=async_res)
+        mock_wait.assert_not_called()
+
+  def test_record_recalculates_uncomputed_pcolls_after_wait(self):
+    from apache_beam.runners.interactive import pipeline_fragment as pf
+    p = beam.Pipeline(InteractiveRunner())
+    pcoll = p | beam.Create([1])
+    rm = RecordingManager(p)
+
+    # Mock wait_for_dependencies to simulate pcoll completing during the wait
+    def mock_wait(pcolls, async_result=None):
+      # Simulate pcoll completing by marking it computed
+      ie.current_env().mark_pcollection_computed({pcoll})
+      return True
+
+    with patch.object(rm, '_wait_for_dependencies', side_effect=mock_wait), \
+         patch.object(rm, '_clear') as mock_clear, \
+         patch.object(pf, 'PipelineFragment') as mock_fragment:
+
+      rm.record([pcoll], max_n=10, max_duration='inf')
+
+      # Since pcoll was marked computed during the wait,
+      # uncomputed_pcolls becomes empty.
+      # So _clear() and PipelineFragment should NOT be called.
+      mock_clear.assert_not_called()
+      mock_fragment.assert_not_called()
+
+  def test_get_pipeline_graph_not_cached(self):
+    p = beam.Pipeline(InteractiveRunner())
+    pcoll1 = p | 'Create1' >> beam.Create([1])
+    rm = RecordingManager(p)
+
+    # First call: graph is created
+    graph1 = rm._get_pipeline_graph()
+    self.assertIsNotNone(graph1)
+    self.assertEqual(graph1._pipeline_instrument.user_pipeline, p)
+
+    # Add a new transform to the pipeline
+    pcoll2 = pcoll1 | 'Map1' >> beam.Map(lambda x: x + 1)
+
+    # Second call: graph should be updated and contain the new Map1 transform
+    graph2 = rm._get_pipeline_graph()
+    self.assertIsNotNone(graph2)
+    self.assertIsNot(graph1, graph2)
+
+    # Verify that the new graph contains the newly added transform
+    transform_names = [
+        t.unique_name
+        for t in graph2._pipeline_proto.components.transforms.values()
+    ]
+    self.assertIn('Map1', transform_names)
+
 
 if __name__ == '__main__':
   unittest.main()