Extend program.TensorBoard API to allow other servers and absl flags

tensorboard/BUILD

@@ -57,6 +57,19 @@ py_library(
     ],
 )
 
+py_test(
+    name = "program_test",
+    size = "small",
+    srcs = ["program_test.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":program",
+        "//tensorboard:expect_tensorflow_installed",
+        "//tensorboard/backend:application",
+        "@org_pocoo_werkzeug",
+    ],
+)
+
 py_library(
     name = "default",
     srcs = ["default.py"],
@@ -116,7 +129,7 @@ py_library(
 py_library(
     name = "expect_futures_installed",
     # This is a dummy rule used as a futures dependency in open-source.
-    # We expect numpy to already be installed on the system, e.g. via
+    # We expect futures to already be installed on the system, e.g. via
     # `pip install futures`
     visibility = ["//visibility:public"],
 )

tensorboard/backend/BUILD

@@ -83,8 +83,6 @@ py_test(
     srcs_version = "PY2AND3",
     deps = [
         ":application",
-        "//tensorboard",
-        "//tensorboard:default",
         "//tensorboard:expect_tensorflow_installed",
         "//tensorboard/backend/event_processing:event_multiplexer",
         "//tensorboard/plugins:base_plugin",

tensorboard/backend/application_test.py

@@ -40,8 +40,6 @@
 from werkzeug import test as werkzeug_test
 from werkzeug import wrappers
 
-from tensorboard import default
-from tensorboard import program as tensorboard
 from tensorboard.backend import application
 from tensorboard.backend.event_processing import plugin_event_multiplexer as event_multiplexer  # pylint: disable=line-too-long
 from tensorboard.plugins import base_plugin
@@ -115,9 +113,7 @@ def is_active(self):
     return self._is_active_value
 
 
-class TensorboardServerTest(tf.test.TestCase):
-  _only_use_meta_graph = False  # Server data contains only a GraphDef
-
+class ApplicationTest(tf.test.TestCase):
   def setUp(self):
     plugins = [
         FakePlugin(
@@ -150,8 +146,7 @@ def testPluginsListing(self):
     self.assertEqual(parsed_object, {'foo': True, 'bar': False})
 
 
-class TensorboardServerBaseUrlTest(tf.test.TestCase):
-  _only_use_meta_graph = False  # Server data contains only a GraphDef
+class ApplicationBaseUrlTest(tf.test.TestCase):
   path_prefix = '/test'
   def setUp(self):
     plugins = [
@@ -191,7 +186,7 @@ def testPluginsListing(self):
     self.assertEqual(parsed_object, {'foo': True, 'bar': False})
 
 
-class TensorboardServerPluginNameTest(tf.test.TestCase):
+class ApplicationPluginNameTest(tf.test.TestCase):
 
   def _test(self, name, should_be_okay):
     temp_dir = tempfile.mkdtemp(prefix=self.get_temp_dir())
@@ -233,8 +228,7 @@ def testComprehensiveName(self):
     self._test('Scalar-Dashboard_3000.1', True)
 
 
-
-class TensorboardServerPluginRouteTest(tf.test.TestCase):
+class ApplicationPluginRouteTest(tf.test.TestCase):
 
   def _test(self, route, should_be_okay):
     temp_dir = tempfile.mkdtemp(prefix=self.get_temp_dir())
@@ -267,11 +261,6 @@ def testSlashlessRoute(self):
     self._test('runaway', False)
 
 
-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 assertPlatformSpecificLogdirParsing(self, pathObj, logdir, expected):
@@ -397,6 +386,7 @@ class TensorBoardPluginsTest(tf.test.TestCase):
 
   def setUp(self):
     self.context = None
+    dummy_assets_zip_provider = lambda: None
     # The application should have added routes for both plugins.
     self.app = application.standard_tensorboard_wsgi(
         FakeFlags(logdir=self.get_temp_dir()),
@@ -414,7 +404,7 @@ def setUp(self):
               routes_mapping={'/bar_route': self._bar_handler},
               construction_callback=self._construction_callback)),
         ],
-        default.get_assets_zip_provider())
+        dummy_assets_zip_provider)
 
   def _foo_handler(self):
     pass
@@ -441,50 +431,7 @@ def testNameToPluginMapping(self):
     self.assertEqual('bar', mapping['bar'].plugin_name)
 
 
-class TensorboardSimpleServerConstructionTest(tf.test.TestCase):
-  """Tests that the default HTTP server is constructed without error.
-
-  Mostly useful for IPv4/IPv6 testing. This test should run with only IPv4, only
-  IPv6, and both IPv4 and IPv6 enabled.
-  """
-
-  class _StubApplication(object):
-    tag = ''
-
-  def testMakeServerBlankHost(self):
-    # Test that we can bind to all interfaces without throwing an error
-    server, url = tensorboard.make_simple_server(
-        self._StubApplication(),
-        host='',
-        port=0)  # Grab any available port
-    self.assertTrue(server)
-    self.assertTrue(url)
-
-  def testSpecifiedHost(self):
-    one_passed = False
-    try:
-      _, url = tensorboard.make_simple_server(
-          self._StubApplication(),
-          host='127.0.0.1',
-          port=0)
-      self.assertStartsWith(actual=url, expected_start='http://127.0.0.1:')
-      one_passed = True
-    except socket.error:
-      # IPv4 is not supported
-      pass
-    try:
-      _, url = tensorboard.make_simple_server(
-          self._StubApplication(),
-          host='::1',
-          port=0)
-      self.assertStartsWith(actual=url, expected_start='http://[::1]:')
-      one_passed = True
-    except socket.error:
-      # IPv6 is not supported
-      pass
-    self.assertTrue(one_passed)  # We expect either IPv4 or IPv6 to be supported
-
-class TensorBoardApplcationConstructionTest(tf.test.TestCase):
+class ApplicationConstructionTest(tf.test.TestCase):
 
   def testExceptions(self):
     logdir = '/fake/foo'

tensorboard/default.py

@@ -51,23 +51,34 @@
 
 logger = logging.getLogger(__name__)
 
-PLUGIN_LOADERS = [
+_PLUGINS = [
     core_plugin.CorePluginLoader(),
-    base_plugin.BasicLoader(beholder_plugin.BeholderPlugin),
-    base_plugin.BasicLoader(scalars_plugin.ScalarsPlugin),
-    base_plugin.BasicLoader(custom_scalars_plugin.CustomScalarsPlugin),
-    base_plugin.BasicLoader(images_plugin.ImagesPlugin),
-    base_plugin.BasicLoader(audio_plugin.AudioPlugin),
-    base_plugin.BasicLoader(graphs_plugin.GraphsPlugin),
-    base_plugin.BasicLoader(distributions_plugin.DistributionsPlugin),
-    base_plugin.BasicLoader(histograms_plugin.HistogramsPlugin),
-    base_plugin.BasicLoader(pr_curves_plugin.PrCurvesPlugin),
-    base_plugin.BasicLoader(projector_plugin.ProjectorPlugin),
-    base_plugin.BasicLoader(text_plugin.TextPlugin),
+    beholder_plugin.BeholderPlugin,
+    scalars_plugin.ScalarsPlugin,
+    custom_scalars_plugin.CustomScalarsPlugin,
+    images_plugin.ImagesPlugin,
+    audio_plugin.AudioPlugin,
+    graphs_plugin.GraphsPlugin,
+    distributions_plugin.DistributionsPlugin,
+    histograms_plugin.HistogramsPlugin,
+    pr_curves_plugin.PrCurvesPlugin,
+    projector_plugin.ProjectorPlugin,
+    text_plugin.TextPlugin,
     profile_plugin.ProfilePluginLoader(),
     debugger_plugin_loader.DebuggerPluginLoader(),
 ]
 
+def get_plugins():
+  """Returns a list specifying TensorBoard's default first-party plugins.
+
+  Plugins are specified in this list either via a TBLoader instance to load the
+  plugin, or the TBPlugin class itself which will be loaded using a BasicLoader.
+
+  This list can be passed to the `tensorboard.program.TensorBoard` API.
+
+  :rtype: list[Union[base_plugin.TBLoader, Type[base_plugin.TBPlugin]]]
+  """
+  return _PLUGINS[:]
 
 def get_assets_zip_provider():
   """Opens stock TensorBoard web assets collection.

tensorboard/main.py

@@ -34,32 +34,25 @@
 os.environ['GCS_READ_CACHE_DISABLED'] = '1'
 # pylint: enable=g-import-not-at-top
 
-import logging
 import sys
 
 from tensorboard import default
 from tensorboard import program
 
-logger = logging.getLogger(__name__)
-
 
 def run_main():
   """Initializes flags and calls main()."""
   program.setup_environment()
-  server = program.TensorBoard(default.PLUGIN_LOADERS,
-                               default.get_assets_zip_provider())
-  server.configure(sys.argv[1:])
+  tensorboard = program.TensorBoard(default.get_plugins(),
+                                    default.get_assets_zip_provider())
   try:
     from absl import app
-    app.run(server.main, sys.argv[:1] + server.unparsed_argv)
+    app.run(tensorboard.main, flags_parser=tensorboard.configure)
     raise AssertionError("absl.app.run() shouldn't return")
   except ImportError:
     pass
-  if server.unparsed_argv:
-    sys.stderr.write('Unknown flags: %s\nPass --help for help.\n' %
-                     (server.unparsed_argv,))
-    sys.exit(1)
-  sys.exit(server.main())
+  tensorboard.configure(sys.argv)
+  sys.exit(tensorboard.main())
 
 
 if __name__ == '__main__':

tensorboard/plugins/core/core_plugin.py

@@ -295,7 +295,8 @@ def define_flags(self, parser):
         type=int,
         default=6006,
         help='''\
-Port to serve TensorBoard on (default: %(default)s)\
+Port to serve TensorBoard on. Pass 0 to request an unused port selected
+by the operating system. (default: %(default)s)\
 ''')
 
     parser.add_argument(