HParams: Update tb_http_client to support the hparams plugin backend take 2 (#6340)

## 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.

I originally attempted this in #6318 but due to incompatibilities with
internal Colab had to revert it in #6336.
The issue was that Colab does not allow `POST` requests and our method
of mapping POST requests to GET requests didn't meet the expectations of
the hparams_plugin backend.

I've improved the tb_http_client in #6337 to now work with the
hparams_plugin backend so this should now be fine

## Technical description of changes
I mostly copied the internal code to the oss implementation and made a
few small adjustments to imports.
The biggest change is that I now utilize the `serializeUnder` parameter
of our custom POST function to ensure the requests meet the expectations
of the hparams backend

## Screenshots of UI changes
### enableHparamsInTimeSeries=false&tensorboardColab=false

![image](https://user-images.githubusercontent.com/78179109/233498859-df60aa0b-43d0-486d-8e4a-5be8fab7d08e.png)

### enableHparamsInTimeSeries=true&tensorboardColab=false

![image](https://user-images.githubusercontent.com/78179109/233498895-01e6e773-faf7-41de-9e34-199925a3e80f.png)

### enableHparamsInTimeSeries=false&tensorboardColab=true

![image](https://user-images.githubusercontent.com/78179109/233498942-fc414b81-256b-4762-837b-56148dc1dbcf.png)

### enableHparamsInTimeSeries=true&tensorboardColab=true

![image](https://user-images.githubusercontent.com/78179109/233498974-c9fd8555-db2a-49e4-9188-b7c19cfaaddf.png)



## Detailed steps to verify changes work correctly (as executed by you)
I ensured all of the above states result in the runs table being
displayed properly and that no errors are thrown by the hparams plugin
backend

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,116 @@ 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,
+        {},
+        'request'
+      )
+      .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`,
+              listSessionRequestParams,
+              {},
+              'request'
+            )
+            .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);
+        })
+      );
   }
 }