text: # Copyright 2017 The TensorFlow Authors. 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.
  # ==============================================================================
  """TensorFlow 2.0 Profiler for both Eager Mode and Graph Mode.
  
  The profiler has two mode:
  - Programmatic Mode: start(), stop() and Profiler class. It will perform
                      when calling start() or create Profiler class and will stop
                      when calling stop() or destroying Profiler class.
  - On-demand Mode: start_profiler_server(). It will perform profiling when
                    receive profiling request.
  
  NOTE: Only one active profiler session is allowed. Use of simultaneous
  Programmatic Mode and On-demand Mode is undefined and will likely fail.
  
  NOTE: The Keras TensorBoard callback will automatically perform sampled
  profiling. Before enabling customized profiling, set the callback flag
  "profile_batches=[]" to disable automatic sampled profiling.
  customized profiling.
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import datetime
  import os
  import threading
  
  from tensorflow.python import pywrap_tensorflow
  from tensorflow.python.eager import context
  from tensorflow.python.framework import c_api_util
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.util import compat
  
  _profiler = None
  _profiler_lock = threading.Lock()
  _run_num = 0
  # This suffix should be kept in sync with kProfileEmptySuffix in
  # tensorflow/core/profiler/rpc/client/capture_profile.cc.
  _EVENT_FILE_SUFFIX = '.profile-empty'
  
  
  class ProfilerAlreadyRunningError(Exception):
    pass
  
  
  class ProfilerNotRunningError(Exception):
    pass
  
  
  def start():
    """Start profiling.
  
    Raises:
      ProfilerAlreadyRunningError: If another profiling session is running.
    """
    global _profiler
    with _profiler_lock:
      if _profiler is not None:
        raise ProfilerAlreadyRunningError('Another profiler is running.')
      context.ensure_initialized()
      profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
      if context.default_execution_mode == context.EAGER_MODE:
        pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
            profiler_context,
            context.context()._handle)  # pylint: disable=protected-access
      _profiler = pywrap_tensorflow.TFE_NewProfiler(profiler_context)
      pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
      if not pywrap_tensorflow.TFE_ProfilerIsOk(_profiler):
        logging.warning('Another profiler session is running which is probably '
                        'created by profiler server. Please avoid using profiler '
                        'server and profiler APIs at the same time.')
  
  
  def stop():
    """Stop current profiling session and return its result.
  
    Returns:
      A binary string of tensorflow.tpu.Trace. User can write the string
      to file for offline analysis by tensorboard.
  
    Raises:
      ProfilerNotRunningError: If there is no active profiling session.
    """
    global _profiler
    global _run_num
    with _profiler_lock:
      if _profiler is None:
        raise ProfilerNotRunningError(
            'Cannot stop profiling. No profiler is running.')
      with c_api_util.tf_buffer() as buffer_:
        pywrap_tensorflow.TFE_ProfilerSerializeToString(
            context.context()._handle,  # pylint: disable=protected-access
            _profiler,
            buffer_)
        result = pywrap_tensorflow.TF_GetBuffer(buffer_)
      pywrap_tensorflow.TFE_DeleteProfiler(_profiler)
      _profiler = None
      _run_num += 1
    return result
  
  
  def maybe_create_event_file(logdir):
    """Create an empty event file if not already exists.
  
    This event file indicates that we have a plugins/profile/ directory in the
    current logdir.
  
    Args:
      logdir: log directory.
    """
    for file_name in gfile.ListDirectory(logdir):
      if file_name.endswith(_EVENT_FILE_SUFFIX):
        return
    # TODO(b/127330388): Use summary_ops_v2.create_file_writer instead.
    event_writer = pywrap_tensorflow.EventsWriter(
        compat.as_bytes(os.path.join(logdir, 'events')))
    event_writer.InitWithSuffix(compat.as_bytes(_EVENT_FILE_SUFFIX))
  
  
  def save(logdir, result):
    """Save profile result to TensorBoard logdir.
  
    Args:
      logdir: log directory read by TensorBoard.
      result: profiling result returned by stop().
    """
    plugin_dir = os.path.join(
        logdir, 'plugins', 'profile',
        datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
    gfile.MakeDirs(plugin_dir)
    maybe_create_event_file(logdir)
    with gfile.Open(os.path.join(plugin_dir, 'local.trace'), 'wb') as f:
      f.write(result)
  
  
  def start_profiler_server(port):
    """Start a profiler grpc server that listens to given port.
  
    The profiler server will keep the program running even the training finishes.
    Please shutdown the server with CTRL-C. It can be used in both eager mode and
    graph mode. The service defined in
    tensorflow/core/profiler/profiler_service.proto. Please use
    tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable
    file following https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace
  
    Args:
      port: port profiler server listens to.
    """
    profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
    if context.default_execution_mode == context.EAGER_MODE:
      pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
          profiler_context,
          context.context()._handle)  # pylint: disable=protected-access
    pywrap_tensorflow.TFE_StartProfilerServer(profiler_context, port)
    pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
  
  
  class Profiler(object):
    """Context-manager eager profiler api.
  
    Example usage:
    ```python
    with Profiler("/path/to/logdir"):
      # do some work
    ```
    """
  
    def __init__(self, logdir):
      self._logdir = logdir
  
    def __enter__(self):
      start()
  
    def __exit__(self, typ, value, tb):
      result = stop()
      save(self._logdir, result)
  