Dynamic work rebalancing support for InMemory reader.

----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=123075083

Author: chamikaramj

google/cloud/dataflow/worker/executor_test.py

@@ -95,7 +95,7 @@ def __init__(self, source):
   def get_progress(self):
     next_progress = super(ProgressRequestRecordingInMemoryReader,
                           self).get_progress()
-    self.progress_record.append(next_progress.percent_complete)
+    self.progress_record.append(next_progress.position.record_index)
     return next_progress
 
 
@@ -332,11 +332,7 @@ def test_in_memory_source_progress_reporting(self):
     ]))
     self.assertEqual(elements, output_buffer)
 
-    expected_progress_record = []
-    len_elements = len(elements)
-    for i in range(len_elements):
-      expected_progress_record.append(float(i + 1) / len_elements)
-
+    expected_progress_record = range(len(elements))
     self.assertEqual(expected_progress_record,
                      source.last_reader.progress_record)
 

google/cloud/dataflow/worker/inmemory.py

@@ -15,9 +15,11 @@
 """In-memory input source."""
 
 import itertools
+import logging
 
 from google.cloud.dataflow import coders
 from google.cloud.dataflow.io import iobase
+from google.cloud.dataflow.io import range_trackers
 
 
 class InMemorySource(iobase.NativeSource):
@@ -53,11 +55,14 @@ class InMemoryReader(iobase.NativeSourceReader):
   """A reader for in-memory source."""
 
   def __init__(self, source):
-    self.source = source
+    self._source = source
 
-    # Index of the next item to be read by the InMemoryReader.
-    # Starts at source.start_index.
-    self.current_index = source.start_index
+    # Index of the last item returned by InMemoryReader.
+    # Initialized to None.
+    self._current_index = None
+
+    self._range_tracker = range_trackers.OffsetRangeTracker(
+        self._source.start_index, self._source.end_index)
 
   def __enter__(self):
     return self
@@ -66,21 +71,49 @@ def __exit__(self, exception_type, exception_value, traceback):
     pass
 
   def __iter__(self):
-    for value in itertools.islice(self.source.elements,
-                                  self.source.start_index,
-                                  self.source.end_index):
-      self.current_index += 1
-      yield self.source.coder.decode(value)
+    for value in itertools.islice(self._source.elements,
+                                  self._source.start_index,
+                                  self._source.end_index):
+      claimed = False
+      if self._current_index is None:
+        claimed = self._range_tracker.try_claim(
+            self._source.start_index)
+      else:
+        claimed = self._range_tracker.try_claim(
+            self._current_index + 1)
+
+      if claimed:
+        if self._current_index is None:
+          self._current_index = self._source.start_index
+        else:
+          self._current_index += 1
+
+        yield self._source.coder.decode(value)
+      else:
+        return
 
   def get_progress(self):
-    if (self.current_index >= self.source.end_index or
-        self.source.start_index >= self.source.end_index):
-      percent_complete = 1
-    elif self.current_index == self.source.start_index:
-      percent_complete = 0
-    else:
-      percent_complete = (
-          float(self.current_index - self.source.start_index) / (
-              self.source.end_index - self.source.start_index))
-
-    return iobase.ReaderProgress(percent_complete=percent_complete)
+    if self._current_index is None:
+      return None
+
+    return iobase.ReaderProgress(
+        position=iobase.ReaderPosition(record_index=self._current_index))
+
+  def request_dynamic_split(self, dynamic_split_request):
+    assert dynamic_split_request is not None
+    progress = dynamic_split_request.progress
+    split_position = progress.position
+    if split_position is None:
+      logging.debug('InMemory reader only supports split requests that are '
+                    'based on positions. Received : %r', dynamic_split_request)
+      return None
+
+    index_position = split_position.record_index
+    if index_position is None:
+      logging.debug('InMemory reader only supports split requests that are '
+                    'based on index positions. Received : %r',
+                    dynamic_split_request)
+      return None
+
+    if self._range_tracker.try_split(index_position):
+      return iobase.DynamicSplitResultWithPosition(split_position)

google/cloud/dataflow/worker/inmemory_test.py

@@ -17,6 +17,7 @@
 import logging
 import unittest
 
+from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.worker import inmemory
 
 
@@ -41,31 +42,115 @@ def test_norange(self):
   def test_in_memory_source_updates_progress_none(self):
     source = inmemory.InMemorySource([], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(1, reader.get_progress().percent_complete)
+      self.assertEqual(None, reader.get_progress())
 
   def test_in_memory_source_updates_progress_one(self):
     source = inmemory.InMemorySource([1], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(0, reader.get_progress().percent_complete)
+      self.assertEqual(None, reader.get_progress())
       i = 0
       for item in reader:
-        i += 1
+        self.assertEqual(i, reader.get_progress().position.record_index)
         self.assertEqual(11, item)
-        self.assertEqual(1, reader.get_progress().percent_complete)
+        i += 1
       self.assertEqual(1, i)
-      self.assertEqual(1, reader.get_progress().percent_complete)
+      self.assertEqual(0, reader.get_progress().position.record_index)
 
   def test_in_memory_source_updates_progress_many(self):
     source = inmemory.InMemorySource([1, 2, 3, 4, 5], coder=FakeCoder())
     with source.reader() as reader:
-      self.assertEqual(0, reader.get_progress().percent_complete)
+      self.assertEqual(None, reader.get_progress())
       i = 0
       for item in reader:
+        self.assertEqual(i, reader.get_progress().position.record_index)
+        self.assertEqual(11 + i, item)
         i += 1
-        self.assertEqual(i + 10, item)
-        self.assertEqual(float(i) / 5, reader.get_progress().percent_complete)
       self.assertEqual(5, i)
-      self.assertEqual(1, reader.get_progress().percent_complete)
+      self.assertEqual(4, reader.get_progress().position.record_index)
+
+  def try_splitting_reader_at(self, reader, split_request, expected_response):
+    actual_response = reader.request_dynamic_split(split_request)
+
+    if expected_response is None:
+      self.assertIsNone(actual_response)
+    else:
+      self.assertIsNotNone(actual_response.stop_position)
+      self.assertIsInstance(actual_response.stop_position,
+                            iobase.ReaderPosition)
+      self.assertIsNotNone(actual_response.stop_position.record_index)
+      self.assertEqual(expected_response.stop_position.record_index,
+                       actual_response.stop_position.record_index)
+
+  def test_in_memory_source_dynamic_split(self):
+    source = inmemory.InMemorySource([10, 20, 30, 40, 50, 60],
+                                     coder=FakeCoder())
+
+    # Unstarted reader
+    with source.reader() as reader:
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=2))),
+          None)
+
+    # Proposed split position out of range
+    with source.reader() as reader:
+      reader_iter = iter(reader)
+      next(reader_iter)
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=-1))),
+          None)
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=10))),
+          None)
+
+    # Already read past proposed split position
+    with source.reader() as reader:
+      reader_iter = iter(reader)
+      next(reader_iter)
+      next(reader_iter)
+      next(reader_iter)
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=1))),
+          None)
+
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=2))),
+          None)
+
+    # Successful split
+    with source.reader() as reader:
+      reader_iter = iter(reader)
+      next(reader_iter)
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=4))),
+          iobase.DynamicSplitResultWithPosition(
+              stop_position=iobase.ReaderPosition(record_index=4)))
+
+      self.try_splitting_reader_at(
+          reader,
+          iobase.DynamicSplitRequest(
+              iobase.ReaderProgress(
+                  position=iobase.ReaderPosition(record_index=2))),
+          iobase.DynamicSplitResultWithPosition(
+              stop_position=iobase.ReaderPosition(record_index=2)))
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)