Function DoCallPyFunc

tensorflow/python/lib/core/py_func.cc:169–282

Calls the registered py function through the trampoline.

// Calls the registered py function through the trampoline.
Status DoCallPyFunc(PyCall* call, bool* out_log_on_error) {
  *out_log_on_error = true;
  PyObject* trampoline = GetPyTrampoline();
  if (trampoline == nullptr) {
    return errors::InvalidArgument(
        "Missing py trampoline. Most likely, it is a link error.");
  }

  // Prepare the argument.
  PyObject* args = nullptr;
  TFE_Context* ctx = nullptr;
  std::unique_ptr<EagerExecutor> new_executor = nullptr;
  EagerExecutor* old_executor = nullptr;
  if (call->eager) {
    // See FuncRegistry._ctx.
    ctx = reinterpret_cast<TFE_Context*>(PyCapsule_GetPointer(
        PyObject_GetAttrString(trampoline, "_ctx"), nullptr));
    CHECK_NE(ctx, nullptr);
    TF_RETURN_IF_ERROR(MakeArgTuple(call, ctx->context, &args));
    new_executor.reset(new EagerExecutor(call->eager_async));
    old_executor = ctx->context->Executor();
    ctx->context->SetExecutorForThread(new_executor.get());
  } else {
    TF_RETURN_IF_ERROR(MakeArgTuple(call, nullptr, &args));
  }
  CHECK(args);

  // Invokes the trampoline.
  PyObject* result = PyEval_CallObject(trampoline, args);
  Py_DECREF(args);
  Status s = Status::OK();
  if (result == nullptr) {
    if (PyErr_Occurred()) {
      if (PyErr_ExceptionMatches(PyExc_ValueError) ||
          PyErr_ExceptionMatches(PyExc_TypeError)) {
        s = errors::InvalidArgument(PyExceptionFetch());
      } else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
        *out_log_on_error = false;
        s = errors::OutOfRange(PyExceptionFetch());
      } else if (PyErr_ExceptionMatches(PyExc_MemoryError)) {
        s = errors::ResourceExhausted(PyExceptionFetch());
      } else if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
        s = errors::Unimplemented(PyExceptionFetch());
      } else {
        // TODO(ebrevdo): Check if exception is an OpError and use the
        // OpError.error_code property to map it back in the Status.
        s = errors::Unknown(PyExceptionFetch());
      }
    } else {
      s = errors::Internal("Failed to run py callback ", call->token,
                           ": see error log.");
    }
  }

  if (new_executor != nullptr) {
    s.Update(new_executor->WaitForAllPendingNodes());
    ctx->context->SetExecutorForThread(old_executor);
  }