Skip to content

Commit

Permalink
Extensible plugin infrastructural changes (#2257)
Browse files Browse the repository at this point in the history
This is the change based on tensorflow/community#90.

The change makes two improvements:
- allow dynamic plugins (as spec in the RFC; uses entry_points to discover) to be loaded
- change plugins_listing contract to include metadata about plugins including their respective FE entry point
  • Loading branch information
stephanwlee authored May 29, 2019
1 parent 39c80de commit af8a478
Show file tree
Hide file tree
Showing 10 changed files with 265 additions and 22 deletions.
24 changes: 24 additions & 0 deletions tensorboard/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,7 @@ py_library(
srcs_version = "PY2AND3",
visibility = ["//visibility:public"],
deps = [
"//tensorboard:expect_pkg_resources_installed",
"//tensorboard/plugins:base_plugin",
"//tensorboard/plugins/audio:audio_plugin",
"//tensorboard/plugins/beholder:beholder_plugin_loader",
Expand All @@ -192,6 +193,21 @@ py_library(
],
)

py_test(
name = "default_test",
size = "small",
srcs = ["default_test.py"],
srcs_version = "PY2AND3",
tags = ["support_notf"],
deps = [
":default",
":test",
"//tensorboard:expect_pkg_resources_installed",
"//tensorboard/plugins:base_plugin",
"@org_pythonhosted_mock",
],
)

py_library(
name = "version",
srcs = ["version.py"],
Expand Down Expand Up @@ -272,6 +288,14 @@ py_library(
visibility = ["//visibility:public"],
)

py_library(
name = "expect_pkg_resources_installed",
# This is a dummy rule used as a pkg-resources dependency in open-source.
# We expect pkg-resources to already be installed on the system, e.g., via
# `pip install setuptools`.
visibility = ["//visibility:public"],
)

py_library(
name = "tf_contrib_ffmpeg",
# This is a dummy rule for the open source world, which indicates
Expand Down
10 changes: 9 additions & 1 deletion tensorboard/backend/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,15 @@ def _serve_plugins_listing(self, request):
response = {}
for plugin in self._plugins:
start = time.time()
response[plugin.plugin_name] = plugin.is_active()
module_path = None
if plugin.es_module_path() is not None:
module_path = (self._path_prefix + DATA_PREFIX + PLUGIN_PREFIX + '/' +
plugin.plugin_name + plugin.es_module_path())

response[plugin.plugin_name] = {
'enabled': plugin.is_active(),
'es_module_path': module_path,
}
elapsed = time.time() - start
logger.info(
'Plugin listing: is_active() for %s took %0.3f seconds',
Expand Down
69 changes: 65 additions & 4 deletions tensorboard/backend/application_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ def __init__(self,
plugin_name,
is_active_value,
routes_mapping,
es_module_path_value=None,
construction_callback=None):
"""Constructs a fake plugin.
Expand All @@ -90,12 +91,15 @@ def __init__(self,
is_active_value: Whether the plugin is active.
routes_mapping: A dictionary mapping from route (string URL path) to the
method called when a user issues a request to that route.
es_module_path_value: An optional string value that indicates a frontend
module entry to the plugin. Must be one of the keys of routes_mapping.
construction_callback: An optional callback called when the plugin is
constructed. The callback is passed the TBContext.
"""
self.plugin_name = plugin_name
self._is_active_value = is_active_value
self._routes_mapping = routes_mapping
self._es_module_path_value = es_module_path_value

if construction_callback:
construction_callback(context)
Expand All @@ -116,6 +120,15 @@ def is_active(self):
"""
return self._is_active_value

def es_module_path(self):
"""Returns a path to plugin frontend entry.
Returns:
A string that corresponds to a key of routes_mapping. For non-dynamic
plugins, it returns None.
"""
return self._es_module_path_value


class ApplicationTest(tb_test.TestCase):
def setUp(self):
Expand All @@ -124,6 +137,15 @@ def setUp(self):
None, plugin_name='foo', is_active_value=True, routes_mapping={}),
FakePlugin(
None, plugin_name='bar', is_active_value=False, routes_mapping={}),
FakePlugin(
None,
plugin_name='baz',
is_active_value=True,
routes_mapping={
'/esmodule': lambda req: None,
},
es_module_path_value='/esmodule'
),
]
app = application.TensorBoardWSGI(plugins)
self.server = werkzeug_test.Client(app, wrappers.BaseResponse)
Expand All @@ -146,8 +168,23 @@ def testRequestNonexistentPage(self):
def testPluginsListing(self):
"""Test the format of the data/plugins_listing endpoint."""
parsed_object = self._get_json('/data/plugins_listing')
# Plugin foo is active. Plugin bar is not.
self.assertEqual(parsed_object, {'foo': True, 'bar': False})
self.assertEqual(
parsed_object,
{
'foo': {
'enabled': True,
'es_module_path': None,
},
'bar': {
'enabled': False,
'es_module_path': None,
},
'baz': {
'enabled': True,
'es_module_path': '/data/plugin/baz/esmodule',
},
}
)


class ApplicationBaseUrlTest(tb_test.TestCase):
Expand All @@ -158,6 +195,15 @@ def setUp(self):
None, plugin_name='foo', is_active_value=True, routes_mapping={}),
FakePlugin(
None, plugin_name='bar', is_active_value=False, routes_mapping={}),
FakePlugin(
None,
plugin_name='baz',
is_active_value=True,
routes_mapping={
'/esmodule': lambda req: None,
},
es_module_path_value='/esmodule'
),
]
app = application.TensorBoardWSGI(plugins, path_prefix=self.path_prefix)
self.server = werkzeug_test.Client(app, wrappers.BaseResponse)
Expand Down Expand Up @@ -186,8 +232,23 @@ def testBaseUrlNonexistentPluginsListing(self):
def testPluginsListing(self):
"""Test the format of the data/plugins_listing endpoint."""
parsed_object = self._get_json(self.path_prefix + '/data/plugins_listing')
# Plugin foo is active. Plugin bar is not.
self.assertEqual(parsed_object, {'foo': True, 'bar': False})
self.assertEqual(
parsed_object,
{
'foo': {
'enabled': True,
'es_module_path': None,
},
'bar': {
'enabled': False,
'es_module_path': None,
},
'baz': {
'enabled': True,
'es_module_path': '/test/data/plugin/baz/esmodule',
},
}
)


class ApplicationPluginNameTest(tb_test.TestCase):
Expand Down
41 changes: 28 additions & 13 deletions tensorboard/components/tf_tensorboard/tf-tensorboard.html
Original file line number Diff line number Diff line change
Expand Up @@ -465,17 +465,18 @@ <h3>There’s no dashboard by the name of “<tt>[[_selectedDashboard]]</tt>.”
value: _.values(tf_tensorboard.dashboardRegistry),
},

_pluginsListing: {
type: Object,
value: () => ({}),
},

/**
* The set of currently active dashboards. `null` if not yet fetched.
*
* TODO(@wchargin): Consider templating in an initial value for
* this property.
*
* The set of currently active dashboards.
* @type {Array<string>?}
*/
_activeDashboards: {
type: Array,
value: null,
computed: '_computeActiveDashboard(_dashboardData, _pluginsListing)'
},

/** @type {tf_tensorboard.ActiveDashboardsLoadState} */
Expand Down Expand Up @@ -812,15 +813,29 @@ <h3>There’s no dashboard by the name of “<tt>[[_selectedDashboard]]</tt>.”
this._lastReloadTime = new Date().toString();
},

_fetchActiveDashboards() {
_computeActiveDashboard() {
return this._dashboardData
.map((d) => d.plugin)
.filter((dashboardName) => {
// TODO(stephanwlee): Remove boolean code path when releasing
// 2.0.
// PluginsListing can be an object whose key is name of the
// plugin and value is a boolean indicating whether if it is
// enabled. This is deprecated but we will maintain backwards
// compatibility for some time.
const maybeMetadata = this._pluginsListing[dashboardName];
if (typeof maybeMetadata === 'boolean') return maybeMetadata;
return maybeMetadata && maybeMetadata.enabled;
});
},

_fetchPluginsListing() {
this._canceller.cancelAll();
const updateActiveDashboards = this._canceller.cancellable(result => {
const updatePluginsListing = this._canceller.cancellable(result => {
if (result.cancelled) {
return;
}
const activePlugins = result.value;
this._activeDashboards = this._dashboardData.map(
d => d.plugin).filter(p => activePlugins[p]);
this._pluginsListing = result.value;
this._activeDashboardsLoadState = tf_tensorboard.ActiveDashboardsLoadState.LOADED;
});
const onFailure = () => {
Expand All @@ -834,7 +849,7 @@ <h3>There’s no dashboard by the name of “<tt>[[_selectedDashboard]]</tt>.”
};
return this._requestManager
.request(tf_backend.getRouter().pluginsListing())
.then(updateActiveDashboards, onFailure);
.then(updatePluginsListing, onFailure);
},

_computeActiveDashboardsNotLoaded(state) {
Expand Down Expand Up @@ -878,7 +893,7 @@ <h3>There’s no dashboard by the name of “<tt>[[_selectedDashboard]]</tt>.”
_reloadData() {
this._refreshing = true;
return Promise.all([
this._fetchActiveDashboards(),
this._fetchPluginsListing(),
tf_backend.environmentStore.refresh(),
tf_backend.runsStore.refresh(),
tf_backend.experimentsStore.refresh(),
Expand Down
24 changes: 23 additions & 1 deletion tensorboard/default.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
# Copyright 2019 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.
Expand Down Expand Up @@ -31,6 +31,8 @@
import logging
import os

import pkg_resources

from tensorboard.compat import tf
from tensorboard.plugins import base_plugin
from tensorboard.plugins.audio import audio_plugin
Expand Down Expand Up @@ -86,4 +88,24 @@ def get_plugins():
:rtype: list[Union[base_plugin.TBLoader, Type[base_plugin.TBPlugin]]]
"""

return _PLUGINS[:]


def get_dynamic_plugins():
"""Returns a list specifying TensorBoard's dynamically loaded plugins.
A dynamic TensorBoard plugin is specified using entry_points [1] and it is
the robust way to integrate plugins into TensorBoard.
This list can be passed to the `tensorboard.program.TensorBoard` API.
Returns:
list of base_plugin.TBLoader or base_plugin.TBPlugin.
[1]: https://packaging.python.org/specifications/entry-points/
"""
return [
entry_point.load()
for entry_point in pkg_resources.iter_entry_points('tensorboard_plugins')
]
74 changes: 74 additions & 0 deletions tensorboard/default_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
# 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.
# ==============================================================================
"""Unit tests for `tensorboard.default`."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

try:
# python version >= 3.3
from unittest import mock # pylint: disable=g-import-not-at-top
except ImportError:
import mock # pylint: disable=g-import-not-at-top,unused-import

import pkg_resources

from tensorboard import default
from tensorboard.plugins import base_plugin
from tensorboard import test


class FakePlugin(base_plugin.TBPlugin):
"""FakePlugin for testing."""

plugin_name = 'fake'


class FakeEntryPoint(pkg_resources.EntryPoint):
"""EntryPoint class that fake loads FakePlugin."""

@classmethod
def create(cls):
"""Creates an instance of FakeEntryPoint.
Returns:
instance of FakeEntryPoint
"""
return cls('foo', 'bar')

def load(self):
"""Returns FakePlugin instead of resolving module.
Returns:
FakePlugin
"""
return FakePlugin


class DefaultTest(test.TestCase):

@mock.patch.object(pkg_resources, 'iter_entry_points')
def test_get_dynamic_plugin(self, mock_iter_entry_points):
mock_iter_entry_points.return_value = [FakeEntryPoint.create()]

actual_plugins = default.get_dynamic_plugins()

mock_iter_entry_points.assert_called_with('tensorboard_plugins')
self.assertEqual(actual_plugins, [FakePlugin])


if __name__ == "__main__":
test.main()
Loading

0 comments on commit af8a478

Please sign in to comment.