Migrate _parse_samples_per_plugin() into flag definition itself (#3637)

* migrate _parse_samples_per_plugin() into flag definition itself

* CR: rename _parse_dict to _parse_samples_per_plugin

Author: nfelt

tensorboard/backend/application.py

@@ -96,16 +96,6 @@
 logger = tb_logging.get_logger()
 
 
-def _parse_samples_per_plugin(flags):
-    result = {}
-    if not flags or not flags.samples_per_plugin:
-        return result
-    for token in flags.samples_per_plugin.split(","):
-        k, v = token.strip().split("=")
-        result[k] = int(v)
-    return result
-
-
 def _apply_tensor_size_guidance(sampling_hints):
     """Apply user per-summary size guidance overrides."""
     tensor_size_guidance = dict(DEFAULT_TENSOR_SIZE_GUIDANCE)
@@ -131,7 +121,7 @@ def standard_tensorboard_wsgi(flags, plugin_loaders, assets_zip_provider):
     multiplexer = None
     reload_interval = flags.reload_interval
     # Regular logdir loading mode.
-    sampling_hints = _parse_samples_per_plugin(flags)
+    sampling_hints = flags.samples_per_plugin
     multiplexer = event_multiplexer.EventMultiplexer(
         size_guidance=DEFAULT_SIZE_GUIDANCE,
         tensor_size_guidance=_apply_tensor_size_guidance(sampling_hints),
@@ -208,7 +198,7 @@ def TensorBoardWSGIApp(
         multiplexer=deprecated_multiplexer,
         assets_zip_provider=assets_zip_provider,
         plugin_name_to_instance=plugin_name_to_instance,
-        sampling_hints=_parse_samples_per_plugin(flags),
+        sampling_hints=flags.samples_per_plugin,
         window_title=flags.window_title,
     )
     tbplugins = []

tensorboard/backend/application_test.py

@@ -58,7 +58,7 @@ def __init__(
         logdir_spec="",
         purge_orphaned_data=True,
         reload_interval=60,
-        samples_per_plugin="",
+        samples_per_plugin=None,
         max_reload_threads=1,
         reload_task="auto",
         window_title="",
@@ -71,7 +71,7 @@ def __init__(
         self.logdir_spec = logdir_spec
         self.purge_orphaned_data = purge_orphaned_data
         self.reload_interval = reload_interval
-        self.samples_per_plugin = samples_per_plugin
+        self.samples_per_plugin = samples_per_plugin or {}
         self.max_reload_threads = max_reload_threads
         self.reload_task = reload_task
         self.window_title = window_title

tensorboard/plugins/core/core_plugin.py

@@ -514,7 +514,7 @@ def define_flags(self, parser):
 
         parser.add_argument(
             "--samples_per_plugin",
-            type=str,
+            type=_parse_samples_per_plugin,
             default="",
             help="""\
 An optional comma separated list of plugin_name=num_samples pairs to
@@ -576,3 +576,13 @@ def _gzip(bytestring):
     with gzip.GzipFile(fileobj=out, mode="wb", compresslevel=3, mtime=0) as f:
         f.write(bytestring)
     return out.getvalue()
+
+
+def _parse_samples_per_plugin(value):
+    """Parses `value` as a string-to-int dict in the form `foo=12,bar=34`."""
+    result = {}
+    for token in value.split(","):
+        if token:
+            k, v = token.strip().split("=")
+            result[k] = int(v)
+    return result