experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Internal change

Browse Source 
PiperOrigin-RevId: 343600096
Change-Id: I6c8dcf7fe14393116e1519f9c0b4087839f2da39
This commit is contained in:
A. Unique TensorFlower 2020-11-20 18:07:04 -08:00 committed by TensorFlower Gardener
parent e984937a4b
commit 2f4d749c04
7 changed files with 127 additions and 232 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/core/platform/status.h
View File

@ -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__)

5
tensorflow/python/BUILD
View File

@ -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",
],
)

4
tensorflow/python/eager/pywrap_tfe_src.cc
View File

@ -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;

17
tensorflow/python/util/stack_trace.cc
View File

@ -39,8 +39,7 @@ const char* GetPythonString(PyObject* o) {
namespace tensorflow {
std::vector<StackFrame> StackTrace::ToStackFrames(
const StackTraceMapper& mapper, const StackTraceFilter& filtered) const {
std::vector<StackFrame> StackTrace::ToStackFrames() const {
std::vector<StackFrame> result;
result.reserve(code_objs_.size());
@ -48,18 +47,8 @@ std::vector<StackFrame> 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<StackFrame> 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;

20
tensorflow/python/util/stack_trace.h
View File

@ -40,17 +40,10 @@ inline void DCheckPyGilStateForStackTrace() {
#endif
}
// Maps filename/line_no combination into a stack frame.
using StackTraceMapper =
std::function<absl::optional<StackFrame>(std::string, int)>;
// Returns "true" for filenames which should be skipped.
using StackTraceFilter = std::function<bool(std::string)>;
// 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<StackFrame> ToStackFrames(
const StackTraceMapper& mapper = {},
const StackTraceFilter& filtered = {}) const;
std::vector<StackFrame> ToStackFrames() const;
private:
absl::InlinedVector<PyCodeObject*, kStackTraceInitialSize> code_objs_;
absl::InlinedVector<int, kStackTraceInitialSize> last_instructions_;
absl::InlinedVector<PyCodeObject*, kDefaultStackTraceInitialSize> code_objs_;
absl::InlinedVector<int, kDefaultStackTraceInitialSize> last_instructions_;
// Python GIL must be acquired beforehand.
ABSL_ATTRIBUTE_HOT

293
tensorflow/python/util/tf_stack.cc
View File

@ -19,20 +19,12 @@ limitations under the License.
#include <algorithm>
#include <vector>
#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<StackFrame>);
PYBIND11_MAKE_OPAQUE(StackTrace);
PYBIND11_MAKE_OPAQUE(std::vector<FrameSummary>);
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<py::str>(
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<py::str>(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<py::object>(code) : py::none();
}
std::string StackFrameToString(const StackFrame& frame) {
return py::str("<FrameSummary file {}, line {} in {}>")
.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<StackFrame const> stack_frames)
: stack_frames_cache_(std::vector<StackFrame>(stack_frames.begin(),
stack_frames.end())) {}
absl::Span<StackFrame const> 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("<FrameSummary file {}, line {} in {}>")
.format(filename, lineno, name);
}
private:
mutable absl::optional<std::vector<StackFrame>> stack_frames_cache_;
StackTrace captured_;
// TODO(cheshire): store those as C++ datastructures instead.
StackTraceMapper stack_trace_mapper_;
StackTraceFilter stack_trace_filter_;
};
std::vector<FrameSummary> 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<FrameSummary> ret;
// 16 is somewhat arbitrary, but TensorFlow stack traces tend to be deep.
ret.reserve(limit < 0 ? 16 : static_cast<size_t>(limit));
for (; f != nullptr && (limit < 0 || ret.size() < static_cast<size_t>(limit));
f = f->f_back) {
const PyCodeObject* co = f->f_code;
int lineno = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(f));
auto filename = py::reinterpret_borrow<py::str>(co->co_filename);
auto name = py::reinterpret_borrow<py::str>(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<py::int_>(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<py::object>(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_<StackFrame>(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_<FrameSummary>(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::tuple<std::string, int, std::string>>()(
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_<StackTraceWrapper>(m, "StackTraceWrapper", py::module_local(true))
py::bind_vector<std::vector<FrameSummary>>(m, "StackSummary",
py::module_local(true))
// TODO(slebedev): upstream negative indexing support into pybind11.
.def(
"__getitem__",
[](const StackTraceWrapper& self, ssize_t index) {
absl::Span<StackFrame const> frames = self.ToFrames();
[](const std::vector<FrameSummary>& self, ssize_t index) {
const size_t eff_index =
index < 0 ? frames.size() + index : static_cast<size_t>(index);
if (eff_index >= frames.size()) {
index < 0 ? self.size() + index : static_cast<size_t>(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<StackFrame const> 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<StackFrame> 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("<Opaque Stack Trace, access to initialize>");
.def("__repr__", [](const std::vector<FrameSummary>& 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<ssize_t>(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<StackFrame> {
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<py::str>(mapped[0])),
py::cast<py::int_>(mapped[1]),
std::string(py::cast<py::str>(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<ssize_t>(limit),
mappers, filters);
});
}
} // namespace tensorflow

13
tensorflow/python/util/tf_stack.py
View File

@ -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