Fixed the cache_manager write implementation.

KevinGG · KevinGG · commit 19aaad2d7ed3 · 2021-06-17T17:00:30.000-07:00
Addtionally:
1. Disabled writing arbitrary typed elements in streaming cache.
2. Temporarily use SafeFastPrimitivesCoder for WindowedValueHolder
   wrapped elements (i.e., all types of elements).
3. Fixed forward the screenshot integration test dependency version.
diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager.py b/sdks/python/apache_beam/runners/interactive/cache_manager.py
@@ -29,6 +29,7 @@
 from apache_beam.io import textio
 from apache_beam.io import tfrecordio
 from apache_beam.transforms import combiners
+from apache_beam.runners.interactive.utils import obfuscate
 
 
 class CacheManager(object):
@@ -165,7 +166,7 @@ def __init__(self, cache_dir=None, cache_format='text'):
       self._cache_dir = cache_dir
     else:
       self._cache_dir = tempfile.mkdtemp(
-          prefix='it-', dir=os.environ.get('TEST_TMPDIR', None))
+          prefix='ib-', dir=os.environ.get('TEST_TMPDIR', None))
     self._versions = collections.defaultdict(lambda: self._CacheVersion())
     self.cache_format = cache_format
 
@@ -208,7 +209,12 @@ def save_pcoder(self, pcoder, *labels):
 
   def load_pcoder(self, *labels):
     saved_pcoder = self._saved_pcoders.get(self._path(*labels), None)
-    return self._default_pcoder if saved_pcoder is None else saved_pcoder
+    # TODO(BEAM-12506): Get rid of the SafeFastPrimitivesCoder for
+    # WindowedValueHolder.
+    if saved_pcoder is None or isinstance(saved_pcoder,
+                                          coders.FastPrimitivesCoder):
+      return self._default_pcoder
+    return saved_pcoder
 
   def read(self, *labels, **args):
     # Return an iterator to an empty list if it doesn't exist.
@@ -224,11 +230,22 @@ def read(self, *labels, **args):
     return reader, version
 
   def write(self, values, *labels):
-    sink = self.sink(labels)._sink
-    path = self._path(*labels)
+    """Imitates how a WriteCache tranform works without running a pipeline.
 
-    init_result = sink.initialize_write()
-    writer = sink.open_writer(init_result, path)
+    For testing and cache manager development, not for production usage because
+    the write is not sharded and does not use Beam execution model.
+    """
+    pcoder = coders.registry.get_coder(type(values[0]))
+    # Save the pcoder for the actual labels.
+    self.save_pcoder(pcoder, *labels)
+    single_shard_labels = [*labels[:-1], '-00000-of-00001']
+    # Save the pcoder for the labels that imitates the sharded cache file name
+    # suffix.
+    self.save_pcoder(pcoder, *single_shard_labels)
+    # Put a '-%05d-of-%05d' suffix to the cache file.
+    sink = self.sink(single_shard_labels)._sink
+    path = self._path(*labels[:-1])
+    writer = sink.open_writer(path, labels[-1])
     for v in values:
       writer.write(v)
     writer.close()
@@ -253,7 +270,7 @@ def cleanup(self):
     self._saved_pcoders = {}
 
   def _glob_path(self, *labels):
-    return self._path(*labels) + '-*-of-*'
+    return self._path(*labels) + '*-*-of-*'
 
   def _path(self, *labels):
     return filesystems.FileSystems.join(self._cache_dir, *labels)
diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager_test.py b/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
@@ -43,36 +43,20 @@ class FileBasedCacheManagerTest(object):
   cache_format = None  # type: str
 
   def setUp(self):
-    self.test_dir = tempfile.mkdtemp()
     self.cache_manager = cache.FileBasedCacheManager(
-        self.test_dir, cache_format=self.cache_format)
+        cache_format=self.cache_format)
 
   def tearDown(self):
-    # The test_dir might have already been removed by cache_manager.cleanup().
-    if os.path.exists(self.test_dir):
-      shutil.rmtree(self.test_dir)
+    self.cache_manager.cleanup()
 
-  def mock_write_cache(self, pcoll_list, prefix, cache_label):
+  def mock_write_cache(self, values, prefix, cache_label):
     """Cache the PCollection where cache.WriteCache would write to."""
-    cache_path = filesystems.FileSystems.join(
-        self.cache_manager._cache_dir, prefix)
-    if not filesystems.FileSystems.exists(cache_path):
-      filesystems.FileSystems.mkdirs(cache_path)
-
     # Pause for 0.1 sec, because the Jenkins test runs so fast that the file
     # writes happen at the same timestamp.
     time.sleep(0.1)
 
-    cache_file = cache_label + '-1-of-2'
     labels = [prefix, cache_label]
-
-    # Usually, the pcoder will be inferred from `pcoll.element_type`
-    pcoder = coders.registry.get_coder(object)
-    # Save a pcoder for reading.
-    self.cache_manager.save_pcoder(pcoder, *labels)
-    # Save a pcoder for the fake write to the file.
-    self.cache_manager.save_pcoder(pcoder, prefix, cache_file)
-    self.cache_manager.write(pcoll_list, prefix, cache_file)
+    self.cache_manager.write(values, *labels)
 
   def test_exists(self):
     """Test that CacheManager can correctly tell if the cache exists or not."""
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
@@ -251,7 +251,7 @@ def __init__(
       self._cache_dir = cache_dir
     else:
       self._cache_dir = tempfile.mkdtemp(
-          prefix='interactive-temp-', dir=os.environ.get('TEST_TMPDIR', None))
+          prefix='ib-', dir=os.environ.get('TEST_TMPDIR', None))
 
     # List of saved pcoders keyed by PCollection path. It is OK to keep this
     # list in memory because once FileBasedCacheManager object is
@@ -343,8 +343,9 @@ def write(self, values, *labels):
         if isinstance(v, (TestStreamFileHeader, TestStreamFileRecord)):
           val = v.SerializeToString()
         else:
-          val = v
-          self.save_pcoder(coders.registry.get_coder(type(val)), *labels)
+          raise TypeError(
+              'Values given to streaming cache should be either '
+              'TestStreamFileHeader or TestStreamFileRecord.')
         f.write(self.load_pcoder(*labels).encode(val) + b'\n')
 
   def clear(self, *labels):
@@ -389,7 +390,12 @@ def save_pcoder(self, pcoder, *labels):
   def load_pcoder(self, *labels):
     saved_pcoder = self._saved_pcoders.get(
         os.path.join(self._cache_dir, *labels), None)
-    return self._default_pcoder if saved_pcoder is None else saved_pcoder
+    # TODO(BEAM-12506): Get rid of the SafeFastPrimitivesCoder for
+    # WindowedValueHolder.
+    if saved_pcoder is None or isinstance(saved_pcoder,
+                                          coders.FastPrimitivesCoder):
+      return self._default_pcoder
+    return saved_pcoder
 
   def cleanup(self):
 
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
@@ -26,9 +26,7 @@
 from apache_beam.portability.api.beam_interactive_api_pb2 import TestStreamFileHeader
 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 FileBasedCacheManager
 from apache_beam.runners.interactive.cache_manager import SafeFastPrimitivesCoder
-from apache_beam.runners.interactive.cache_manager import WriteCache
 from apache_beam.runners.interactive.caching.cacheable import CacheKey
 from apache_beam.runners.interactive.caching.streaming_cache import StreamingCache
 from apache_beam.runners.interactive.testing.test_cache_manager import FileRecordsBuilder
@@ -446,27 +444,9 @@ def test_always_default_coder_for_test_stream_records(self):
     self.assertIs(
         type(cache.load_pcoder(CACHED_NUMBERS)), type(cache._default_pcoder))
 
-  def test_load_saved_coder_for_non_record_types(self):
+  def test_streaming_cache_does_not_write_non_record_or_header_types(self):
     cache = StreamingCache(cache_dir=None)
-    cache.write('some value', 'a key')
-    self.assertIs(
-        type(cache.load_pcoder('a key')), type(coders.registry.get_coder(str)))
-
-  def test_wrap_file_based_cache_respect_load_saved_coders(self):
-    file_based_cache = FileBasedCacheManager(tempfile.mkdtemp())
-    pipeline = beam.Pipeline()
-    pcoll = pipeline | beam.Create([1, 2, 3])
-    _ = pcoll | WriteCache(file_based_cache, 'a key')
-    self.assertIs(
-        type(file_based_cache.load_pcoder('full', 'a key')),
-        type(coders.registry.get_coder(int)))
-    streaming_cache = StreamingCache(
-        cache_dir=file_based_cache._cache_dir,
-        saved_pcoders=file_based_cache._saved_pcoders)
-    streaming_cache.write('some value', 'another key')
-    self.assertIs(
-        type(streaming_cache.load_pcoder('full', 'a key')),
-        type(coders.registry.get_coder(int)))
+    self.assertRaises(TypeError, cache.write, 'some value', 'a key')
 
 
 if __name__ == '__main__':
diff --git a/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py b/sdks/python/apache_beam/runners/interactive/pipeline_instrument_test.py
@@ -342,7 +342,7 @@ def test_instrument_example_unbounded_pipeline_to_read_cache(self):
       if not isinstance(pcoll, beam.pvalue.PCollection):
         continue
       cache_key = self.cache_key_of(name, pcoll)
-      self._mock_write_cache(p_original, [b''], cache_key)
+      self._mock_write_cache(p_original, [], cache_key)
 
     # Instrument the original pipeline to create the pipeline the user will see.
     instrumenter = instr.build_pipeline_instrument(p_original)
@@ -498,7 +498,7 @@ def test_instrument_mixed_streaming_batch(self):
     ib.watch(locals())
 
     self._mock_write_cache(
-        p_original, [b''], self.cache_key_of('source_2', source_2))
+        p_original, [], self.cache_key_of('source_2', source_2))
     ie.current_env().mark_pcollection_computed([source_2])
 
     # Instrument the original pipeline to create the pipeline the user will see.
@@ -710,7 +710,7 @@ def test_instrument_example_unbounded_pipeline_to_multiple_read_cache(self):
       if not isinstance(pcoll, beam.pvalue.PCollection):
         continue
       cache_key = self.cache_key_of(name, pcoll)
-      self._mock_write_cache(p_original, [b''], cache_key)
+      self._mock_write_cache(p_original, [], cache_key)
 
     # Instrument the original pipeline to create the pipeline the user will see.
     instrumenter = instr.build_pipeline_instrument(p_original)
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
@@ -218,7 +218,7 @@ def get_version():
     # headless chrome based integration tests
     'selenium>=3.141.0,<4',
     'needle>=0.5.0,<1',
-    'chromedriver-binary>=88,<89',
+    'chromedriver-binary>=91,<92',
     # use a fixed major version of PIL for different python versions
     'pillow>=7.1.1,<8',
 ]

Original file line number	Diff line number	Diff line change
`@@ -218,7 +218,7 @@ def get_version():`
`218`	`218`	`# headless chrome based integration tests`
`219`	`219`	`'selenium>=3.141.0,<4',`
`220`	`220`	`'needle>=0.5.0,<1',`
`221`		`- 'chromedriver-binary>=88,<89',`
	`221`	`+ 'chromedriver-binary>=91,<92',`
`222`	`222`	`# use a fixed major version of PIL for different python versions`
`223`	`223`	`'pillow>=7.1.1,<8',`
`224`	`224`	`]`