1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""TensorFlow 2.0 Profiler for both Eager Mode and Graph Mode.
16
17The profiler has two mode:
18- Programmatic Mode: start(), stop() and Profiler class. It will perform
19                    when calling start() or create Profiler class and will stop
20                    when calling stop() or destroying Profiler class.
21- On-demand Mode: start_profiler_server(). It will perform profiling when
22                  receive profiling request.
23
24NOTE: Only one active profiler session is allowed. Use of simultaneous
25Programmatic Mode and On-demand Mode is undefined and will likely fail.
26
27NOTE: The Keras TensorBoard callback will automatically perform sampled
28profiling. Before enabling customized profiling, set the callback flag
29"profile_batches=[]" to disable automatic sampled profiling.
30customized profiling.
31"""
32
33from __future__ import absolute_import
34from __future__ import division
35from __future__ import print_function
36
37import datetime
38import os
39import threading
40
41from tensorflow.python import _pywrap_events_writer
42from tensorflow.python.eager import context
43from tensorflow.python.framework import errors
44from tensorflow.python.platform import gfile
45from tensorflow.python.platform import tf_logging as logging
46from tensorflow.python.profiler.internal import _pywrap_profiler
47from tensorflow.python.util import compat
48from tensorflow.python.util.deprecation import deprecated
49
50_profiler = None
51_profiler_lock = threading.Lock()
52_run_num = 0
53# This suffix should be kept in sync with kProfileEmptySuffix in
54# tensorflow/core/profiler/rpc/client/capture_profile.cc.
55_EVENT_FILE_SUFFIX = '.profile-empty'
56
57
58class ProfilerAlreadyRunningError(Exception):
59  pass
60
61
62class ProfilerNotRunningError(Exception):
63  pass
64
65
66@deprecated('2020-07-01', 'use `tf.profiler.experimental.start` instead.')
67def start():
68  """Start profiling.
69
70  Raises:
71    ProfilerAlreadyRunningError: If another profiling session is running.
72  """
73  global _profiler
74  with _profiler_lock:
75    if _profiler is not None:
76      raise ProfilerAlreadyRunningError('Another profiler is running.')
77    if context.default_execution_mode == context.EAGER_MODE:
78      context.ensure_initialized()
79    _profiler = _pywrap_profiler.ProfilerSession()
80    try:
81      _profiler.start('', {})
82    except errors.AlreadyExistsError:
83      logging.warning('Another profiler session is running which is probably '
84                      'created by profiler server. Please avoid using profiler '
85                      'server and profiler APIs at the same time.')
86      raise ProfilerAlreadyRunningError('Another profiler is running.')
87
88
89@deprecated('2020-07-01', 'use `tf.profiler.experimental.stop` instead.')
90def stop():
91  """Stop current profiling session and return its result.
92
93  Returns:
94    A binary string of tensorflow.tpu.Trace. User can write the string
95    to file for offline analysis by tensorboard.
96
97  Raises:
98    ProfilerNotRunningError: If there is no active profiling session.
99  """
100  global _profiler
101  global _run_num
102  with _profiler_lock:
103    if _profiler is None:
104      raise ProfilerNotRunningError(
105          'Cannot stop profiling. No profiler is running.')
106    if context.default_execution_mode == context.EAGER_MODE:
107      context.context().executor.wait()
108    result = _profiler.stop()
109    _profiler = None
110    _run_num += 1
111  return result
112
113
114@deprecated(
115    '2020-07-01',
116    '`tf.python.eager.profiler` has deprecated, use `tf.profiler` instead.'
117)
118def maybe_create_event_file(logdir):
119  """Create an empty event file if not already exists.
120
121  This event file indicates that we have a plugins/profile/ directory in the
122  current logdir.
123
124  Args:
125    logdir: log directory.
126  """
127  for file_name in gfile.ListDirectory(logdir):
128    if file_name.endswith(_EVENT_FILE_SUFFIX):
129      return
130  # TODO(b/127330388): Use summary_ops_v2.create_file_writer instead.
131  event_writer = _pywrap_events_writer.EventsWriter(
132      compat.as_bytes(os.path.join(logdir, 'events')))
133  event_writer.InitWithSuffix(compat.as_bytes(_EVENT_FILE_SUFFIX))
134
135
136@deprecated(
137    '2020-07-01',
138    '`tf.python.eager.profiler` has deprecated, use `tf.profiler` instead.'
139)
140def save(logdir, result):
141  """Save profile result to TensorBoard logdir.
142
143  Args:
144    logdir: log directory read by TensorBoard.
145    result: profiling result returned by stop().
146  """
147  plugin_dir = os.path.join(
148      logdir, 'plugins', 'profile',
149      datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
150  gfile.MakeDirs(plugin_dir)
151  maybe_create_event_file(logdir)
152  with gfile.Open(os.path.join(plugin_dir, 'local.trace'), 'wb') as f:
153    f.write(result)
154
155
156@deprecated('2020-07-01', 'use `tf.profiler.experimental.server.start`.')
157def start_profiler_server(port):
158  """Start a profiler grpc server that listens to given port.
159
160  The profiler server will keep the program running even the training finishes.
161  Please shutdown the server with CTRL-C. It can be used in both eager mode and
162  graph mode. The service defined in
163  tensorflow/core/profiler/profiler_service.proto. Please use
164  tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable
165  file following https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace
166
167  Args:
168    port: port profiler server listens to.
169  """
170  if context.default_execution_mode == context.EAGER_MODE:
171    context.ensure_initialized()
172  _pywrap_profiler.start_server(port)
173
174
175@deprecated('2020-07-01', 'use `tf.profiler.experimental.Profile` instead.')
176class Profiler(object):
177  """Context-manager eager profiler api.
178
179  Example usage:
180  ```python
181  with Profiler("/path/to/logdir"):
182    # do some work
183  ```
184  """
185
186  def __init__(self, logdir):
187    self._logdir = logdir
188
189  def __enter__(self):
190    start()
191
192  def __exit__(self, typ, value, tb):
193    result = stop()
194    save(self._logdir, result)
195