beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From pabl...@apache.org
Subject [beam] branch master updated: Add ElementLimiters to all Cache Managers.
Date Fri, 07 Aug 2020 19:25:20 GMT
This is an automated email from the ASF dual-hosted git repository.

pabloem 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 d30c912  Add ElementLimiters to all Cache Managers.
     new af2f071  Merge pull request #12412 from [BEAM-10603] Add ElementLimiters to all Cache
Managers.
d30c912 is described below

commit d30c912cb31ce88aafa3310ba9b4bf75f061d521
Author: Sam Rohde <srohde@google.com>
AuthorDate: Mon Jul 27 14:56:06 2020 -0700

    Add ElementLimiters to all Cache Managers.
    
    Change-Id: I093e1dc1f99c31b1e0bb868fc3bcfc507bd24c8d
---
 .../runners/interactive/cache_manager.py           |  47 ++++++++-
 .../runners/interactive/cache_manager_test.py      |  28 ++++++
 .../runners/interactive/caching/streaming_cache.py |  51 +++++++---
 .../interactive/caching/streaming_cache_test.py    | 110 +++++++++++++++++++++
 .../interactive/options/capture_limiters.py        |   4 +-
 .../interactive/options/capture_limiters_test.py   |   2 +-
 .../interactive/testing/test_cache_manager.py      |  18 +++-
 7 files changed, 240 insertions(+), 20 deletions(-)

diff --git a/sdks/python/apache_beam/runners/interactive/cache_manager.py b/sdks/python/apache_beam/runners/interactive/cache_manager.py
index 9894535..48f1fc5 100644
--- a/sdks/python/apache_beam/runners/interactive/cache_manager.py
+++ b/sdks/python/apache_beam/runners/interactive/cache_manager.py
@@ -67,13 +67,17 @@ class CacheManager(object):
     """Returns the latest version number of the PCollection cache."""
     raise NotImplementedError
 
-  def read(self, *labels):
-    # type (*str) -> Tuple[str, Generator[Any]]
+  def read(self, *labels, **args):
+    # type (*str, Dict[str, Any]) -> Tuple[str, Generator[Any]]
 
     """Return the PCollection as a list as well as the version number.
 
     Args:
       *labels: List of labels for PCollection instance.
+      **args: Dict of additional arguments. Currently only supports 'limiters'
+        as a list of ElementLimiters, and 'tail' as a boolean. Limiters limits
+        the amount of elements read and duration with respect to processing
+        time.
 
     Returns:
       A tuple containing an iterator for the items in the PCollection and the
@@ -97,6 +101,17 @@ class CacheManager(object):
     """
     raise NotImplementedError
 
+  def clear(self, *labels):
+    # type (*str) -> Boolean
+
+    """Clears the cache entry of the given labels and returns True on success.
+
+    Args:
+      value: An encodable (with corresponding PCoder) value
+      *labels: List of labels for PCollection instance
+    """
+    raise NotImplementedError
+
   def source(self, *labels):
     # type (*str) -> ptransform.PTransform
 
@@ -196,17 +211,35 @@ class FileBasedCacheManager(CacheManager):
         self._default_pcoder if self._default_pcoder is not None else
         self._saved_pcoders[self._path(*labels)])
 
-  def read(self, *labels):
+  def read(self, *labels, **args):
     # Return an iterator to an empty list if it doesn't exist.
     if not self.exists(*labels):
       return iter([]), -1
 
+    limiters = args.pop('limiters', [])
+
     # Otherwise, return a generator to the cached PCollection.
     source = self.source(*labels)._source
     range_tracker = source.get_range_tracker(None, None)
     reader = source.read(range_tracker)
     version = self._latest_version(*labels)
-    return reader, version
+
+    # The return type is a generator, so in order to implement the limiter for
+    # the FileBasedCacheManager we wrap the original generator with the logic
+    # to limit yielded elements.
+    def limit_reader(r):
+      for e in r:
+        # Update the limiters and break early out of reading from cache if any
+        # are triggered.
+        for l in limiters:
+          l.update(e)
+
+        if any(l.is_triggered() for l in limiters):
+          break
+
+        yield e
+
+    return limit_reader(reader), version
 
   def write(self, values, *labels):
     sink = self.sink(labels)._sink
@@ -218,6 +251,12 @@ class FileBasedCacheManager(CacheManager):
       writer.write(v)
     writer.close()
 
+  def clear(self, *labels):
+    if self.exists(*labels):
+      filesystems.FileSystems.delete(self._match(*labels))
+      return True
+    return False
+
   def source(self, *labels):
     return self._reader_class(
         self._glob_path(*labels), coder=self.load_pcoder(*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
index 7868e90..e7dc936 100644
--- a/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
+++ b/sdks/python/apache_beam/runners/interactive/cache_manager_test.py
@@ -30,6 +30,7 @@ import unittest
 from apache_beam import coders
 from apache_beam.io import filesystems
 from apache_beam.runners.interactive import cache_manager as cache
+from apache_beam.runners.interactive.options.capture_limiters import CountLimiter
 
 
 class FileBasedCacheManagerTest(object):
@@ -91,6 +92,18 @@ class FileBasedCacheManagerTest(object):
     self.mock_write_cache(cache_version_one, prefix, cache_label)
     self.assertTrue(self.cache_manager.exists(prefix, cache_label))
 
+  def test_clear(self):
+    """Test that CacheManager can correctly tell if the cache exists or not."""
+    prefix = 'full'
+    cache_label = 'some-cache-label'
+    cache_version_one = ['cache', 'version', 'one']
+
+    self.assertFalse(self.cache_manager.exists(prefix, cache_label))
+    self.mock_write_cache(cache_version_one, prefix, cache_label)
+    self.assertTrue(self.cache_manager.exists(prefix, cache_label))
+    self.assertTrue(self.cache_manager.clear(prefix, cache_label))
+    self.assertFalse(self.cache_manager.exists(prefix, cache_label))
+
   def test_read_basic(self):
     """Test the condition where the cache is read once after written once."""
     prefix = 'full'
@@ -180,6 +193,21 @@ class FileBasedCacheManagerTest(object):
     self.assertTrue(
         self.cache_manager.is_latest_version(version, prefix, cache_label))
 
+  def test_read_with_count_limiter(self):
+    """Test the condition where the cache is read once after written once."""
+    prefix = 'full'
+    cache_label = 'some-cache-label'
+    cache_version_one = ['cache', 'version', 'one']
+
+    self.mock_write_cache(cache_version_one, prefix, cache_label)
+    reader, version = self.cache_manager.read(
+        prefix, cache_label, limiters=[CountLimiter(2)])
+    pcoll_list = list(reader)
+    self.assertListEqual(pcoll_list, ['cache', 'version'])
+    self.assertEqual(version, 0)
+    self.assertTrue(
+        self.cache_manager.is_latest_version(version, prefix, cache_label))
+
 
 class TextFileBasedCacheManagerTest(
     FileBasedCacheManagerTest,
diff --git a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
index b2204cf..77f976d 100644
--- a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
+++ b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache.py
@@ -153,13 +153,23 @@ class StreamingCacheSource:
       cache_dir,
       labels,
       is_cache_complete=None,
-      coder=SafeFastPrimitivesCoder()):
+      coder=None,
+      limiters=None):
+    if not coder:
+      coder = SafeFastPrimitivesCoder()
+
+    if not is_cache_complete:
+      is_cache_complete = lambda _: True
+
+    if not limiters:
+      limiters = []
+
     self._cache_dir = cache_dir
     self._coder = coder
     self._labels = labels
     self._path = os.path.join(self._cache_dir, *self._labels)
-    self._is_cache_complete = (
-        is_cache_complete if is_cache_complete else lambda _: True)
+    self._is_cache_complete = is_cache_complete
+    self._limiters = limiters
 
     from apache_beam.runners.interactive.pipeline_instrument import CacheKey
     self._pipeline_id = CacheKey.from_str(labels[-1]).pipeline_id
@@ -193,7 +203,8 @@ class StreamingCacheSource:
 
       # Check if we are at EOF or if we have an incomplete line.
       if not line or (line and line[-1] != b'\n'[0]):
-        if not tail:
+        # Read at least the first line to get the header.
+        if not tail and pos != 0:
           break
 
         # Complete reading only when the cache is complete.
@@ -210,10 +221,16 @@ class StreamingCacheSource:
         proto_cls = TestStreamFileHeader if pos == 0 else TestStreamFileRecord
         msg = self._try_parse_as(proto_cls, to_decode)
         if msg:
-          yield msg
+          for l in self._limiters:
+            l.update(msg)
+
+          if any(l.is_triggered() for l in self._limiters):
+            break
         else:
           break
 
+        yield msg
+
   def _try_parse_as(self, proto_cls, to_decode):
     try:
       msg = proto_cls()
@@ -285,7 +302,7 @@ class StreamingCache(CacheManager):
     return os.path.exists(path)
 
   # TODO(srohde): Modify this to return the correct version.
-  def read(self, *labels):
+  def read(self, *labels, **args):
     """Returns a generator to read all records from file.
 
     Does not tail.
@@ -293,8 +310,12 @@ class StreamingCache(CacheManager):
     if not self.exists(*labels):
       return iter([]), -1
 
+    limiters = args.pop('limiters', [])
+    tail = args.pop('tail', False)
+
     reader = StreamingCacheSource(
-        self._cache_dir, labels, self._is_cache_complete).read(tail=False)
+        self._cache_dir, labels, self._is_cache_complete,
+        limiters=limiters).read(tail=tail)
 
     # Return an empty iterator if there is nothing in the file yet. This can
     # only happen when tail is False.
@@ -304,7 +325,7 @@ class StreamingCache(CacheManager):
       return iter([]), -1
     return StreamingCache.Reader([header], [reader]).read(), 1
 
-  def read_multiple(self, labels):
+  def read_multiple(self, labels, limiters=None, tail=True):
     """Returns a generator to read all records from file.
 
     Does tail until the cache is complete. This is because it is used in the
@@ -312,9 +333,9 @@ class StreamingCache(CacheManager):
     pipeline runtime which needs to block.
     """
     readers = [
-        StreamingCacheSource(self._cache_dir, l,
-                             self._is_cache_complete).read(tail=True)
-        for l in labels
+        StreamingCacheSource(
+            self._cache_dir, l, self._is_cache_complete,
+            limiters=limiters).read(tail=tail) for l in labels
     ]
     headers = [next(r) for r in readers]
     return StreamingCache.Reader(headers, readers).read()
@@ -334,6 +355,14 @@ class StreamingCache(CacheManager):
           val = v
         f.write(self._default_pcoder.encode(val) + b'\n')
 
+  def clear(self, *labels):
+    directory = os.path.join(self._cache_dir, *labels[:-1])
+    filepath = os.path.join(directory, labels[-1])
+    if os.path.exists(filepath):
+      os.remove(filepath)
+      return True
+    return False
+
   def source(self, *labels):
     """Returns the StreamingCacheManager source.
 
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
index a56b851..23390cc 100644
--- a/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
+++ b/sdks/python/apache_beam/runners/interactive/caching/streaming_cache_test.py
@@ -28,6 +28,8 @@ from apache_beam.portability.api.beam_interactive_api_pb2 import TestStreamFileR
 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.streaming_cache import StreamingCache
+from apache_beam.runners.interactive.options.capture_limiters import CountLimiter
+from apache_beam.runners.interactive.options.capture_limiters import ProcessingTimeLimiter
 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
@@ -64,6 +66,14 @@ class StreamingCacheTest(unittest.TestCase):
     # Assert that an empty reader returns an empty list.
     self.assertFalse([e for e in reader])
 
+  def test_clear(self):
+    cache = StreamingCache(cache_dir=None)
+    self.assertFalse(cache.exists('my_label'))
+    cache.write([TestStreamFileRecord()], 'my_label')
+    self.assertTrue(cache.exists('my_label'))
+    self.assertTrue(cache.clear('my_label'))
+    self.assertFalse(cache.exists('my_label'))
+
   def test_single_reader(self):
     """Tests that we expect to see all the correctly emitted TestStreamPayloads.
     """
@@ -403,6 +413,106 @@ class StreamingCacheTest(unittest.TestCase):
 
     self.assertListEqual(actual_events, expected_events)
 
+  def test_single_reader_with_count_limiter(self):
+    """Tests that we expect to see all the correctly emitted TestStreamPayloads.
+    """
+    CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', ''))
+
+    values = (FileRecordsBuilder(tag=CACHED_PCOLLECTION_KEY)
+              .add_element(element=0, event_time_secs=0)
+              .advance_processing_time(1)
+              .add_element(element=1, event_time_secs=1)
+              .advance_processing_time(1)
+              .add_element(element=2, event_time_secs=2)
+              .build()) # yapf: disable
+
+    cache = StreamingCache(cache_dir=None)
+    cache.write(values, CACHED_PCOLLECTION_KEY)
+
+    reader, _ = cache.read(CACHED_PCOLLECTION_KEY, limiters=[CountLimiter(2)])
+    coder = coders.FastPrimitivesCoder()
+    events = list(reader)
+
+    # Units here are in microseconds.
+    # These are a slice of the original values such that we only get two
+    # elements.
+    expected = [
+        TestStreamPayload.Event(
+            element_event=TestStreamPayload.Event.AddElements(
+                elements=[
+                    TestStreamPayload.TimestampedElement(
+                        encoded_element=coder.encode(0), timestamp=0)
+                ],
+                tag=CACHED_PCOLLECTION_KEY)),
+        TestStreamPayload.Event(
+            processing_time_event=TestStreamPayload.Event.AdvanceProcessingTime(
+                advance_duration=1 * 10**6)),
+        TestStreamPayload.Event(
+            element_event=TestStreamPayload.Event.AddElements(
+                elements=[
+                    TestStreamPayload.TimestampedElement(
+                        encoded_element=coder.encode(1), timestamp=1 * 10**6)
+                ],
+                tag=CACHED_PCOLLECTION_KEY)),
+        TestStreamPayload.Event(
+            processing_time_event=TestStreamPayload.Event.AdvanceProcessingTime(
+                advance_duration=1 * 10**6)),
+    ]
+    self.assertSequenceEqual(events, expected)
+
+  def test_single_reader_with_processing_time_limiter(self):
+    """Tests that we expect to see all the correctly emitted TestStreamPayloads.
+    """
+    CACHED_PCOLLECTION_KEY = repr(CacheKey('arbitrary_key', '', '', ''))
+
+    values = (FileRecordsBuilder(tag=CACHED_PCOLLECTION_KEY)
+              .advance_processing_time(1e-6)
+              .add_element(element=0, event_time_secs=0)
+              .advance_processing_time(1)
+              .add_element(element=1, event_time_secs=1)
+              .advance_processing_time(1)
+              .add_element(element=2, event_time_secs=2)
+              .advance_processing_time(1)
+              .add_element(element=3, event_time_secs=2)
+              .advance_processing_time(1)
+              .add_element(element=4, event_time_secs=2)
+              .build()) # yapf: disable
+
+    cache = StreamingCache(cache_dir=None)
+    cache.write(values, CACHED_PCOLLECTION_KEY)
+
+    reader, _ = cache.read(
+        CACHED_PCOLLECTION_KEY, limiters=[ProcessingTimeLimiter(2)])
+    coder = coders.FastPrimitivesCoder()
+    events = list(reader)
+
+    # Units here are in microseconds.
+    # Expects that the elements are a slice of the original values where all
+    # processing time is less than the duration.
+    expected = [
+        TestStreamPayload.Event(
+            processing_time_event=TestStreamPayload.Event.AdvanceProcessingTime(
+                advance_duration=1)),
+        TestStreamPayload.Event(
+            element_event=TestStreamPayload.Event.AddElements(
+                elements=[
+                    TestStreamPayload.TimestampedElement(
+                        encoded_element=coder.encode(0), timestamp=0)
+                ],
+                tag=CACHED_PCOLLECTION_KEY)),
+        TestStreamPayload.Event(
+            processing_time_event=TestStreamPayload.Event.AdvanceProcessingTime(
+                advance_duration=1 * 10**6)),
+        TestStreamPayload.Event(
+            element_event=TestStreamPayload.Event.AddElements(
+                elements=[
+                    TestStreamPayload.TimestampedElement(
+                        encoded_element=coder.encode(1), timestamp=1 * 10**6)
+                ],
+                tag=CACHED_PCOLLECTION_KEY)),
+    ]
+    self.assertSequenceEqual(events, expected)
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/sdks/python/apache_beam/runners/interactive/options/capture_limiters.py b/sdks/python/apache_beam/runners/interactive/options/capture_limiters.py
index 2c84f80..c9888ab 100644
--- a/sdks/python/apache_beam/runners/interactive/options/capture_limiters.py
+++ b/sdks/python/apache_beam/runners/interactive/options/capture_limiters.py
@@ -109,12 +109,14 @@ class CountLimiter(ElementLimiter):
       self._count += 1
 
   def is_triggered(self):
-    return self._count >= self._max_count
+    return self._count > self._max_count
 
 
 class ProcessingTimeLimiter(ElementLimiter):
   """Limits by how long the ProcessingTime passed in the element stream.
 
+  Reads all elements from the timespan [start, start + duration).
+
   This measures the duration from the first element in the stream. Each
   subsequent element has a delta "advance_duration" that moves the internal
   clock forward. This triggers when the duration from the internal clock and
diff --git a/sdks/python/apache_beam/runners/interactive/options/capture_limiters_test.py
b/sdks/python/apache_beam/runners/interactive/options/capture_limiters_test.py
index 850c56e2c..347cb8e 100644
--- a/sdks/python/apache_beam/runners/interactive/options/capture_limiters_test.py
+++ b/sdks/python/apache_beam/runners/interactive/options/capture_limiters_test.py
@@ -28,7 +28,7 @@ class CaptureLimitersTest(unittest.TestCase):
   def test_count_limiter(self):
     limiter = CountLimiter(5)
 
-    for e in range(4):
+    for e in range(5):
       limiter.update(e)
 
     self.assertFalse(limiter.is_triggered())
diff --git a/sdks/python/apache_beam/runners/interactive/testing/test_cache_manager.py b/sdks/python/apache_beam/runners/interactive/testing/test_cache_manager.py
index f39f016..098f249 100644
--- a/sdks/python/apache_beam/runners/interactive/testing/test_cache_manager.py
+++ b/sdks/python/apache_beam/runners/interactive/testing/test_cache_manager.py
@@ -45,11 +45,23 @@ class InMemoryCache(CacheManager):
   def _latest_version(self, *labels):
     return True
 
-  def read(self, *labels):
+  def read(self, *labels, **args):
     if not self.exists(*labels):
       return itertools.chain([]), -1
-    ret = itertools.chain(self._cached[self._key(*labels)])
-    return ret, None
+
+    limiters = args.pop('limiters', [])
+
+    def limit_reader(r):
+      for e in r:
+        for l in limiters:
+          l.update(e)
+
+        if any(l.is_triggered() for l in limiters):
+          break
+
+        yield e
+
+    return limit_reader(itertools.chain(self._cached[self._key(*labels)])), None
 
   def write(self, value, *labels):
     if not self.exists(*labels):


Mime
View raw message