Hparams: Adds sort argument to DataProvider.read_hyperparameters(). (#6518)

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

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

Author: bmd3k

tensorboard/data/provider.py

@@ -396,7 +396,9 @@ def list_hyperparameters(self, ctx=None, *, experiment_ids):
         """
         return []
 
-    def read_hyperparameters(self, ctx=None, *, experiment_ids, filters=None):
+    def read_hyperparameters(
+        self, ctx=None, *, experiment_ids, filters=None, sort=None
+    ):
         """Read hyperparameter values.
 
         Args:
@@ -405,6 +407,8 @@ def read_hyperparameters(self, ctx=None, *, experiment_ids, filters=None):
             experiments.
           filters: A Collection[HyperparameterFilter] that constrain the
             returned session groups based on hyperparameter value.
+          sort: A Sequence[HyperparameterSort] that specify how the results
+            should be sorted.
 
         Returns:
           A Collection[HyperparameterSessionGroup] describing the groups and
@@ -709,6 +713,30 @@ class HyperparameterFilter:
     ]
 
 
+class HyperparameterSortDirection(enum.Enum):
+    """Describes which direction to sort a value."""
+
+    # Sort values ascending.
+    ASCENDING = "ascending"
+    # Sort values descending.
+    DESCENDING = "descending"
+
+
+@dataclasses.dataclass(frozen=True)
+class HyperparameterSort:
+    """A sort criterium based on hyperparameter value.
+
+    Attributes:
+      hyperparameter_name: A string identifier for the hyperparameter to use for
+        the sort. It corresponds to the hyperparameter_name field in the
+        Hyperparameter class.
+      sort_direction: The direction to sort.
+    """
+
+    hyperparameter_name: str
+    sort_direction: HyperparameterSortDirection
+
+
 class _TimeSeries:
     """Metadata about time series data for a particular run and tag.
 

tensorboard/plugins/hparams/backend_context.py

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

tensorboard/plugins/hparams/list_session_groups.py

@@ -108,10 +108,12 @@ 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)
+        sort = _build_data_provider_sort(self._request.col_params)
         response = self._backend_context.session_groups_from_data_provider(
             self._request_context,
             self._experiment_id,
             filters,
+            sort,
         )
 
         session_groups = []
@@ -856,3 +858,37 @@ def _build_data_provider_filter(col_param):
         filter_type=filter_type,
         filter=fltr,
     )
+
+
+def _build_data_provider_sort(col_params):
+    """Builds HyperparameterSorts from ColParams."""
+    sort = []
+    for col_param in col_params:
+        sort_item = _build_data_provider_sort_item(col_param)
+        if sort_item is None:
+            continue
+        sort.append(sort_item)
+    return sort
+
+
+def _build_data_provider_sort_item(col_param):
+    """Builds HyperparameterSort from ColParam.
+
+    Args:
+      col_param: ColParam that possibly contains sort information.
+
+    Returns:
+      None if col_param does not specify sort information.
+    """
+    if col_param.order == api_pb2.ORDER_UNSPECIFIED:
+        return None
+
+    sort_direction = (
+        provider.HyperparameterSortDirection.ASCENDING
+        if col_param.order == api_pb2.ORDER_ASC
+        else provider.HyperparameterSortDirection.DESCENDING
+    )
+    return provider.HyperparameterSort(
+        hyperparameter_name=col_param.hparam,
+        sort_direction=sort_direction,
+    )

tensorboard/plugins/hparams/list_session_groups_test.py

@@ -1312,6 +1312,26 @@ def test_experiment_without_any_hparams(self):
         response = self._run_handler(request)
         self.assertProtoEquals("", response)
 
+    def test_experiment_from_data_provider_sends_empty_filter_and_sort_from_col_params(
+        self,
+    ):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        # The request specifies col params but without filter or sort information.
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+            }
+            col_params: {
+              hparam: 'hparam2'
+            }
+        """
+        self._run_handler(request)
+        self.assertEquals(
+            self._get_read_hyperparameters_call_filters(),
+            [],
+        )
+        self.assertEquals(self._get_read_hyperparameters_call_sort(), [])
+
     def test_experiment_from_data_provider_sends_regex_filter(self):
         self._mock_tb_context.data_provider.list_tensors.side_effect = None
         request = """
@@ -1471,6 +1491,37 @@ def test_experiment_from_data_provider_sends_discrete_filter(self):
             ],
         )
 
+    def test_experiment_from_data_provider_sends_sort(self):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        request = """
+            col_params: {
+              hparam: 'hparam1'
+              order: ORDER_ASC
+            }
+            col_params: {
+              hparam: 'hparam2'
+              order: ORDER_UNSPECIFIED
+            }
+            col_params: {
+              hparam: 'hparam3'
+              order: ORDER_DESC
+            }
+        """
+        self._run_handler(request)
+        self.assertEquals(
+            self._get_read_hyperparameters_call_sort(),
+            [
+                provider.HyperparameterSort(
+                    hyperparameter_name="hparam1",
+                    sort_direction=provider.HyperparameterSortDirection.ASCENDING,
+                ),
+                provider.HyperparameterSort(
+                    hyperparameter_name="hparam3",
+                    sort_direction=provider.HyperparameterSortDirection.DESCENDING,
+                ),
+            ],
+        )
+
     def test_experiment_from_data_provider_with_no_sessions_or_hparam_values(
         self,
     ):
@@ -1913,6 +1964,12 @@ def _get_read_hyperparameters_call_filters(self):
         )
         return call_args[1]["filters"]
 
+    def _get_read_hyperparameters_call_sort(self):
+        call_args = (
+            self._mock_tb_context.data_provider.read_hyperparameters.call_args
+        )
+        return call_args[1]["sort"]
+
 
 def _reduce_session_group_to_names(session_group):
     return [session.name for session in session_group.sessions]