Adds support for reading custom sources using DataflowPipelineRunner.

Adds support for performing custom source split operations and reading sub-sources generated by split operations.

Generalizes WorkItem executing by introducing an interface 'Executor'. Adds an executor for performing custom source split operations.

Adds a local runner based integration test for custom sources.

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

Author: chamikaramj

google/cloud/dataflow/internal/apiclient.py

@@ -21,13 +21,16 @@
 import re
 import time
 
+
 from google.cloud.dataflow import utils
 from google.cloud.dataflow import version
+from google.cloud.dataflow.internal import pickler
 from google.cloud.dataflow.internal.auth import get_service_credentials
 from google.cloud.dataflow.internal.json_value import to_json_value
 from google.cloud.dataflow.io import iobase
 from google.cloud.dataflow.transforms import cy_combiners
 from google.cloud.dataflow.utils import dependency
+from google.cloud.dataflow.utils import names
 from google.cloud.dataflow.utils import retry
 from google.cloud.dataflow.utils.names import PropertyNames
 from google.cloud.dataflow.utils.options import GoogleCloudOptions
@@ -665,6 +668,7 @@ def report_status(self,
                     completed,
                     progress,
                     dynamic_split_result_to_report=None,
+                    source_operation_response=None,
                     exception_details=None):
     """Reports status for a work item (success or failure).
 
@@ -688,6 +692,9 @@ def report_status(self,
         handling the work item.
       dynamic_split_result_to_report: A successful dynamic split result that
         should be sent to the Dataflow service along with the status report.
+      source_operation_response: Response to a source operation request from
+        the service. This will be sent to the service along with the status
+        report.
       exception_details: A string representation of the stack trace for an
         exception raised while executing the work item. The string is the
         output of the standard traceback.format_exc() function.
@@ -746,9 +753,13 @@ def report_status(self,
       status.message = exception_details
       work_item_status.errors.append(status)
 
+    if source_operation_response is not None:
+      work_item_status.sourceOperationResponse = source_operation_response
+
     # Look through the work item for metrics to send.
-    for counter in work_item.map_task.itercounters():
-      append_counter(work_item_status, counter, tentative=not completed)
+    if work_item.map_task:
+      for counter in work_item.map_task.itercounters():
+        append_counter(work_item_status, counter, tentative=not completed)
 
     report_request = dataflow.ReportWorkItemStatusRequest()
     report_request.currentWorkerTime = worker_info.formatted_current_time
@@ -871,3 +882,44 @@ def set_mean(accumulator, metric_update):
     cy_combiners.AllCombineFn: ('and', set_scalar),
     cy_combiners.AnyCombineFn: ('or', set_scalar),
 }
+
+
+def splits_to_split_response(bundles):
+  """Generates a response to a custom source split request.
+
+  Args:
+    bundles: a set of bundles generated by a BoundedSource.split() invocation.
+  Returns:
+   a SourceOperationResponse object.
+  """
+  derived_sources = []
+  for bundle in bundles:
+    derived_source = dataflow.DerivedSource()
+    derived_source.derivationMode = (
+        dataflow.DerivedSource.DerivationModeValueValuesEnum
+        .SOURCE_DERIVATION_MODE_INDEPENDENT)
+    derived_source.source = dataflow.Source()
+    derived_source.source.doesNotNeedSplitting = True
+
+    derived_source.source.spec = dataflow.Source.SpecValue()
+    derived_source.source.spec.additionalProperties.append(
+        dataflow.Source.SpecValue.AdditionalProperty(
+            key=names.SERIALIZED_SOURCE_KEY,
+            value=to_json_value(pickler.dumps(
+                (bundle.source, bundle.start_position, bundle.stop_position)),
+                                with_type=True)))
+    derived_source.source.spec.additionalProperties.append(
+        dataflow.Source.SpecValue.AdditionalProperty(key='@type',
+                                                     value=to_json_value(
+                                                         names.SOURCE_TYPE)))
+    derived_sources.append(derived_source)
+
+  split_response = dataflow.SourceSplitResponse()
+  split_response.bundles = derived_sources
+  split_response.outcome = (
+      dataflow.SourceSplitResponse.OutcomeValueValuesEnum
+      .SOURCE_SPLIT_OUTCOME_SPLITTING_HAPPENED)
+
+  response = dataflow.SourceOperationResponse()
+  response.split = split_response
+  return response

google/cloud/dataflow/runners/dataflow_runner.py

@@ -34,6 +34,7 @@
 from google.cloud.dataflow.runners.runner import PipelineState
 from google.cloud.dataflow.runners.runner import PValueCache
 from google.cloud.dataflow.typehints import typehints
+from google.cloud.dataflow.utils import names
 from google.cloud.dataflow.utils.names import PropertyNames
 from google.cloud.dataflow.utils.names import TransformNames
 from google.cloud.dataflow.utils.options import StandardOptions
@@ -456,12 +457,17 @@ def run_Read(self, transform_node):
     # TODO(mairbek): refactor if-else tree to use registerable functions.
     # Initialize the source specific properties.
 
-    if isinstance(transform.source, iobase.BoundedSource):
-      raise ValueError('DataflowPipelineRunner does not support reading '
-                       'BoundedSource implementations yet. Please use a source '
-                       'provided by Dataflow SDK or use DirectPipelineRunner.')
-
-    if transform.source.format == 'text':
+    if not hasattr(transform.source, 'format'):
+      # If a format is not set, we assume the source to be a custom source.
+      source_dict = dict()
+      spec_dict = dict()
+
+      spec_dict[names.SERIALIZED_SOURCE_KEY] = pickler.dumps(transform.source)
+      spec_dict['@type'] = names.SOURCE_TYPE
+      source_dict['spec'] = spec_dict
+      step.add_property(PropertyNames.SOURCE_STEP_INPUT,
+                        source_dict)
+    elif transform.source.format == 'text':
       step.add_property(PropertyNames.FILE_PATTERN, transform.source.path)
     elif transform.source.format == 'bigquery':
       # TODO(silviuc): Add table validation if transform.source.validate.
@@ -494,15 +500,22 @@ def run_Read(self, transform_node):
       if transform.source.id_label:
         step.add_property(PropertyNames.PUBSUB_ID_LABEL,
                           transform.source.id_label)
-    elif transform.source.format == 'custom':
-      # TODO(silviuc): Implement custom sources.
-      raise NotImplementedError
     else:
       raise ValueError(
           'Source %r has unexpected format %s.' % (
               transform.source, transform.source.format))
-    step.add_property(PropertyNames.FORMAT, transform.source.format)
-    step.encoding = self._get_cloud_encoding(transform.source.coder)
+
+    if not hasattr(transform.source, 'format'):
+      step.add_property(PropertyNames.FORMAT, names.SOURCE_FORMAT)
+    else:
+      step.add_property(PropertyNames.FORMAT, transform.source.format)
+
+    if isinstance(transform.source, iobase.BoundedSource):
+      coder = transform.source.default_output_coder()
+    else:
+      coder = transform.source.coder
+
+    step.encoding = self._get_cloud_encoding(coder)
     step.add_property(
         PropertyNames.OUTPUT_INFO,
         [{PropertyNames.USER_NAME: (

google/cloud/dataflow/utils/names.py

@@ -19,6 +19,11 @@
 PICKLED_MAIN_SESSION_FILE = 'pickled_main_session'
 DATAFLOW_SDK_TARBALL_FILE = 'dataflow_python_sdk.tar'
 
+# String constants related to sources framework
+SOURCE_FORMAT = 'custom_source'
+SOURCE_TYPE = 'CustomSourcesType'
+SERIALIZED_SOURCE_KEY = 'serialized_source'
+
 
 class TransformNames(object):
   """Transform strings as they are expected in the CloudWorkflow protos."""
@@ -61,6 +66,7 @@ class PropertyNames(object):
   PUBSUB_ID_LABEL = 'pubsub_id_label'
   SERIALIZED_FN = 'serialized_fn'
   SHARD_NAME_TEMPLATE = 'shard_template'
+  SOURCE_STEP_INPUT = 'custom_source_step_input'
   STEP_NAME = 'step_name'
   USER_FN = 'user_fn'
   USER_NAME = 'user_name'

google/cloud/dataflow/worker/batchworker.py

@@ -201,6 +201,7 @@ def cloud_time_to_timestamp(self, cloud_time_string):
   def report_status(self,
                     completed=False,
                     progress=None,
+                    source_operation_response=None,
                     exception_details=None):
     """Reports to the service status of a work item (completion or progress).
 
@@ -209,6 +210,7 @@ def report_status(self,
         either because it succeeded or because it failed. False if this is a
         progress report.
       progress: Progress of processing the work_item.
+      source_operation_response: Response to a custom source operation
       exception_details: A string representation of the stack trace for an
         exception raised while executing the work item. The string is the
         output of the standard traceback.format_exc() function.
@@ -226,7 +228,8 @@ def report_status(self,
         completed,
         progress if not completed else None,
         self.dynamic_split_result_to_report if not completed else None,
-        exception_details)
+        source_operation_response=source_operation_response,
+        exception_details=exception_details)
 
     # Resetting dynamic_split_result_to_report after reporting status
     # successfully.
@@ -368,6 +371,7 @@ def report_completion_status(
       self,
       current_work_item,
       progress_reporter,
+      source_operation_response=None,
       exception_details=None):
     """Reports to the service a work item completion (successful or failed).
 
@@ -383,6 +387,7 @@ def report_completion_status(
       current_work_item: A WorkItem instance describing the work.
       progress_reporter: A ProgressReporter configured to process work item
         current_work_item.
+      source_operation_response: Response to a custom source operation.
       exception_details: A string representation of the stack trace for an
         exception raised while executing the work item. The string is the
         output of the standard traceback.format_exc() function.
@@ -395,8 +400,10 @@ def report_completion_status(
                  'successfully' if exception_details is None
                  else 'with exception')
 
-    progress_reporter.report_status(completed=True,
-                                    exception_details=exception_details)
+    progress_reporter.report_status(
+        completed=True,
+        source_operation_response=source_operation_response,
+        exception_details=exception_details)
 
   @staticmethod
   def log_memory_usage_if_needed(worker_id, force=False):
@@ -416,12 +423,21 @@ def log_memory_usage_if_needed(worker_id, force=False):
   def shutdown(self):
     self._shutdown = True
 
+  def get_executor_for_work_item(self, work_item):
+    if work_item.map_task is not None:
+      return executor.MapTaskExecutor(work_item.map_task)
+    elif work_item.source_operation_split_task is not None:
+      return executor.CustomSourceSplitExecutor(
+          work_item.source_operation_split_task)
+    else:
+      raise ValueError('Unknown type of work item : %s', work_item)
+
   def do_work(self, work_item, deferred_exception_details=None):
     """Executes worker operations and adds any failures to the report status."""
     logging.info('Executing %s', work_item)
     BatchWorker.log_memory_usage_if_needed(self.worker_id, force=True)
 
-    work_executor = executor.MapTaskExecutor()
+    work_executor = self.get_executor_for_work_item(work_item)
     progress_reporter = ProgressReporter(
         work_item, work_executor, self, self.client)
 
@@ -441,7 +457,7 @@ def do_work(self, work_item, deferred_exception_details=None):
     exception_details = None
     try:
       progress_reporter.start_reporting_progress()
-      work_executor.execute(work_item.map_task)
+      work_executor.execute()
     except Exception:  # pylint: disable=broad-except
       exception_details = traceback.format_exc()
       logging.error('An exception was raised when trying to execute the '
@@ -464,8 +480,14 @@ def do_work(self, work_item, deferred_exception_details=None):
           exception_details = traceback.format_exc()
 
       with work_item.lock:
-        self.report_completion_status(work_item, progress_reporter,
-                                      exception_details=exception_details)
+        source_split_response = None
+        if isinstance(work_executor, executor.CustomSourceSplitExecutor):
+          source_split_response = work_executor.response
+
+        self.report_completion_status(
+            work_item, progress_reporter,
+            source_operation_response=source_split_response,
+            exception_details=exception_details)
         work_item.done = True
 
   def status_server(self):
@@ -559,9 +581,13 @@ def run(self):
           time.sleep(1.0 * (1 - 0.5 * random.random()))
           continue
 
+        stage_name = None
+        if work_item.map_task:
+          stage_name = work_item.map_task.stage_name
+
         with logger.PerThreadLoggingContext(
             work_item_id=work_item.proto.id,
-            stage_name=work_item.map_task.stage_name):
+            stage_name=stage_name):
           # TODO(silviuc): Add more detailed timing and profiling support.
           start_time = time.time()
 

google/cloud/dataflow/worker/batchworker_test.py

@@ -87,9 +87,9 @@ def test_worker_starts_and_stops_progress_reporter(
     worker.do_work(mock_work_item)
 
     mock_report_status.assert_called_with(
-        completed=True, exception_details=None)
+        completed=True, source_operation_response=None, exception_details=None)
     mock_start.assert_called_once_with()
-    mock_execute.assert_called_once_with(mock.ANY)
+    mock_execute.assert_called_once_with()
     mock_stop.assert_called_once_with()
 
   @patch.object(executor.MapTaskExecutor, 'execute')
@@ -103,7 +103,9 @@ def test_worker_fails_for_deferred_exceptions(
     worker.do_work(mock_work_item, deferred_exception_details='deferred_exc')
 
     mock_report_status.assert_called_with(
-        completed=True, exception_details='deferred_exc')
+        completed=True,
+        source_operation_response=None,
+        exception_details='deferred_exc')
     assert not mock_stop.called
     assert not mock_start.called
     assert not mock_execute.called
@@ -121,10 +123,11 @@ def __eq__(self, other):
 
     mock_report_status.assert_called_with(
         completed=True,
+        source_operation_response=None,
         exception_details=AnyStringWith(expected_exception))
 
     mock_start.assert_called_once_with()
-    mock_execute.assert_called_once_with(mock.ANY)
+    mock_execute.assert_called_once_with()
     mock_stop.assert_called_once_with()
 
   @patch.object(executor.MapTaskExecutor, 'execute')
@@ -167,8 +170,8 @@ class ProgressReporterTest(unittest.TestCase):
   @patch.object(batchworker.ProgressReporter, 'process_report_status_response')
   def test_progress_reporter_reports_progress(
       self, mock_report_response, mock_next_progress):  # pylint: disable=unused-argument
-    work_item = workitem.BatchWorkItem(
-        proto=mock.MagicMock(), map_task=mock.MagicMock())
+    work_item = workitem.BatchWorkItem(proto=mock.MagicMock())
+    work_item.map_task = mock.MagicMock()
     mock_work_executor = mock.MagicMock()
     mock_batch_worker = mock.MagicMock()
     mock_client = mock.MagicMock()
@@ -181,7 +184,8 @@ def test_progress_reporter_reports_progress(
     time.sleep(10)
     progress_reporter.stop_reporting_progress()
     mock_client.report_status.assert_called_with(
-        mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY)
+        mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY,
+        exception_details=mock.ANY, source_operation_response=mock.ANY)
 
   @patch.object(batchworker.ProgressReporter, 'next_progress_report_interval')
   @patch.object(batchworker.ProgressReporter, 'process_report_status_response')
@@ -200,7 +204,7 @@ def test_progress_reporter_sends_last_update(
     progress_reporter.stop_reporting_progress()
     mock_client.report_status.assert_called_with(
         mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock_split_result,
-        mock.ANY)
+        exception_details=mock.ANY, source_operation_response=mock.ANY)
 
 
 if __name__ == '__main__':