From 4e6905a35a18c633164f850edad10cb8caf651e0 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 1 Apr 2020 14:32:11 -0700 Subject: [PATCH] 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 | 2 +- .../profiler/internal/profiler_wrapper.cc | 21 +++++++++- tensorflow/python/profiler/profiler_v2.py | 42 ++++++++++++++++--- .../python/profiler/profiler_v2_test.py | 29 +++++++++++++ ...rflow.profiler.experimental.-profile.pbtxt | 2 +- ...filer.experimental.-profiler-options.pbtxt | 23 ++++++++++ .../v2/tensorflow.profiler.experimental.pbtxt | 6 ++- 7 files changed, 115 insertions(+), 10 deletions(-) create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profiler-options.pbtxt diff --git a/tensorflow/python/eager/profiler.py b/tensorflow/python/eager/profiler.py index 91761986107..beb1e6fcf78 100644 --- a/tensorflow/python/eager/profiler.py +++ b/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 ' diff --git a/tensorflow/python/profiler/internal/profiler_wrapper.cc b/tensorflow/python/profiler/internal/profiler_wrapper.cc index eaf7d09105b..b76f28c2a74 100644 --- a/tensorflow/python/profiler/internal/profiler_wrapper.cc +++ b/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(kw.first); + if (key == "host_tracer_level") { + options.host_tracer_level = py::cast(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(kw.second) > 0; + VLOG(1) << "enable_python_tracer set to " + << options.enable_python_tracer; + } + } + return options; + } + std::unique_ptr session_; tensorflow::string logdir_; }; diff --git a/tensorflow/python/profiler/profiler_v2.py b/tensorflow/python/profiler/profiler_v2.py index 3c15f70aa04..28d22bad584 100644 --- a/tensorflow/python/profiler/profiler_v2.py +++ b/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() diff --git a/tensorflow/python/profiler/profiler_v2_test.py b/tensorflow/python/profiler/profiler_v2_test.py index de35cc78a75..50d29c0532f 100644 --- a/tensorflow/python/profiler/profiler_v2_test.py +++ b/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() diff --git a/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt index c777d3705d9..f20dc86f2d5 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profile.pbtxt @@ -4,6 +4,6 @@ tf_class { is_instance: "" 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\'], " } } diff --git a/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profiler-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profiler-options.pbtxt new file mode 100644 index 00000000000..7f5248180a4 --- /dev/null +++ b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.-profiler-options.pbtxt @@ -0,0 +1,23 @@ +path: "tensorflow.profiler.experimental.ProfilerOptions" +tf_class { + is_instance: "" + is_instance: "" + is_instance: "" + member { + name: "host_tracer_level" + mtype: "" + } + member { + name: "python_tracer_level" + mtype: "" + } + member_method { + name: "__init__" + } + member_method { + name: "count" + } + member_method { + name: "index" + } +} diff --git a/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt index ed55fbef167..3841e39adcf 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.profiler.experimental.pbtxt @@ -4,6 +4,10 @@ tf_module { name: "Profile" mtype: "" } + member { + name: "ProfilerOptions" + mtype: "" + } member { name: "Trace" mtype: "" @@ -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"