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

Copy stack traces to instantiated functions.

Browse Source 
Otherwise, instantiated functions have no Python stack traces, and one has to
get them manually through FunctionDefinitionLibrary.

More importantly, it makes it impossible to assign different stack traces in
cases where it is actually required: namely, inlining.

PiperOrigin-RevId: 346435458
Change-Id: I41a3188e453566fbae6d29b261eefe4d24cf6453
This commit is contained in:
George Karpenkov 2020-12-08 16:27:46 -08:00 committed by TensorFlower Gardener
parent 2a0fd3f6b1
commit 20d69d0834
13 changed files with 94 additions and 50 deletions

6
tensorflow/c/c_api_function.cc
View File

@ -213,7 +213,11 @@ TF_Function* TF_GraphToFunctionWithControlOutputs(
TF_DeleteFunction(tf_function);
return nullptr;
}
tf_function->graph_with_debug_info = &fn_body->graph;
for (const Node* n : fn_body->graph.nodes()) {
tf_function->stack_traces[n->name()] = n->GetStackTrace();
}
return tf_function;
}

4
tensorflow/c/c_api_internal.h
View File

@ -157,9 +157,7 @@ struct TF_DeviceList {
struct TF_Function {
tensorflow::FunctionDef fdef;
// Graph with nodes with debug stack traces.
const tensorflow::Graph* graph_with_debug_info = nullptr;
tensorflow::StackTracesMap stack_traces;
};
struct TF_ApiDefMap {

4
tensorflow/c/eager/c_api.cc
View File

@ -749,8 +749,8 @@ void TFE_ContextAddFunctionDef(TFE_Context* ctx,
void TFE_ContextAddFunction(TFE_Context* ctx, TF_Function* function,
TF_Status* status) {
status->status = tensorflow::unwrap(ctx)->AddFunctionDefWithDebugInfo(
function->fdef, function->graph_with_debug_info);
status->status = tensorflow::unwrap(ctx)->AddFunctionDefWithStackTraces(
function->fdef, function->stack_traces);
}
void TFE_ContextRemoveFunction(TFE_Context* ctx, const char* name,

10
tensorflow/c/eager/immediate_execution_context.h
View File

@ -111,11 +111,11 @@ class ImmediateExecutionContext : public AbstractContext {
// already exists.
virtual Status AddFunctionDef(const FunctionDef& fdef) = 0;
// Same as `AddFunctionDef`, and additionally saves a pointer to the Graph
// which has nodes containing stack traces for the nodes in `fdef`. Assumes
// `graph` is alive while the function is alive.
virtual Status AddFunctionDefWithDebugInfo(const FunctionDef& fdef,
const Graph* graph) = 0;
// Same as `AddFunctionDef`, but additionally saves the `stack_traces` under
// the key of the function definition name (to be retrieved during function
// instantiation).
virtual Status AddFunctionDefWithStackTraces(
const FunctionDef& fdef, const StackTracesMap& stack_traces) = 0;
// Find and return a added function by its name.
virtual const FunctionDef* FindFunctionDef(const string& name) const = 0;

11
tensorflow/core/common_runtime/eager/context.cc
View File

@ -705,10 +705,10 @@ Status EagerContext::RegisterExistingFunctionsOnRemoteWorkers(
return Status::OK();
}
Status EagerContext::AddFunctionDefWithDebugInfo(
const FunctionDef& fdef, const Graph* graph_with_debug_info) {
Status EagerContext::AddFunctionDefWithStackTraces(
const FunctionDef& fdef, const StackTracesMap& stack_traces) {
return AddFunctionDef(fdef, FunctionDefLibrary(),
/* add_to_local_only=*/false, graph_with_debug_info);
/* add_to_local_only=*/false, stack_traces);
}
Status EagerContext::AddFunctionDef(const FunctionDef& fdef) {
@ -719,7 +719,7 @@ Status EagerContext::AddFunctionDef(const FunctionDef& fdef) {
Status EagerContext::AddFunctionDef(const FunctionDef& fdef,
const FunctionDefLibrary& library,
const bool add_to_local_only,
const Graph* graph_with_debug_info) {
const StackTracesMap& stack_traces) {
bool is_first_ref = false;
{
mutex_lock l(cache_mu_);
@ -753,8 +753,7 @@ Status EagerContext::AddFunctionDef(const FunctionDef& fdef,
is_first_ref = registered_function->RefCountIsOne();
}
if (is_first_ref) {
TF_RETURN_IF_ERROR(
func_lib_def_.AddFunctionDef(fdef, graph_with_debug_info));
TF_RETURN_IF_ERROR(func_lib_def_.AddFunctionDef(fdef, stack_traces));
TF_RETURN_IF_ERROR(func_lib_def_.AddLibrary(library));
if (!add_to_local_only) {
return MaybeRegisterFunctionRemotely(fdef);

6
tensorflow/core/common_runtime/eager/context.h
View File

@ -234,8 +234,8 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
// entry to the KernelAndDevice cache for it if it's not exist.
Status AddFunctionDef(const FunctionDef& fdef) override;
Status AddFunctionDefWithDebugInfo(
const FunctionDef& fdef, const Graph* graph_with_debug_info) override;
Status AddFunctionDefWithStackTraces(
const FunctionDef& fdef, const StackTracesMap& stack_traces) override;
// `library` contains all FunctionDefs and GradientDefs to expand `fdef`. Add
// it to the local FunctionLibraryDefinition as well, but no need to add it
@ -244,7 +244,7 @@ class EagerContext : public ImmediateExecutionContext, public core::RefCounted {
Status AddFunctionDef(const FunctionDef& fdef,
const FunctionDefLibrary& library,
const bool add_to_local_only = false,
const Graph* graph_with_debug_info = nullptr);
const StackTracesMap& stack_traces = {});
const FunctionDef* GetFunctionDef(const string& function_name);

12
tensorflow/core/common_runtime/function_def_utils.cc
View File

@ -35,14 +35,22 @@ Status FunctionDefToBodyHelper(
InstantiationResult result;
TF_RETURN_IF_ERROR(InstantiateFunction(fdef, attrs, get_func_sig, &result));
std::unique_ptr<Graph> graph(new Graph(lib_def));
auto graph = absl::make_unique<Graph>(lib_def);
graph->SetConstructionContext(ConstructionContext::kFunctionDef);
GraphConstructorOptions opts;
opts.allow_internal_ops = true;
opts.expect_device_spec = false;
TF_RETURN_IF_ERROR(ConvertNodeDefsToGraph(opts, result.nodes, graph.get()));
const StackTracesMap& stack_traces =
lib_def->GetStackTraces(fdef.signature().name());
for (Node* n : graph->nodes()) {
auto it = stack_traces.find(n->name());
if (n && it != stack_traces.end()) {
n->SetStackTrace(it->second);
}
}
// Call BuildControlFlowInfo to validate that this function body has
// well-formed control flow.
std::vector<ControlFlowInfo> dummy;

29
tensorflow/core/common_runtime/function_test.cc
View File

@ -393,6 +393,35 @@ TEST_F(FunctionLibraryRuntimeTest, XTimesTwo) {
test::ExpectTensorEqual<float>(y, test::AsTensor<float>({2, 4, 6, 8}));
}
TEST_F(FunctionLibraryRuntimeTest, InstantiationStackTraceCopying) {
class DummyStackTrace : public AbstractStackTrace {
absl::Span<StackFrame const> ToFrames() const override { return {}; }
std::string ToString(const TracePrintingOptions& opts) const override {
return "DummyStackTrace";
}
};
FunctionDef func = test::function::XTimesTwo();
Init({});
StackTracesMap stack_traces;
stack_traces["two"] = std::make_shared<DummyStackTrace>();
TF_CHECK_OK(lib_def_->AddFunctionDef(func, stack_traces));
FunctionLibraryRuntime::Handle handle;
TF_CHECK_OK(Instantiate(flr0_, "XTimesTwo", {{"T", DT_FLOAT}}, {}, &handle));
const FunctionBody* func_body = flr0_->GetFunctionBody(handle);
for (const Node* node : func_body->graph->nodes()) {
if (node->name() == "two") {
EXPECT_EQ(node->GetStackTrace()->ToString({}), "DummyStackTrace");
}
}
TF_CHECK_OK(flr0_->ReleaseHandle(handle));
}
TEST_F(FunctionLibraryRuntimeTest, XTimesTwo_MultiDeviceBacked) {
Init({test::function::XTimesTwo()});
auto x = test::AsTensor<float>({1, 2, 3, 4});

18
tensorflow/core/framework/function.cc
View File

@ -1174,13 +1174,13 @@ Status FunctionCallFrame::SetRetval(int index, const Tensor& val) {
FunctionLibraryDefinition::FunctionDefAndOpRegistration::
FunctionDefAndOpRegistration(const FunctionDef& fdef_in,
const Graph* graph_with_debug_info)
const StackTracesMap& stack_traces)
: fdef(fdef_in),
// Exact shape inference for functions is handled by ShapeRefiner.
// Here we pass a dummy shape inference function for legacy code paths.
op_registration_data(fdef.signature(), shape_inference::UnknownShape,
true /* is_function */),
graph_with_debug_info(graph_with_debug_info) {}
stack_traces(stack_traces) {}
FunctionLibraryDefinition::FunctionLibraryDefinition(
const FunctionLibraryDefinition& other)
@ -1233,14 +1233,14 @@ FunctionLibraryDefinition::FindHelper(const string& func) const {
}
Status FunctionLibraryDefinition::AddFunctionDef(
const FunctionDef& fdef, const Graph* graph_with_debug_info) {
const FunctionDef& fdef, const StackTracesMap& stack_traces) {
mutex_lock l(mu_);
bool added;
return AddFunctionDefHelper(fdef, graph_with_debug_info, &added);
return AddFunctionDefHelper(fdef, stack_traces, &added);
}
Status FunctionLibraryDefinition::AddFunctionDefHelper(
const FunctionDef& fdef, const Graph* graph_with_debug_info, bool* added) {
const FunctionDef& fdef, const StackTracesMap& stack_traces, bool* added) {
*added = false;
std::shared_ptr<FunctionDefAndOpRegistration>& entry =
function_defs_[fdef.signature().name()];
@ -1260,8 +1260,7 @@ Status FunctionLibraryDefinition::AddFunctionDefHelper(
"Cannot add function '", fdef.signature().name(),
"' because an op with the same name already exists.");
}
entry = std::make_shared<FunctionDefAndOpRegistration>(fdef,
graph_with_debug_info);
entry = std::make_shared<FunctionDefAndOpRegistration>(fdef, stack_traces);
*added = true;
return Status::OK();
}
@ -1403,7 +1402,7 @@ Status FunctionLibraryDefinition::AddLibrary(
Status s;
bool added;
for (const FunctionDef& fdef : lib_def.function()) {
s = AddFunctionDefHelper(fdef, /*graph_with_debug_info=*/nullptr, &added);
s = AddFunctionDefHelper(fdef, /*stack_traces=*/{}, &added);
if (!s.ok()) {
Remove(funcs, funcs_with_grads);
return s;
@ -1430,8 +1429,7 @@ Status FunctionLibraryDefinition::ReplaceFunction(const string& func,
mutex_lock l(mu_);
bool added;
TF_RETURN_IF_ERROR(RemoveFunctionHelper(func));
TF_RETURN_IF_ERROR(
AddFunctionDefHelper(fdef, /*graph_with_debug_info=*/nullptr, &added));
TF_RETURN_IF_ERROR(AddFunctionDefHelper(fdef, /*stack_traces=*/{}, &added));
return Status::OK();
}

21
tensorflow/core/framework/function.h
View File

@ -348,6 +348,10 @@ class AbstractStackTrace {
virtual std::string ToString(const TracePrintingOptions& opts) const = 0;
};
using StackTracesMap =
std::unordered_map<std::string,
std::shared_ptr<tensorflow::AbstractStackTrace>>;
// Helper to maintain a map between function names in a given
// FunctionDefLibrary and function definitions.
//
@ -397,7 +401,7 @@ class FunctionLibraryDefinition : public OpRegistryInterface {
// Associates `graph` with a function `func_name`. Lifetime assumption:
// `graph` has to outlive all instantiated graphs.
Status AddFunctionDef(const FunctionDef& fdef,
const Graph* graph_with_debug_info = nullptr)
const StackTracesMap& stack_traces = {})
TF_LOCKS_EXCLUDED(mu_);
// Adds gradient definition 'grad' to this function library.
@ -509,10 +513,14 @@ class FunctionLibraryDefinition : public OpRegistryInterface {
// Returns graph with debug stack traces for the given function, or `nullptr`
// if none found.
const Graph* GetGraphWithDebugInfo(const std::string& func_name) const {
const StackTracesMap& GetStackTraces(const std::string& func_name) const {
tf_shared_lock l(mu_);
std::shared_ptr<FunctionDefAndOpRegistration> entry = FindHelper(func_name);
return entry ? entry->graph_with_debug_info : nullptr;
if (entry) {
return entry->stack_traces;
}
static const auto* empty_map = new StackTracesMap;
return *empty_map;
}
private:
@ -520,12 +528,11 @@ class FunctionLibraryDefinition : public OpRegistryInterface {
struct FunctionDefAndOpRegistration {
explicit FunctionDefAndOpRegistration(
const FunctionDef& fdef_in,
const Graph* graph_with_debug_info = nullptr);
const FunctionDef& fdef_in, const StackTracesMap& stack_traces = {});
const FunctionDef fdef;
const OpRegistrationData op_registration_data;
const Graph* graph_with_debug_info;
const StackTracesMap stack_traces;
};
std::shared_ptr<FunctionDefAndOpRegistration> FindHelper(
@ -539,7 +546,7 @@ class FunctionLibraryDefinition : public OpRegistryInterface {
// Same as AddFunctionDef/AddGradientDef except these methods set
// `added` to true if the `fdef`/`grad` were actually added to this.
Status AddFunctionDefHelper(const FunctionDef& fdef,
const Graph* graph_with_debug_info, bool* added)
const StackTracesMap& stack_traces, bool* added)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
Status AddGradientDefHelper(const GradientDef& grad, bool* added)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);

20
tensorflow/python/BUILD
View File

@ -5600,25 +5600,25 @@ tf_python_pybind_extension(
hdrs = [
"//tensorflow/c:headers",
"//tensorflow/c/eager:headers",
# Using header directly is required to avoid ODR violations.
"util/stack_trace.h",
],
# TODO(b/138203821): change to "util._tf_stack" once the bug is fixed.
module_name = "_tf_stack",
deps = [
":stack_trace",
"//tensorflow/c:pywrap_required_hdrs",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/framework:pywrap_required_hdrs",
"//tensorflow/core/platform:path",
"//third_party/python_runtime:headers", # buildcleaner: keep
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"@pybind11",
],
"//third_party/python_runtime:headers", # buildcleaner: keep
"//tensorflow/c:pywrap_required_hdrs",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/framework:pywrap_required_hdrs",
"//tensorflow/core/platform:path",
] + if_static([
":stack_trace",
]),
)
tf_py_test(

1
tensorflow/python/saved_model/load_test.py
View File

@ -2097,6 +2097,7 @@ class SingleCycleTests(test.TestCase, parameterized.TestCase):
# allocations at a lower level.
@test_util.assert_no_new_pyobjects_executing_eagerly
def test_functions_cleaned(self):
self.skipTest("TODO(b/175152958): The test is leaking function definitions")
if sys.version_info.major < 3:
self.skipTest("Not working in Python 2")
root = module.Module()

2
tensorflow/python/util/tf_stack.cc
View File

@ -19,7 +19,7 @@ limitations under the License.
// We store the retrieved stack trace within the Node object directly. Then
// whenever the graph is instantiated/copies, we copy the stack trace with it.
// Since the graph instantiation goes through the protobuf roundtrip, we store
// the original Graph with stack traces attached in FunctionLibraryDefinition.
// the original stack traces mapping attached in FunctionLibraryDefinition.
#include <Python.h>
#include <frameobject.h>