[BEAM-10708] Clean up pipeline instrument (#14176)

Ning Kang · TheNeuralBit · web-flow · commit d5595449d14d · 2021-03-17T09:07:51.000-07:00
* [BEAM-10708] Clean up pipeline instrument 1. Moved cacheable and cachekey to dataclasses. 2. Cleaned up logic of finding the user defined pipeline. 3. Fixed missing add_derived_pipeline in background_caching_job module. * add python_version<3.7 to the dataclasses dep comment in setup * Add python_version<"3.7" to dataclasses requirement Co-authored-by: Brian Hulette <hulettbh@gmail.com>
diff --git a/sdks/python/apache_beam/runners/interactive/background_caching_job.py b/sdks/python/apache_beam/runners/interactive/background_caching_job.py
@@ -141,6 +141,7 @@ def attempt_to_run_background_caching_job(
     from apache_beam.runners.interactive import pipeline_instrument as instr
     runner_pipeline = beam.pipeline.Pipeline.from_runner_api(
         user_pipeline.to_runner_api(use_fake_coders=True), runner, options)
+    ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline)
     background_caching_job_result = beam.pipeline.Pipeline.from_runner_api(
         instr.build_pipeline_instrument(
             runner_pipeline).background_caching_pipeline_proto(),
diff --git a/sdks/python/apache_beam/runners/interactive/caching/cacheable.py b/sdks/python/apache_beam/runners/interactive/caching/cacheable.py
@@ -0,0 +1,78 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Module for dataclasses to hold metadata of cacheable PCollections in the user
+code scope.
+
+For internal use only; no backwards-compatibility guarantees.
+"""
+
+# pytype: skip-file
+
+from __future__ import absolute_import
+
+from dataclasses import dataclass
+
+import apache_beam as beam
+from apache_beam.runners.interactive.utils import obfuscate
+
+
+@dataclass
+class Cacheable:
+  pcoll_id: str
+  var: str
+  version: str
+  pcoll: beam.pvalue.PCollection
+  producer_version: str
+
+  def __hash__(self):
+    return hash((
+        self.pcoll_id,
+        self.var,
+        self.version,
+        self.pcoll,
+        self.producer_version))
+
+  def to_key(self):
+    return CacheKey(
+        self.var,
+        self.version,
+        self.producer_version,
+        str(id(self.pcoll.pipeline)))
+
+
+@dataclass
+class CacheKey:
+  var: str
+  version: str
+  producer_version: str
+  pipeline_id: str
+
+  def __post_init__(self):
+    # Normalize arbitrary variable name to a fixed length hex str.
+    self.var = obfuscate(self.var)[:10]
+
+  @staticmethod
+  def from_str(r):
+    r_split = r.split('-')
+    ck = CacheKey(*r_split)
+    ck.var = r_split[0]
+    return ck
+
+  def __repr__(self):
+    return '-'.join(
+        [self.var, self.version, self.producer_version, self.pipeline_id])
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
@@ -39,6 +39,7 @@
 from apache_beam.portability.api.beam_runner_api_pb2 import TestStreamPayload
 from apache_beam.runners.interactive.cache_manager import CacheManager
 from apache_beam.runners.interactive.cache_manager import SafeFastPrimitivesCoder
+from apache_beam.runners.interactive.caching.cacheable import CacheKey
 from apache_beam.testing.test_stream import OutputFormat
 from apache_beam.testing.test_stream import ReverseTestStream
 from apache_beam.utils import timestamp
@@ -156,8 +157,6 @@ def __init__(self, cache_dir, labels, is_cache_complete=None, coder=None):
     self._labels = labels
     self._path = os.path.join(self._cache_dir, *self._labels)
     self._is_cache_complete = is_cache_complete
-
-    from apache_beam.runners.interactive.pipeline_instrument import CacheKey
     self._pipeline_id = CacheKey.from_str(labels[-1]).pipeline_id
 
   def _wait_until_file_exists(self, timeout_secs=30):
@@ -168,7 +167,6 @@ def _wait_until_file_exists(self, timeout_secs=30):
     while not os.path.exists(self._path):
       time.sleep(1)
       if time.time() - start > timeout_secs:
-        from apache_beam.runners.interactive.pipeline_instrument import CacheKey
         pcollection_var = CacheKey.from_str(self._labels[-1]).var
         raise RuntimeError(
             'Timed out waiting for cache file for PCollection `{}` to be '
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
@@ -27,8 +27,8 @@
 from apache_beam.portability.api.beam_interactive_api_pb2 import TestStreamFileRecord
 from apache_beam.portability.api.beam_runner_api_pb2 import TestStreamPayload
 from apache_beam.runners.interactive.cache_manager import SafeFastPrimitivesCoder
+from apache_beam.runners.interactive.caching.cacheable import CacheKey
 from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache
-from apache_beam.runners.interactive.pipeline_instrument import CacheKey
 from apache_beam.runners.interactive.testing.test_cache_manager import FileRecordsBuilder
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.test_stream import TestStream
diff --git a/sdks/python/apache_beam/runners/interactive/interactive_runner.py b/sdks/python/apache_beam/runners/interactive/interactive_runner.py
@@ -190,7 +190,7 @@ def visit_transform(self, transform_node):
 
     if not self._skip_display:
       a_pipeline_graph = pipeline_graph.PipelineGraph(
-          pipeline_instrument.original_pipeline,
+          pipeline_instrument.original_pipeline_proto,
           render_option=self._render_option)
       a_pipeline_graph.display_graph()
 
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py b/sdks/python/apache_beam/runners/interactive/pipeline_instrument.py
@@ -32,65 +32,15 @@
 from apache_beam.runners.interactive import interactive_environment as ie
 from apache_beam.runners.interactive import pipeline_fragment as pf
 from apache_beam.runners.interactive import background_caching_job
-from apache_beam.runners.interactive.utils import obfuscate
+from apache_beam.runners.interactive.caching.cacheable import Cacheable
+from apache_beam.runners.interactive.caching.cacheable import CacheKey
 from apache_beam.testing import test_stream
 from apache_beam.transforms.window import WindowedValue
 
 READ_CACHE = "_ReadCache_"
 WRITE_CACHE = "_WriteCache_"
 
 
-# TODO: turn this into a dataclass object when we finally get off of Python2.
-class Cacheable:
-  def __init__(self, pcoll_id, var, version, pcoll, producer_version):
-    self.pcoll_id = pcoll_id
-    self.var = var
-    self.version = version
-    self.pcoll = pcoll
-    self.producer_version = producer_version
-
-  def __eq__(self, other):
-    return (
-        self.pcoll_id == other.pcoll_id and self.var == other.var and
-        self.version == other.version and self.pcoll == other.pcoll and
-        self.producer_version == other.producer_version)
-
-  def __hash__(self):
-    return hash((
-        self.pcoll_id,
-        self.var,
-        self.version,
-        self.pcoll,
-        self.producer_version))
-
-  def to_key(self):
-    return CacheKey(
-        self.var,
-        self.version,
-        self.producer_version,
-        str(id(self.pcoll.pipeline)))
-
-
-# TODO: turn this into a dataclass object when we finally get off of Python2.
-class CacheKey:
-  def __init__(self, var, version, producer_version, pipeline_id):
-    # Makes sure that the variable name is obfuscated and only first 10
-    # characters taken so that the CacheKey has a constant length.
-    self.var = obfuscate(var)[:10]
-    self.version = version
-    self.producer_version = producer_version
-    self.pipeline_id = pipeline_id
-
-  @staticmethod
-  def from_str(r):
-    split = r.split('-')
-    return CacheKey(split[0], split[1], split[2], split[3])
-
-  def __repr__(self):
-    return '-'.join(
-        [self.var, self.version, self.producer_version, self.pipeline_id])
-
-
 class PipelineInstrument(object):
   """A pipeline instrument for pipeline to be executed by interactive runner.
 
@@ -103,36 +53,34 @@ class PipelineInstrument(object):
   """
   def __init__(self, pipeline, options=None):
     self._pipeline = pipeline
-    # The cache manager per user-defined pipeline is lazily initiated the first
-    # time accessed. It is owned by interactive_environment module. This
-    # shortcut reference will be initialized when the user pipeline associated
-    # to the given pipeline is identified.
-    self._cache_manager = None
-
-    # Invoke a round trip through the runner API. This makes sure the Pipeline
-    # proto is stable. The snapshot of pipeline will not be mutated within this
-    # module and can be used to recover original pipeline if needed.
-    self._pipeline_snap = beam.pipeline.Pipeline.from_runner_api(
-        pipeline.to_runner_api(use_fake_coders=True), pipeline.runner, options)
-    ie.current_env().add_derived_pipeline(self._pipeline, self._pipeline_snap)
+
+    self._user_pipeline = ie.current_env().user_pipeline(pipeline)
+    if not self._user_pipeline:
+      self._user_pipeline = pipeline
+    self._cache_manager = ie.current_env().get_cache_manager(
+        self._user_pipeline, create_if_absent=True)
+    # Check if the user defined pipeline contains any source to cache.
+    # If so, during the check, the cache manager is converted into a
+    # streaming cache manager, thus re-assign.
+    if background_caching_job.has_source_to_cache(self._user_pipeline):
+      self._cache_manager = ie.current_env().get_cache_manager(
+          self._user_pipeline)
 
     self._background_caching_pipeline = beam.pipeline.Pipeline.from_runner_api(
         pipeline.to_runner_api(use_fake_coders=True), pipeline.runner, options)
     ie.current_env().add_derived_pipeline(
         self._pipeline, self._background_caching_pipeline)
 
     # Snapshot of original pipeline information.
-    (self._original_pipeline_proto,
-     self._original_context) = self._pipeline_snap.to_runner_api(
-         return_context=True, use_fake_coders=True)
+    (self._original_pipeline_proto, context) = self._pipeline.to_runner_api(
+        return_context=True, use_fake_coders=True)
 
     # All compute-once-against-original-pipeline fields.
     self._unbounded_sources = unbounded_sources(
         self._background_caching_pipeline)
     # TODO(BEAM-7760): once cache scope changed, this is not needed to manage
     # relationships across pipelines, runners, and jobs.
-    self._pcolls_to_pcoll_id = pcolls_to_pcoll_id(
-        self._pipeline_snap, self._original_context)
+    self._pcolls_to_pcoll_id = pcolls_to_pcoll_id(self._pipeline, context)
 
     # A mapping from PCollection id to python id() value in user defined
     # pipeline instance.
@@ -149,11 +97,6 @@ def __init__(self, pipeline, options=None):
     # (Dict[str, AppliedPTransform]).
     self._cached_pcoll_read = {}
 
-    # Reference to the user defined pipeline instance based on the given
-    # pipeline. The class never mutates it.
-    # Note: the original pipeline is not the user pipeline.
-    self._user_pipeline = None
-
     # A dict from PCollections in the runner pipeline instance to their
     # corresponding PCollections in the user pipeline instance. Populated
     # after preprocess().
@@ -421,15 +364,9 @@ def pcolls_to_pcoll_id(self):
 
   @property
   def original_pipeline_proto(self):
-    """Returns the portable proto representation of the pipeline before
-    instrumentation."""
+    """Returns a snapshot of the pipeline proto before instrumentation."""
     return self._original_pipeline_proto
 
-  @property
-  def original_pipeline(self):
-    """Returns a snapshot of the pipeline before instrumentation."""
-    return self._pipeline_snap
-
   @property
   def user_pipeline(self):
     """Returns a reference to the pipeline instance defined by the user. If a
@@ -571,29 +508,11 @@ def _process(self, pcoll):
           cacheable_key = self._pin._cacheable_key(pcoll)
           user_pcoll = self._pin.cacheables[cacheable_key].pcoll
           if (cacheable_key in self._pin.cacheables and user_pcoll != pcoll):
-            if not self._pin._user_pipeline:
-              # Retrieve a reference to the user defined pipeline instance.
-              self._pin._user_pipeline = user_pcoll.pipeline
-              # Retrieve a reference to the cache manager for the user defined
-              # pipeline instance.
-              self._pin._cache_manager = ie.current_env().get_cache_manager(
-                  self._pin._user_pipeline, create_if_absent=True)
-              # Check if the user defined pipeline contains any source to cache.
-              # If so, during the check, the cache manager is converted into a
-              # streaming cache manager, thus re-assign the reference.
-              if background_caching_job.has_source_to_cache(
-                  self._pin._user_pipeline):
-                self._pin._cache_manager = ie.current_env().get_cache_manager(
-                    self._pin._user_pipeline)
             self._pin._runner_pcoll_to_user_pcoll[pcoll] = user_pcoll
             self._pin.cacheables[cacheable_key].pcoll = pcoll
 
     v = PreprocessVisitor(self)
     self._pipeline.visit(v)
-    if not self._user_pipeline:
-      self._user_pipeline = self._pipeline
-      self._cache_manager = ie.current_env().get_cache_manager(
-          self._user_pipeline, create_if_absent=True)
 
   def _write_cache(
       self,
@@ -679,7 +598,6 @@ def _read_cache(self, pipeline, pcoll, is_unbounded_source_output):
     key = self.cache_key(pcoll)
     # Can only read from cache when the cache with expected key exists and its
     # computation has been completed.
-
     is_cached = self._cache_manager.exists('full', key)
     is_computed = (
         pcoll in self._runner_pcoll_to_user_pcoll and
@@ -886,7 +804,6 @@ def cacheables(pcolls_to_pcoll_id):
   for watching in ie.current_env().watching():
     for key, val in watching:
       if isinstance(val, beam.pvalue.PCollection):
-
         pcoll_id = pcolls_to_pcoll_id.get(str(val), None)
         # It's highly possible that PCollection str is not unique across
         # multiple pipelines, further check during instrument is needed.
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py b/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
@@ -279,6 +279,7 @@ def test_instrument_example_pipeline_to_read_cache(self):
     self._mock_write_cache(p_origin, [b'1', b'4', b'9'], second_pcoll_cache_key)
     # Mark the completeness of PCollections from the original(user) pipeline.
     ie.current_env().mark_pcollection_computed((init_pcoll, second_pcoll))
+    ie.current_env().add_derived_pipeline(p_origin, p_copy)
     instr.build_pipeline_instrument(p_copy)
 
     cached_init_pcoll = (
@@ -315,6 +316,7 @@ def test_find_out_correct_user_pipeline(self):
         user_pipeline.to_runner_api(use_fake_coders=True),
         user_pipeline.runner,
         options=None)
+    ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline)
     # This is a totally irrelevant user pipeline in the watched scope.
     irrelevant_user_pipeline = beam.Pipeline(
         interactive_runner.InteractiveRunner())
@@ -506,6 +508,7 @@ def test_instrument_mixed_streaming_batch(self):
         p_original.to_runner_api(),
         runner=interactive_runner.InteractiveRunner(),
         options=options)
+    ie.current_env().add_derived_pipeline(p_original, p_copy)
     instrumenter = instr.build_pipeline_instrument(p_copy)
     actual_pipeline = beam.Pipeline.from_runner_api(
         proto=instrumenter.instrumented_pipeline_proto(),
@@ -766,6 +769,7 @@ def test_pipeline_pruned_when_input_pcoll_is_cached(self):
         user_pipeline.to_runner_api(use_fake_coders=True),
         user_pipeline.runner,
         None)
+    ie.current_env().add_derived_pipeline(user_pipeline, runner_pipeline)
 
     # Mock as if init_pcoll is cached.
     init_pcoll_cache_key = self.cache_key_of('init_pcoll', init_pcoll)
diff --git a/sdks/python/container/license_scripts/dep_urls_py.yaml b/sdks/python/container/license_scripts/dep_urls_py.yaml
@@ -47,6 +47,8 @@ pip_dependencies:
     license: "https://raw.githubusercontent.com/certifi/python-certifi/master/LICENSE"
   cython:
     license: "https://raw.githubusercontent.com/cython/cython/master/LICENSE.txt"
+  dataclasses:
+    license: "https://raw.githubusercontent.com/ericvsmith/dataclasses/master/LICENSE.txt"
   enum34:
     # The original repo is down. This license taken from somebody's clone:
     # https://github.com/jamespharaoh/python-enum34/blob/master/enum/LICENSE
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
@@ -131,6 +131,10 @@ def get_version():
     # Avro 1.9.2 for python3 was broken. The issue was fixed in version 1.9.2.1
     'avro-python3>=1.8.1,!=1.9.2,<1.10.0',
     'crcmod>=1.7,<2.0',
+    # dataclasses backport for python_version<3.7. No version bound because this
+    # is Python standard since Python 3.7 and each Python version is compatible
+    # with a specific dataclasses version.
+    'dataclasses;python_version<"3.7"',
     # Dill doesn't have forwards-compatibility guarantees within minor version.
     # Pickles created with a new version of dill may not unpickle using older
     # version of dill. It is best to use the same version of dill on client and
