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 { PythonTraceEntryPythonTraceEntry48 PythonTraceEntry(uint64 start, uint64 end, PyCodeObject* code, 49 PyCFunctionObject* func) 50 : start_time_ns(start), 51 end_time_ns(end), 52 code_object(code), 53 function_object(func) { 54 Py_XINCREF(code_object); 55 Py_XINCREF(function_object); 56 } ~PythonTraceEntryPythonTraceEntry57 ~PythonTraceEntry() { 58 Py_XDECREF(code_object); 59 Py_XDECREF(function_object); 60 } PythonTraceEntryPythonTraceEntry61 PythonTraceEntry(PythonTraceEntry&& other) { 62 start_time_ns = other.start_time_ns; 63 end_time_ns = other.end_time_ns; 64 code_object = other.code_object; 65 function_object = other.function_object; 66 other.code_object = nullptr; 67 other.function_object = nullptr; 68 } 69 70 std::string Name() const; 71 72 uint64 start_time_ns; 73 uint64 end_time_ns; 74 PyCodeObject* code_object; 75 PyCFunctionObject* function_object; 76 77 PythonTraceEntry(const PythonTraceEntry& other) = delete; 78 void operator=(const PythonTraceEntry&) = delete; 79 void operator=(PythonTraceEntry&&) = delete; 80 }; 81 82 struct PerThreadEvents { 83 std::deque<PythonTraceEntry> completed; 84 std::stack<PythonTraceEntry> active; 85 }; 86 87 class PythonHookContext { 88 public: 89 void Start(const PythonHooksOptions& option); 90 void Stop(); 91 void Finalize(XSpace* space); 92 void ProfileFast(PyFrameObject* frame, int what, PyObject* arg); 93 94 private: 95 void CollectData(XPlane* raw_plane); 96 static void EnableTraceMe(bool enable); 97 98 void SetProfilerInAllThreads(); 99 static void ClearProfilerInAllThreads(); 100 101 void operator=(const PythonHookContext&) = delete; 102 void operator=(PythonHookContext&&) = delete; 103 104 absl::flat_hash_map<int64, PerThreadEvents> entries_; 105 uint64 start_timestamp_ns_; 106 PythonHooksOptions options_; 107 // In end to end mode, Python get uninitialized before Stop()/Finalize(), we 108 // need to buffer the result. 109 absl::optional<XPlane> end_to_end_xplane_; 110 }; 111 112 // Singleton for tracing python function calls. 113 class PythonHooks { 114 public: 115 static PythonHooks* GetSingleton(); 116 Start(const PythonHooksOptions & option)117 void Start(const PythonHooksOptions& option) { 118 if (active_context_) return; 119 active_context_ = std::make_unique<PythonHookContext>(); 120 active_context_->Start(option); 121 } 122 Stop()123 std::unique_ptr<PythonHookContext> Stop() { 124 if (e2e_context_) { 125 auto* e2e_context = e2e_context_; 126 e2e_context_ = nullptr; 127 return absl::WrapUnique(e2e_context); 128 } 129 130 if (!active_context_) return nullptr; 131 active_context_->Stop(); 132 std::unique_ptr<PythonHookContext> output = std::move(active_context_); 133 active_context_.reset(); 134 return output; 135 } 136 137 void ProfileSlow(const py::object& frame, const string& event, 138 const py::object& arg); 139 ProfileFast(PyFrameObject * frame,int what,PyObject * arg)140 void ProfileFast(PyFrameObject* frame, int what, PyObject* arg) { 141 if (TF_PREDICT_TRUE(active_context_)) { 142 active_context_->ProfileFast(frame, what, arg); 143 } 144 } 145 set_e2e_context(PythonHookContext * e2e_context)146 static void set_e2e_context(PythonHookContext* e2e_context) { 147 e2e_context_ = e2e_context; 148 } 149 e2e_context()150 static PythonHookContext* e2e_context() { return e2e_context_; } 151 152 private: 153 // active_context_ are accessed when GIL is held, therefore no race 154 // conditions. 155 std::unique_ptr<PythonHookContext> active_context_; 156 static PythonHookContext* e2e_context_; 157 }; 158 159 } // namespace profiler 160 } // namespace tensorflow 161 162 #endif // TENSORFLOW_PYTHON_PROFILER_INTERNAL_PYTHON_HOOKS_H_ 163