Hparams: Adds filters to DataProvider.read_hyperparameters(). (#6506)

Adds support for specifying filters in
DataProvider.read_hyperparameters(), represented as keyword argument
`filters: Collection[HyperparameterFilter]`.

Updates list_session_groups.py to translate ColParams from the HTTP
request into HyperparameterFilter objects and include it in the request
to the DataProvider implementation.

Author: bmd3k

tensorboard/data/provider.py

@@ -396,13 +396,15 @@ def list_hyperparameters(self, ctx=None, *, experiment_ids):
         """
         return []
 
-    def read_hyperparameters(self, ctx=None, *, experiment_ids):
+    def read_hyperparameters(self, ctx=None, *, experiment_ids, filters=None):
         """Read hyperparameter values.
 
         Args:
           ctx: A TensorBoard `RequestContext` value.
           experiment_ids: A Collection[string] of IDs of the enclosing
             experiments.
+          filters: A Collection[HyperparameterFilter] that constrain the
+            returned session groups based on hyperparameter value.
 
         Returns:
           A Collection[HyperparameterSessionGroup] describing the groups and
@@ -668,6 +670,45 @@ class HyperparameterSessionGroup:
     hyperparameter_values: Collection[HyperparameterValue]
 
 
+class HyperparameterFilterType(enum.Enum):
+    """Describes how to represent filter values."""
+
+    # A regular expression string. Normally represented as str.
+    REGEX = "regex"
+    # A range of numeric values. Normally represented as Tuple[float, float].
+    INTERVAL = "interval"
+    # A finite set of values. Normally represented as Collection[float|str|bool].
+    DISCRETE = "discrete"
+
+
+@dataclasses.dataclass(frozen=True)
+class HyperparameterFilter:
+    """A constraint based on hyperparameter value.
+
+    Attributes:
+      hyperparameter_name: A string identifier for the hyperparameter to use for
+        the filter. It corresponds to the hyperparameter_name field in the
+        Hyperparameter class.
+      filter_type: A HyperparameterFilterType describing how we represent the
+        filter values in the 'filter' attribute.
+      filter: A representation of the set of the filter values.
+
+        If filter_type is REGEX, a str containing the regular expression.
+        If filter_type is INTERVAL, a Tuple[float, float] describing the min and
+          max values of the filter interval.
+        If filter_type is DISCRETE a Collection[float|str|bool] describing the
+          finite set of filter values.
+    """
+
+    hyperparameter_name: str
+    filter_type: HyperparameterFilterType
+    filter: Union[
+        str,
+        Tuple[float, float],
+        Collection[Union[float, str, bool]],
+    ]
+
+
 class _TimeSeries:
     """Metadata about time series data for a particular run and tag.
 

tensorboard/plugins/hparams/backend_context.py

@@ -192,10 +192,10 @@ def hparams_from_data_provider(self, ctx, experiment_id):
             ctx, experiment_ids=[experiment_id]
         )
 
-    def session_groups_from_data_provider(self, ctx, experiment_id):
+    def session_groups_from_data_provider(self, ctx, experiment_id, filters):
         """Calls DataProvider.read_hyperparameters() and returns the result."""
         return self._tb_context.data_provider.read_hyperparameters(
-            ctx, experiment_ids=[experiment_id]
+            ctx, experiment_ids=[experiment_id], filters=filters
         )
 
     def _find_experiment_tag(self, hparams_run_to_tag_to_content):

tensorboard/plugins/hparams/list_session_groups.py

@@ -107,8 +107,11 @@ def _session_groups_from_tags(self):
 
     def _session_groups_from_data_provider(self):
         """Constructs lists of SessionGroups based on DataProvider results."""
+        filters = _build_data_provider_filters(self._request.col_params)
         response = self._backend_context.session_groups_from_data_provider(
-            self._request_context, self._experiment_id
+            self._request_context,
+            self._experiment_id,
+            filters,
         )
 
         session_groups = []
@@ -811,3 +814,45 @@ def _measurements(session_group, metric_name):
         if not metric_value:
             continue
         yield _Measurement(metric_value, session_index)
+
+
+def _build_data_provider_filters(col_params):
+    """Builds HyperparameterFilters from ColParams."""
+    filters = []
+    for col_param in col_params:
+        fltr = _build_data_provider_filter(col_param)
+        if fltr is None:
+            continue
+        filters.append(fltr)
+    return filters
+
+
+def _build_data_provider_filter(col_param):
+    """Builds HyperparameterFilter from ColParam.
+
+    Args:
+      col_param: ColParam that possibly contains filter information.
+
+    Returns:
+      None if col_param does not specify filter information.
+    """
+    if col_param.HasField("filter_regexp"):
+        filter_type = provider.HyperparameterFilterType.REGEX
+        fltr = col_param.filter_regexp
+    elif col_param.HasField("filter_interval"):
+        filter_type = provider.HyperparameterFilterType.INTERVAL
+        fltr = (
+            col_param.filter_interval.min_value,
+            col_param.filter_interval.max_value,
+        )
+    elif col_param.HasField("filter_discrete"):
+        filter_type = provider.HyperparameterFilterType.DISCRETE
+        fltr = [_value_to_python(b) for b in col_param.filter_discrete.values]
+    else:
+        return None
+
+    return provider.HyperparameterFilter(
+        hyperparameter_name=col_param.hparam,
+        filter_type=filter_type,
+        filter=fltr,
+    )

tensorboard/plugins/hparams/list_session_groups_test.py

@@ -363,9 +363,8 @@ def _mock_read_scalars(
 
     def _mock_read_hyperparameters(
         self,
-        ctx,
-        *,
-        experiment_ids,
+        *args,
+        **kwargs,
     ):
         return self._hyperparameters
 
@@ -1313,6 +1312,165 @@ def test_experiment_without_any_hparams(self):
         response = self._run_handler(request)
         self.assertProtoEquals("", response)
 
+    def test_experiment_from_data_provider_sends_regex_filter(self):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+              filter_regexp: 'v.*ue'
+            }
+        """
+        self._run_handler(request)
+        self.assertEquals(
+            self._get_read_hyperparameters_call_filters(),
+            [
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam1",
+                    filter_type=provider.HyperparameterFilterType.REGEX,
+                    filter="v.*ue",
+                )
+            ],
+        )
+
+    def test_experiment_from_data_provider_sends_interval_filter(self):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+              filter_interval: {
+                min_value: 0.1
+                max_value: 0.2
+              }
+            }
+            col_params: {
+              hparam: 'hparam2'
+              filter_interval: {
+                min_value: 0.1
+                max_value: Infinity
+              }
+            }
+            col_params: {
+              hparam: 'hparam3'
+              filter_interval: {
+                min_value: -Infinity
+                max_value: 0.2
+              }
+            }
+            col_params: {
+              hparam: 'hparam4'
+              filter_interval: {
+              }
+            }
+        """
+        self._run_handler(request)
+        self.assertEquals(
+            self._get_read_hyperparameters_call_filters(),
+            [
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam1",
+                    filter_type=provider.HyperparameterFilterType.INTERVAL,
+                    filter=(0.1, 0.2),
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam2",
+                    filter_type=provider.HyperparameterFilterType.INTERVAL,
+                    filter=(0.1, float("inf")),
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam3",
+                    filter_type=provider.HyperparameterFilterType.INTERVAL,
+                    filter=(float("-inf"), 0.2),
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam4",
+                    filter_type=provider.HyperparameterFilterType.INTERVAL,
+                    filter=(0.0, 0.0),
+                ),
+            ],
+        )
+
+    def test_experiment_from_data_provider_sends_discrete_filter(self):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+              filter_discrete: {
+                values: {
+                  bool_value: true
+                }
+                values: {
+                  bool_value: false
+                }
+              }
+            }
+            col_params: {
+              hparam: 'hparam2'
+              filter_discrete: {
+                values: {
+                  number_value: 2.0
+                }
+              }
+            }
+            col_params: {
+              hparam: 'hparam3'
+              filter_discrete: {
+                values: {
+                  string_value: '3_string'
+                }
+              }
+            }
+            col_params: {
+              hparam: 'hparam4'
+              filter_discrete: {
+                values: {
+                  string_value: '4_string'
+                }
+                values: {
+                  bool_value: true
+                }
+                values: {
+                  number_value: 4.0
+                }
+              }
+            }
+            col_params: {
+              hparam: 'hparam5'
+              filter_discrete: {}
+            }
+        """
+        self._run_handler(request)
+
+        self.assertEquals(
+            self._get_read_hyperparameters_call_filters(),
+            [
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam1",
+                    filter_type=provider.HyperparameterFilterType.DISCRETE,
+                    filter=[True, False],
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam2",
+                    filter_type=provider.HyperparameterFilterType.DISCRETE,
+                    filter=[2.0],
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam3",
+                    filter_type=provider.HyperparameterFilterType.DISCRETE,
+                    filter=["3_string"],
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam4",
+                    filter_type=provider.HyperparameterFilterType.DISCRETE,
+                    filter=["4_string", True, 4.0],
+                ),
+                provider.HyperparameterFilter(
+                    hyperparameter_name="hparam5",
+                    filter_type=provider.HyperparameterFilterType.DISCRETE,
+                    filter=[],
+                ),
+            ],
+        )
+
     def test_experiment_from_data_provider_with_no_sessions_or_hparam_values(
         self,
     ):
@@ -1749,6 +1907,12 @@ def _serialized_plugin_data(self, data_oneof_field, text_protobuffer):
         getattr(plugin_data, data_oneof_field).CopyFrom(protobuffer)
         return metadata.create_summary_metadata(plugin_data).plugin_data.content
 
+    def _get_read_hyperparameters_call_filters(self):
+        call_args = (
+            self._mock_tb_context.data_provider.read_hyperparameters.call_args
+        )
+        return call_args[1]["filters"]
+
 
 def _reduce_session_group_to_names(session_group):
     return [session.name for session in session_group.sessions]