Hparams: Limit number of Hparams retrieved by default. (#6777)

Limit the number of hparams retrieved by default to 1000 for display in
Time Series. Print retrieval information: # of columns loaded and
whether there could be more.

This is to address performance issues with loading 1000s of hparams from
the backend and trying to render them all.

In a future PR we will add ability to "Load All" if user wants to
nonetheless load them all.

Screenshot (with default set to 2 instead of 1000):

![image](https://github.com/tensorflow/tensorboard/assets/17152369/b0f2c4e6-bda0-4ed7-bead-9f1ef6eaaca7)

Author: bmd3k

tensorboard/webapp/hparams/_redux/BUILD

@@ -111,6 +111,7 @@ tf_ng_module(
     deps = [
         ":hparams_actions",
         ":hparams_data_source",
+        ":hparams_selectors",
         ":types",
         "//tensorboard/webapp:app_state",
         "//tensorboard/webapp/app_routing:types",

tensorboard/webapp/hparams/_redux/hparams_data_source.ts

@@ -87,11 +87,15 @@ export class HparamsDataSource {
     return experimentIds.map((eid, index) => `${index}:${eid}`).join(',');
   }
 
-  fetchExperimentInfo(experimentIds: string[]): Observable<HparamSpec[]> {
+  fetchExperimentInfo(
+    experimentIds: string[],
+    hparamsLimit: number
+  ): Observable<HparamSpec[]> {
     const formattedExperimentIds = this.formatExperimentIds(experimentIds);
 
     const experimentRequest: BackendHparamsExperimentRequest = {
       experimentName: formattedExperimentIds,
+      hparamsLimit,
       // The hparams feature generates its own metric data and does not require
       // the backend to calculate it.
       includeMetrics: false,
@@ -132,7 +136,7 @@ export class HparamsDataSource {
     const colParams: BackendListSessionGroupRequest['colParams'] = [];
 
     for (const hparamSpec of hparamSpecs) {
-      colParams.push({hparam: hparamSpec.name});
+      colParams.push({hparam: hparamSpec.name, includeInResult: true});
     }
 
     const listSessionRequestParams: BackendListSessionGroupRequest = {

tensorboard/webapp/hparams/_redux/hparams_data_source_test.ts

@@ -28,6 +28,7 @@ import {
   BackendListSessionGroupResponse,
   RunStatus,
 } from '../types';
+import {buildHparamSpec} from './testing';
 
 describe('HparamsDataSource Test', () => {
   let httpMock: HttpTestingController;
@@ -46,7 +47,7 @@ describe('HparamsDataSource Test', () => {
   describe('fetchExperimentInfo', () => {
     it('uses /experiment when a single experiment id is provided', () => {
       const returnValue = jasmine.createSpy();
-      dataSource.fetchExperimentInfo(['eid']).subscribe(returnValue);
+      dataSource.fetchExperimentInfo(['eid'], 0).subscribe(returnValue);
       httpMock
         .expectOne('/experiment/eid/data/plugin/hparams/experiment')
         .flush(createHparamsExperimentResponse());
@@ -55,7 +56,9 @@ describe('HparamsDataSource Test', () => {
 
     it('uses /compare when a multiple experiment ids are provided', () => {
       const returnValue = jasmine.createSpy();
-      dataSource.fetchExperimentInfo(['eid1', 'eid2']).subscribe(returnValue);
+      dataSource
+        .fetchExperimentInfo(['eid1', 'eid2'], 0)
+        .subscribe(returnValue);
       httpMock
         .expectOne('/compare/0:eid1,1:eid2/data/plugin/hparams/experiment')
         .flush(createHparamsExperimentResponse());
@@ -64,7 +67,7 @@ describe('HparamsDataSource Test', () => {
 
     it('maps interval and discrete domains to domain', () => {
       const returnValue = jasmine.createSpy();
-      dataSource.fetchExperimentInfo(['eid']).subscribe(returnValue);
+      dataSource.fetchExperimentInfo(['eid'], 0).subscribe(returnValue);
       httpMock
         .expectOne('/experiment/eid/data/plugin/hparams/experiment')
         .flush(createHparamsExperimentResponse());
@@ -95,7 +98,7 @@ describe('HparamsDataSource Test', () => {
 
     it('treats missing domains as discrete domains', () => {
       const returnValue = jasmine.createSpy();
-      dataSource.fetchExperimentInfo(['eid']).subscribe(returnValue);
+      dataSource.fetchExperimentInfo(['eid'], 0).subscribe(returnValue);
       httpMock
         .expectOne('/experiment/eid/data/plugin/hparams/experiment')
         .flush(createHparamsExperimentNoDomainResponse());
@@ -122,6 +125,15 @@ describe('HparamsDataSource Test', () => {
         },
       ]);
     });
+
+    it('sets hparamsLimit and includeMetrics', () => {
+      dataSource.fetchExperimentInfo(['eid'], 100).subscribe();
+      const actual = httpMock.expectOne(
+        '/experiment/eid/data/plugin/hparams/experiment'
+      );
+      expect(actual.request.body.hparamsLimit).toEqual(100);
+      expect(actual.request.body.includeMetrics).toBeFalse();
+    });
   });
 
   describe('fetchSessionGroups', () => {
@@ -176,6 +188,25 @@ describe('HparamsDataSource Test', () => {
       expect(sessionGroups[0].sessions[0].name).toEqual('eid1/run_name_1');
       expect(sessionGroups[1].sessions[0].name).toEqual('eid2/run_name_2');
     });
+
+    it('adds hparams as colParams', () => {
+      dataSource
+        .fetchSessionGroups(
+          ['eid'],
+          [
+            buildHparamSpec({name: 'hparam1'}),
+            buildHparamSpec({name: 'hparam2'}),
+          ]
+        )
+        .subscribe();
+      const actual = httpMock.expectOne(
+        '/experiment/eid/data/plugin/hparams/session_groups'
+      );
+      expect(actual.request.body.colParams).toEqual([
+        {hparam: 'hparam1', includeInResult: true},
+        {hparam: 'hparam2', includeInResult: true},
+      ]);
+    });
   });
 });
 

tensorboard/webapp/hparams/_redux/hparams_effects.ts

@@ -39,6 +39,7 @@ import * as hparamsActions from './hparams_actions';
 import {HparamsDataSource} from './hparams_data_source';
 import {HparamSpec, SessionGroup} from '../types';
 import {RouteKind} from '../../app_routing/types';
+import {getNumDashboardHparamsToLoad} from './hparams_selectors';
 
 /**
  * Effects for fetching the hparams data from the backend.
@@ -70,40 +71,51 @@ export class HparamsEffects {
   /** @export */
   loadHparamsData$ = createEffect(() => {
     return merge(this.navigated$, this.loadHparamsOnReload$).pipe(
-      withLatestFrom(this.store.select(getActiveRoute)),
+      withLatestFrom(
+        this.store.select(getActiveRoute),
+        this.store.select(getNumDashboardHparamsToLoad)
+      ),
       filter(
         ([, activeRoute]) =>
           activeRoute?.routeKind === RouteKind.EXPERIMENT ||
           activeRoute?.routeKind === RouteKind.COMPARE_EXPERIMENT
       ),
       throttleTime(10),
-      switchMap(([experimentIds]) =>
-        this.loadHparamsForExperiments(experimentIds)
+      switchMap(([experimentIds, , numHparamsToLoad]) =>
+        this.loadHparamsForExperiments(experimentIds, numHparamsToLoad)
       ),
       map((resp) => hparamsActions.hparamsFetchSessionGroupsSucceeded(resp))
     );
   });
 
-  private loadHparamsForExperiments(experimentIds: string[]): Observable<{
+  private loadHparamsForExperiments(
+    experimentIds: string[],
+    hparamsLimit: number
+  ): Observable<{
     hparamSpecs: HparamSpec[];
     sessionGroups: SessionGroup[];
   }> {
-    return this.dataSource.fetchExperimentInfo(experimentIds).pipe(
-      switchMap((hparamSpecs) => {
-        return this.dataSource
-          .fetchSessionGroups(experimentIds, hparamSpecs)
-          .pipe(
-            catchError((error) => {
-              // HParam plugin return 400 when there are no hparams
-              // for an experiment.
-              if (error instanceof HttpErrorResponse && error.status === 400) {
-                return of([] as SessionGroup[]);
-              }
-              return throwError(() => error);
-            }),
-            map((sessionGroups) => ({hparamSpecs, sessionGroups}))
-          );
-      })
-    );
+    return this.dataSource
+      .fetchExperimentInfo(experimentIds, hparamsLimit)
+      .pipe(
+        switchMap((hparamSpecs) => {
+          return this.dataSource
+            .fetchSessionGroups(experimentIds, hparamSpecs)
+            .pipe(
+              catchError((error) => {
+                // HParam plugin return 400 when there are no hparams
+                // for an experiment.
+                if (
+                  error instanceof HttpErrorResponse &&
+                  error.status === 400
+                ) {
+                  return of([] as SessionGroup[]);
+                }
+                return throwError(() => error);
+              }),
+              map((sessionGroups) => ({hparamSpecs, sessionGroups}))
+            );
+        })
+      );
   }
 }

tensorboard/webapp/hparams/_redux/hparams_effects_test.ts

@@ -27,6 +27,7 @@ import {buildHparamSpec} from './testing';
 import {HparamSpec, RunStatus, SessionGroup} from '../types';
 import * as selectors from '../../selectors';
 import * as hparamsActions from './hparams_actions';
+import * as hparamsSelectors from './hparams_selectors';
 import * as appRoutingActions from '../../app_routing/actions';
 import * as coreActions from '../../core/actions';
 import {RouteKind} from '../../app_routing/types';
@@ -112,6 +113,10 @@ describe('hparams effects', () => {
       store.overrideSelector(selectors.getExperimentIdsFromRoute, [
         'expFromRoute',
       ]);
+      store.overrideSelector(
+        hparamsSelectors.getNumDashboardHparamsToLoad,
+        1111
+      );
       store.refreshState();
     });
 
@@ -130,9 +135,10 @@ describe('hparams effects', () => {
 
     it('fetches data after navigation', () => {
       action.next(appRoutingActions.navigated({} as any));
-      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith([
-        'expFromRoute',
-      ]);
+      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith(
+        ['expFromRoute'],
+        1111
+      );
       expect(dataSource.fetchSessionGroups).toHaveBeenCalledWith(
         ['expFromRoute'],
         [buildHparamSpec({name: 'h1'})]
@@ -158,9 +164,10 @@ describe('hparams effects', () => {
 
     it('fetches data on reload', () => {
       action.next(coreActions.reload());
-      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith([
-        'expFromRoute',
-      ]);
+      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith(
+        ['expFromRoute'],
+        1111
+      );
       expect(dataSource.fetchSessionGroups).toHaveBeenCalledWith(
         ['expFromRoute'],
         [buildHparamSpec({name: 'h1'})]
@@ -175,9 +182,10 @@ describe('hparams effects', () => {
 
     it('fetches data on manualReload', () => {
       action.next(coreActions.manualReload());
-      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith([
-        'expFromRoute',
-      ]);
+      expect(dataSource.fetchExperimentInfo).toHaveBeenCalledWith(
+        ['expFromRoute'],
+        1111
+      );
       expect(dataSource.fetchSessionGroups).toHaveBeenCalledWith(
         ['expFromRoute'],
         [buildHparamSpec({name: 'h1'})]

tensorboard/webapp/hparams/_redux/hparams_reducers.ts

@@ -19,6 +19,8 @@ import {Side} from '../../widgets/data_table/types';
 import * as actions from './hparams_actions';
 import {HparamsState} from './types';
 
+const INITIAL_NUM_HPARAMS_TO_LOAD = 1000;
+
 const initialState: HparamsState = {
   dashboardHparamSpecs: [],
   dashboardSessionGroups: [],
@@ -27,6 +29,8 @@ const initialState: HparamsState = {
     metrics: new Map(),
   },
   dashboardDisplayedHparamColumns: [],
+  numDashboardHparamsToLoad: INITIAL_NUM_HPARAMS_TO_LOAD,
+  numDashboardHparamsLoaded: 0,
 };
 
 const reducer: ActionReducer<HparamsState, Action> = createReducer(
@@ -49,6 +53,7 @@ const reducer: ActionReducer<HparamsState, Action> = createReducer(
       ...state,
       dashboardHparamSpecs: nextDashboardHparamSpecs,
       dashboardSessionGroups: nextDashboardSessionGroups,
+      numDashboardHparamsLoaded: nextDashboardHparamSpecs.length,
     };
   }),
   on(actions.dashboardHparamFilterAdded, (state, action) => {

tensorboard/webapp/hparams/_redux/hparams_reducers_test.ts

@@ -144,6 +144,25 @@ describe('hparams/_redux/hparams_reducers_test', () => {
 
       expect(state2.dashboardSessionGroups).toEqual([mockSessionGroup]);
     });
+
+    it('calculates numDashboardHparamsLoaded from action.hparamSpecs', () => {
+      const state = buildHparamsState({
+        numDashboardHparamsLoaded: 0,
+      });
+      const state2 = reducers(
+        state,
+        actions.hparamsFetchSessionGroupsSucceeded({
+          hparamSpecs: [
+            buildHparamSpec(),
+            buildHparamSpec(),
+            buildHparamSpec(),
+          ],
+          sessionGroups: [],
+        })
+      );
+
+      expect(state2.numDashboardHparamsLoaded).toEqual(3);
+    });
   });
 
   describe('dashboardHparamFilterAdded', () => {

tensorboard/webapp/hparams/_redux/hparams_selectors.ts

@@ -71,3 +71,17 @@ export const getDashboardMetricsFilterMap = createSelector(
     return state.dashboardFilters.metrics;
   }
 );
+
+export const getNumDashboardHparamsToLoad = createSelector(
+  getHparamsState,
+  (state) => {
+    return state.numDashboardHparamsToLoad;
+  }
+);
+
+export const getNumDashboardHparamsLoaded = createSelector(
+  getHparamsState,
+  (state) => {
+    return state.numDashboardHparamsLoaded;
+  }
+);

tensorboard/webapp/hparams/_redux/hparams_selectors_test.ts

@@ -163,4 +163,28 @@ describe('hparams/_redux/hparams_selectors_test', () => {
       ]);
     });
   });
+
+  describe('#getNumDashboardHparamsToLoad', () => {
+    it('returns dashboard specs', () => {
+      const state = buildStateFromHparamsState(
+        buildHparamsState({
+          numDashboardHparamsToLoad: 5,
+        })
+      );
+
+      expect(selectors.getNumDashboardHparamsToLoad(state)).toEqual(5);
+    });
+  });
+
+  describe('#getNumDashboardHparamsToLoad', () => {
+    it('returns dashboard specs', () => {
+      const state = buildStateFromHparamsState(
+        buildHparamsState({
+          numDashboardHparamsLoaded: 22,
+        })
+      );
+
+      expect(selectors.getNumDashboardHparamsLoaded(state)).toEqual(22);
+    });
+  });
 });

tensorboard/webapp/hparams/_redux/testing.ts

@@ -39,6 +39,8 @@ export function buildHparamsState(
     },
     dashboardDisplayedHparamColumns:
       overrides.dashboardDisplayedHparamColumns ?? [],
+    numDashboardHparamsLoaded: overrides.numDashboardHparamsLoaded ?? 0,
+    numDashboardHparamsToLoad: overrides.numDashboardHparamsToLoad ?? 0,
   } as HparamsState;
 }