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 |
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 | |
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 |
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 |
80 | ~StackTrace() { Clear(); } |
81 | |
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 |
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 |
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. |
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. |
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 | |