HParams: Port runs_data_source fetchHparamsMetadata to oss (#6318)

## Motivation for features / changes
As part of the effort to surface hparams data in the time series
dashboard we need to start actually fetching it. The logic existed
internally but for some reason was not previously available in OSS.

## Technical description of changes
I just copied the internal code to the oss implementation and made a few
small adjustments to imports

## Screenshots of UI changes
None

## Detailed steps to verify changes work correctly (as executed by you)
Tests should pass
Patch #6317, enable the feature flag, and ensure the hparams appear in
the runs table

## Alternate designs / implementations considered
Make a more detailed stub?

Author: rileyajones

tensorboard/webapp/runs/data_source/BUILD

@@ -48,7 +48,9 @@ tf_ts_library(
         "runs_data_source_test.ts",
     ],
     deps = [
+        ":backend_types",
         ":data_source",
+        ":testing",
         "//tensorboard/webapp/angular:expect_angular_core_testing",
         "//tensorboard/webapp/webapp_data_source:http_client_testing",
         "@npm//@types/jasmine",

tensorboard/webapp/runs/data_source/runs_data_source.ts

@@ -13,21 +13,71 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 import {Injectable} from '@angular/core';
-import {Observable, of} from 'rxjs';
-import {map} from 'rxjs/operators';
-import {TBHttpClient} from '../../webapp_data_source/tb_http_client';
+import {Observable, of, throwError} from 'rxjs';
+import {catchError, map, mergeMap} from 'rxjs/operators';
 import {
+  HttpErrorResponse,
+  TBHttpClient,
+} from '../../webapp_data_source/tb_http_client';
+import * as backendTypes from './runs_backend_types';
+import {
+  Domain,
+  DomainType,
   HparamsAndMetadata,
+  HparamSpec,
+  HparamValue,
+  MetricSpec,
   Run,
   RunsDataSource,
+  RunToHparamsAndMetrics,
 } from './runs_data_source_types';
 
+const HPARAMS_HTTP_PATH_PREFIX = 'data/plugin/hparams';
+
 type BackendGetRunsResponse = string[];
 
 function runToRunId(run: string, experimentId: string) {
   return `${experimentId}/${run}`;
 }
 
+function transformBackendHparamSpec(
+  hparamInfo: backendTypes.HparamSpec
+): HparamSpec {
+  let domain: Domain;
+  if (backendTypes.isDiscreteDomainHparamSpec(hparamInfo)) {
+    domain = {type: DomainType.DISCRETE, values: hparamInfo.domainDiscrete};
+  } else if (backendTypes.isIntervalDomainHparamSpec(hparamInfo)) {
+    domain = {...hparamInfo.domainInterval, type: DomainType.INTERVAL};
+  } else {
+    domain = {
+      type: DomainType.INTERVAL,
+      minValue: -Infinity,
+      maxValue: Infinity,
+    };
+  }
+  return {
+    description: hparamInfo.description,
+    displayName: hparamInfo.displayName,
+    name: hparamInfo.name,
+    type: hparamInfo.type,
+    domain,
+  };
+}
+
+function transformBackendMetricSpec(
+  metricInfo: backendTypes.MetricSpec
+): MetricSpec {
+  const {name, ...otherSpec} = metricInfo;
+  return {
+    ...otherSpec,
+    tag: name.tag,
+  };
+}
+
+declare interface GetExperimentHparamRequestPayload {
+  experimentName: string;
+}
+
 @Injectable()
 export class TBRunsDataSource implements RunsDataSource {
   constructor(private readonly http: TBHttpClient) {}
@@ -48,11 +98,112 @@ export class TBRunsDataSource implements RunsDataSource {
   }
 
   fetchHparamsMetadata(experimentId: string): Observable<HparamsAndMetadata> {
-    // Return a stub implementation.
-    return of({
-      hparamSpecs: [],
-      metricSpecs: [],
-      runToHparamsAndMetrics: {},
-    });
+    const requestPayload: GetExperimentHparamRequestPayload = {
+      experimentName: experimentId,
+    };
+    return this.http
+      .post<backendTypes.BackendHparamsExperimentResponse>(
+        `/experiment/${experimentId}/${HPARAMS_HTTP_PATH_PREFIX}/experiment`,
+        requestPayload
+      )
+      .pipe(
+        map((response) => {
+          const colParams: backendTypes.BackendListSessionGroupRequest['colParams'] =
+            [];
+
+          for (const hparamInfo of response.hparamInfos) {
+            colParams.push({hparam: hparamInfo.name});
+          }
+          for (const metricInfo of response.metricInfos) {
+            colParams.push({metric: metricInfo.name});
+          }
+
+          const listSessionRequestParams: backendTypes.BackendListSessionGroupRequest =
+            {
+              experimentName: experimentId,
+              allowedStatuses: [
+                backendTypes.RunStatus.STATUS_FAILURE,
+                backendTypes.RunStatus.STATUS_RUNNING,
+                backendTypes.RunStatus.STATUS_SUCCESS,
+                backendTypes.RunStatus.STATUS_UNKNOWN,
+              ],
+              colParams,
+              startIndex: 0,
+              // arbitrary large number so it does not get clipped.
+              sliceSize: 1e6,
+            };
+
+          return {
+            experimentHparamsInfo: response,
+            listSessionRequestParams,
+          };
+        }),
+        mergeMap(({experimentHparamsInfo, listSessionRequestParams}) => {
+          return this.http
+            .post<backendTypes.BackendListSessionGroupResponse>(
+              `/experiment/${experimentId}/${HPARAMS_HTTP_PATH_PREFIX}/session_groups`,
+              JSON.stringify(listSessionRequestParams)
+            )
+            .pipe(
+              map((sessionGroupsList) => {
+                return {experimentHparamsInfo, sessionGroupsList};
+              })
+            );
+        }),
+        map(({experimentHparamsInfo, sessionGroupsList}) => {
+          const runToHparamsAndMetrics: RunToHparamsAndMetrics = {};
+
+          // Reorganize the sessionGroup/session into run to <hparams,
+          // metrics>.
+          for (const sessionGroup of sessionGroupsList.sessionGroups) {
+            const hparams: HparamValue[] = Object.entries(
+              sessionGroup.hparams
+            ).map((keyValue) => {
+              const [hparam, value] = keyValue;
+              return {name: hparam, value};
+            });
+
+            for (const session of sessionGroup.sessions) {
+              for (const metricValue of session.metricValues) {
+                const runName = metricValue.name.group
+                  ? `${session.name}/${metricValue.name.group}`
+                  : session.name;
+                const runId = `${experimentId}/${runName}`;
+                const hparamsAndMetrics = runToHparamsAndMetrics[runId] || {
+                  metrics: [],
+                  hparams,
+                };
+                hparamsAndMetrics.metrics.push({
+                  tag: metricValue.name.tag,
+                  trainingStep: metricValue.trainingStep,
+                  value: metricValue.value,
+                });
+                runToHparamsAndMetrics[runId] = hparamsAndMetrics;
+              }
+            }
+          }
+          return {
+            hparamSpecs: experimentHparamsInfo.hparamInfos.map(
+              transformBackendHparamSpec
+            ),
+            metricSpecs: experimentHparamsInfo.metricInfos.map(
+              transformBackendMetricSpec
+            ),
+            runToHparamsAndMetrics,
+          };
+        }),
+        catchError((error) => {
+          // HParams plugin return 400 when there are no hparams for an
+          // experiment.
+          if (error instanceof HttpErrorResponse && error.status === 400) {
+            return of({
+              hparamSpecs: [],
+              metricSpecs: [],
+              runToHparamsAndMetrics: {},
+            });
+          }
+          return throwError(error);
+        })
+      );
   }
 }