diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index b93992246ab..21802a1d819 100644 --- a/tensorflow/python/BUILD +++ b/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( diff --git a/tensorflow/python/lib/core/py_util.cc b/tensorflow/python/lib/core/py_util.cc index 739cab46b10..a78f0a12f21 100644 --- a/tensorflow/python/lib/core/py_util.cc +++ b/tensorflow/python/lib/core/py_util.cc @@ -18,10 +18,6 @@ limitations under the License. // Place `` before to avoid build failure in macOS. #include -// The empty line above is on purpose as otherwise clang-format will -// automatically move before . -#include - #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/strcat.h" diff --git a/tensorflow/python/lib/core/py_util.h b/tensorflow/python/lib/core/py_util.h index a9f39d39461..af1b21699e6 100644 --- a/tensorflow/python/lib/core/py_util.h +++ b/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 + +#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_ diff --git a/tensorflow/python/util/stack_trace.cc b/tensorflow/python/util/stack_trace.cc new file mode 100644 index 00000000000..cf574f6f292 --- /dev/null +++ b/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 \n"; + // TODO(kkb): Add source code line. See tf_stack.cc's + // FrameSummary::line() function. + } + return result.str(); +} +} // namespace tensorflow diff --git a/tensorflow/python/util/stack_trace.h b/tensorflow/python/util/stack_trace.h new file mode 100644 index 00000000000..0b9a737bf7e --- /dev/null +++ b/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 +#include + +#include +#include +#include + +#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 code_objs_; + std::array 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_