Hparams: Generate metric_infos from data provider session groups. (#6539)

In calls to the /experiment endpoint for hparams, we now generate
metric_infos based on the session_groups returned by the DataProvider.
We are able to reuse most of the logic that generates metric_infos for
the classic tensor-based hparams.

The end result is that metrics are listed in the left pane of the
hparams dashboard.

We have not yet integrated generating metric values into the
/session_groups operation so the values do not yet appear in the table
view.

Author: bmd3k

tensorboard/plugins/hparams/backend_context.py

@@ -98,7 +98,9 @@ def experiment_from_metadata(
             return experiment_from_runs
 
         experiment_from_data_provider_hparams = (
-            self._experiment_from_data_provider_hparams(data_provider_hparams)
+            self._experiment_from_data_provider_hparams(
+                ctx, experiment_id, data_provider_hparams
+            )
         )
         return (
             experiment_from_data_provider_hparams
@@ -224,7 +226,7 @@ def _compute_experiment_from_runs(
         """
         hparam_infos = self._compute_hparam_infos(hparams_run_to_tag_to_content)
         if hparam_infos:
-            metric_infos = self._compute_metric_infos(
+            metric_infos = self._compute_metric_infos_from_runs(
                 ctx, experiment_id, hparams_run_to_tag_to_content
             )
         else:
@@ -316,6 +318,8 @@ def _compute_hparam_info_from_values(self, name, values):
 
     def _experiment_from_data_provider_hparams(
         self,
+        ctx,
+        experiment_id,
         data_provider_hparams,
     ):
         """Returns an experiment protobuffer based on data provider hparams.
@@ -334,18 +338,24 @@ def _experiment_from_data_provider_hparams(
             # until all internal implementations of DataProvider can be
             # migrated to use new return value of provider.ListHyperparametersResult.
             hyperparameters = data_provider_hparams
+            session_groups = []
         else:
             # Is instance of provider.ListHyperparametersResult
             hyperparameters = data_provider_hparams.hyperparameters
-
-        if not hyperparameters:
-            return None
+            session_groups = data_provider_hparams.session_groups
 
         hparam_infos = [
             self._convert_data_provider_hparam(dp_hparam)
             for dp_hparam in hyperparameters
         ]
-        return api_pb2.Experiment(hparam_infos=hparam_infos)
+        metric_infos = (
+            self.compute_metric_infos_from_data_provider_session_groups(
+                ctx, experiment_id, session_groups
+            )
+        )
+        return api_pb2.Experiment(
+            hparam_infos=hparam_infos, metric_infos=metric_infos
+        )
 
     def _convert_data_provider_hparam(self, dp_hparam):
         """Builds an HParamInfo message from data provider Hyperparameter.
@@ -374,19 +384,37 @@ def _convert_data_provider_hparam(self, dp_hparam):
             hparam_info.domain_discrete.extend(dp_hparam.domain)
         return hparam_info
 
-    def _compute_metric_infos(
+    def _compute_metric_infos_from_runs(
         self, ctx, experiment_id, hparams_run_to_tag_to_content
     ):
+        session_runs = set(
+            run
+            for run, tags in hparams_run_to_tag_to_content.items()
+            if metadata.SESSION_START_INFO_TAG in tags
+        )
         return (
             api_pb2.MetricInfo(name=api_pb2.MetricName(group=group, tag=tag))
             for tag, group in self._compute_metric_names(
-                ctx, experiment_id, hparams_run_to_tag_to_content
+                ctx, experiment_id, session_runs
             )
         )
 
-    def _compute_metric_names(
-        self, ctx, experiment_id, hparams_run_to_tag_to_content
+    def compute_metric_infos_from_data_provider_session_groups(
+        self, ctx, experiment_id, session_groups
     ):
+        session_runs = set(
+            f"{s.experiment_id}/{s.run}"
+            for sg in session_groups
+            for s in sg.sessions
+        )
+        return [
+            api_pb2.MetricInfo(name=api_pb2.MetricName(group=group, tag=tag))
+            for tag, group in self._compute_metric_names(
+                ctx, experiment_id, session_runs
+            )
+        ]
+
+    def _compute_metric_names(self, ctx, experiment_id, session_runs):
         """Computes the list of metric names from all the scalar (run, tag)
         pairs.
 
@@ -412,11 +440,6 @@ def _compute_metric_names(
           A python list containing pairs. Each pair is a (tag, group) pair
           representing a metric name used in some session.
         """
-        session_runs = set(
-            run
-            for run, tags in hparams_run_to_tag_to_content.items()
-            if metadata.SESSION_START_INFO_TAG in tags
-        )
         metric_names_set = set()
         scalars_run_to_tag_to_content = self.scalars_metadata(
             ctx, experiment_id

tensorboard/plugins/hparams/backend_context_test.py

@@ -561,6 +561,59 @@ def test_experiment_from_data_provider_discrete_string_hparam(self):
         """
         self.assertProtoEquals(expected_exp, actual_exp)
 
+    def test_experiment_from_data_provider_session_group(self):
+        self._mock_tb_context.data_provider.list_tensors.side_effect = None
+        # The sessions chosen here mimic those returned in the implementation
+        # of _mock_list_tensors. These work nicely with the scalars returned
+        # in _mock_list_scalars and generate the same set of metric_infos as
+        # the tensor-based tests in this file.
+        self._hyperparameters = provider.ListHyperparametersResult(
+            hyperparameters=[],
+            session_groups=[
+                provider.HyperparameterSessionGroup(
+                    root=provider.HyperparameterSessionRun(
+                        experiment_id="exp", run=""
+                    ),
+                    sessions=[
+                        provider.HyperparameterSessionRun(
+                            experiment_id="exp", run="session_1"
+                        ),
+                        provider.HyperparameterSessionRun(
+                            experiment_id="exp", run="session_2"
+                        ),
+                    ],
+                    hyperparameter_values=[],
+                ),
+                provider.HyperparameterSessionGroup(
+                    root=provider.HyperparameterSessionRun(
+                        experiment_id="exp", run=""
+                    ),
+                    sessions=[
+                        provider.HyperparameterSessionRun(
+                            experiment_id="exp", run="session_3"
+                        ),
+                    ],
+                    hyperparameter_values=[],
+                ),
+            ],
+        )
+        actual_exp = self._experiment_from_metadata()
+        expected_exp = """
+            metric_infos: {
+              name: {group: '', tag: 'accuracy'}
+            }
+            metric_infos: {
+              name: {group: '', tag: 'loss'}
+            }
+            metric_infos: {
+              name: {group: 'eval', tag: 'loss'}
+            }
+            metric_infos: {
+              name: {group: 'train', tag: 'loss'}
+            }
+        """
+        self.assertProtoEquals(expected_exp, actual_exp)
+
     def test_experiment_from_data_provider_old_response_type(self):
         self._hyperparameters = [
             provider.Hyperparameter(