Hparams: Add include_in_result field and implement support for hparams_to_include (#6559)

This PR extracts the filtering logic for hparams only (metrics filtering
will be done later) from bmd3k's
[commit](bmd3k@ab75cd1).

Googlers, see comments in cl/559852837 for more context.

Test internally at cl/563163418.

#hparams

Author: yatbear

tensorboard/data/provider.py

@@ -397,7 +397,13 @@ def list_hyperparameters(self, ctx=None, *, experiment_ids):
         return ListHyperparametersResult(hyperparameters=[], session_groups=[])
 
     def read_hyperparameters(
-        self, ctx=None, *, experiment_ids, filters=None, sort=None
+        self,
+        ctx=None,
+        *,
+        experiment_ids,
+        filters=None,
+        sort=None,
+        hparams_to_include=None,
     ):
         """Read hyperparameter values.
 
@@ -409,6 +415,10 @@ def read_hyperparameters(
             returned session groups based on hyperparameter value.
           sort: A Sequence[HyperparameterSort] that specify how the results
             should be sorted.
+          hparams_to_include: An optional Collection[str] of the full names of
+            hyperparameters to include in the results. This collection will be
+            augmented to include all the hyperparameters specified in `filters`
+            and `sort`. If None, all hyperparameters will be returned.
 
         Returns:
           A Sequence[HyperparameterSessionGroup] describing the groups and

tensorboard/plugins/hparams/api.proto

@@ -322,9 +322,9 @@ message ListSessionGroupsRequest {
   optional bool include_metrics = 8;
 }
 
-// Defines parmeters for a ListSessionGroupsRequest for a specific column.
+// Defines parameters for a ListSessionGroupsRequest for a specific column.
 // See the comment for "ListSessionGroupsRequest" above for more details.
-// NEXT_TAG: 9
+// NEXT_TAG: 10
 message ColParams {
   oneof name {
     MetricName metric = 1;
@@ -373,6 +373,12 @@ message ColParams {
   // Specifies whether to exclude session groups whose column value is missing
   // from the response.
   bool exclude_missing_values = 8;
+
+  // Specifies whether to include the values for the field in the operation
+  // result. Defaults to True.
+  // Note: Hparams and metrics that do not appear in any `ColParams` in a
+  // request will also not be included in the result.
+  optional bool include_in_result = 9;
 }
 
 enum SortOrder {

tensorboard/plugins/hparams/backend_context.py

@@ -200,11 +200,15 @@ def hparams_from_data_provider(self, ctx, experiment_id):
         )
 
     def session_groups_from_data_provider(
-        self, ctx, experiment_id, filters, sort
+        self, ctx, experiment_id, filters, sort, hparams_to_include
     ):
         """Calls DataProvider.read_hyperparameters() and returns the result."""
         return self._tb_context.data_provider.read_hyperparameters(
-            ctx, experiment_ids=[experiment_id], filters=filters, sort=sort
+            ctx,
+            experiment_ids=[experiment_id],
+            filters=filters,
+            sort=sort,
+            hparams_to_include=hparams_to_include,
         )
 
     def _find_experiment_tag(

tensorboard/plugins/hparams/list_session_groups.py

@@ -114,17 +114,29 @@ def _session_groups_from_tags(self):
         )
         session_groups = self._filter(session_groups, filters)
         self._sort(session_groups, extractors)
+
+        if _specifies_include(self._request.col_params):
+            _reduce_to_hparams_to_include(
+                session_groups, self._request.col_params
+            )
+
         return session_groups
 
     def _session_groups_from_data_provider(self):
         """Constructs lists of SessionGroups based on DataProvider results."""
         filters = _build_data_provider_filters(self._request.col_params)
         sort = _build_data_provider_sort(self._request.col_params)
+        hparams_to_include = (
+            _get_hparams_to_include(self._request.col_params)
+            if _specifies_include(self._request.col_params)
+            else None
+        )
         response = self._backend_context.session_groups_from_data_provider(
             self._request_context,
             self._experiment_id,
             filters,
             sort,
+            hparams_to_include,
         )
 
         metric_infos = (
@@ -968,3 +980,62 @@ def _build_data_provider_sort_item(col_param):
         hyperparameter_name=col_param.hparam,
         sort_direction=sort_direction,
     )
+
+
+def _specifies_include(col_params):
+    """Determines whether any `ColParam` contains the `include_in_result` field.
+
+    In the case where none of the col_params contains the field, we should assume
+    that all fields should be included in the response.
+    """
+    return any(
+        col_param.HasField("include_in_result") for col_param in col_params
+    )
+
+
+def _get_hparams_to_include(col_params):
+    """Generates the list of hparams to include in the response.
+
+    The determination is based on the `include_in_result` field in ColParam. If
+    a ColParam either has `include_in_result: True` or does not specify the
+    field at all, then it should be included in the result.
+
+    Args:
+      col_params: A collection of `ColParams` protos.
+
+    Returns:
+      A list of names of hyperparameters to include in the response.
+    """
+    hparams_to_include = []
+    for col_param in col_params:
+        if (
+            col_param.HasField("include_in_result")
+            and not col_param.include_in_result
+        ):
+            # Explicitly set to exclude this hparam.
+            continue
+        if col_param.hparam:
+            hparams_to_include.append(col_param.hparam)
+    return hparams_to_include
+
+
+def _reduce_to_hparams_to_include(session_groups, col_params):
+    """Removes hparams from session_groups that should not be included.
+
+    Args:
+      session_groups: A collection of `SessionGroup` protos, which will be
+        modified in place.
+      col_params: A collection of `ColParams` protos.
+    """
+    hparams_to_include = _get_hparams_to_include(col_params)
+
+    for session_group in session_groups:
+        new_hparams = {
+            hparam: value
+            for (hparam, value) in session_group.hparams.items()
+            if hparam in hparams_to_include
+        }
+
+        session_group.ClearField("hparams")
+        for (hparam, value) in new_hparams.items():
+            session_group.hparams[hparam].CopyFrom(value)

tensorboard/plugins/hparams/list_session_groups_test.py

@@ -616,6 +616,43 @@ def test_some_allowed_statuses(self):
             ],
         )
 
+    def test_include_in_result(self):
+        request = """
+            start_index: 0
+            slice_size: 3
+            allowed_statuses: [
+              STATUS_UNKNOWN,
+              STATUS_SUCCESS,
+              STATUS_FAILURE,
+              STATUS_RUNNING
+            ]
+            aggregation_type: AGGREGATION_AVG
+            col_params {
+              hparam: "bool_hparam"
+              include_in_result: True
+            }
+            col_params {
+              hparam: "initial_temp"
+            }
+            col_params {
+              hparam: "string_hparam"
+              include_in_result: False
+            }
+        """
+        response = self._run_handler(request)
+
+        # Each of the session groups and sessions have two hparams and three metrics to include.
+        # Only check the first two session groups.
+        self.assertCountEqual(
+            response.session_groups[0].hparams, ["bool_hparam", "initial_temp"]
+        )
+        self.assertLen(response.session_groups[0].metric_values, 3)
+
+        self.assertCountEqual(
+            response.session_groups[0].hparams, ["bool_hparam", "initial_temp"]
+        )
+        self.assertLen(response.session_groups[0].metric_values, 3)
+
     def test_some_allowed_statuses_empty_groups(self):
         request = """
             start_index: 0
@@ -2405,6 +2442,37 @@ def test_experiment_from_data_provider_include_metrics(
                 response.session_groups[0].sessions[0].metric_values, 2
             )
 
+    def test_experiment_from_data_provider_sends_hparams_include_in_result(
+        self,
+    ):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+              include_in_result: True
+            }
+            col_params: {
+              hparam: 'hparam2'
+              include_in_result: False
+            }
+            col_params: {
+              hparam: 'hparam3'
+            }
+            col_params: {
+              hparam: 'hparam4'
+              include_in_result: True
+            }
+            col_params: {
+              metric: {tag: 'metric1'}
+              include_in_result: True
+            }
+        """
+        self._run_handler(request)
+        self.assertCountEqual(
+            self._get_read_hyperparameters_call_hparams_to_include(),
+            ["hparam1", "hparam3", "hparam4"],
+        )
+
     def _run_handler(self, request):
         request_proto = api_pb2.ListSessionGroupsRequest()
         text_format.Merge(request, request_proto)
@@ -2456,6 +2524,12 @@ def _get_read_hyperparameters_call_sort(self):
         )
         return call_args[1]["sort"]
 
+    def _get_read_hyperparameters_call_hparams_to_include(self):
+        call_args = (
+            self._mock_tb_context.data_provider.read_hyperparameters.call_args
+        )
+        return call_args[1]["hparams_to_include"]
+
 
 def _reduce_session_group_to_names(session_group):
     return [session.name for session in session_group.sessions]