[beam] branch master updated: Implicitly watch and track anonymous pipeline and PCollections

[yi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yichi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new e28ed48  Implicitly watch and track anonymous pipeline and PCollections
     new 1563588  Merge pull request #15551 from KevinGG/anonymous
e28ed48 is described below

commit e28ed48decaca52b6b1c2cc987695e13a027b1d8
Author: KevinGG <ka...@gmail.com>
AuthorDate: Wed Sep 22 11:25:14 2021 -0700

    Implicitly watch and track anonymous pipeline and PCollections
---
 .../apache_beam/runners/interactive/interactive_beam.py    | 13 ++++++++++++-
 .../runners/interactive/interactive_beam_test.py           |  1 -
 .../runners/interactive/interactive_environment.py         | 14 +++++++++-----
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/sdks/python/apache_beam/runners/interactive/interactive_beam.py b/sdks/python/apache_beam/runners/interactive/interactive_beam.py
index fc55bf1..3786193 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_beam.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_beam.py
@@ -471,7 +471,13 @@ def show(
   assert len(pcolls) > 0, (
       'Need at least 1 PCollection to show data visualization.')
 
-  user_pipeline = ie.current_env().user_pipeline(next(iter(pcolls)).pipeline)
+  pcoll_pipeline = next(iter(pcolls)).pipeline
+  user_pipeline = ie.current_env().user_pipeline(pcoll_pipeline)
+  # Possibly showing a PCollection defined in a local scope that is not
+  # explicitly watched. Ad hoc watch it though it's a little late.
+  if not user_pipeline:
+    watch({'anonymous_pipeline_{}'.format(id(pcoll_pipeline)): pcoll_pipeline})
+    user_pipeline = pcoll_pipeline
 
   if isinstance(n, str):
     assert n == 'inf', (
@@ -607,6 +613,11 @@ def collect(pcoll, n='inf', duration='inf', include_window_info=False):
     duration = float('inf')
 
   user_pipeline = ie.current_env().user_pipeline(pcoll.pipeline)
+  # Possibly collecting a PCollection defined in a local scope that is not
+  # explicitly watched. Ad hoc watch it though it's a little late.
+  if not user_pipeline:
+    watch({'anonymous_pipeline_{}'.format(id(pcoll.pipeline)): pcoll.pipeline})
+    user_pipeline = pcoll.pipeline
   recording_manager = ie.current_env().get_recording_manager(
       user_pipeline, create_if_absent=True)
 
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py b/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
index b093aa2..da948fd 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_beam_test.py
@@ -188,7 +188,6 @@ class InteractiveBeamTest(unittest.TestCase):
     p2 = beam.Pipeline(ir.InteractiveRunner())
 
     ib.watch(locals())
-    ie.current_env().track_user_pipelines()
 
     # Get the descriptions. This test is simple as there isn't much logic in the
     # method.
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_environment.py b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
index e22dee7..fe10ab1 100644
--- a/sdks/python/apache_beam/runners/interactive/interactive_environment.py
+++ b/sdks/python/apache_beam/runners/interactive/interactive_environment.py
@@ -286,16 +286,21 @@ class InteractiveEnvironment(object):
   def _track_user_pipelines(self, watchable):
     """Tracks user pipelines from the given watchable."""
 
+    pipelines = set()
     if isinstance(watchable, beam.Pipeline):
-      self._tracked_user_pipelines.add_user_pipeline(watchable)
+      pipelines.add(watchable)
     elif isinstance(watchable, dict):
       for v in watchable.values():
         if isinstance(v, beam.Pipeline):
-          self._tracked_user_pipelines.add_user_pipeline(v)
+          pipelines.add(v)
     elif isinstance(watchable, Iterable):
       for v in watchable:
         if isinstance(v, beam.Pipeline):
-          self._tracked_user_pipelines.add_user_pipeline(v)
+          pipelines.add(v)
+    for p in pipelines:
+      self._tracked_user_pipelines.add_user_pipeline(p)
+      _ = self.get_cache_manager(p, create_if_absent=True)
+      _ = self.get_recording_manager(p, create_if_absent=True)
 
   def watch(self, watchable):
     """Watches a watchable.
@@ -306,12 +311,11 @@ class InteractiveEnvironment(object):
     matter since they are different instances. Duplicated variables are also
     allowed when watching.
     """
-    self._track_user_pipelines(watchable)
-
     if isinstance(watchable, dict):
       self._watching_dict_list.append(watchable.items())
     else:
       self._watching_set.add(watchable)
+    self._track_user_pipelines(watchable)
 
   def watching(self):
     """Analyzes and returns a list of pair lists referring to variable names and