profiler session allow more options to be specified at python level. such as host trace level and python tracer enable.

This change is totally backward compatible, because profiler_v2_test have passed.

PiperOrigin-RevId: 304263623
Change-Id: Idf23e8d89d285b02d00dd2c67c8acfed61fb4ee5

tensorflow/python/eager/profiler.py

@@ -78,7 +78,7 @@ def start():
       context.ensure_initialized()
     _profiler = _pywrap_profiler.ProfilerSession()
     try:
-      _profiler.start('')
+      _profiler.start('', {})
     except errors.AlreadyExistsError:
       logging.warning('Another profiler session is running which is probably '
                       'created by profiler server. Please avoid using profiler '

tensorflow/python/profiler/internal/profiler_wrapper.cc

@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "absl/memory/memory.h"
 #include "include/pybind11/pybind11.h"
+#include "include/pybind11/pytypes.h"
 #include "tensorflow/core/platform/host_info.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/profiler/convert/xplane_to_profile_response.h"
@@ -48,8 +49,8 @@ tensorflow::ProfileRequest MakeProfileRequest(
 
 class ProfilerSessionWrapper {
  public:
-  void Start(const char* logdir) {
-    session_ = tensorflow::ProfilerSession::Create();
+  void Start(const char* logdir, const py::dict& options) {
+    session_ = tensorflow::ProfilerSession::Create(GetOptions(options));
     logdir_ = logdir;
     tensorflow::MaybeRaiseRegisteredFromStatus(session_->Status());
   }
@@ -92,6 +93,22 @@ class ProfilerSessionWrapper {
   }
 
  private:
+  tensorflow::profiler::ProfilerOptions GetOptions(const py::dict& opts) {
+    tensorflow::profiler::ProfilerOptions options;
+    for (const auto& kw : opts) {
+      std::string key = py::cast<std::string>(kw.first);
+      if (key == "host_tracer_level") {
+        options.host_tracer_level = py::cast<int>(kw.second);
+        VLOG(1) << "host_tracer_level set to " << options.host_tracer_level;
+      } else if (key == "python_tracer_level") {
+        options.enable_python_tracer = py::cast<int>(kw.second) > 0;
+        VLOG(1) << "enable_python_tracer set to "
+                << options.enable_python_tracer;
+      }
+    }
+    return options;
+  }
+
   std::unique_ptr<tensorflow::ProfilerSession> session_;
   tensorflow::string logdir_;
 };

tensorflow/python/profiler/profiler_v2.py

@@ -34,6 +34,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import threading
 
 from tensorflow.python.framework import errors
@@ -45,19 +46,39 @@ _profiler = None
 _profiler_lock = threading.Lock()
 
 
+@tf_export('profiler.experimental.ProfilerOptions', v1=[])
+class ProfilerOptions(
+    collections.namedtuple('ProfilerOptions',
+                           ['host_tracer_level', 'python_tracer_level'])):
+  """Options to control profiler behaviors.
+
+  A `tf.profiler.ProfilerOptions` hold the knobs to control tf.profiler's
+  behavior.
+
+  Fields:
+    host_tracer_level: for adjust TraceMe levels. i.e. 1 => critical,
+                       2 => info, 3 => verbose. [default to 2]
+    python_tracer_level: for enable python function call tracing, 1 => enable.
+                         0 => disable [default to 0]
+  """
+  pass
+
+
 @tf_export('profiler.experimental.start', v1=[])
-def start(logdir):
+def start(logdir, options=None):
   """Starts profiling.
 
   Args:
     logdir: A log directory read by TensorBoard to export the profile results.
+    options: namedtuple of ProfilerOptions for miscellaneous profiler options.
 
   Raises:
     AlreadyExistsError: If another profiling session is running.
 
   Example usage:
   ```python
-  tf.profiler.experimental.start('logdir_path')
+  tf.profiler.experimental.start(
+      'logdir_path', tf.profiler.ProfilerOptions(host_tracer_level=2))
   # do your training here.
   tf.profiler.experimental.stop()
   ```
@@ -74,7 +95,10 @@ def start(logdir):
                                       'Another profiler is running.')
     _profiler = _pywrap_profiler.ProfilerSession()
     try:
-      _profiler.start(logdir)
+      # support for namedtuple in pybind11 is missing, we change it to
+      # dict type first.
+      opts = dict(options._asdict()) if options is not None else {}
+      _profiler.start(logdir, opts)
     except errors.AlreadyExistsError:
       logging.warning('Another profiler session is running which is probably '
                       'created by profiler server. Please avoid using profiler '
@@ -154,11 +178,19 @@ class Profile(object):
   ```
   """
 
-  def __init__(self, logdir):
+  def __init__(self, logdir, options=None):
+    """Creates a context manager object for profiler API.
+
+    Args:
+      logdir: profile data will save to this directory.
+      options: An optional tf.profiler.ProfilerOptions can be provided to fine
+        tune the profiler's behavior.
+    """
     self._logdir = logdir
+    self._options = options
 
   def __enter__(self):
-    start(self._logdir)
+    start(self._logdir, self._options)
 
   def __exit__(self, typ, value, tb):
     stop()

tensorflow/python/profiler/profiler_v2_test.py

@@ -86,6 +86,35 @@ class ProfilerTest(test_util.TensorFlowTestCase):
     trace_file = os.path.join(profile_dir, run, hostname + '.trace.json.gz')
     self.assertTrue(gfile.Exists(trace_file))
 
+  def test_profile_with_options(self):
+    logdir = self.get_temp_dir()
+    options = profiler.ProfilerOptions(
+        host_tracer_level=3, python_tracer_level=1)
+    profiler.start(logdir, options)
+    with traceme.TraceMe('three_times_five'):
+      three = constant_op.constant(3)
+      five = constant_op.constant(5)
+      product = three * five
+    self.assertAllEqual(15, product)
+
+    profiler.stop()
+    file_list = gfile.ListDirectory(logdir)
+    self.assertEqual(len(file_list), 2)
+
+  def test_context_manager_with_options(self):
+    logdir = self.get_temp_dir()
+    options = profiler.ProfilerOptions(
+        host_tracer_level=3, python_tracer_level=1)
+    with profiler.Profile(logdir, options):
+      with traceme.TraceMe('three_times_five'):
+        three = constant_op.constant(3)
+        five = constant_op.constant(5)
+        product = three * five
+      self.assertAllEqual(15, product)
+
+    file_list = gfile.ListDirectory(logdir)
+    self.assertEqual(len(file_list), 2)
+
 
 if __name__ == '__main__':
   test.main()

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt

@@ -4,6 +4,6 @@ tf_class {
   is_instance: "<type \'object\'>"
   member_method {
     name: "__init__"
-    argspec: "args=[\'self\', \'logdir\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'self\', \'logdir\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
 }

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profiler-options.pbtxt

@@ -0,0 +1,23 @@
+path: "tensorflow.profiler.experimental.ProfilerOptions"
+tf_class {
+  is_instance: "<class \'tensorflow.python.profiler.profiler_v2.ProfilerOptions\'>"
+  is_instance: "<class \'tensorflow.python.profiler.profiler_v2.ProfilerOptions\'>"
+  is_instance: "<type \'tuple\'>"
+  member {
+    name: "host_tracer_level"
+    mtype: "<type \'property\'>"
+  }
+  member {
+    name: "python_tracer_level"
+    mtype: "<type \'property\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "count"
+  }
+  member_method {
+    name: "index"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt

@@ -4,6 +4,10 @@ tf_module {
     name: "Profile"
     mtype: "<type \'type\'>"
   }
+  member {
+    name: "ProfilerOptions"
+    mtype: "<type \'type\'>"
+  }
   member {
     name: "Trace"
     mtype: "<type \'type\'>"
@@ -18,7 +22,7 @@ tf_module {
   }
   member_method {
     name: "start"
-    argspec: "args=[\'logdir\'], varargs=None, keywords=None, defaults=None"
+    argspec: "args=[\'logdir\', \'options\'], varargs=None, keywords=None, defaults=[\'None\'], "
   }
   member_method {
     name: "stop"