Add support for plugin assets to the tensorboard EventAccumulator.

Change: 150151834

tensorflow/tensorboard/backend/event_processing/BUILD

@@ -91,6 +91,7 @@ py_library(
     deps = [
         ":directory_watcher",
         ":event_file_loader",
+        ":plugin_asset_util",
         ":reservoir",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:framework",
@@ -146,6 +147,30 @@ py_test(
     ],
 )
 
+py_library(
+    name = "plugin_asset_util",
+    srcs = ["plugin_asset_util.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        "//tensorflow/python:platform",
+    ],
+)
+
+py_test(
+    name = "plugin_asset_util_test",
+    size = "small",
+    srcs = ["plugin_asset_util_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":event_multiplexer",
+        ":plugin_asset_util",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:summary",
+    ],
+)
+
 py_library(
     name = "event_file_inspector",
     srcs = ["event_file_inspector.py"],

tensorflow/tensorboard/backend/event_processing/event_accumulator.py

@@ -33,6 +33,7 @@ from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.util import compat
 from tensorflow.tensorboard.backend.event_processing import directory_watcher
 from tensorflow.tensorboard.backend.event_processing import event_file_loader
+from tensorflow.tensorboard.backend.event_processing import plugin_asset_util
 from tensorflow.tensorboard.backend.event_processing import reservoir
 
 namedtuple = collections.namedtuple
@@ -210,6 +211,7 @@ class EventAccumulator(object):
     self._tensors = reservoir.Reservoir(size=sizes[TENSORS])
 
     self._generator_mutex = threading.Lock()
+    self.path = path
     self._generator = _GeneratorFromPath(path)
 
     self._compression_bps = compression_bps
@@ -237,6 +239,33 @@ class EventAccumulator(object):
         self._ProcessEvent(event)
     return self
 
+  def PluginAssets(self, plugin_name):
+    """Return a list of all plugin assets for the given plugin.
+
+    Args:
+      plugin_name: The string name of a plugin to retrieve assets for.
+
+    Returns:
+      A list of string plugin asset names, or empty list if none are available.
+      If the plugin was not registered, an empty list is returned.
+    """
+    return plugin_asset_util.ListAssets(self.path, plugin_name)
+
+  def RetrievePluginAsset(self, plugin_name, asset_name):
+    """Return the contents of a given plugin asset.
+
+    Args:
+      plugin_name: The string name of a plugin.
+      asset_name: The string name of an asset.
+
+    Returns:
+      The string contents of the plugin asset.
+
+    Raises:
+      KeyError: If the asset is not available.
+    """
+    return plugin_asset_util.RetrieveAsset(self.path, plugin_name, asset_name)
+
   def FirstEventTimestamp(self):
     """Returns the timestamp in seconds of the first event.
 

tensorflow/tensorboard/backend/event_processing/event_multiplexer.py

@@ -200,6 +200,39 @@ class EventMultiplexer(object):
     logging.info('Finished with EventMultiplexer.Reload()')
     return self
 
+  def PluginAssets(self, plugin_name):
+    """Get index of runs and assets for a given plugin.
+
+    Args:
+      plugin_name: Name of the plugin we are checking for.
+
+    Returns:
+      A dictionary that maps from run_name to a list of plugin
+        assets for that run.
+    """
+    with self._accumulators_mutex:
+      # To avoid nested locks, we construct a copy of the run-accumulator map
+      items = list(six.iteritems(self._accumulators))
+
+    return {run: accum.PluginAssets(plugin_name) for run, accum in items}
+
+  def RetrievePluginAsset(self, run, plugin_name, asset_name):
+    """Return the contents for a specific plugin asset from a run.
+
+    Args:
+      run: The string name of the run.
+      plugin_name: The string name of a plugin.
+      asset_name: The string name of an asset.
+
+    Returns:
+      The string contents of the plugin asset.
+
+    Raises:
+      KeyError: If the asset is not available.
+    """
+    accumulator = self._GetAccumulator(run)
+    return accumulator.RetrievePluginAsset(plugin_name, asset_name)
+
   def FirstEventTimestamp(self, run):
     """Return the timestamp of the first event of the given run.
 

tensorflow/tensorboard/backend/event_processing/plugin_asset_util.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Load plugin assets from disk."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os.path
+
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.platform import gfile
+
+_PLUGINS_DIR = "plugins"
+
+
+def _IsDirectory(parent, item):
+  """Helper that returns if parent/item is a directory."""
+  return gfile.IsDirectory(os.path.join(parent, item))
+
+
+def ListPlugins(logdir):
+  """List all the plugins that have registered assets in logdir.
+
+  If the plugins_dir does not exist, it returns an empty list. This maintains
+  compatibility with old directories that have no plugins written.
+
+  Args:
+    logdir: A directory that was created by a TensorFlow events writer.
+
+  Returns:
+    a list of plugin names, as strings
+  """
+  plugins_dir = os.path.join(logdir, _PLUGINS_DIR)
+  if not gfile.IsDirectory(plugins_dir):
+    return []
+  entries = gfile.ListDirectory(plugins_dir)
+  return [x for x in entries if _IsDirectory(plugins_dir, x)]
+
+
+def ListAssets(logdir, plugin_name):
+  """List all the assets that are available for given plugin in a logdir.
+
+  Args:
+    logdir: A directory that was created by a TensorFlow summary.FileWriter.
+    plugin_name: A string name of a plugin to list assets for.
+
+  Returns:
+    A string list of available plugin assets. If the plugin subdirectory does
+    not exist (either because the logdir doesn't exist, or because the plugin
+    didn't register) an empty list is returned.
+  """
+  plugin_dir = os.path.join(logdir, _PLUGINS_DIR, plugin_name)
+  if not gfile.IsDirectory(plugin_dir):
+    return []
+  entries = gfile.ListDirectory(plugin_dir)
+  return [x for x in entries if not _IsDirectory(plugin_dir, x)]
+
+
+def RetrieveAsset(logdir, plugin_name, asset_name):
+  """Retrieve a particular plugin asset from a logdir.
+
+  Args:
+    logdir: A directory that was created by a TensorFlow summary.FileWriter.
+    plugin_name: The plugin we want an asset from.
+    asset_name: The name of the requested asset.
+
+  Returns:
+    string contents of the plugin asset.
+
+  Raises:
+    KeyError: if the asset does not exist.
+  """
+
+  asset_path = os.path.join(logdir, _PLUGINS_DIR, plugin_name, asset_name)
+  try:
+    with gfile.Open(asset_path, "r") as f:
+      return f.read()
+  except errors_impl.NotFoundError:
+    raise KeyError("Asset path %s not found" % asset_path)
+  except errors_impl.OpError as e:
+    raise KeyError("Couldn't read asset path: %s, OpError %s" % (asset_path, e))

tensorflow/tensorboard/backend/event_processing/plugin_asset_util_test.py

+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os.path
+
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+from tensorflow.python.summary import plugin_asset
+from tensorflow.python.summary.writer import writer
+from tensorflow.tensorboard.backend.event_processing import event_multiplexer
+from tensorflow.tensorboard.backend.event_processing import plugin_asset_util
+
+
+class GenericContentPlugin(plugin_asset.PluginAsset):
+
+  def __init__(self):
+    self.contents = "hello world"
+
+  def assets(self):
+    return {"contents.txt": self.contents}
+
+
+class PluginAlpha(GenericContentPlugin):
+  plugin_name = "Alpha"
+
+
+class PluginBeta(GenericContentPlugin):
+  plugin_name = "Beta"
+
+
+class PluginGamma(GenericContentPlugin):
+  plugin_name = "Gamma"
+
+
+class PluginAssetUtilitiesTest(test.TestCase):
+
+  def testNonExistentDirectory(self):
+    tempdir = self.get_temp_dir()
+    fake_dir = os.path.join(tempdir, "nonexistent_dir")
+    self.assertEqual([], plugin_asset_util.ListPlugins(fake_dir))
+    self.assertEqual([], plugin_asset_util.ListAssets(fake_dir, "fake_plugin"))
+    with self.assertRaises(KeyError):
+      plugin_asset_util.RetrieveAsset(fake_dir, "fake_plugin", "fake_asset")
+
+  def testSimplePluginCase(self):
+    tempdir = self.get_temp_dir()
+    with ops.Graph().as_default() as g:
+      plugin_asset.get_plugin_asset(PluginAlpha)
+      fw = writer.FileWriter(tempdir)
+      fw.add_graph(g)
+    self.assertEqual(["Alpha"], plugin_asset_util.ListPlugins(tempdir))
+    assets = plugin_asset_util.ListAssets(tempdir, "Alpha")
+    self.assertEqual(["contents.txt"], assets)
+    contents = plugin_asset_util.RetrieveAsset(tempdir, "Alpha", "contents.txt")
+    self.assertEqual("hello world", contents)
+
+  def testEventMultiplexerIntegration(self):
+    tempdir = self.get_temp_dir()
+    with ops.Graph().as_default() as g:
+      plugin_instance = plugin_asset.get_plugin_asset(PluginAlpha)
+      plugin_instance.contents = "graph one"
+      plugin_asset.get_plugin_asset(PluginBeta)
+
+      fw = writer.FileWriter(os.path.join(tempdir, "one"))
+      fw.add_graph(g)
+      fw.close()
+
+    with ops.Graph().as_default() as g:
+      plugin_instance = plugin_asset.get_plugin_asset(PluginAlpha)
+      plugin_instance.contents = "graph two"
+      fw = writer.FileWriter(os.path.join(tempdir, "two"))
+      fw.add_graph(g)
+      fw.close()
+
+    multiplexer = event_multiplexer.EventMultiplexer()
+    multiplexer.AddRunsFromDirectory(tempdir)
+
+    self.assertEqual(
+        multiplexer.PluginAssets("Alpha"),
+        {"one": ["contents.txt"], "two": ["contents.txt"]})
+    self.assertEqual(
+        multiplexer.RetrievePluginAsset("one", "Alpha", "contents.txt"),
+        "graph one")
+    self.assertEqual(
+        multiplexer.RetrievePluginAsset("one", "Beta", "contents.txt"),
+        "hello world")
+    self.assertEqual(
+        multiplexer.RetrievePluginAsset("two", "Alpha", "contents.txt"),
+        "graph two")
+
+    self.assertEqual(
+        multiplexer.PluginAssets("Beta"),
+        {"one": ["contents.txt"], "two": []})
+    self.assertEqual(multiplexer.PluginAssets("Gamma"), {"one": [], "two": []})
+
+
+if __name__ == "__main__":
+  test.main()