data: add histograms plugin integration (#2981)

Summary:
This patch teaches the histograms plugin about the generic data APIs and
their newfound tensor support, guarded behind the `--generic_data=true`
feature flag. The distributions plugin tags along for free.

Testing:
After this change, setting `--generic_data` to either `false` or `true`
does not change the visual appearance of the distributions or histograms
dashboards nor the outputs of either the `/data/tags` routes or the
`/data/histograms` or `/data/distributions` routes. (Replacing `e.numpy`
with `(e.numpy + 100)` in `histograms_plugin.py` does cause these checks
to fail, so everything is wired up properly.)

Unit tests updated for the histograms plugin following the form of those
in the scalars plugin. Tests for the distributions plugin tests are left
unchanged, as it never actually touches its context or multiplexer.

wchargin-branch: generic-histograms

Author: wchargin

tensorboard/plugins/distribution/BUILD

@@ -15,6 +15,7 @@ py_library(
     visibility = ["//visibility:public"],
     deps = [
         ":compressor",
+        "//tensorboard:plugin_util",
         "//tensorboard/backend:http_util",
         "//tensorboard/plugins:base_plugin",
         "//tensorboard/plugins/histogram:histograms_plugin",

tensorboard/plugins/distribution/distributions_plugin.py

@@ -24,6 +24,7 @@
 
 from werkzeug import wrappers
 
+from tensorboard import plugin_util
 from tensorboard.backend import http_util
 from tensorboard.plugins import base_plugin
 from tensorboard.plugins.distribution import compressor
@@ -52,7 +53,6 @@ def __init__(self, context):
       context: A base_plugin.TBContext instance.
     """
     self._histograms_plugin = histograms_plugin.HistogramsPlugin(context)
-    self._multiplexer = context.multiplexer
 
   def get_plugin_apps(self):
     return {
@@ -73,14 +73,14 @@ def frontend_metadata(self):
         element_name='tf-distribution-dashboard',
     )
 
-  def distributions_impl(self, tag, run):
+  def distributions_impl(self, tag, run, experiment):
     """Result of the form `(body, mime_type)`.
 
     Raises:
       tensorboard.errors.PublicError: On invalid request.
     """
     (histograms, mime_type) = self._histograms_plugin.histograms_impl(
-        tag, run, downsample_to=self.SAMPLE_SIZE)
+        tag, run, experiment=experiment, downsample_to=self.SAMPLE_SIZE)
     return ([self._compress(histogram) for histogram in histograms],
             mime_type)
 
@@ -89,18 +89,20 @@ def _compress(self, histogram):
     converted_buckets = compressor.compress_histogram(buckets)
     return [wall_time, step, converted_buckets]
 
-  def index_impl(self):
-    return self._histograms_plugin.index_impl()
+  def index_impl(self, experiment):
+    return self._histograms_plugin.index_impl(experiment=experiment)
 
   @wrappers.Request.application
   def tags_route(self, request):
-    index = self.index_impl()
+    experiment = plugin_util.experiment_id(request.environ)
+    index = self.index_impl(experiment=experiment)
     return http_util.Respond(request, index, 'application/json')
 
   @wrappers.Request.application
   def distributions_route(self, request):
     """Given a tag and single run, return an array of compressed histograms."""
+    experiment = plugin_util.experiment_id(request.environ)
     tag = request.args.get('tag')
     run = request.args.get('run')
-    (body, mime_type) = self.distributions_impl(tag, run)
+    (body, mime_type) = self.distributions_impl(tag, run, experiment=experiment)
     return http_util.Respond(request, body, mime_type)

tensorboard/plugins/distribution/distributions_plugin_test.py

@@ -121,14 +121,16 @@ def test_index(self):
                 'description': self._HTML_DESCRIPTION,
             },
         },
-    }, self.plugin.index_impl())
+    }, self.plugin.index_impl(experiment='exp'))
 
   def _test_distributions(self, run_name, tag_name, should_work=True):
     self.set_up_with_runs([self._RUN_WITH_SCALARS,
                            self._RUN_WITH_LEGACY_DISTRIBUTION,
                            self._RUN_WITH_DISTRIBUTION])
     if should_work:
-      (data, mime_type) = self.plugin.distributions_impl(tag_name, run_name)
+      (data, mime_type) = self.plugin.distributions_impl(
+          tag_name, run_name, experiment='exp'
+      )
       self.assertEqual('application/json', mime_type)
       self.assertEqual(len(data), self._STEPS)
       for i in xrange(self._STEPS):
@@ -138,7 +140,9 @@ def _test_distributions(self, run_name, tag_name, should_work=True):
         self.assertEqual(bps, compressor.NORMAL_HISTOGRAM_BPS)
     else:
       with self.assertRaises(errors.NotFoundError):
-        self.plugin.distributions_impl(self._DISTRIBUTION_TAG, run_name)
+        self.plugin.distributions_impl(
+            self._DISTRIBUTION_TAG, run_name, experiment='exp'
+        )
 
   def test_distributions_with_scalars(self):
     self._test_distributions(self._RUN_WITH_SCALARS, self._DISTRIBUTION_TAG,

tensorboard/plugins/histogram/BUILD

@@ -21,6 +21,7 @@ py_library(
         "//tensorboard:expect_numpy_installed",
         "//tensorboard:plugin_util",
         "//tensorboard/backend:http_util",
+        "//tensorboard/data:provider",
         "//tensorboard/plugins:base_plugin",
         "//tensorboard/util:tensor_util",
         "@org_pocoo_werkzeug",
@@ -42,6 +43,7 @@ py_test(
         "//tensorboard/backend:application",
         "//tensorboard/backend/event_processing:event_accumulator",
         "//tensorboard/backend/event_processing:event_multiplexer",
+        "//tensorboard/data:provider",
         "//tensorboard/plugins:base_plugin",
         "//tensorboard/util:test_util",
         "@org_pocoo_werkzeug",

tensorboard/plugins/histogram/histograms_plugin.py

@@ -33,6 +33,7 @@
 from tensorboard import plugin_util
 from tensorboard.backend import http_util
 from tensorboard.compat import tf
+from tensorboard.data import provider
 from tensorboard.plugins import base_plugin
 from tensorboard.plugins.histogram import metadata
 from tensorboard.util import tensor_util
@@ -59,8 +60,12 @@ def __init__(self, context):
     Args:
       context: A base_plugin.TBContext instance.
     """
-    self._db_connection_provider = context.db_connection_provider
     self._multiplexer = context.multiplexer
+    self._db_connection_provider = context.db_connection_provider
+    if context.flags and context.flags.generic_data == 'true':
+      self._data_provider = context.data_provider
+    else:
+      self._data_provider = None
 
   def get_plugin_apps(self):
     return {
@@ -70,6 +75,11 @@ def get_plugin_apps(self):
 
   def is_active(self):
     """This plugin is active iff any run has at least one histograms tag."""
+    if self._data_provider:
+      # We don't have an experiment ID, and modifying the backend core
+      # to provide one would break backward compatibility. Hack for now.
+      return True
+
     if self._db_connection_provider:
       # The plugin is active if one relevant tag can be found in the database.
       db = self._db_connection_provider()
@@ -82,10 +92,28 @@ def is_active(self):
       ''', (metadata.PLUGIN_NAME,))
       return bool(list(cursor))
 
-    return bool(self._multiplexer) and any(self.index_impl().values())
+    if self._multiplexer:
+      return any(self.index_impl(experiment='').values())
+
+    return False
 
-  def index_impl(self):
+  def index_impl(self, experiment):
     """Return {runName: {tagName: {displayName: ..., description: ...}}}."""
+    if self._data_provider:
+      mapping = self._data_provider.list_tensors(
+          experiment_id=experiment,
+          plugin_name=metadata.PLUGIN_NAME,
+      )
+      result = {run: {} for run in mapping}
+      for (run, tag_to_content) in six.iteritems(mapping):
+        for (tag, metadatum) in six.iteritems(tag_to_content):
+          description = plugin_util.markdown_to_safe_html(metadatum.description)
+          result[run][tag] = {
+              'displayName': metadatum.display_name,
+              'description': description,
+          }
+      return result
+
     if self._db_connection_provider:
       # Read tags from the database.
       db = self._db_connection_provider()
@@ -128,7 +156,7 @@ def index_impl(self):
   def frontend_metadata(self):
     return base_plugin.FrontendMetadata(element_name='tf-histogram-dashboard')
 
-  def histograms_impl(self, tag, run, downsample_to=None):
+  def histograms_impl(self, tag, run, experiment, downsample_to=None):
     """Result of the form `(body, mime_type)`.
 
     At most `downsample_to` events will be returned. If this value is
@@ -137,7 +165,33 @@ def histograms_impl(self, tag, run, downsample_to=None):
     Raises:
       tensorboard.errors.PublicError: On invalid request.
     """
-    if self._db_connection_provider:
+    if self._data_provider:
+      # Downsample reads to 500 histograms per time series, which is
+      # the default size guidance for histograms under the multiplexer
+      # loading logic.
+      SAMPLE_COUNT = downsample_to if downsample_to is not None else 500
+      all_histograms = self._data_provider.read_tensors(
+          experiment_id=experiment,
+          plugin_name=metadata.PLUGIN_NAME,
+          downsample=SAMPLE_COUNT,
+          run_tag_filter=provider.RunTagFilter(runs=[run], tags=[tag]),
+      )
+      histograms = all_histograms.get(run, {}).get(tag, None)
+      if histograms is None:
+        raise errors.NotFoundError(
+            "No histogram tag %r for run %r" % (tag, run)
+        )
+      # Downsample again, even though the data provider is supposed to,
+      # because the multiplexer provider currently doesn't. (For
+      # well-behaved data providers, this is a no-op.)
+      if downsample_to is not None:
+        rng = random.Random(0)
+        histograms = _downsample(rng, histograms, downsample_to)
+      events = [
+          (e.wall_time, e.step, e.numpy.tolist())
+          for e in histograms
+      ]
+    elif self._db_connection_provider:
       # Serve data from the database.
       db = self._db_connection_provider()
       cursor = db.cursor()
@@ -205,11 +259,9 @@ def histograms_impl(self, tag, run, downsample_to=None):
         raise errors.NotFoundError(
             'No histogram tag %r for run %r' % (tag, run)
         )
-      if downsample_to is not None and len(tensor_events) > downsample_to:
-        rand_indices = random.Random(0).sample(
-            six.moves.xrange(len(tensor_events)), downsample_to)
-        indices = sorted(rand_indices)
-        tensor_events = [tensor_events[i] for i in indices]
+      if downsample_to is not None:
+        rng = random.Random(0)
+        tensor_events = _downsample(rng, tensor_events, downsample_to)
       events = [[e.wall_time, e.step, tensor_util.make_ndarray(e.tensor_proto).tolist()]
                 for e in tensor_events]
     return (events, 'application/json')
@@ -228,14 +280,42 @@ def _get_values(self, data_blob, dtype_enum, shape_string):
 
   @wrappers.Request.application
   def tags_route(self, request):
-    index = self.index_impl()
+    experiment = plugin_util.experiment_id(request.environ)
+    index = self.index_impl(experiment=experiment)
     return http_util.Respond(request, index, 'application/json')
 
   @wrappers.Request.application
   def histograms_route(self, request):
     """Given a tag and single run, return array of histogram values."""
+    experiment = plugin_util.experiment_id(request.environ)
     tag = request.args.get('tag')
     run = request.args.get('run')
     (body, mime_type) = self.histograms_impl(
-        tag, run, downsample_to=self.SAMPLE_SIZE)
+        tag, run, experiment=experiment, downsample_to=self.SAMPLE_SIZE)
     return http_util.Respond(request, body, mime_type)
+
+
+def _downsample(rng, xs, k):
+  """Uniformly choose a maximal at-most-`k`-subsequence of `xs`.
+
+  If `k` is larger than `xs`, then the contents of `xs` itself will be
+  returned.
+
+  This differs from `random.sample` in that it returns a subsequence
+  (i.e., order is preserved) and that it permits `k > len(xs)`.
+
+  Args:
+    rng: A `random` interface.
+    xs: A sequence (`collections.abc.Sequence`).
+    k: A non-negative integer.
+
+  Returns:
+    A new list whose elements are a subsequence of `xs` of length
+    `min(k, len(xs))`, uniformly selected among such subsequences.
+  """
+
+  if k > len(xs):
+    return list(xs)
+  indices = rng.sample(six.moves.xrange(len(xs)), k)
+  indices.sort()
+  return [xs[i] for i in indices]