Use the new SummaryMetadata field of the Value proto to pass per-plugin information.

This will let TensorBoard figure out which plugins should handle which types of summaries.

PiperOrigin-RevId: 159202469

5 files changed, 96 insertions, 8 deletions

diff --git a/tensorflow/core/debug/BUILD b/tensorflow/core/debug/BUILD
index 2fc49d4412..30824a9072 100644
--- a/tensorflow/core/debug/BUILD
+++ b/tensorflow/core/debug/BUILD
@@ -123,6 +123,7 @@ tf_cuda_library(
     linkstatic = 1,
     deps = [
         ":debug_service_proto_cc",
+        ":debugger_event_metadata_proto_cc",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
@@ -144,6 +145,7 @@ tf_cuda_library(
         ":debug_graph_utils",
         ":debug_io_utils",
         ":debug_service_proto_cc",
+        ":debugger_event_metadata_proto_cc",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
@@ -194,6 +196,7 @@ tf_cc_test(
     deps = [
         ":debug_grpc_testlib",
         ":debug_io_utils",
+        ":debugger_event_metadata_proto_cc",
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
@@ -204,6 +207,7 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/platform/default/build_config:platformlib",
     ],
 )
 
@@ -247,6 +251,12 @@ tf_cc_test(
     ],
 )
 
+tf_proto_library_cc(
+    name = "debugger_event_metadata_proto",
+    srcs = ["debugger_event_metadata.proto"],
+    cc_api_version = 2,
+)
+
 # TODO(cais): Add the following back in when tfdbg is supported on Android.
 # filegroup(
 #     name = "android_srcs",
diff --git a/tensorflow/core/debug/debug_grpc_testlib.cc b/tensorflow/core/debug/debug_grpc_testlib.cc
index c19842a2f6..7317aa0372 100644
--- a/tensorflow/core/debug/debug_grpc_testlib.cc
+++ b/tensorflow/core/debug/debug_grpc_testlib.cc
@@ -16,10 +16,12 @@ limitations under the License.
 #include "tensorflow/core/debug/debug_grpc_testlib.h"
 
 #include "tensorflow/core/debug/debug_graph_utils.h"
+#include "tensorflow/core/debug/debugger_event_metadata.pb.h"
 #include "tensorflow/core/framework/summary.pb.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/tracing.h"
 
 namespace tensorflow {
@@ -44,8 +46,6 @@ namespace test {
           tensorflow::str_util::Split(val.node_name(), ':');
 
       const string node_name = name_items[0];
-      int32 output_slot = 0;
-      tensorflow::strings::safe_strto32(name_items[1], &output_slot);
       const string debug_op = name_items[2];
 
       const TensorProto& tensor_proto = val.tensor();
@@ -54,9 +54,24 @@ namespace test {
         return ::grpc::Status::CANCELLED;
       }
 
-      device_names.push_back(val.tag());
+      // Obtain the device name, which is encoded in JSON.
+      third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
+      for (int i = 0; i < val.metadata().plugin_data_size(); i++) {
+        if (val.metadata().plugin_data(i).plugin_name() != "debugger") {
+          // This plugin data was meant for another plugin.
+          continue;
+        }
+        auto status = tensorflow::protobuf::util::JsonStringToMessage(
+            val.metadata().plugin_data(i).content(), &metadata);
+        if (status.ok()) {
+          // The device name has been determined.
+          break;
+        }
+      }
+
+      device_names.push_back(metadata.device());
       node_names.push_back(node_name);
-      output_slots.push_back(output_slot);
+      output_slots.push_back(metadata.output_slot());
       debug_ops.push_back(debug_op);
       debug_tensors.push_back(tensor);
     }
diff --git a/tensorflow/core/debug/debug_io_utils.cc b/tensorflow/core/debug/debug_io_utils.cc
index 88dffb934c..69fc367789 100644
--- a/tensorflow/core/debug/debug_io_utils.cc
+++ b/tensorflow/core/debug/debug_io_utils.cc
@@ -26,11 +26,13 @@ limitations under the License.
 #pragma comment(lib,"Ws2_32.lib")
 #endif
 
+#include "tensorflow/core/debug/debugger_event_metadata.pb.h"
 #include "tensorflow/core/framework/summary.pb.h"
 #include "tensorflow/core/lib/hash/hash.h"
 #include "tensorflow/core/lib/io/path.h"
 #include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
+#include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/util/event.pb.h"
 
 #define GRPC_OSS_UNIMPLEMENTED_ERROR \
@@ -55,7 +57,32 @@ Event WrapTensorAsEvent(const DebugNodeKey& debug_node_key,
   // "DebugIdentity", the debug node_name in the Summary proto will be
   // "foo/node_a:0:DebugIdentity".
   summ_val->set_node_name(debug_node_key.debug_node_name);
-  summ_val->set_tag(debug_node_key.device_name);
+
+  // Tag by the node name. This allows TensorBoard to quickly fetch data per op.
+  summ_val->set_tag(debug_node_key.node_name);
+
+  // Store data within debugger metadata to be stored for each event.
+  third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
+  metadata.set_device(debug_node_key.device_name);
+  metadata.set_output_slot(debug_node_key.output_slot);
+
+  // Encode the data in JSON.
+  string json_output;
+  tensorflow::protobuf::util::JsonPrintOptions json_options;
+  json_options.always_print_primitive_fields = true;
+  auto status = tensorflow::protobuf::util::MessageToJsonString(
+      metadata, &json_output, json_options);
+  if (status.ok()) {
+    // Store summary metadata. Set the plugin to use this data as "debugger".
+    SummaryMetadata::PluginData* plugin_data =
+        summ_val->mutable_metadata()->add_plugin_data();
+    plugin_data->set_plugin_name("debugger");
+    plugin_data->set_content(json_output);
+  } else {
+    LOG(WARNING) << "Failed to convert DebuggerEventMetadata proto to JSON. "
+                 << "The debug_node_name is " << debug_node_key.debug_node_name
+                 << ".";
+  }
 
   if (tensor.dtype() == DT_STRING) {
     // Treat DT_STRING specially, so that tensor_util.MakeNdarray can convert
diff --git a/tensorflow/core/debug/debug_io_utils_test.cc b/tensorflow/core/debug/debug_io_utils_test.cc
index 35c95fb98c..08ef4001bc 100644
--- a/tensorflow/core/debug/debug_io_utils_test.cc
+++ b/tensorflow/core/debug/debug_io_utils_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/core/debug/debug_io_utils.h"
 
+#include "tensorflow/core/debug/debugger_event_metadata.pb.h"
 #include "tensorflow/core/framework/summary.pb.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/lib/core/notification.h"
@@ -124,10 +125,18 @@ TEST_F(DebugIOUtilsTest, DumpStringTensorToFileSunnyDay) {
 
   ASSERT_GE(wall_time, event.wall_time());
   ASSERT_EQ(1, event.summary().value().size());
-  ASSERT_EQ(kDebugNodeKey.device_name, event.summary().value(0).tag());
+  ASSERT_EQ(kDebugNodeKey.node_name, event.summary().value(0).tag());
   ASSERT_EQ(kDebugNodeKey.debug_node_name,
             event.summary().value(0).node_name());
 
+  // Determine and validate some information from the metadata.
+  third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
+  auto status = tensorflow::protobuf::util::JsonStringToMessage(
+      event.summary().value(0).metadata().plugin_data(0).content(), &metadata);
+  ASSERT_TRUE(status.ok());
+  ASSERT_EQ(kDebugNodeKey.device_name, metadata.device());
+  ASSERT_EQ(kDebugNodeKey.output_slot, metadata.output_slot());
+
   Tensor b_prime(DT_STRING);
   ASSERT_TRUE(b_prime.FromProto(event.summary().value(0).tensor()));
 
@@ -229,10 +238,19 @@ TEST_F(DebugIOUtilsTest, PublishTensorToMultipleFileURLs) {
 
     ASSERT_GE(wall_time, event.wall_time());
     ASSERT_EQ(1, event.summary().value().size());
-    ASSERT_EQ(kDebugNodeKey.device_name, event.summary().value(0).tag());
+    ASSERT_EQ(kDebugNodeKey.node_name, event.summary().value(0).tag());
     ASSERT_EQ(kDebugNodeKey.debug_node_name,
               event.summary().value(0).node_name());
 
+    // Determine and validate some information from the metadata.
+    third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
+    auto status = tensorflow::protobuf::util::JsonStringToMessage(
+        event.summary().value(0).metadata().plugin_data(0).content(),
+        &metadata);
+    ASSERT_TRUE(status.ok());
+    ASSERT_EQ(kDebugNodeKey.device_name, metadata.device());
+    ASSERT_EQ(kDebugNodeKey.output_slot, metadata.output_slot());
+
     Tensor a_prime(DT_FLOAT);
     ASSERT_TRUE(a_prime.FromProto(event.summary().value(0).tensor()));
 
@@ -333,10 +351,19 @@ TEST_F(DebugIOUtilsTest, PublishTensorConcurrentlyToPartiallyOverlappingPaths) {
 
       ASSERT_GE(wall_time, event.wall_time());
       ASSERT_EQ(1, event.summary().value().size());
-      ASSERT_EQ(kDebugNodeKey.device_name, event.summary().value(0).tag());
+      ASSERT_EQ(kDebugNodeKey.node_name, event.summary().value(0).tag());
       ASSERT_EQ(kDebugNodeKey.debug_node_name,
                 event.summary().value(0).node_name());
 
+      // Determine and validate some information from the metadata.
+      third_party::tensorflow::core::debug::DebuggerEventMetadata metadata;
+      auto status = tensorflow::protobuf::util::JsonStringToMessage(
+          event.summary().value(0).metadata().plugin_data(0).content(),
+          &metadata);
+      ASSERT_TRUE(status.ok());
+      ASSERT_EQ(kDebugNodeKey.device_name, metadata.device());
+      ASSERT_EQ(kDebugNodeKey.output_slot, metadata.output_slot());
+
       Tensor a_prime(DT_FLOAT);
       ASSERT_TRUE(a_prime.FromProto(event.summary().value(0).tensor()));
 
diff --git a/tensorflow/core/debug/debugger_event_metadata.proto b/tensorflow/core/debug/debugger_event_metadata.proto
new file mode 100644
index 0000000000..44ef305f5a
--- /dev/null
+++ b/tensorflow/core/debug/debugger_event_metadata.proto
@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package third_party.tensorflow.core.debug;
+
+// Encapsulates per-event data related to debugging.
+message DebuggerEventMetadata {
+  string device = 1;
+  int32 output_slot = 2;
+};