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 #ifndef TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_ 16 #define TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_ 17 18 #include <memory> 19 #include <stack> 20 #include <vector> 21 22 #include "absl/container/flat_hash_map.h" 23 #include "absl/memory/memory.h" 24 #include "pybind11/cast.h" 25 #include "pybind11/pybind11.h" 26 #include "pybind11/pytypes.h" 27 #include "tensorflow/core/platform/macros.h" 28 #include "tensorflow/core/platform/types.h" 29 #include "tensorflow/core/profiler/protobuf/xplane.pb.h" 30 31 namespace tensorflow { 32 namespace profiler { 33 34 namespace py = ::pybind11; 35 36 struct PythonHooksOptions { 37 bool enable_trace_python_function = false; 38 bool enable_python_traceme = true; 39 bool end_to_end_mode = false; 40 // Incomplete events are defined as those python calls which we only see 41 // either start or end, but not both. If we want to include them in the final 42 // result, profiler start, end time are used respectively to the absent 43 // timestamps. 44 bool include_incomplete_events = true; 45 }; 46 47 struct PythonTraceEntry { 48 // Capture the source/line information for a PyCodeObject object. 49 // In eager mode, keeping a reference to PyCodeObject leaks device memory. PythonTraceEntryPythonTraceEntry50 PythonTraceEntry(uint64 start, uint64 end, PyCodeObject* py_code_object) 51 : start_time_ns(start), 52 end_time_ns(end), 53 co_filename(py_code_object->co_filename), 54 co_name(py_code_object->co_name), 55 co_firstlineno(py_code_object->co_firstlineno) { 56 Py_XINCREF(co_filename); 57 Py_XINCREF(co_name); 58 } 59 // Capture the source/line information for a PyCFunctionObject object. 60 // In eager mode, keeping a reference to PyCFunctionObject leaks device 61 // memory. PythonTraceEntryPythonTraceEntry62 PythonTraceEntry(uint64 start, uint64 end, PyCFunctionObject* py_c_function) 63 : start_time_ns(start), 64 end_time_ns(end), 65 method_def(py_c_function->m_ml), 66 m_module(py_c_function->m_module) { 67 Py_XINCREF(m_module); 68 } 69 ~PythonTraceEntryPythonTraceEntry70 ~PythonTraceEntry() { 71 Py_XDECREF(co_filename); 72 Py_XDECREF(co_name); 73 Py_XDECREF(m_module); 74 } 75 PythonTraceEntryPythonTraceEntry76 PythonTraceEntry(PythonTraceEntry&& other) { 77 start_time_ns = other.start_time_ns; 78 end_time_ns = other.end_time_ns; 79 co_firstlineno = other.co_firstlineno; 80 co_filename = other.co_filename; 81 co_name = other.co_name; 82 method_def = other.method_def; 83 m_module = other.m_module; 84 other.co_filename = nullptr; 85 other.co_name = nullptr; 86 other.method_def = nullptr; 87 other.m_module = nullptr; 88 } 89 90 std::string Name() const; 91 92 uint64 start_time_ns; 93 uint64 end_time_ns; 94 PyObject* co_filename = nullptr; 95 PyObject* co_name = nullptr; 96 int co_firstlineno = 0; 97 PyMethodDef* method_def = nullptr; 98 PyObject* m_module = nullptr; 99 100 PythonTraceEntry(const PythonTraceEntry& other) = delete; 101 void operator=(const PythonTraceEntry&) = delete; 102 void operator=(PythonTraceEntry&&) = delete; 103 }; 104 105 struct PerThreadEvents { 106 std::deque<PythonTraceEntry> completed; 107 std::stack<PythonTraceEntry> active; 108 }; 109 110 class PythonHooks; 111 112 class PythonHookContext { 113 public: 114 void Finalize(XSpace* space); 115 116 friend class ::tensorflow::profiler::PythonHooks; 117 118 private: 119 void Start(const PythonHooksOptions& option); 120 void Stop(); 121 void ProfileFast(PyFrameObject* frame, int what, PyObject* arg); 122 void CollectData(XPlane* raw_plane); 123 static void EnableTraceMe(bool enable); 124 125 static void SetProfilerInAllThreads(); 126 static void ClearProfilerInAllThreads(); 127 128 void operator=(const PythonHookContext&) = delete; 129 void operator=(PythonHookContext&&) = delete; 130 131 absl::flat_hash_map<int64_t, PerThreadEvents> entries_; 132 uint64 start_timestamp_ns_; 133 PythonHooksOptions options_; 134 // In end to end mode, Python get uninitialized before Stop()/Finalize(), we 135 // need to buffer the result. 136 absl::optional<XPlane> end_to_end_xplane_; 137 }; 138 139 // Singleton for tracing python function calls. 140 class PythonHooks { 141 public: 142 static PythonHooks* GetSingleton(); 143 Start(const PythonHooksOptions & option)144 void Start(const PythonHooksOptions& option) { 145 if (active_context_) return; 146 active_context_ = std::make_unique<PythonHookContext>(); 147 active_context_->Start(option); 148 } 149 Stop()150 std::unique_ptr<PythonHookContext> Stop() { 151 if (e2e_context_) { 152 auto* e2e_context = e2e_context_; 153 e2e_context_ = nullptr; 154 return absl::WrapUnique(e2e_context); 155 } 156 157 if (!active_context_) return nullptr; 158 active_context_->Stop(); 159 std::unique_ptr<PythonHookContext> output = std::move(active_context_); 160 active_context_.reset(); 161 return output; 162 } 163 164 friend class ::tensorflow::profiler::PythonHookContext; 165 166 private: 167 void ProfileSlow(const py::object& frame, const string& event, 168 const py::object& arg); 169 ProfileFast(PyFrameObject * frame,int what,PyObject * arg)170 void ProfileFast(PyFrameObject* frame, int what, PyObject* arg) { 171 if (TF_PREDICT_TRUE(active_context_)) { 172 active_context_->ProfileFast(frame, what, arg); 173 } 174 } 175 set_e2e_context(PythonHookContext * e2e_context)176 static void set_e2e_context(PythonHookContext* e2e_context) { 177 e2e_context_ = e2e_context; 178 } 179 e2e_context()180 static PythonHookContext* e2e_context() { return e2e_context_; } 181 182 static int ProfileFunction(PyObject* obj, PyFrameObject* frame, int what, 183 PyObject* arg); 184 185 // active_context_ are accessed when GIL is held, therefore no race 186 // conditions. 187 std::unique_ptr<PythonHookContext> active_context_; 188 static PythonHookContext* e2e_context_; 189 }; 190 191 } // namespace profiler 192 } // namespace tensorflow 193 194 #endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_ 195