Modify EventAccumulator and TensorBoard to handle MetaGraph in Events.

Change: 135824634

1 files changed, 24 insertions, 8 deletions

diff --git a/tensorflow/tensorboard/backend/server_test.py b/tensorflow/tensorboard/backend/server_test.py
index 6877fe0b31..73e9f8bafd 100644
--- a/tensorflow/tensorboard/backend/server_test.py
+++ b/tensorflow/tensorboard/backend/server_test.py
@@ -37,6 +37,7 @@ import tensorflow as tf
 
 from google.protobuf import text_format
 from tensorflow.contrib.tensorboard.plugins.projector.projector_config_pb2 import ProjectorConfig
+from tensorflow.core.protobuf import meta_graph_pb2
 from tensorflow.python.platform import resource_loader
 from tensorflow.python.summary import event_multiplexer
 from tensorflow.tensorboard.backend import server
@@ -44,6 +45,7 @@ from tensorflow.tensorboard.plugins import REGISTERED_PLUGINS
 
 
 class TensorboardServerTest(tf.test.TestCase):
+  _only_use_meta_graph = False  # Server data contains only a GraphDef
 
   # Number of scalar-containing events to make.
   _SCALAR_COUNT = 99
@@ -114,13 +116,16 @@ class TensorboardServerTest(tf.test.TestCase):
     self.assertTrue(isinstance(run_json['run1']['firstEventTimestamp'],
                                numbers.Number))
     del run_json['run1']['firstEventTimestamp']
-    self.assertEqual(run_json, {'run1': {'compressedHistograms': ['histogram'],
-                                         'scalars': ['simple_values'],
-                                         'histograms': ['histogram'],
-                                         'images': ['image'],
-                                         'audio': ['audio'],
-                                         'graph': True,
-                                         'run_metadata': ['test run']}})
+    self.assertEqual(run_json, {'run1': {
+        'compressedHistograms': ['histogram'],
+        'scalars': ['simple_values'],
+        'histograms': ['histogram'],
+        'images': ['image'],
+        'audio': ['audio'],
+        # if only_use_meta_graph, the graph is extracted from the metagraph
+        'graph': True,
+        'meta_graph': self._only_use_meta_graph,
+        'run_metadata': ['test run']}})
 
   def testApplicationPaths_getCached(self):
     """Test the format of the /data/runs endpoint."""
@@ -344,7 +349,13 @@ class TensorboardServerTest(tf.test.TestCase):
     node2 = graph_def.node.add()
     node2.name = 'b'
     node2.attr['very_large_attr'].s = b'a' * 2048  # 2 KB attribute
-    writer.add_graph(graph_def)
+
+    meta_graph_def = meta_graph_pb2.MetaGraphDef(graph_def=graph_def)
+
+    if self._only_use_meta_graph:
+      writer.add_meta_graph(meta_graph_def)
+    else:
+      writer.add_graph(graph_def)
 
     # Add a simple run metadata event.
     run_metadata = tf.RunMetadata()
@@ -412,6 +423,11 @@ class TensorboardServerTest(tf.test.TestCase):
       saver.save(sess, checkpoint_path)
 
 
+class TensorboardServerUsingMetagraphOnlyTest(TensorboardServerTest):
+  # Tests new ability to use only the MetaGraphDef
+  _only_use_meta_graph = True  # Server data contains only a MetaGraphDef
+
+
 class ParseEventFilesSpecTest(tf.test.TestCase):
 
   def testRunName(self):