diff --git a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc index 17cc4d3282a..d29aa799479 100644 --- a/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc +++ b/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc @@ -1109,9 +1109,10 @@ mlir::Location ImporterBase::GetLocation(const NodeDef& node_def) { for (int i = 0, e = original_nodes.size(); i != e; ++i) { auto node_name = original_nodes[i]; auto func_name = (i < original_funcs.size()) ? original_funcs[i] : ""; - // Use the catenation of function and node names as the lookup key. This - // is to match the utility of generating the GraphDebugInfo. - node_call_sites.push_back(node_name_to_call_site(func_name + node_name)); + // Use the catenation of function and node names as the lookup key. + // This matches the way that the key is formed on the python side. + std::string key = node_name + "@" + func_name; + node_call_sites.push_back(node_name_to_call_site(key)); } return mlir::FusedLoc::get(node_call_sites, context_); } diff --git a/tensorflow/core/protobuf/graph_debug_info.proto b/tensorflow/core/protobuf/graph_debug_info.proto index 45f930cfaca..a3e3b18d72a 100644 --- a/tensorflow/core/protobuf/graph_debug_info.proto +++ b/tensorflow/core/protobuf/graph_debug_info.proto @@ -1,6 +1,7 @@ syntax = "proto3"; package tensorflow; + option cc_enable_arenas = true; option java_outer_classname = "GraphDebugInfoProtos"; option java_multiple_files = true; @@ -37,5 +38,14 @@ message GraphDebugInfo { repeated string files = 1; // This maps a node name to a stack trace in the source code. + // The map key is a mangling of the containing function and op name with + // syntax: + // op.name '@' func_name + // For ops in the top-level graph, the func_name is the empty string. + // Note that op names are restricted to a small number of characters which + // exclude '@', making it impossible to collide keys of this form. Function + // names accept a much wider set of characters. + // It would be preferable to avoid mangling and use a tuple key of (op.name, + // func_name), but this is not supported with protocol buffers. map traces = 2; } diff --git a/tensorflow/lite/python/lite_test.py b/tensorflow/lite/python/lite_test.py index 1d11cb75e9e..1ee16c386d4 100644 --- a/tensorflow/lite/python/lite_test.py +++ b/tensorflow/lite/python/lite_test.py @@ -1077,7 +1077,7 @@ class FromSessionTest(TestModels, parameterized.TestCase): # Check the add node in the inlined function is included. func = sess.graph.as_graph_def().library.function[0].signature.name - self.assertIn((func + 'add'), converter._debug_info.traces) + self.assertIn(('add@' + func), converter._debug_info.traces) class FromFrozenGraphFile(test_util.TensorFlowTestCase): diff --git a/tensorflow/python/framework/error_interpolation.py b/tensorflow/python/framework/error_interpolation.py index f4036265e46..0c6feb64008 100644 --- a/tensorflow/python/framework/error_interpolation.py +++ b/tensorflow/python/framework/error_interpolation.py @@ -38,12 +38,31 @@ _INTERPOLATION_PATTERN = re.compile(_INTERPOLATION_REGEX, re.DOTALL) _ParseTag = collections.namedtuple("_ParseTag", ["type", "name"]) -_BAD_FILE_SUBSTRINGS = [ - os.path.join("tensorflow", "python"), - os.path.join("tensorflow", "contrib"), - os.path.join("tensorflow_estimator", "python"), - os.path.join("tensorflow_estimator", "contrib"), - "/context_lib.py")) + + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/framework/meta_graph_test.py b/tensorflow/python/framework/meta_graph_test.py index 948956e8bbd..eff613b4204 100644 --- a/tensorflow/python/framework/meta_graph_test.py +++ b/tensorflow/python/framework/meta_graph_test.py @@ -741,11 +741,11 @@ class ScopedMetaGraphTest(test.TestCase): biases1 = resource_variable_ops.ResourceVariable( [0.1] * 3, name="biases") nn_ops.relu(math_ops.matmul(images, weights1) + biases1, name="relu") - operations = [] + func_named_operations = [] for op in graph1.get_operations(): - operations.append(("", op)) + func_named_operations.append(("", op)) debug_info_def = error_interpolation.create_graph_debug_info_def( - operations=operations) + func_named_operations) # The unique file names in all the stack traces should be larger or equal # than 1. diff --git a/tensorflow/python/saved_model/BUILD b/tensorflow/python/saved_model/BUILD index 7fd56e118e8..6357cc7b544 100644 --- a/tensorflow/python/saved_model/BUILD +++ b/tensorflow/python/saved_model/BUILD @@ -323,9 +323,12 @@ tf_py_test( additional_deps = [ ":loader", ":save", + ":save_options", ":signature_constants", ":tag_constants", "@absl_py//absl/testing:parameterized", + "//tensorflow/core:protos_all_py", + "//tensorflow/python:error_interpolation", "//tensorflow/python/eager:def_function", "//tensorflow/python/eager:test", ], diff --git a/tensorflow/python/saved_model/constants.py b/tensorflow/python/saved_model/constants.py index 90511a409ed..7c4d91b003d 100644 --- a/tensorflow/python/saved_model/constants.py +++ b/tensorflow/python/saved_model/constants.py @@ -86,6 +86,25 @@ tf_export( "saved_model.constants.SAVED_MODEL_FILENAME_PBTXT" ]).export_constant(__name__, "SAVED_MODEL_FILENAME_PBTXT") +# Subdirectory where debugging related files are written. +DEBUG_DIRECTORY = "debug" +tf_export( + "saved_model.DEBUG_DIRECTORY", + v1=[ + "saved_model.DEBUG_DIRECTORY", + "saved_model.constants.DEBUG_DIRECTORY", + ]).export_constant(__name__, "DEBUG_DIRECTORY") + +# File name for GraphDebugInfo protocol buffer which corresponds to the +# SavedModel. +DEBUG_INFO_FILENAME_PB = "saved_model_debug_info.pb" +tf_export( + "saved_model.DEBUG_INFO_FILENAME_PB", + v1=[ + "saved_model.DEBUG_INFO_FILENAME_PB", + "saved_model.constants.DEBUG_INFO_FILENAME_PB" + ]).export_constant(__name__, "DEBUG_INFO_FILENAME_PB") + # File name for json format of SavedModel. # Not exported while keras_saved_model is in contrib. SAVED_MODEL_FILENAME_JSON = "saved_model.json" diff --git a/tensorflow/python/saved_model/save.py b/tensorflow/python/saved_model/save.py index b1b69f1ff32..18eb69c1d72 100644 --- a/tensorflow/python/saved_model/save.py +++ b/tensorflow/python/saved_model/save.py @@ -31,6 +31,7 @@ from tensorflow.python.eager import def_function from tensorflow.python.eager import function as defun from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import error_interpolation from tensorflow.python.framework import meta_graph from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_util @@ -546,7 +547,8 @@ def _fill_meta_graph_def(meta_graph_def, saveable_view, signature_functions, namespace_whitelist: List of strings containing whitelisted op namespaces. Returns: - An _AssetInfo, which contains information to help creating the SavedModel. + A tuple of (_AssetInfo, Graph) containing the captured assets and + exported Graph generated from tracing the saveable_view. """ # List objects from the eager context to make sure Optimizers give us the # right Graph-dependent variables. @@ -700,6 +702,28 @@ def _write_object_proto(obj, proto, asset_file_def_index): proto.user_object.CopyFrom(registered_type_proto) +def _export_debug_info(exported_graph): + """Exports debug information from a graph. + + Args: + exported_graph: A Graph that has been created by tracing a saveable view. + + Returns: + Corresponding GraphDebugInfo with traces for ops in all functions of the + exported_graph. + """ + exported_operations = [] + for fn_name in exported_graph._functions: # pylint: disable=protected-access + fn = exported_graph._get_function(fn_name) # pylint: disable=protected-access + if not isinstance(fn, defun._EagerDefinedFunction): # pylint: disable=protected-access + continue + + fn_graph = fn.graph + for fn_op in fn_graph.get_operations(): + exported_operations.append((fn_name, fn_op)) + return error_interpolation.create_graph_debug_info_def(exported_operations) + + @tf_export("saved_model.save", v1=["saved_model.save", "saved_model.experimental.save"]) def save(obj, export_dir, signatures=None, options=None): @@ -907,6 +931,16 @@ def save(obj, export_dir, signatures=None, options=None): saveable_view, asset_info.asset_index) meta_graph_def.object_graph_def.CopyFrom(object_graph_proto) file_io.atomic_write_string_to_file(path, saved_model.SerializeToString()) + + # Save debug info, if requested. + if options.save_debug_info: + graph_debug_info = _export_debug_info(exported_graph) + file_io.atomic_write_string_to_file( + os.path.join( + utils_impl.get_or_create_debug_dir(export_dir), + constants.DEBUG_INFO_FILENAME_PB), + graph_debug_info.SerializeToString()) + # Clean reference cycles so repeated export()s don't make work for the garbage # collector. Before this point we need to keep references to captured # constants in the saved graph. diff --git a/tensorflow/python/saved_model/save_options.py b/tensorflow/python/saved_model/save_options.py index fc70c1fcd77..50a8d74dc9e 100644 --- a/tensorflow/python/saved_model/save_options.py +++ b/tensorflow/python/saved_model/save_options.py @@ -33,9 +33,9 @@ class SaveOptions(object): """ # Define object attributes in __slots__ for improved memory and performance. - __slots__ = ("namespace_whitelist",) + __slots__ = ("namespace_whitelist", "save_debug_info") - def __init__(self, namespace_whitelist=None): + def __init__(self, namespace_whitelist=None, save_debug_info=False): """Creates an object that stores options for SavedModel saving. Args: @@ -43,9 +43,14 @@ class SaveOptions(object): when saving a model. Saving an object that uses namespaced ops must explicitly add all namespaces to the whitelist. The namespaced ops must be registered into the framework when loading the SavedModel. + save_debug_info: Boolean indicating whether debug information is saved. + If True, then a debug/saved_model_debug_info.pb file will be written + with the contents of a GraphDebugInfo binary protocol buffer containing + stack trace information for all ops and functions that are saved. """ self.namespace_whitelist = _validate_namespace_whitelist( namespace_whitelist) + self.save_debug_info = save_debug_info def _validate_namespace_whitelist(namespace_whitelist): diff --git a/tensorflow/python/saved_model/save_test.py b/tensorflow/python/saved_model/save_test.py index 542e7130273..200050fbf54 100644 --- a/tensorflow/python/saved_model/save_test.py +++ b/tensorflow/python/saved_model/save_test.py @@ -24,6 +24,7 @@ import sys from google.protobuf import text_format from tensorflow.core.framework import graph_pb2 +from tensorflow.core.protobuf import graph_debug_info_pb2 from tensorflow.python.client import session as session_lib from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import backprop @@ -48,6 +49,7 @@ from tensorflow.python.ops import variables from tensorflow.python.saved_model import loader from tensorflow.python.saved_model import loader_impl from tensorflow.python.saved_model import save +from tensorflow.python.saved_model import save_options from tensorflow.python.saved_model import signature_constants from tensorflow.python.saved_model import tag_constants from tensorflow.python.training.tracking import tracking @@ -415,6 +417,48 @@ class SavingOptionsTest(test.TestCase): save._verify_ops(graph_def, []) save._verify_ops(graph_def, ["Test"]) + def test_save_debug_info_enabled(self): + root = tracking.AutoTrackable() + root.f = def_function.function( + lambda x: math_ops.mul(2., x, name="DEBUG_INFO_OP"), + input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) + save_dir = os.path.join(self.get_temp_dir(), "saved_model") + save.save( + root, + save_dir, + root.f, + options=save_options.SaveOptions(save_debug_info=True)) + debug_info_file_name = os.path.join(save_dir, "debug", + "saved_model_debug_info.pb") + self.assertTrue(os.path.exists(debug_info_file_name)) + debug_info = graph_debug_info_pb2.GraphDebugInfo() + with open(debug_info_file_name, "rb") as f: + debug_info.ParseFromString(f.read()) + + # Verify that there is a trace for DEBUG_INFO_OP just to ensure that + # function debug info tracing is nominally functioning. + found_op = False + for key in debug_info.traces.keys(): + if key.startswith("DEBUG_INFO_OP@"): + found_op = True + break + self.assertTrue(found_op, "Did not find DEBUG_INFO_OP in trace") + + def test_save_debug_info_disabled(self): + root = tracking.AutoTrackable() + root.f = def_function.function( + lambda x: math_ops.mul(2., x, name="DEBUG_INFO_OP"), + input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]) + save_dir = os.path.join(self.get_temp_dir(), "saved_model") + save.save( + root, + save_dir, + root.f, + options=save_options.SaveOptions(save_debug_info=False)) + debug_info_file_name = os.path.join(save_dir, "debug", + "saved_model_debug_info.pb") + self.assertFalse(os.path.exists(debug_info_file_name)) + class AssetTests(test.TestCase): diff --git a/tensorflow/python/saved_model/utils_impl.py b/tensorflow/python/saved_model/utils_impl.py index 3dd7d6c7ae4..70e507a3c13 100644 --- a/tensorflow/python/saved_model/utils_impl.py +++ b/tensorflow/python/saved_model/utils_impl.py @@ -244,3 +244,19 @@ def get_assets_dir(export_dir): return os.path.join( compat.as_text(export_dir), compat.as_text(constants.ASSETS_DIRECTORY)) + + +def get_or_create_debug_dir(export_dir): + """Returns path to the debug sub-directory, creating if it does not exist.""" + debug_dir = get_debug_dir(export_dir) + + if not file_io.file_exists(debug_dir): + file_io.recursive_create_dir(debug_dir) + + return debug_dir + + +def get_debug_dir(export_dir): + """Returns path to the debug sub-directory in the SavedModel.""" + return os.path.join( + compat.as_text(export_dir), compat.as_text(constants.DEBUG_DIRECTORY)) diff --git a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-save-options.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-save-options.pbtxt index 174043aec78..ea31605ba1f 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-save-options.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.-save-options.pbtxt @@ -6,8 +6,12 @@ tf_class { name: "namespace_whitelist" mtype: "" } + member { + name: "save_debug_info" + mtype: "" + } member_method { name: "__init__" - argspec: "args=[\'self\', \'namespace_whitelist\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'namespace_whitelist\', \'save_debug_info\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], " } } diff --git a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.constants.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.constants.pbtxt index 20e10aa094f..4e16707ba2a 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.constants.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.constants.pbtxt @@ -8,6 +8,14 @@ tf_module { name: "ASSETS_KEY" mtype: "" } + member { + name: "DEBUG_DIRECTORY" + mtype: "" + } + member { + name: "DEBUG_INFO_FILENAME_PB" + mtype: "" + } member { name: "LEGACY_INIT_OP_KEY" mtype: "" diff --git a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt index f3558109ce8..24e468bbf86 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.saved_model.pbtxt @@ -28,6 +28,14 @@ tf_module { name: "CLASSIFY_OUTPUT_SCORES" mtype: "" } + member { + name: "DEBUG_DIRECTORY" + mtype: "" + } + member { + name: "DEBUG_INFO_FILENAME_PB" + mtype: "" + } member { name: "DEFAULT_SERVING_SIGNATURE_DEF_KEY" mtype: "" diff --git a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-save-options.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-save-options.pbtxt index 174043aec78..ea31605ba1f 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-save-options.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.-save-options.pbtxt @@ -6,8 +6,12 @@ tf_class { name: "namespace_whitelist" mtype: "" } + member { + name: "save_debug_info" + mtype: "" + } member_method { name: "__init__" - argspec: "args=[\'self\', \'namespace_whitelist\'], varargs=None, keywords=None, defaults=[\'None\'], " + argspec: "args=[\'self\', \'namespace_whitelist\', \'save_debug_info\'], varargs=None, keywords=None, defaults=[\'None\', \'False\'], " } } diff --git a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt index 94fa0eaad53..a78ca5c349a 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.saved_model.pbtxt @@ -24,6 +24,14 @@ tf_module { name: "CLASSIFY_OUTPUT_SCORES" mtype: "" } + member { + name: "DEBUG_DIRECTORY" + mtype: "" + } + member { + name: "DEBUG_INFO_FILENAME_PB" + mtype: "" + } member { name: "DEFAULT_SERVING_SIGNATURE_DEF_KEY" mtype: ""