PyCodeObject have references to PyFrameObject, therefore it might prevent some variable from garbage collected if we take its references. but we can only increase references for the objects that we do needed.
PiperOrigin-RevId: 362075013
Change-Id: I763c3bad8bcecf5160e1be1c0df8250a8eae29c9
diff --git a/tensorflow/python/profiler/internal/python_hooks.cc b/tensorflow/python/profiler/internal/python_hooks.cc
index aca1de1..7aa2847 100644
--- a/tensorflow/python/profiler/internal/python_hooks.cc
+++ b/tensorflow/python/profiler/internal/python_hooks.cc
@@ -40,17 +40,18 @@
setprofile(callback);
}
-std::string GetEventName(PyCodeObject* py_code) {
- string filename(py::reinterpret_borrow<py::str>(py_code->co_filename));
+std::string GetEventName(PyObject* co_filename, PyObject* co_name,
+ int co_firstlineno) {
+ string filename(py::reinterpret_borrow<py::str>(co_filename));
string function;
- if (py_code->co_name == nullptr) {
+ if (co_name == nullptr) {
function = "<unknown>";
} else {
- function = py::reinterpret_borrow<py::str>(py_code->co_name);
+ function = py::reinterpret_borrow<py::str>(co_name);
}
- return absl::StrCat("$", io::Basename(filename), ":", py_code->co_firstlineno,
- " ", function);
+ return absl::StrCat("$", io::Basename(filename), ":", co_firstlineno, " ",
+ function);
}
string GetEventName(PyCFunctionObject* py_cfunc) {
@@ -85,8 +86,8 @@
std::string PythonTraceEntry::Name() const {
std::string event_name;
- if (code_object) {
- return GetEventName(code_object);
+ if (co_filename) {
+ return GetEventName(co_filename, co_name, co_firstlineno);
} else if (function_object) {
return GetEventName(function_object);
}
@@ -252,7 +253,8 @@
switch (what) {
case PyTrace_CALL: {
PyCodeObject* f_code = frame->f_code;
- thread_traces.active.emplace(now, 0, f_code, nullptr);
+ thread_traces.active.emplace(now, 0, f_code->co_filename, f_code->co_name,
+ f_code->co_firstlineno);
break;
}
case PyTrace_RETURN:
@@ -264,8 +266,9 @@
thread_traces.active.pop();
} else if (options_.include_incomplete_events) {
PyCodeObject* f_code = frame->f_code;
- thread_traces.completed.emplace_back(start_timestamp_ns_, now, f_code,
- nullptr);
+ thread_traces.completed.emplace_back(
+ start_timestamp_ns_, now, f_code->co_filename, f_code->co_name,
+ f_code->co_firstlineno);
}
break;
}
@@ -273,7 +276,7 @@
if (PyCFunction_Check(arg)) {
// Python stack does not have a filename/line_no for native calls.
auto* func = reinterpret_cast<PyCFunctionObject*>(arg);
- entries_[thread_id].active.emplace(now, 0, nullptr, func);
+ entries_[thread_id].active.emplace(now, 0, func);
}
break;
}
@@ -290,7 +293,7 @@
// Python stack does not have a filename/line_no for native calls.
auto* func = reinterpret_cast<PyCFunctionObject*>(arg);
entries_[thread_id].completed.emplace_back(start_timestamp_ns_, now,
- nullptr, func);
+ func);
}
}
break;
diff --git a/tensorflow/python/profiler/internal/python_hooks.h b/tensorflow/python/profiler/internal/python_hooks.h
index e3d11db..76c1034 100644
--- a/tensorflow/python/profiler/internal/python_hooks.h
+++ b/tensorflow/python/profiler/internal/python_hooks.h
@@ -45,25 +45,36 @@
};
struct PythonTraceEntry {
- PythonTraceEntry(uint64 start, uint64 end, PyCodeObject* code,
- PyCFunctionObject* func)
+ PythonTraceEntry(uint64 start, uint64 end, PyObject* filename, PyObject* name,
+ int firstlineno)
: start_time_ns(start),
end_time_ns(end),
- code_object(code),
- function_object(func) {
- Py_XINCREF(code_object);
+ co_filename(filename),
+ co_name(name),
+ co_firstlineno(firstlineno) {
+ Py_XINCREF(co_filename);
+ Py_XINCREF(co_name);
+ }
+ PythonTraceEntry(uint64 start, uint64 end, PyCFunctionObject* func)
+ : start_time_ns(start), end_time_ns(end), function_object(func) {
Py_XINCREF(function_object);
}
+
~PythonTraceEntry() {
- Py_XDECREF(code_object);
+ Py_XDECREF(co_filename);
+ Py_XDECREF(co_name);
Py_XDECREF(function_object);
}
+
PythonTraceEntry(PythonTraceEntry&& other) {
start_time_ns = other.start_time_ns;
end_time_ns = other.end_time_ns;
- code_object = other.code_object;
+ co_firstlineno = other.co_firstlineno;
+ co_filename = other.co_filename;
+ co_name = other.co_name;
function_object = other.function_object;
- other.code_object = nullptr;
+ other.co_filename = nullptr;
+ other.co_name = nullptr;
other.function_object = nullptr;
}
@@ -71,8 +82,10 @@
uint64 start_time_ns;
uint64 end_time_ns;
- PyCodeObject* code_object;
- PyCFunctionObject* function_object;
+ PyObject* co_filename = nullptr;
+ PyObject* co_name = nullptr;
+ int co_firstlineno = 0;
+ PyCFunctionObject* function_object = nullptr;
PythonTraceEntry(const PythonTraceEntry& other) = delete;
void operator=(const PythonTraceEntry&) = delete;
