[NFC][Stack Traces] Do not build a functor where we can use the translation map directly

PiperOrigin-RevId: 356528885
Change-Id: I4482fd75dd10fee71b72b0537cbf96a4762dde92
This commit is contained in:
Mihai Maruseac 2021-02-09 10:07:32 -08:00 committed by TensorFlower Gardener
parent fb194caf50
commit f11e226a34
5 changed files with 31 additions and 20 deletions

View File

@ -482,7 +482,6 @@ cc_library(
],
deps = [
"//tensorflow/core/platform:stack_frame",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],

View File

@ -19,22 +19,20 @@ limitations under the License.
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/strings/match.h"
#include "absl/types/optional.h"
#include "tensorflow/core/platform/stack_frame.h"
namespace tensorflow {
using SourceLoc = std::tuple<std::string, int>;
// Maps filename/line_no combination into a stack frame.
using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
using StackTraceMap =
std::function<absl::optional<StackFrame>(std::pair<const char*, int>)>;
// Returns "true" on filenames which should be skipped.
using StackTraceFilter = std::function<bool(const char*)>;
using ToStackFramesFunctor = std::vector<StackFrame>(int, const SourceMap&,
using ToStackFramesFunctor = std::vector<StackFrame>(int, const StackTraceMap&,
const StackTraceFilter&,
bool, int);
@ -56,12 +54,11 @@ class ManagedStackTrace {
: id_(id), to_stack_frames_(to_stack_frames) {}
// Returns stack trace as a vector of `StackFrame`s.
std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
const StackTraceFilter& filtered = {},
bool reverse_traversal = false,
int limit = -1) const {
return to_stack_frames_(id_, source_map, filtered, reverse_traversal,
limit);
return to_stack_frames_(id_, mapper, filtered, reverse_traversal, limit);
}
private:

View File

@ -42,7 +42,7 @@ const char* GetPythonString(PyObject* o) {
namespace tensorflow {
std::vector<StackFrame> StackTrace::ToStackFrames(
const SourceMap& source_map, const StackTraceFilter& filtered,
const StackTraceMap& mapper, const StackTraceFilter& filtered,
bool reverse_traversal, int limit) const {
DCheckPyGilStateForStackTrace();
std::vector<StackFrame> result;
@ -61,9 +61,11 @@ std::vector<StackFrame> StackTrace::ToStackFrames(
continue;
}
const auto it = source_map.find(std::make_tuple(file_name, line_number));
if (it != source_map.end()) {
result.push_back(it->second);
absl::optional<StackFrame> mapped =
mapper ? mapper(std::make_pair(file_name, line_number)) : absl::nullopt;
if (mapped) {
result.push_back(*mapped);
} else {
result.emplace_back(StackFrame{file_name, line_number,
GetPythonString(code_obj.first->co_name)});

View File

@ -88,12 +88,12 @@ class StackTrace final {
}
// Returns a structured representation of the captured stack trace.
// `source_map` provides a custom mapping for translating stack frames,
// `filter` returns `true` for the stack frames which should be omitted.
// `mapper` provides a custom mapping for translating stack frames, `filter`
// returns `true` for the stack frames which should be omitted.
//
// `reverse_traversal` changes the traversal order of the stack trace, and
// `limit` bounds the number of returned frames (after filtering).
std::vector<StackFrame> ToStackFrames(const SourceMap& source_map = {},
std::vector<StackFrame> ToStackFrames(const StackTraceMap& mapper = {},
const StackTraceFilter& filtered = {},
bool reverse_traversal = false,
int limit = -1) const;
@ -149,11 +149,11 @@ extern StackTraceManager* const stack_trace_manager;
// Converts the ManagedStackTrace (identified by ID) to a vector of stack
// frames.
inline std::vector<StackFrame> ManagedStackTraceToStackFrames(
int id, const SourceMap& source_map, const StackTraceFilter& filtered,
int id, const StackTraceMap& mapper, const StackTraceFilter& filtered,
bool reverse_traversal, int limit) {
PyGILState_STATE gstate = PyGILState_Ensure();
std::vector<StackFrame> result = stack_trace_manager->Get(id)->ToStackFrames(
source_map, filtered, reverse_traversal, limit);
mapper, filtered, reverse_traversal, limit);
PyGILState_Release(gstate);
return result;
}

View File

@ -54,6 +54,10 @@ namespace {
namespace py = pybind11;
using SourceLoc = std::tuple<std::string, int>;
using SourceMap = absl::flat_hash_map<SourceLoc, StackFrame>;
using StringSet = absl::flat_hash_set<std::string>;
// Python wrapper for a SourceMap.
@ -145,7 +149,8 @@ class StackTraceWrapper : public AbstractStackTrace {
PyGILState_STATE state = PyGILState_Ensure();
stack_frames_cache_ = captured_.ToStackFrames(
*source_map_, [&](const char* f) { return StackTraceFiltering(f); });
[&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
[&](const char* f) { return StackTraceFiltering(f); });
stack_frames_cache_->pop_back(); // Drop last stack frame.
PyGILState_Release(state);
return *stack_frames_cache_;
@ -158,7 +163,7 @@ class StackTraceWrapper : public AbstractStackTrace {
PyGILState_STATE state = PyGILState_Ensure();
std::vector<StackFrame> last_frame = captured_.ToStackFrames(
*source_map_,
[&](std::pair<const char*, int> p) { return StackTraceMapping(p); },
[&](const char* file_name) {
return StackTraceFiltering(file_name) ||
IsInternalFrameForFilename(file_name);
@ -221,6 +226,14 @@ class StackTraceWrapper : public AbstractStackTrace {
});
}
absl::optional<StackFrame> StackTraceMapping(SourceLoc loc) const {
if (source_map_->contains(loc)) {
return source_map_->at(loc);
}
return absl::nullopt;
}
bool StackTraceFiltering(const char* file_name) const {
return filter_->contains(file_name);
}