From 2f4d749c04de49b23ad3baec962b6f4537c4e035 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 20 Nov 2020 18:07:04 -0800 Subject: [PATCH] Internal change PiperOrigin-RevId: 343600096 Change-Id: I6c8dcf7fe14393116e1519f9c0b4087839f2da39 --- tensorflow/core/platform/status.h | 7 - tensorflow/python/BUILD | 5 - tensorflow/python/eager/pywrap_tfe_src.cc | 4 +- tensorflow/python/util/stack_trace.cc | 17 +- tensorflow/python/util/stack_trace.h | 20 +- tensorflow/python/util/tf_stack.cc | 293 ++++++++-------------- tensorflow/python/util/tf_stack.py | 13 +- 7 files changed, 127 insertions(+), 232 deletions(-) diff --git a/tensorflow/core/platform/status.h b/tensorflow/core/platform/status.h index 87d4eb26fd1..fc570caf6b1 100644 --- a/tensorflow/core/platform/status.h +++ b/tensorflow/core/platform/status.h @@ -34,13 +34,6 @@ struct StackFrame { std::string file_name; int line_number; std::string function_name; - - bool operator==(const StackFrame& other) const { - return line_number == other.line_number && - function_name == other.function_name && file_name == other.file_name; - } - - bool operator!=(const StackFrame& other) const { return !(*this == other); } }; #if defined(__clang__) diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index 6473e7824a8..31c205c638e 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -5690,12 +5690,7 @@ pybind_extension( # TODO(b/138203821): change to "util._tf_stack" once the bug is fixed. module_name = "_tf_stack", deps = [ - ":stack_trace", "//third_party/python_runtime:headers", # buildcleaner: keep - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/hash", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", "@pybind11", ], ) diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc index 15d8ec9434b..d41f5d736c8 100644 --- a/tensorflow/python/eager/pywrap_tfe_src.cc +++ b/tensorflow/python/eager/pywrap_tfe_src.cc @@ -866,7 +866,7 @@ void TFE_Py_ExecuteCancelable(TFE_Context* ctx, const char* device_name, if (!out_status->status.ok()) return; tensorflow::unwrap(op)->SetStackTrace(tensorflow::GetStackTrace( - tensorflow::StackTrace::kStackTraceInitialSize)); + tensorflow::StackTrace::kDefaultStackTraceInitialSize)); for (int i = 0; i < inputs->size() && out_status->status.ok(); ++i) { TFE_OpAddInput(op, inputs->at(i), out_status); @@ -3631,7 +3631,7 @@ PyObject* TFE_Py_FastPathExecute_C(PyObject* args) { } tensorflow::unwrap(op)->SetStackTrace(tensorflow::GetStackTrace( - tensorflow::StackTrace::kStackTraceInitialSize)); + tensorflow::StackTrace::kDefaultStackTraceInitialSize)); const tensorflow::OpDef* op_def = tensorflow::unwrap(op)->OpDef(); if (op_def == nullptr) return nullptr; diff --git a/tensorflow/python/util/stack_trace.cc b/tensorflow/python/util/stack_trace.cc index 4197dbe6c1c..b3bb9fb027c 100644 --- a/tensorflow/python/util/stack_trace.cc +++ b/tensorflow/python/util/stack_trace.cc @@ -39,8 +39,7 @@ const char* GetPythonString(PyObject* o) { namespace tensorflow { -std::vector StackTrace::ToStackFrames( - const StackTraceMapper& mapper, const StackTraceFilter& filtered) const { +std::vector StackTrace::ToStackFrames() const { std::vector result; result.reserve(code_objs_.size()); @@ -48,18 +47,8 @@ std::vector StackTrace::ToStackFrames( const char* file_name = GetPythonString(code_objs_[i]->co_filename); const int line_number = PyCode_Addr2Line(code_objs_[i], last_instructions_[i]); - - if (!result.empty() && filtered && filtered(file_name)) { - continue; // Never filter the innermost frame. - } - - if (absl::optional mapped = - mapper ? mapper(file_name, line_number) : absl::nullopt) { - result.push_back(*mapped); - } else { - result.emplace_back(StackFrame{file_name, line_number, - GetPythonString(code_objs_[i]->co_name)}); - } + result.emplace_back(StackFrame{file_name, line_number, + GetPythonString(code_objs_[i]->co_name)}); } return result; diff --git a/tensorflow/python/util/stack_trace.h b/tensorflow/python/util/stack_trace.h index c12a841a1f3..2734a028e82 100644 --- a/tensorflow/python/util/stack_trace.h +++ b/tensorflow/python/util/stack_trace.h @@ -40,17 +40,10 @@ inline void DCheckPyGilStateForStackTrace() { #endif } -// Maps filename/line_no combination into a stack frame. -using StackTraceMapper = - std::function(std::string, int)>; - -// Returns "true" for filenames which should be skipped. -using StackTraceFilter = std::function; - // A class for capturing Python stack trace. class StackTrace final { public: - static constexpr int kStackTraceInitialSize = 10; + static constexpr int kDefaultStackTraceInitialSize = 10; StackTrace() {} @@ -100,16 +93,11 @@ class StackTrace final { } // Returns a structured representation of the captured stack trace. - // `mapper` provides a custom mapping for translating stack frames, `filter` - // returns `true` for the stack frames which should be omitted, and if - // `drop_last` is set, the last stack frame is dropped. - std::vector ToStackFrames( - const StackTraceMapper& mapper = {}, - const StackTraceFilter& filtered = {}) const; + std::vector ToStackFrames() const; private: - absl::InlinedVector code_objs_; - absl::InlinedVector last_instructions_; + absl::InlinedVector code_objs_; + absl::InlinedVector last_instructions_; // Python GIL must be acquired beforehand. ABSL_ATTRIBUTE_HOT diff --git a/tensorflow/python/util/tf_stack.cc b/tensorflow/python/util/tf_stack.cc index 07566a5f290..b3a797568a7 100644 --- a/tensorflow/python/util/tf_stack.cc +++ b/tensorflow/python/util/tf_stack.cc @@ -19,20 +19,12 @@ limitations under the License. #include #include -#include "absl/algorithm/container.h" -#include "absl/hash/hash.h" -#include "absl/strings/str_join.h" -#include "absl/types/span.h" #include "pybind11/pybind11.h" -#include "pybind11/stl.h" #include "pybind11/stl_bind.h" -#include "tensorflow/python/util/stack_trace.h" -struct StackFrame; // Forward declaration. -struct StackTrace; +struct FrameSummary; // Forward declaration. -PYBIND11_MAKE_OPAQUE(std::vector); -PYBIND11_MAKE_OPAQUE(StackTrace); +PYBIND11_MAKE_OPAQUE(std::vector); namespace tensorflow { @@ -40,16 +32,21 @@ namespace { namespace py = pybind11; -py::object LineContents(const StackFrame& frame) { - static const auto* linecache = - new py::module(py::module::import("linecache")); - const auto& checkcache = linecache->attr("checkcache"); - const auto& getline = linecache->attr("getline"); - checkcache(py::str(frame.file_name)); - const auto& code = py::cast( - getline(py::str(frame.file_name), py::int_(frame.line_number)) - .attr("strip")()); - ssize_t size = 0; +struct FrameSummary { + py::str filename; + int lineno; + py::str name; + py::object globals; + + py::object line() const { + static const auto* linecache = + new py::module(py::module::import("linecache")); + const auto& checkcache = linecache->attr("checkcache"); + const auto& getline = linecache->attr("getline"); + checkcache(filename); + const auto& code = + py::cast(getline(filename, lineno, globals).attr("strip")()); + ssize_t size = 0; #if PY_MAJOR_VERSION == 3 if (PyUnicode_AsUTF8AndSize(code.ptr(), &size) == nullptr) { throw py::error_already_set(); @@ -58,199 +55,131 @@ py::object LineContents(const StackFrame& frame) { size = PyString_Size(code.ptr()); #endif return size > 0 ? static_cast(code) : py::none(); -} - -std::string StackFrameToString(const StackFrame& frame) { - return py::str("") - .format(py::str(frame.file_name), py::int_(frame.line_number), - py::str(frame.function_name)); -} - -class StackTraceWrapper { - public: - StackTraceWrapper(StackTrace&& captured, - const StackTraceMapper& stack_trace_mapper, - const StackTraceFilter& stack_trace_filter) - : captured_(std::move(captured)), - stack_trace_mapper_(stack_trace_mapper), - stack_trace_filter_(stack_trace_filter) {} - - explicit StackTraceWrapper(absl::Span stack_frames) - : stack_frames_cache_(std::vector(stack_frames.begin(), - stack_frames.end())) {} - - absl::Span ToFrames() const { - GenerateCache(); - return *stack_frames_cache_; } - std::string ToString() const { - GenerateCache(); - return absl::StrJoin(*stack_frames_cache_, "\n", - [&](std::string* out, const StackFrame& frame) { - absl::StrAppend(out, StackFrameToString(frame)); - }); + bool operator==(const FrameSummary& other) const { + return filename == other.filename && lineno == other.lineno && + name == other.name && globals == other.globals; } - bool IsCacheGenerated() const { return stack_frames_cache_.has_value(); } + bool operator!=(const FrameSummary& other) const { return !(*this == other); } - void GenerateCache() const { - if (stack_frames_cache_) { - return; - } - stack_frames_cache_ = - captured_.ToStackFrames(stack_trace_mapper_, stack_trace_filter_); - stack_frames_cache_->pop_back(); // Drop last stack frame. + py::str toString() const { + return py::str("") + .format(filename, lineno, name); } - - private: - mutable absl::optional> stack_frames_cache_; - StackTrace captured_; - - // TODO(cheshire): store those as C++ datastructures instead. - StackTraceMapper stack_trace_mapper_; - StackTraceFilter stack_trace_filter_; }; +std::vector ExtractStack(ssize_t limit, const py::list& mappers, + const py::list& filters) { + const py::dict& source_map = + mappers.size() == 0 + ? py::dict() + : mappers[mappers.size() - 1].attr("get_effective_source_map")(); + const py::set& filtered_filenames = + filters.size() == 0 + ? py::set() + : filters[filters.size() - 1].attr("get_filtered_filenames")(); + + const auto* tstate = PyThreadState_GET(); + // Drop extract_stack() wrapper-function frame from the result. + const PyFrameObject* f = tstate->frame->f_back; // TODO(slebedev): INCREF? + + std::vector ret; + // 16 is somewhat arbitrary, but TensorFlow stack traces tend to be deep. + ret.reserve(limit < 0 ? 16 : static_cast(limit)); + for (; f != nullptr && (limit < 0 || ret.size() < static_cast(limit)); + f = f->f_back) { + const PyCodeObject* co = f->f_code; + int lineno = PyFrame_GetLineNumber(const_cast(f)); + auto filename = py::reinterpret_borrow(co->co_filename); + auto name = py::reinterpret_borrow(co->co_name); + + // TODO(slebedev): consider moving the mappers/filters to C++ as well. + if (source_map.size() > 0) { + const auto& key = py::make_tuple(filename, lineno); + if (source_map.contains(key)) { + const py::tuple& mapped = source_map[key]; + filename = mapped[0]; + lineno = py::cast(mapped[1]); + name = mapped[2]; + } + } + + if (!ret.empty() && // Never filter the innermost frame. + filtered_filenames.size() > 0 && + PySet_Contains(filtered_filenames.ptr(), filename.ptr())) { + continue; + } + + const auto& globals = py::reinterpret_borrow(f->f_globals); + ret.push_back({std::move(filename), lineno, std::move(name), globals}); + } + + std::reverse(ret.begin(), ret.end()); + return ret; +} + } // namespace PYBIND11_MODULE(_tf_stack, m) { - py::class_(m, "StackFrame") - .def_property_readonly( - "filename", - [](const StackFrame& self) { return py::str(self.file_name); }) - .def_property_readonly( - "lineno", - [](const StackFrame& self) { return py::int_(self.line_number); }) - .def_property_readonly( - "name", - [](const StackFrame& self) { return py::str(self.function_name); }) - .def_property_readonly( - "line", - [](const StackFrame& self) { return py::str(LineContents(self)); }) + py::class_(m, "FrameSummary") + .def_readonly("filename", &FrameSummary::filename) + .def_readonly("lineno", &FrameSummary::lineno) + .def_readonly("name", &FrameSummary::name) + .def_property_readonly("line", &FrameSummary::line) // For compatibility with the traceback module. - .def("__eq__", &StackFrame::operator==) - .def("__ne__", &StackFrame::operator!=) + .def("__eq__", &FrameSummary::operator==) + .def("__ne__", &FrameSummary::operator!=) .def("__hash__", - [](const StackFrame& self) { - return absl::Hash>()( - std::make_tuple(self.file_name, self.line_number, - self.function_name)); + [](const FrameSummary& self) { + return py::hash( + py::make_tuple(self.filename, self.lineno, self.name)); }) .def("__getitem__", - [](const StackFrame& self, const py::object& index) -> py::object { - return py::make_tuple( - py::str(self.file_name), py::int_(self.line_number), - py::str(self.function_name), LineContents(self))[index]; + [](const FrameSummary& self, const py::object& index) -> py::object { + return py::make_tuple(self.filename, self.lineno, self.name, + self.line())[index]; }) .def("__iter__", - [](const StackFrame& self) { - return py::iter(py::make_tuple( - py::str(self.file_name), py::int_(self.line_number), - py::str(self.function_name), LineContents(self)) - - ); + [](const FrameSummary& self) { + return py::iter(py::make_tuple(self.filename, self.lineno, + self.name, self.line())); }) - .def("__repr__", - [](const StackFrame& self) { return StackFrameToString(self); }) - .def("__len__", [](const StackFrame&) { return 4; }); + .def("__repr__", [](const FrameSummary& self) { return self.toString(); }) + .def("__len__", [](const FrameSummary&) { return 4; }); - py::class_(m, "StackTraceWrapper", py::module_local(true)) + py::bind_vector>(m, "StackSummary", + py::module_local(true)) // TODO(slebedev): upstream negative indexing support into pybind11. .def( "__getitem__", - [](const StackTraceWrapper& self, ssize_t index) { - absl::Span frames = self.ToFrames(); + [](const std::vector& self, ssize_t index) { const size_t eff_index = - index < 0 ? frames.size() + index : static_cast(index); - if (eff_index >= frames.size()) { + index < 0 ? self.size() + index : static_cast(index); + if (eff_index > self.size()) { throw py::index_error(); } - return frames[eff_index]; + return self[eff_index]; }, py::return_value_policy::reference_internal) - .def( - "__getitem__", - [](const StackTraceWrapper& self, py::slice slice) { - absl::Span frames = self.ToFrames(); - py::ssize_t start, stop, step, slicelength; - if (!slice.compute(frames.size(), &start, &stop, &step, - &slicelength)) { - throw py::error_already_set(); - } - if (step == 1) { - return StackTraceWrapper{frames.subspan(start, slicelength)}; - } - std::vector out; - out.reserve(slicelength); - for (int i = start; i < stop; i += step) { - out.push_back(frames[i]); - } - return StackTraceWrapper{out}; - }, - py::return_value_policy::reference_internal) - .def("__len__", - [](const StackTraceWrapper& self) { return self.ToFrames().size(); }) - .def("__eq__", - [](const StackTraceWrapper& self, const StackTraceWrapper& other) { - return self.ToFrames() == other.ToFrames(); - }) - .def("__hash__", - [](const StackTraceWrapper& self) { - return py::hash(py::str(self.ToString())); - }) - .def("__repr__", [](const StackTraceWrapper& self) { - if (self.IsCacheGenerated()) { - return py::str(""); + .def("__repr__", [](const std::vector& self) { + py::list frames; + for (const auto& frame : self) { + frames.append(frame.toString()); } - return py::str(self.ToString()); + // "\n".join(frames) + return py::cast("\n").attr("join")(frames); }); - m.def( - "extract_stack", - [](const py::object& limit, const py::list& mappers, - const py::list& filters) { - // In Python 3.X ``traceback.extract_stack`` allows ``limit`` to - // either be None or -1. - int casted_limit = limit.is_none() ? -1 : py::cast(limit); - - // Raise limit by one since we are dropping the last frame. - if (casted_limit != -1) casted_limit++; - - const py::dict& source_map = mappers.empty() - ? py::dict() - : mappers[mappers.size() - 1].attr( - "get_effective_source_map")(); - const py::set& filtered_filenames = - filters.empty() - ? py::set() - : filters[filters.size() - 1].attr("get_filtered_filenames")(); - - auto mapper = [=](std::string filename, - int line_no) -> absl::optional { - if (source_map.empty()) { - return absl::nullopt; - } - const auto& key = - py::make_tuple(py::str(filename), py::int_(line_no)); - if (source_map.contains(key)) { - const py::tuple& mapped = source_map[key]; - return StackFrame{std::string(py::cast(mapped[0])), - py::cast(mapped[1]), - std::string(py::cast(mapped[2]))}; - } - - return absl::nullopt; - }; - - auto filter = [=](std::string filename) -> bool { - return filtered_filenames.contains(py::str(filename)); - }; - return StackTraceWrapper{StackTrace::Capture(casted_limit), mapper, - filter}; - }, - py::return_value_policy::move); + m.def("extract_stack", [](const py::object& limit, const py::list& mappers, + const py::list& filters) { + // In Python 3.X ``traceback.extract_stack`` allows ``limit`` to + // either be None or -1. + return ExtractStack(limit.is_none() ? -1 : py::cast(limit), + mappers, filters); + }); } } // namespace tensorflow diff --git a/tensorflow/python/util/tf_stack.py b/tensorflow/python/util/tf_stack.py index e0875590365..628cd4e1854 100644 --- a/tensorflow/python/util/tf_stack.py +++ b/tensorflow/python/util/tf_stack.py @@ -141,15 +141,16 @@ def extract_stack(limit=-1): limit: A limit on the number of frames to return. Returns: - A sequence of StackFrame objects (filename, lineno, name, line) + A sequence of FrameSummary objects (filename, lineno, name, line) corresponding to the call stack of the current thread. """ # N.B ExtractStack in tf_stack.cc will drop this frame prior to # traversing the stack. thread_key = _get_thread_key() - return _tf_stack.extract_stack(limit, _source_mapper_stacks[thread_key], - _source_filter_stacks[thread_key]) + return _tf_stack.extract_stack( + limit, + _source_mapper_stacks[thread_key], + _source_filter_stacks[thread_key]) - -StackSummary = _tf_stack.StackTraceWrapper -FrameSummary = _tf_stack.StackFrame +StackSummary = _tf_stack.StackSummary +FrameSummary = _tf_stack.FrameSummary