Implement fast deferred-decoding Python stack trace class.

PiperOrigin-RevId: 315634209
Change-Id: Iac2a77dd98cabafefe54cabec1bdde025950eb8d

tensorflow/python/BUILD

@@ -1070,6 +1070,7 @@ cc_library(
     deps = [
         "//tensorflow/core:lib",
         "//tensorflow/core:script_ops_op_lib",
+        "//tensorflow/core/platform:logging",
         "//third_party/python_runtime:headers",
     ],
 )
@@ -5623,6 +5624,19 @@ tf_py_test(
     ],
 )
 
+cc_library(
+    name = "stack_trace",
+    srcs = ["util/stack_trace.cc"],
+    hdrs = ["util/stack_trace.h"],
+    deps = [
+        ":py_util",
+        "//third_party/python_runtime:headers",  # buildcleaner: keep
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:optional",
+    ],
+)
+
 py_library(
     name = "util",
     srcs = glob(

tensorflow/python/lib/core/py_util.cc

@@ -18,10 +18,6 @@ limitations under the License.
 // Place `<locale>` before <Python.h> to avoid build failure in macOS.
 #include <locale>
 
-// The empty line above is on purpose as otherwise clang-format will
-// automatically move <Python.h> before <locale>.
-#include <Python.h>
-
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 

tensorflow/python/lib/core/py_util.h

@@ -16,12 +16,24 @@ limitations under the License.
 #ifndef TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_
 #define TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_
 
+#include <Python.h>
+
+#include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/types.h"
 
 namespace tensorflow {
+
 // Fetch the exception message as a string. An exception must be set
 // (PyErr_Occurred() must be true).
 string PyExceptionFetch();
-}  // end namespace tensorflow
+
+// Assert that Python GIL is held.
+inline void DCheckPyGilState() {
+#if PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION >= 4
+  DCHECK(PyGILState_Check());
+#endif
+}
+
+}  // namespace tensorflow
 
 #endif  // TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_

tensorflow/python/util/stack_trace.cc

@@ -0,0 +1,52 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#include "tensorflow/python/util/stack_trace.h"
+
+namespace {
+
+// Returns C string from a Python string object. Handles Python2/3 strings.
+// TODO(kkb): This is a generic Python utility function. factor out as a
+// utility.
+const char* GetPythonString(PyObject* o) {
+#if PY_MAJOR_VERSION >= 3
+  if (PyBytes_Check(o)) {
+    return PyBytes_AsString(o);
+  } else {
+    return PyUnicode_AsUTF8(o);
+  }
+#else
+  return PyBytes_AsString(o);
+#endif
+}
+}  // namespace
+
+namespace tensorflow {
+std::string StackTrace::ToString() const {
+  DCheckPyGilState();
+
+  std::ostringstream result;
+  for (int i = size_ - 1; i >= 0; --i) {
+    result << "  File \"" << PyUnicode_AsUTF8(code_objs_[i]->co_filename)
+           << "\", line "
+           << PyCode_Addr2Line(code_objs_[i], last_instructions_[i]) << ", in "
+           << GetPythonString(code_objs_[i]->co_name)
+           << "\n    <source line unimplemented>\n";
+    // TODO(kkb): Add source code line.  See tf_stack.cc's
+    // FrameSummary::line() function.
+  }
+  return result.str();
+}
+}  // namespace tensorflow

tensorflow/python/util/stack_trace.h

@@ -0,0 +1,108 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
+#define TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
+
+#include <Python.h>
+#include <frameobject.h>
+
+#include <array>
+#include <sstream>
+#include <string>
+
+#include "absl/base/attributes.h"
+#include "absl/base/optimization.h"
+#include "tensorflow/python/lib/core/py_util.h"
+
+namespace tensorflow {
+
+// A class for capturing Python stack trace.
+class StackTrace final {
+ public:
+  static constexpr int kMaxDepth = 10;
+
+  StackTrace() : size_(0) {}
+
+  // Returns `StackTrace` object that captures the current Python stack trace.
+  // Python GIL must be acquired beforehand.
+  ABSL_MUST_USE_RESULT
+  ABSL_ATTRIBUTE_HOT
+  static StackTrace Capture() {
+    DCheckPyGilState();
+
+    StackTrace result;
+    const PyFrameObject* frame = PyThreadState_GET()->frame;
+    int i = 0;
+    for (; i < kMaxDepth && frame != nullptr; frame = frame->f_back, ++i) {
+      PyCodeObject* code_obj = frame->f_code;
+      DCHECK(frame->f_trace == nullptr);
+      DCHECK(code_obj != nullptr);
+
+      Py_INCREF(code_obj);
+      result.code_objs_[i] = code_obj;
+      result.last_instructions_[i] = frame->f_lasti;
+    }
+    result.size_ = i;
+    return result;
+  }
+
+  // Python GIL must be acquired beforehand.
+  ABSL_ATTRIBUTE_HOT
+  ~StackTrace() { Clear(); }
+
+  StackTrace(StackTrace&& other) {
+    code_objs_ = other.code_objs_;
+    last_instructions_ = other.last_instructions_;
+    size_ = other.size_;
+    other.size_ = 0;
+  }
+
+  // Python GIL must be acquired beforehand.
+  ABSL_ATTRIBUTE_HOT
+  StackTrace& operator=(StackTrace&& other) {
+    Clear();
+
+    code_objs_ = other.code_objs_;
+    last_instructions_ = other.last_instructions_;
+    size_ = other.size_;
+    other.size_ = 0;
+    return *this;
+  }
+
+  // Returns string representation of the captured stack trace.
+  std::string ToString() const;
+
+  // TODO(kkb): Implement structured stack trace object getter.
+
+ private:
+  std::array<PyCodeObject*, kMaxDepth> code_objs_;
+  std::array<int, kMaxDepth> last_instructions_;
+  int size_;
+
+  // Python GIL must be acquired beforehand.
+  ABSL_ATTRIBUTE_HOT
+  void Clear() {
+    DCheckPyGilState();
+    for (int i = 0; i < size_; ++i) Py_DECREF(code_objs_[i]);
+  }
+
+  StackTrace(const StackTrace&) = delete;
+  StackTrace& operator=(const StackTrace&) = delete;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_