experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

reduce public interface for python hooks.

Browse Source 
PiperOrigin-RevId: 358447444
Change-Id: I949c08d8aa4da929667b42766a9289a549c87ff8
This commit is contained in:
A. Unique TensorFlower 2021-02-19 11:32:17 -08:00 committed by TensorFlower Gardener
parent 03b926d9f9
commit 9b7ff60faa
2 changed files with 23 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
tensorflow/python/profiler/internal/python_hooks.cc
View File

@ -30,13 +30,6 @@ namespace py = ::pybind11;
namespace {
template <typename T>
int ProfileFunction(PyObject* obj, PyFrameObject* frame, int what,
PyObject* arg) {
T::GetSingleton()->ProfileFast(frame, what, arg);
return 0;
}
void SysSetProfileNone() {
py::object setprofile = py::module::import("sys").attr("setprofile");
setprofile(py::none());
@ -214,6 +207,12 @@ void PythonHookContext::Finalize(XSpace* space) {
}
}
/*static*/ int PythonHooks::ProfileFunction(PyObject* obj, PyFrameObject* frame,
int what, PyObject* arg) {
GetSingleton()->ProfileFast(frame, what, arg);
return 0;
}
void PythonHooks::ProfileSlow(const py::object& frame, const string& event,
const py::object& arg) {
int what;
@ -301,7 +300,7 @@ void PythonHookContext::ProfileFast(PyFrameObject* frame, int what,
}
}
void PythonHookContext::SetProfilerInAllThreads() {
/*static*/ void PythonHookContext::SetProfilerInAllThreads() {
// We also want any new threads started to use our profiler.
// NOTE: threading does not provide a C API equivalent to
// `threading.setprofile` so we are forced to go via Python to setup the
@ -315,7 +314,7 @@ void PythonHookContext::SetProfilerInAllThreads() {
const py::object& arg) {
singleton->ProfileSlow(frame, event, arg);
SysSetProfileNone();
PyEval_SetProfile(ProfileFunction<PythonHooks>, nullptr);
PyEval_SetProfile(&PythonHooks::ProfileFunction, nullptr);
});
ThreadingSetProfile(callback);
@ -327,7 +326,7 @@ void PythonHookContext::SetProfilerInAllThreads() {
while (next_thread != nullptr) {
VLOG(1) << "Setting profiler in " << next_thread->thread_id;
PyThreadState_Swap(next_thread);
PyEval_SetProfile(ProfileFunction<PythonHooks>, nullptr);
PyEval_SetProfile(&PythonHooks::ProfileFunction, nullptr);
next_thread = next_thread->next;
}
PyThreadState_Swap(curr_thread);

19
tensorflow/python/profiler/internal/python_hooks.h
View File

@ -84,18 +84,22 @@ struct PerThreadEvents {
std::stack<PythonTraceEntry> active;
};
class PythonHooks;
class PythonHookContext {
public:
void Start(const PythonHooksOptions& option);
void Stop();
void Finalize(XSpace* space);
void ProfileFast(PyFrameObject* frame, int what, PyObject* arg);
friend class ::tensorflow::profiler::PythonHooks;
private:
void Start(const PythonHooksOptions& option);
void Stop();
void ProfileFast(PyFrameObject* frame, int what, PyObject* arg);
void CollectData(XPlane* raw_plane);
static void EnableTraceMe(bool enable);
void SetProfilerInAllThreads();
static void SetProfilerInAllThreads();
static void ClearProfilerInAllThreads();
void operator=(const PythonHookContext&) = delete;
@ -134,6 +138,9 @@ class PythonHooks {
return output;
}
friend class ::tensorflow::profiler::PythonHookContext;
private:
void ProfileSlow(const py::object& frame, const string& event,
const py::object& arg);
@ -149,7 +156,9 @@ class PythonHooks {
static PythonHookContext* e2e_context() { return e2e_context_; }
private:
static int ProfileFunction(PyObject* obj, PyFrameObject* frame, int what,
PyObject* arg);
// active_context_ are accessed when GIL is held, therefore no race
// conditions.
std::unique_ptr<PythonHookContext> active_context_;