1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #ifndef TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
17 #define TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
18
19 #include <Python.h>
20 #include <frameobject.h>
21
22 #include <array>
23 #include <limits>
24 #include <sstream>
25 #include <string>
26
27 #include "absl/base/attributes.h"
28 #include "absl/base/optimization.h"
29 #include "absl/container/flat_hash_map.h"
30 #include "absl/container/flat_hash_set.h"
31 #include "absl/container/inlined_vector.h"
32 #include "absl/types/optional.h"
33 #include "tensorflow/core/platform/status.h"
34 #include "tensorflow/core/util/managed_stack_trace.h"
35
36 namespace tensorflow {
37
38 // Assert that Python GIL is held.
39 // TODO(cheshire): Fix duplication vs. py_util.h
DCheckPyGilStateForStackTrace()40 inline void DCheckPyGilStateForStackTrace() {
41 #if PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION >= 4
42 DCHECK(PyGILState_Check());
43 #endif
44 }
45
46 // A class for capturing Python stack trace.
47 class StackTrace final {
48 public:
49 static constexpr int kStackTraceInitialSize = 30;
50
StackTrace()51 StackTrace() {}
52
53 // Returns `StackTrace` object that captures the current Python stack trace.
54 // `limit` determines how many stack frames at most are returned: set to -1
55 // for "no limit".
56 // Python GIL must be acquired beforehand.
57 ABSL_MUST_USE_RESULT
58 ABSL_ATTRIBUTE_HOT
Capture(int limit)59 static StackTrace Capture(int limit) {
60 DCheckPyGilStateForStackTrace();
61 if (limit == -1) limit = std::numeric_limits<int>::max();
62
63 StackTrace result;
64 const PyFrameObject* frame = PyThreadState_GET()->frame;
65 int i = 0;
66 for (; i < limit && frame != nullptr; frame = frame->f_back, ++i) {
67 PyCodeObject* code_obj = frame->f_code;
68 DCHECK(code_obj != nullptr);
69
70 Py_INCREF(code_obj);
71 int line_number =
72 PyFrame_GetLineNumber(const_cast<PyFrameObject*>(frame));
73 result.code_objs_.push_back(std::make_pair(code_obj, line_number));
74 }
75 return result;
76 }
77
78 // Python GIL must be acquired beforehand.
79 ABSL_ATTRIBUTE_HOT
~StackTrace()80 ~StackTrace() { Clear(); }
81
StackTrace(StackTrace && other)82 StackTrace(StackTrace&& other) { std::swap(code_objs_, other.code_objs_); }
83
84 // Python GIL must be acquired beforehand.
85 ABSL_ATTRIBUTE_HOT
86 StackTrace& operator=(StackTrace&& other) {
87 Clear();
88 std::swap(code_objs_, other.code_objs_);
89 return *this;
90 }
91
92 // Returns a structured representation of the captured stack trace.
93 // `source_map` provides a custom mapping for translating stack frames,
94 // `filter` returns `true` for the stack frames which should be omitted.
95 //
96 // `reverse_traversal` changes the traversal order of the stack trace, and
97 // `limit` bounds the number of returned frames (after filtering).
98 std::vector<StackFrame> ToStackFrames(const SourceMap& source_map,
99 const StackTraceFilter& filtered,
100 bool reverse_traversal = false,
101 int limit = -1) const;
102
103 // Python GIL must be acquired beforehand.
104 ABSL_ATTRIBUTE_HOT
Clear()105 void Clear() {
106 if (!code_objs_.empty()) DCheckPyGilStateForStackTrace();
107 for (const auto& p : code_objs_) Py_DECREF(p.first);
108 code_objs_.clear();
109 }
110
111 private:
112 absl::InlinedVector<std::pair<PyCodeObject*, int>, kStackTraceInitialSize>
113 code_objs_;
114
115 StackTrace(const StackTrace&) = delete;
116 StackTrace& operator=(const StackTrace&) = delete;
117 };
118
119 // A class that manages Python stack traces in a circular buffer. Users can
120 // insert stack trace entries and retrieve them by ids.
121 class StackTraceManager {
122 public:
123 static constexpr int kStackTraceCircularBufferSize = 1024;
124
125 // Captures the current Python stack trace and returns an id.
126 // Python GIL must be acquired beforehand.
127 ABSL_MUST_USE_RESULT
128 ABSL_ATTRIBUTE_HOT
Capture(int limit)129 int Capture(int limit) {
130 DCheckPyGilStateForStackTrace();
131 const int id = next_id_++;
132 const int index = id & (kStackTraceCircularBufferSize - 1);
133 stack_traces_[index] = StackTrace::Capture(limit);
134 return id;
135 }
136
137 // Retrieve captured Python stack trace by id. Returns `nullptr` if the
138 // requested stack trace is evicted from the circular buffer.
139 // Python GIL must be acquired beforehand.
140 ABSL_MUST_USE_RESULT
141 StackTrace* Get(int id);
142
143 private:
144 int next_id_ = 0;
145 std::array<StackTrace, kStackTraceCircularBufferSize> stack_traces_;
146 };
147
148 // Singleton StackTraceManager.
149 extern StackTraceManager* const stack_trace_manager;
150
151 // Converts the ManagedStackTrace (identified by ID) to a vector of stack
152 // frames.
ManagedStackTraceToStackFrames(int id,const SourceMap & source_map,const StackTraceFilter & filtered,bool reverse_traversal,int limit)153 inline std::vector<StackFrame> ManagedStackTraceToStackFrames(
154 int id, const SourceMap& source_map, const StackTraceFilter& filtered,
155 bool reverse_traversal, int limit) {
156 PyGILState_STATE gstate = PyGILState_Ensure();
157 StackTrace* stack_trace = stack_trace_manager->Get(id);
158 if (!stack_trace) {
159 // Must have evicted the stack trace by now. Do best effort.
160 return {};
161 }
162
163 std::vector<StackFrame> result = stack_trace->ToStackFrames(
164 source_map, filtered, reverse_traversal, limit);
165 PyGILState_Release(gstate);
166 return result;
167 }
168
169 // Returns Python stack trace object that can be converted to string.
170 // Note that the actual stack trace is kept in a circular buffer for string
171 // conversion could fail if it's evicted before.
172 // Python GIL must be acquired beforehand.
GetStackTrace(int limit)173 inline ManagedStackTrace GetStackTrace(int limit) {
174 DCheckPyGilStateForStackTrace();
175 return ManagedStackTrace(stack_trace_manager->Capture(limit),
176 &ManagedStackTraceToStackFrames);
177 }
178
179 } // namespace tensorflow
180
181 #endif // TENSORFLOW_PYTHON_UTIL_STACK_TRACE_H_
182