Run batchworker do worker under cProfile

Captures all defined command line options with a value and sends to
the worker as sdk_pipeline_options.

Batchworker supports two new options:
--profile : Flag to enable per work item profiling
--profile_location gs:[] : gcs location to save profile
results.

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

Author: aaltay

google/cloud/dataflow/internal/apiclient.py

@@ -262,6 +262,17 @@ def __init__(self, packages, options, environment_version):
       pool.dataDisks.append(disk)
     self.proto.workerPools.append(pool)
 
+    sdk_pipeline_options = options.get_all_options()
+    if sdk_pipeline_options:
+      self.proto.sdkPipelineOptions = (
+          dataflow.Environment.SdkPipelineOptionsValue())
+
+      for k, v in sdk_pipeline_options.iteritems():
+        if v is not None:
+          self.proto.sdkPipelineOptions.additionalProperties.append(
+              dataflow.Environment.SdkPipelineOptionsValue.AdditionalProperty(
+                  key=k, value=to_json_value(v)))
+
 
 class Job(object):
   """Wrapper for a dataflow Job protobuf."""

google/cloud/dataflow/pipeline_test.py

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

google/cloud/dataflow/utils/options.py

@@ -19,11 +19,6 @@
 
 import argparse
 
-# Raw (unparsed) options. They are also added by other modules that want to
-# contribute modules other than the ones defined in this file. See add_option(),
-# below.
-OPTIONS = []
-
 
 class PipelineOptions(object):
   """Pipeline options class used as container for command line options.
@@ -86,6 +81,21 @@ def _add_argparse_args(cls, parser):
     # Override this in subclasses to provide options.
     pass
 
+  def get_all_options(self):
+    """Returns a dictionary of all defined arguments.
+
+    Returns a dictionary of all defined arguments (arguments that are defined in
+    any subclass of PipelineOptions) into a dictionary.
+
+    Returns:
+      Dictionary of all args and values.
+    """
+    parser = argparse.ArgumentParser()
+    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)
+
   def view_as(self, cls):
     view = cls(self._flags)
     view._all_options = self._all_options
@@ -300,6 +310,18 @@ def _add_argparse_args(cls, parser):
                         help='Debug file to write the workflow specification.')
 
 
+class ProfilingOptions(PipelineOptions):
+
+  @classmethod
+  def _add_argparse_args(cls, parser):
+    parser.add_argument('--profile',
+                        action='store_true',
+                        help='Enable work item profiling')
+    parser.add_argument('--profile_location',
+                        default=None,
+                        help='GCS path for saving profiler data.')
+
+
 class SetupOptions(PipelineOptions):
 
   @classmethod

google/cloud/dataflow/utils/pipeline_options_test.py

@@ -0,0 +1,58 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for the pipeline options module."""
+
+import logging
+import unittest
+
+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:
+      options = PipelineOptions(flags=case['flags'])
+      self.assertDictContainsSubset(case['expected'], options.get_all_options())
+      self.assertEqual(options.view_as(MockOptions).mock_flag,
+                       case['expected']['mock_flag'])
+
+if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
+  unittest.main()

google/cloud/dataflow/utils/profiler.py

@@ -0,0 +1,66 @@
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A profiler context manager based on cProfile.Profile objects."""
+
+import cProfile
+import logging
+import os
+import pstats
+import StringIO
+import tempfile
+import time
+
+
+from google.cloud.dataflow.utils.dependency import _dependency_file_copy
+
+
+class Profile(object):
+  """cProfile wrapper context for saving and logging profiler results."""
+
+  SORTBY = 'cumulative'
+
+  def __init__(self, profile_id, profile_location=None, log_results=False):
+    self.stats = None
+    self.profile_id = str(profile_id)
+    self.profile_location = profile_location
+    self.log_results = log_results
+
+  def __enter__(self):
+    logging.info('Start profiling: %s', self.profile_id)
+    self.profile = cProfile.Profile()
+    self.profile.enable()
+    return self
+
+  def __exit__(self, *args):
+    self.profile.disable()
+    logging.info('Stop profiling: %s', self.profile_id)
+
+    if self.profile_location:
+      dump_location = os.path.join(
+          self.profile_location, 'profile',
+          ('%s-%s' % (time.strftime('%Y-%m-%d_%H_%M_%S'), self.profile_id)))
+      fd, filename = tempfile.mkstemp()
+      self.profile.dump_stats(filename)
+      logging.info('Copying profiler data to: [%s]', dump_location)
+      _dependency_file_copy(filename, dump_location)  # pylint: disable=protected-access
+      os.close(fd)
+      os.remove(filename)
+
+    if self.log_results:
+      s = StringIO.StringIO()
+      self.stats = pstats.Stats(
+          self.profile, stream=s).sort_stats(Profile.SORTBY)
+      self.stats.print_stats()
+      logging.info('Profiler data: [%s]', s.getvalue())

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 profiler
 from google.cloud.dataflow.utils import retry
 from google.cloud.dataflow.worker import executor
 from google.cloud.dataflow.worker import logger
@@ -69,7 +70,7 @@ class BatchWorker(object):
   STATUS_HTTP_PORT = 0  # A value of 0 will pick a random unused port.
   MEMORY_USAGE_REPORTING_INTERVAL_SECS = 5 * 60
 
-  def __init__(self, properties):
+  def __init__(self, properties, sdk_pipeline_options):
     """Initializes a worker object from command line arguments."""
     self.project_id = properties['project_id']
     self.job_id = properties['job_id']
@@ -109,6 +110,11 @@ def __init__(self, properties):
     # the currently set work item; does not send progress updates otherwise.
     self.report_progress = False
 
+    # If 'True' each work item will be profiled with cProfile. Results will
+    # be logged and also saved to profile_location if set.
+    self.work_item_profiling = sdk_pipeline_options.get('profile', False)
+    self.profile_location = sdk_pipeline_options.get('profile_location', None)
+
   @property
   def current_work_item(self):
     with self.lock:
@@ -518,7 +524,15 @@ def run(self):
             # failed.  The progress reporting_thread will take care of sending
             # updates and updating in the workitem object the reporting indexes
             # and duration for the lease.
-            self.do_work(work_item)
+
+            if self.work_item_profiling:
+              with profiler.Profile(
+                  profile_id=work_item.proto.id,
+                  profile_location=self.profile_location, log_results=True):
+                self.do_work(work_item)
+            else:
+              self.do_work(work_item)
+
           logging.info('Completed work item: %s in %.9f seconds',
                        work_item.proto.id, time.time() - start_time)
 

google/cloud/dataflow/worker/start.py

@@ -17,6 +17,7 @@
 # thread-safety issue with datetime.datetime.strptime if this module is not
 # already imported.
 import _strptime  # pylint: disable=unused-import
+import json
 import logging
 import random
 import re
@@ -49,6 +50,11 @@ def main():
 
   logging.info('Worker started with properties: %s', properties)
 
+  sdk_pipeline_options = json.loads(
+      properties.get('sdk_pipeline_options', '{}'))
+  logging.info('Worker started with sdk_pipeline_options: %s',
+               sdk_pipeline_options)
+
   if unused_args:
     logging.warning('Unrecognized arguments %s', unused_args)
 
@@ -64,7 +70,7 @@ def main():
     streamingworker.StreamingWorker(properties).run()
   else:
     logging.info('Starting batch worker.')
-    batchworker.BatchWorker(properties).run()
+    batchworker.BatchWorker(properties, sdk_pipeline_options).run()
 
 
 if __name__ == '__main__':