wehub-resource-sync
133 lines
4.1 KiB
C++
133 lines
4.1 KiB
C++
/* 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"
|
|
|
|
#include <Python.h>
|
|
|
|
#include <limits>
|
|
#include <memory>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "absl/base/attributes.h"
|
|
#include "absl/container/flat_hash_map.h"
|
|
#include "absl/log/check.h"
|
|
#include "tensorflow/core/platform/stack_frame.h"
|
|
#include "tensorflow/core/util/managed_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 (PyBytes_Check(o)) {
|
|
return PyBytes_AsString(o);
|
|
} else {
|
|
return PyUnicode_AsUTF8(o);
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
namespace tensorflow {
|
|
|
|
ABSL_MUST_USE_RESULT
|
|
ABSL_ATTRIBUTE_HOT
|
|
std::shared_ptr<StackTrace> StackTrace::Capture(int limit) {
|
|
DCHECK(PyGILState_Check());
|
|
if (limit == -1) limit = std::numeric_limits<int>::max();
|
|
|
|
StackTrace result;
|
|
#if PY_VERSION_HEX >= 0x030B0000
|
|
PyFrameObject* oldframe;
|
|
PyFrameObject* frame = PyThreadState_GetFrame(PyThreadState_GET());
|
|
for (int i = 0; i < limit && frame != nullptr; oldframe = frame,
|
|
frame = PyFrame_GetBack(frame), Py_DECREF(oldframe), ++i) {
|
|
PyCodeObject* code_obj = PyFrame_GetCode(frame);
|
|
DCHECK(code_obj != nullptr);
|
|
|
|
int line_number = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(frame));
|
|
result.code_objs_.push_back(std::make_pair(code_obj, line_number));
|
|
}
|
|
Py_XDECREF(frame);
|
|
#else
|
|
const PyFrameObject* frame = PyThreadState_GET()->frame;
|
|
for (int i = 0; i < limit && frame != nullptr; frame = frame->f_back, ++i) {
|
|
PyCodeObject* code_obj = frame->f_code;
|
|
Py_XINCREF(code_obj);
|
|
DCHECK(code_obj != nullptr);
|
|
|
|
int line_number = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(frame));
|
|
result.code_objs_.push_back(std::make_pair(code_obj, line_number));
|
|
}
|
|
#endif
|
|
|
|
static absl::flat_hash_map<uint64_t, std::shared_ptr<StackTrace>>* cache =
|
|
new absl::flat_hash_map<uint64_t, std::shared_ptr<StackTrace>>();
|
|
#ifdef Py_GIL_DISABLED
|
|
static absl::Mutex mu(absl::kConstInit);
|
|
|
|
absl::MutexLock lock(&mu);
|
|
#endif // Py_GIL_DISABLED
|
|
uint64_t hash_code = result.hash();
|
|
if (!cache->contains(hash_code)) {
|
|
cache->insert(std::make_pair(
|
|
hash_code, std::make_shared<StackTrace>(std::move(result))));
|
|
}
|
|
|
|
return cache->at(hash_code);
|
|
}
|
|
|
|
std::vector<StackFrame> StackTrace::ToStackFrames(
|
|
const SourceMap& source_map, const StackTraceFilter& filtered,
|
|
bool reverse_traversal, int limit) const {
|
|
auto gil_state = PyGILState_Ensure();
|
|
std::vector<StackFrame> result;
|
|
result.reserve(code_objs_.size());
|
|
|
|
if (limit == -1) limit = std::numeric_limits<int>::max();
|
|
|
|
for (int i = 0; i < code_objs_.size(); i++) {
|
|
int idx = reverse_traversal ? i : code_objs_.size() - 1 - i;
|
|
|
|
const std::pair<PyCodeObject*, int>& code_obj = code_objs_[idx];
|
|
const char* file_name = GetPythonString(code_obj.first->co_filename);
|
|
const int line_number = code_obj.second;
|
|
|
|
if (filtered && filtered(file_name)) {
|
|
continue;
|
|
}
|
|
|
|
const auto it = source_map.find(SourceLoc{file_name, line_number});
|
|
if (it != source_map.end()) {
|
|
result.push_back(it->second);
|
|
} else {
|
|
result.emplace_back(StackFrame{file_name, line_number,
|
|
GetPythonString(code_obj.first->co_name)});
|
|
}
|
|
|
|
if (result.size() == limit) {
|
|
break;
|
|
}
|
|
}
|
|
|
|
PyGILState_Release(gil_state);
|
|
return result;
|
|
}
|
|
|
|
} // namespace tensorflow
|