Make sdk pipeline options available in the DoFn context

There is also a bug fix in get_all_options() to apply the overrides.
----Release Notes----
Adds a new API DoFnProcessContext.get_sdk_pipeline_options() that returns
the pipeline options.
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=118083641

Author: aaltay

google/cloud/dataflow/pipeline_test.py

@@ -255,10 +255,10 @@ def test_defaults(self):
   def test_dir(self):
     options = Breakfast()
     self.assertEquals(
-        ['get_all_options', 'slices', 'style', 'view_as'],
+        ['from_dictionary', 'get_all_options', 'slices', 'style', 'view_as'],
         [attr for attr in dir(options) if not attr.startswith('_')])
     self.assertEquals(
-        ['get_all_options', 'style', 'view_as'],
+        ['from_dictionary', 'get_all_options', 'style', 'view_as'],
         [attr for attr in dir(options.view_as(Eggs))
          if not attr.startswith('_')])
 

google/cloud/dataflow/runners/common.py

@@ -171,9 +171,14 @@ def existing_windows(self):
 
 class DoFnState(object):
   """Keeps track of state that DoFns want, currently, user counters.
+
+  Attributes:
+    pipeline_options: a PipelineOptions object associated with this DoFn.
+    step_name: name of the step as a string.
   """
 
-  def __init__(self):
+  def __init__(self, pipeline_options):
+    self.pipeline_options = pipeline_options
     self.step_name = ''
     self._user_counters = {}
 

google/cloud/dataflow/runners/direct_runner.py

@@ -77,8 +77,10 @@ def func_wrapper(self, pvalue, *args, **kwargs):
   @skip_if_cached
   def run_ParDo(self, transform_node):
     transform = transform_node.transform
+    options = transform_node.inputs[0].pipeline.options
     # TODO(gildea): what is the appropriate object to attach the state to?
-    context = DoFnProcessContext(label=transform.label, state=DoFnState())
+    context = DoFnProcessContext(label=transform.label,
+                                 state=DoFnState(options))
 
     # Construct the list of values from side-input PCollections that we'll
     # substitute into the arguments for DoFn methods.
@@ -105,7 +107,6 @@ def get_side_input_value(si):
 
     # TODO(robertwb): Do this type checking inside DoFnRunner to get it on
     # remote workers as well?
-    options = transform_node.inputs[0].pipeline.options
     if options is not None and options.view_as(TypeOptions).runtime_type_check:
       transform.dofn = TypeCheckWrapperDoFn(
           transform.dofn, transform.get_type_hints())

google/cloud/dataflow/transforms/core.py

@@ -61,6 +61,7 @@ class DoFnProcessContext(object):
       (in process method only; always None in start_bundle and finish_bundle)
     windows: windows of the element
       (in process method only; always None in start_bundle and finish_bundle)
+    pipeline_options: PipelineOptions object used for creating the pipeline.
     state: a DoFnState object, which holds the runner's internal state
       for this element.  For example, aggregator state is here.
       Not used by the pipeline code.
@@ -106,6 +107,10 @@ def aggregate_to(self, aggregator, input_value):
     """
     self.state.counter_for(aggregator).update(input_value)
 
+  @property
+  def pipeline_options(self):
+    return self.state.pipeline_options
+
 
 class DoFn(WithTypeHints):
   """A function object used by a transform with custom processing.

google/cloud/dataflow/utils/options.py

@@ -81,6 +81,26 @@ def _add_argparse_args(cls, parser):
     # Override this in subclasses to provide options.
     pass
 
+  @classmethod
+  def from_dictionary(cls, options):
+    """Returns a PipelineOptions from a dictionary of arguments.
+
+    Args:
+      options: Dictinary of argument value pairs.
+
+    Returns:
+      A PipelineOptions object representing the given arguments.
+    """
+    flags = []
+    for k, v in options.iteritems():
+      if isinstance(v, bool):
+        if v:
+          flags.append('--%s' % k)
+      else:
+        flags.append('--%s=%s' % (k, v))
+
+    return cls(flags)
+
   def get_all_options(self):
     """Returns a dictionary of all defined arguments.
 
@@ -94,7 +114,14 @@ def get_all_options(self):
     for cls in PipelineOptions.__subclasses__():
       cls._add_argparse_args(parser)  # pylint: disable=protected-access
     known_args, _ = parser.parse_known_args(self._flags)
-    return vars(known_args)
+    result = vars(known_args)
+
+    # Apply the overrides if any
+    for k in result:
+      if k in self._all_options:
+        result[k] = self._all_options[k]
+
+    return result
 
   def view_as(self, cls):
     view = cls(self._flags)

google/cloud/dataflow/utils/pipeline_options_test.py

@@ -20,38 +20,84 @@
 from google.cloud.dataflow.utils.options import PipelineOptions
 
 
-class SetupTest(unittest.TestCase):
-
-  def test_get_unknown_args(self):
-
-    # Used for testing newly added flags.
-    class MockOptions(PipelineOptions):
-
-      @classmethod
-      def _add_argparse_args(cls, parser):
-        parser.add_argument('--mock_flag',
-                            action='store_true',
-                            help='Enable work item profiling')
-
-    test_cases = [
-        {'flags': ['--num_workers', '5'],
-         'expected': {'num_workers': 5, 'mock_flag': False}},
-        {
-            'flags': [
-                '--profile', '--profile_location', 'gs://bucket/', 'ignored'],
-            'expected': {
-                'profile': True, 'profile_location': 'gs://bucket/',
-                'mock_flag': False}
-        },
-        {'flags': ['--num_workers', '5', '--mock_flag'],
-         'expected': {'num_workers': 5, 'mock_flag': True}},
-    ]
-
-    for case in test_cases:
+class PipelineOptionsTest(unittest.TestCase):
+
+  TEST_CASES = [
+      {'flags': ['--num_workers', '5'],
+       'expected': {'num_workers': 5, 'mock_flag': False, 'mock_option': None}},
+      {
+          'flags': [
+              '--profile', '--profile_location', 'gs://bucket/', 'ignored'],
+          'expected': {
+              'profile': True, 'profile_location': 'gs://bucket/',
+              'mock_flag': False, 'mock_option': None}
+      },
+      {'flags': ['--num_workers', '5', '--mock_flag'],
+       'expected': {'num_workers': 5, 'mock_flag': True, 'mock_option': None}},
+      {'flags': ['--mock_option', 'abc'],
+       'expected': {'mock_flag': False, 'mock_option': 'abc'}},
+      {'flags': ['--mock_option', ' abc def '],
+       'expected': {'mock_flag': False, 'mock_option': ' abc def '}},
+      {'flags': ['--mock_option= abc xyz '],
+       'expected': {'mock_flag': False, 'mock_option': ' abc xyz '}},
+      {'flags': ['--mock_option=gs://my bucket/my folder/my file'],
+       'expected': {'mock_flag': False,
+                    'mock_option': 'gs://my bucket/my folder/my file'}},
+  ]
+
+  # Used for testing newly added flags.
+  class MockOptions(PipelineOptions):
+
+    @classmethod
+    def _add_argparse_args(cls, parser):
+      parser.add_argument('--mock_flag', action='store_true', help='mock flag')
+      parser.add_argument('--mock_option', help='mock option')
+      parser.add_argument('--option with space', help='mock option with space')
+
+  def test_get_all_options(self):
+    for case in PipelineOptionsTest.TEST_CASES:
       options = PipelineOptions(flags=case['flags'])
       self.assertDictContainsSubset(case['expected'], options.get_all_options())
-      self.assertEqual(options.view_as(MockOptions).mock_flag,
+      self.assertEqual(options.view_as(
+          PipelineOptionsTest.MockOptions).mock_flag,
                        case['expected']['mock_flag'])
+      self.assertEqual(options.view_as(
+          PipelineOptionsTest.MockOptions).mock_option,
+                       case['expected']['mock_option'])
+
+  def test_from_dictionary(self):
+    for case in PipelineOptionsTest.TEST_CASES:
+      options = PipelineOptions(flags=case['flags'])
+      all_options_dict = options.get_all_options()
+      options_from_dict = PipelineOptions.from_dictionary(all_options_dict)
+      self.assertEqual(options_from_dict.view_as(
+          PipelineOptionsTest.MockOptions).mock_flag,
+                       case['expected']['mock_flag'])
+      self.assertEqual(options.view_as(
+          PipelineOptionsTest.MockOptions).mock_option,
+                       case['expected']['mock_option'])
+
+  def test_option_with_spcae(self):
+    options = PipelineOptions(flags=['--option with space= value with space'])
+    self.assertEqual(
+        getattr(options.view_as(PipelineOptionsTest.MockOptions),
+                'option with space'), ' value with space')
+    options_from_dict = PipelineOptions.from_dictionary(
+        options.get_all_options())
+    self.assertEqual(
+        getattr(options_from_dict.view_as(PipelineOptionsTest.MockOptions),
+                'option with space'), ' value with space')
+
+  def test_override_options(self):
+    base_flags = ['--num_workers', '5']
+    options = PipelineOptions(base_flags)
+    self.assertEqual(options.get_all_options()['num_workers'], 5)
+    self.assertEqual(options.get_all_options()['mock_flag'], False)
+
+    options.view_as(PipelineOptionsTest.MockOptions).mock_flag = True
+    self.assertEqual(options.get_all_options()['num_workers'], 5)
+    self.assertEqual(options.get_all_options()['mock_flag'], True)
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)

google/cloud/dataflow/worker/batchworker.py

@@ -49,6 +49,7 @@
 from google.cloud.dataflow.internal import auth
 from google.cloud.dataflow.internal import pickler
 from google.cloud.dataflow.utils import names
+from google.cloud.dataflow.utils import options
 from google.cloud.dataflow.utils import profiler
 from google.cloud.dataflow.utils import retry
 from google.cloud.dataflow.worker import executor
@@ -76,6 +77,8 @@ def __init__(self, properties, sdk_pipeline_options):
     self.job_id = properties['job_id']
     self.worker_id = properties['worker_id']
     self.service_path = properties['service_path']
+    self.pipeline_options = options.PipelineOptions.from_dictionary(
+        sdk_pipeline_options)
     self.capabilities = [self.worker_id, 'remote_source', 'custom_source']
     self.work_types = ['map_task', 'seq_map_task', 'remote_source_task']
     # The following properties are passed to the worker when its container
@@ -402,7 +405,8 @@ def do_work(self, work_item):
         self.dynamic_split_result_to_report = None
 
         self.set_current_work_item_and_executor(work_item,
-                                                executor.MapTaskExecutor())
+                                                executor.MapTaskExecutor(
+                                                    self.pipeline_options))
       self.report_progress = True
       self.current_executor.execute(work_item.map_task)
     except Exception:  # pylint: disable=broad-except

google/cloud/dataflow/worker/executor.py

@@ -383,9 +383,9 @@ def process(self, o):
 class DoOperation(Operation):
   """A Do operation that will execute a custom DoFn for each input element."""
 
-  def __init__(self, spec):
+  def __init__(self, spec, pipeline_options):
     super(DoOperation, self).__init__(spec)
-    self.state = common.DoFnState()
+    self.state = common.DoFnState(pipeline_options)
 
   def _read_side_inputs(self, tags_and_types):
     """Generator reading side inputs in the order prescribed by tags_and_types.
@@ -728,7 +728,8 @@ class MapTaskExecutor(object):
   multiple_read_instruction_error_msg = (
       'Found more than one \'read instruction\' in a single \'map task\'')
 
-  def __init__(self):
+  def __init__(self, pipeline_options=None):
+    self.pipeline_options = pipeline_options
     self._ops = []
     self._read_operation = None
 
@@ -775,7 +776,7 @@ def execute(self, map_task, test_shuffle_source=None, test_shuffle_sink=None):
       elif isinstance(spec, maptask.WorkerPartialGroupByKey):
         op = create_pgbk_op(spec)
       elif isinstance(spec, maptask.WorkerDoFn):
-        op = DoOperation(spec)
+        op = DoOperation(spec, self.pipeline_options)
       elif isinstance(spec, maptask.WorkerGroupingShuffleRead):
         op = GroupedShuffleReadOperation(
             spec, shuffle_source=test_shuffle_source)