Refactor the PluginAsset class so that PluginAssets are not directly responsible for serialization.

In the current PluginAsset api, the PluginAsset provides a serialize_to_directory method in which it directly writes contents to disk. This means that we as framework maintainers don't have the flexibility to change the serialization strategy in different contexts, e.g. providing good ways to serialize in contexts where we are writing to a db rather than to disk. Also, it presents a trivial landmine where users may use standard python file APIs rather than gfile and thus provide implementations that work externally but break within g3.

After the change, the PluginAsset instead provides an 'assets' method, which provides asset names and asset contents. How this gets written out is now an implementation detail handled by the tf.summary.FileWriter.

We haven't yet exposed the PluginAsset class as part of the TensorFlow API (it's hidden) so this isn't an API break.
Change: 149985564

4 files changed, 32 insertions, 31 deletions

diff --git a/tensorflow/python/summary/plugin_asset.py b/tensorflow/python/summary/plugin_asset.py
index 18b6bd4fc4..998fb30fa4 100644
--- a/tensorflow/python/summary/plugin_asset.py
+++ b/tensorflow/python/summary/plugin_asset.py
@@ -14,15 +14,16 @@
 # ==============================================================================
 """TensorBoard Plugin asset abstract class.
 
-TensorBoard plugins may need to write arbitrary assets to disk, such as
+TensorBoard plugins may need to provide arbitrary assets, such as
 configuration information for specific outputs, or vocabulary files, or sprite
 images, etc.
 
 This module contains methods that allow plugin assets to be specified at graph
 construction time. Plugin authors define a PluginAsset which is treated as a
-singleton on a per-graph basis. The PluginAsset has a serialize_to_directory
-method which writes its assets to disk within a special plugin directory
-that the tf.summary.FileWriter creates.
+singleton on a per-graph basis. The PluginAsset has an assets method which
+returns a dictionary of asset contents. The tf.summary.FileWriter
+(or any other Summary writer) will serialize these assets in such a way that
+TensorBoard can retrieve them.
 """
 
 from __future__ import absolute_import
@@ -111,30 +112,30 @@ class PluginAsset(object):
 
   Plugin authors are expected to extend the PluginAsset class, so that it:
   - has a unique plugin_name
-  - provides a serialize_to_directory method that dumps its assets in that dir
-  - takes no constructor arguments
+  - provides an assets method that returns an {asset_name: asset_contents}
+    dictionary. For now, asset_contents are strings, although we may add
+    StringIO support later.
 
   LifeCycle of a PluginAsset instance:
   - It is constructed when get_plugin_asset is called on the class for
-  the first time.
+    the first time.
   - It is configured by code that follows the calls to get_plugin_asset
   - When the containing graph is serialized by the tf.summary.FileWriter, the
-  writer calls serialize_to_directory and the PluginAsset instance dumps its
-  contents to disk.
+    writer calls assets and the PluginAsset instance provides its contents to be
+    written to disk.
   """
   __metaclass__ = abc.ABCMeta
 
   plugin_name = None
 
   @abc.abstractmethod
-  def serialize_to_directory(self, directory):
-    """Serialize the assets in this PluginAsset to given directory.
+  def assets(self):
+    """Provide all of the assets contained by the PluginAsset instance.
 
-    The directory will be specific to this plugin (as determined by the
-    plugin_key property on the class). This method will be called when the graph
-    containing this PluginAsset is given to a tf.summary.FileWriter.
+    The assets method should return a dictionary structured as
+    {asset_name: asset_contents}. asset_contents is a string.
 
-    Args:
-      directory: The directory path (as string) that serialize should write to.
+    This method will be called by the tf.summary.FileWriter when it is time to
+    write the assets out to disk.
     """
     raise NotImplementedError()
diff --git a/tensorflow/python/summary/plugin_asset_test.py b/tensorflow/python/summary/plugin_asset_test.py
index 5f80a1a4db..d2edcc053c 100644
--- a/tensorflow/python/summary/plugin_asset_test.py
+++ b/tensorflow/python/summary/plugin_asset_test.py
@@ -26,8 +26,8 @@ from tensorflow.python.summary import plugin_asset
 class _UnnamedPluginAsset(plugin_asset.PluginAsset):
   """An example asset with a dummy serialize method provided, but no name."""
 
-  def serialize_to_directory(self, unused_directory):
-    pass
+  def assets(self):
+    return {}
 
 
 class _ExamplePluginAsset(_UnnamedPluginAsset):
diff --git a/tensorflow/python/summary/writer/writer.py b/tensorflow/python/summary/writer/writer.py
index dec3ee96d5..2dad7b88d3 100644
--- a/tensorflow/python/summary/writer/writer.py
+++ b/tensorflow/python/summary/writer/writer.py
@@ -164,7 +164,7 @@ class SummaryToEventTransformer(object):
 
       # Serialize the graph with additional info.
       true_graph_def = graph.as_graph_def(add_shapes=True)
-      self._write_tensorboard_metadata(graph)
+      self._write_plugin_assets(graph)
     elif (isinstance(graph, graph_pb2.GraphDef) or
           isinstance(graph_def, graph_pb2.GraphDef)):
       # The user passed a `GraphDef`.
@@ -185,13 +185,18 @@ class SummaryToEventTransformer(object):
     # Finally, add the graph_def to the summary writer.
     self._add_graph_def(true_graph_def, global_step)
 
-  def _write_tensorboard_metadata(self, graph):
-    assets = plugin_asset.get_all_plugin_assets(graph)
+  def _write_plugin_assets(self, graph):
+    plugin_assets = plugin_asset.get_all_plugin_assets(graph)
     logdir = self.event_writer.get_logdir()
-    for asset in assets:
-      plugin_dir = os.path.join(logdir, _PLUGINS_DIR, asset.plugin_name)
+    for asset_container in plugin_assets:
+      plugin_name = asset_container.plugin_name
+      plugin_dir = os.path.join(logdir, _PLUGINS_DIR, plugin_name)
       gfile.MakeDirs(plugin_dir)
-      asset.serialize_to_directory(plugin_dir)
+      assets = asset_container.assets()
+      for (asset_name, content) in assets.items():
+        asset_path = os.path.join(plugin_dir, asset_name)
+        with gfile.Open(asset_path, "w") as f:
+          f.write(content)
 
   def add_meta_graph(self, meta_graph_def, global_step=None):
     """Adds a `MetaGraphDef` to the event file.
diff --git a/tensorflow/python/summary/writer/writer_test.py b/tensorflow/python/summary/writer/writer_test.py
index f9e9a8ff33..50797483ee 100644
--- a/tensorflow/python/summary/writer/writer_test.py
+++ b/tensorflow/python/summary/writer/writer_test.py
@@ -359,13 +359,8 @@ class SummaryWriterCacheTest(test.TestCase):
 class ExamplePluginAsset(plugin_asset.PluginAsset):
   plugin_name = "example"
 
-  def serialize_to_directory(self, directory):
-    foo = os.path.join(directory, "foo.txt")
-    bar = os.path.join(directory, "bar.txt")
-    with gfile.Open(foo, "w") as f:
-      f.write("foo!")
-    with gfile.Open(bar, "w") as f:
-      f.write("bar!")
+  def assets(self):
+    return {"foo.txt": "foo!", "bar.txt": "bar!"}
 
 
 class PluginAssetsTest(test.TestCase):