# Copyright 2019 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.
# ==============================================================================
"""Executor for eager execution."""

from tensorflow.python import pywrap_tfe


class Executor(object):
  """A class for handling eager execution.

  The default behavior for asynchronous execution is to serialize all ops on
  a single thread. Having different `Executor` objects in different threads
  enables executing ops asynchronously in parallel:

  ```python
  def thread_function():
    executor = executor.Executor(enable_async=True):
    context.set_executor(executor)

  a = threading.Thread(target=thread_function)
  a.start()
  b = threading.Thread(target=thread_function)
  b.start()
  ```
  """

  __slots__ = ["_handle"]

  def __init__(self, handle):
    self._handle = handle

  def __del__(self):
    try:
      self.wait()
      pywrap_tfe.TFE_DeleteExecutor(self._handle)
    except TypeError:
      # Suppress some exceptions, mainly for the case when we're running on
      # module deletion. Things that can go wrong include the pywrap module
      # already being unloaded, self._handle. no longer being
      # valid, and so on. Printing warnings in these cases is silly
      # (exceptions raised from __del__ are printed as warnings to stderr).
      pass  # 'NoneType' object is not callable when the handle has been
      # partially unloaded.

  def is_async(self):
    return pywrap_tfe.TFE_ExecutorIsAsync(self._handle)

  def handle(self):
    return self._handle

  def wait(self):
    """Waits for ops dispatched in this executor to finish."""
    pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)

  def clear_error(self):
    """Clears errors raised in this executor during execution."""
    pywrap_tfe.TFE_ExecutorClearError(self._handle)


def new_executor(enable_async,
                 enable_streaming_enqueue=True,
                 in_flight_nodes_limit=0):
  handle = pywrap_tfe.TFE_NewExecutor(enable_async, enable_streaming_enqueue,
                                      in_flight_nodes_limit)
  return Executor(handle)