Group colours by regex of the experiment name. (#6846)

## Motivation for features / changes

Users want to color runs by the regex string of the corresponding to
runs experiment names.

## Technical description of changes

- Added new GroupBy type, REGEX_BY_EXP.
- Added a dropdown in dialog window, so users could select between regex
for run name or experiment name.

## Screenshots of UI changes (or N/A)

N/A since internal change. 

## Detailed steps to verify changes work correctly (as executed by you)

- Run tensorboard.corp server.
- Click on color grouping icon and select `Regex`.
- Select `Experiment Name` in dropdown.

## Alternate designs / implementations considered (or N/A)
N/A

Author: AnuarTB

tensorboard/webapp/runs/actions/runs_actions.ts

@@ -40,6 +40,7 @@ export const fetchRunsSucceeded = createAction(
     experimentIds: string[];
     runsForAllExperiments: Run[];
     newRuns: ExperimentIdToRuns;
+    expNameByExpId?: Record<string, string>;
   }>()
 );
 
@@ -79,7 +80,11 @@ export const runColorChanged = createAction(
 
 export const runGroupByChanged = createAction(
   '[Runs] Run Group By Changed',
-  props<{experimentIds: string[]; groupBy: GroupBy}>()
+  props<{
+    experimentIds: string[];
+    groupBy: GroupBy;
+    expNameByExpId?: Record<string, string>;
+  }>()
 );
 
 /**

tensorboard/webapp/runs/effects/runs_effects.ts

@@ -36,6 +36,7 @@ import {
   getExperimentIdsFromRoute,
   getRuns,
   getRunsLoadState,
+  getDashboardExperimentNames,
 } from '../../selectors';
 import {DataLoadState, LoadState} from '../../types/data';
 import * as actions from '../actions';
@@ -217,12 +218,15 @@ export class RunsEffects {
         }
         return {newRuns, runsForAllExperiments};
       }),
-      tap(({newRuns, runsForAllExperiments}) => {
+      withLatestFrom(this.store.select(getDashboardExperimentNames)),
+      tap(([runsData, expNameByExpId]) => {
+        const {newRuns, runsForAllExperiments} = runsData;
         this.store.dispatch(
           actions.fetchRunsSucceeded({
             experimentIds,
             newRuns,
             runsForAllExperiments,
+            expNameByExpId,
           })
         );
       }),

tensorboard/webapp/runs/effects/runs_effects_test.ts

@@ -29,6 +29,7 @@ import * as coreActions from '../../core/actions';
 import * as hparamsActions from '../../hparams/_redux/hparams_actions';
 import {
   getActiveRoute,
+  getDashboardExperimentNames,
   getExperimentIdsFromRoute,
   getRuns,
   getRunsLoadState,
@@ -131,6 +132,10 @@ describe('runs_effects', () => {
             buildCompareRoute(['exp1:123', 'exp2:456'])
           );
           store.overrideSelector(getExperimentIdsFromRoute, ['123', '456']);
+          store.overrideSelector(getDashboardExperimentNames, {
+            456: 'exp2',
+            123: 'exp1',
+          });
           const createFooRuns = () => [
             createRun({
               id: 'foo/runA',
@@ -166,6 +171,10 @@ describe('runs_effects', () => {
                   runs: createFooRuns(),
                 },
               },
+              expNameByExpId: {
+                456: 'exp2',
+                123: 'exp1',
+              },
             }),
           ]);
         });
@@ -211,6 +220,10 @@ describe('runs_effects', () => {
             buildCompareRoute(['exp1:123', ' exp2:456'])
           );
           store.overrideSelector(getExperimentIdsFromRoute, ['123', '456']);
+          store.overrideSelector(getDashboardExperimentNames, {
+            456: 'exp2',
+            123: 'exp1',
+          });
           store.refreshState();
 
           action.next(specAction());
@@ -244,6 +257,10 @@ describe('runs_effects', () => {
                   runs: createBarRuns(),
                 },
               },
+              expNameByExpId: {
+                456: 'exp2',
+                123: 'exp1',
+              },
             }),
           ]);
         });
@@ -310,6 +327,10 @@ describe('runs_effects', () => {
           buildCompareRoute(['exp1:123', ' exp2:456'])
         );
         store.overrideSelector(getExperimentIdsFromRoute, ['123', '456']);
+        store.overrideSelector(getDashboardExperimentNames, {
+          456: 'exp1',
+          123: 'exp2',
+        });
         store.refreshState();
 
         action.next(buildNavigatedAction());
@@ -328,13 +349,18 @@ describe('runs_effects', () => {
                 runs: createBarRuns(),
               },
             },
+            expNameByExpId: {
+              456: 'exp1',
+              123: 'exp2',
+            },
           }),
         ]);
       });
 
       it('ignores a navigation to same route and experiments (hash changes)', () => {
         store.overrideSelector(getActiveRoute, buildRoute());
         store.overrideSelector(getExperimentIdsFromRoute, ['123']);
+        store.overrideSelector(getDashboardExperimentNames, {123: 'exp1'});
         const createFooRuns = () => [
           createRun({
             id: 'foo/runA',
@@ -363,6 +389,7 @@ describe('runs_effects', () => {
             experimentIds: ['123'],
             runsForAllExperiments: [...createFooRuns()],
             newRuns: {},
+            expNameByExpId: {123: 'exp1'},
           }),
         ]);
 
@@ -391,6 +418,7 @@ describe('runs_effects', () => {
             })
           );
         store.overrideSelector(getExperimentIdsFromRoute, ['foo']);
+        store.overrideSelector(getDashboardExperimentNames, {foo: 'exp1'});
         store.refreshState();
 
         action.next(buildNavigatedAction());
@@ -404,6 +432,7 @@ describe('runs_effects', () => {
             experimentIds: ['foo'],
             runsForAllExperiments: [...createFooRuns()],
             newRuns: {},
+            expNameByExpId: {foo: 'exp1'},
           }),
         ]);
       });
@@ -457,6 +486,10 @@ describe('runs_effects', () => {
       // Emulate navigation to a new experiment route.
       store.overrideSelector(getActiveRoute, buildExperimentRouteFromId('456'));
       store.overrideSelector(getExperimentIdsFromRoute, ['456']);
+      store.overrideSelector(getDashboardExperimentNames, {
+        456: 'exp1',
+        123: 'exp2',
+      });
       // Force selectors to re-evaluate with a change in store.
       store.refreshState();
 
@@ -480,13 +513,15 @@ describe('runs_effects', () => {
           newRuns: {
             456: {runs: createBarRuns()},
           },
+          expNameByExpId: {456: 'exp1', 123: 'exp2'},
         }),
         actions.fetchRunsSucceeded({
           experimentIds: ['123'],
           runsForAllExperiments: createFooRuns(),
           newRuns: {
             123: {runs: createFooRuns()},
           },
+          expNameByExpId: {456: 'exp1', 123: 'exp2'},
         }),
       ]);
     });
@@ -552,6 +587,10 @@ describe('runs_effects', () => {
         });
 
         store.overrideSelector(getExperimentIdsFromRoute, ['foo']);
+        store.overrideSelector(getDashboardExperimentNames, {
+          foo: 'exp1',
+          bar: 'exp2',
+        });
         selectSpy
           .withArgs(getRuns, {experimentId: 'foo'})
           .and.returnValue(runsSubject);
@@ -611,6 +650,7 @@ describe('runs_effects', () => {
                 runs: createFooAfterRuns(),
               },
             },
+            expNameByExpId: {foo: 'exp1', bar: 'exp2'},
           }),
           actions.fetchRunsSucceeded({
             experimentIds: ['foo', 'bar'],
@@ -623,6 +663,7 @@ describe('runs_effects', () => {
                 runs: createBarRuns(),
               },
             },
+            expNameByExpId: {foo: 'exp1', bar: 'exp2'},
           }),
         ]);
       });

tensorboard/webapp/runs/store/BUILD

@@ -64,6 +64,7 @@ tf_ts_library(
     ],
     deps = [
         "//tensorboard/webapp/app_routing:namespaced_state_reducer_helper",
+        "//tensorboard/webapp/experiments:types",
         "//tensorboard/webapp/runs:types",
         "//tensorboard/webapp/runs/data_source",
         "//tensorboard/webapp/types",

tensorboard/webapp/runs/store/runs_reducers.ts

@@ -118,7 +118,8 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
     let {colorGroupRegexString, userSetGroupByKey} = state;
     if (groupBy) {
       const regexString =
-        groupBy.key === GroupByKey.REGEX
+        groupBy.key === GroupByKey.REGEX ||
+        groupBy.key === GroupByKey.REGEX_BY_EXP
           ? groupBy.regexString
           : state.colorGroupRegexString;
       colorGroupRegexString = regexString;
@@ -203,48 +204,56 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
     }
     return {...state, runsLoadState: nextRunsLoadState};
   }),
-  on(runsActions.fetchRunsSucceeded, (state, {runsForAllExperiments}) => {
-    const groupKeyToColorId = new Map(state.groupKeyToColorId);
-    const defaultRunColorIdForGroupBy = new Map(
-      state.defaultRunColorIdForGroupBy
-    );
+  on(
+    runsActions.fetchRunsSucceeded,
+    (state, {runsForAllExperiments, expNameByExpId}) => {
+      const groupKeyToColorId = new Map(state.groupKeyToColorId);
+      const defaultRunColorIdForGroupBy = new Map(
+        state.defaultRunColorIdForGroupBy
+      );
 
-    let groupBy = state.initialGroupBy;
-    if (state.userSetGroupByKey !== null) {
-      groupBy = createGroupBy(
-        state.userSetGroupByKey,
-        state.colorGroupRegexString
+      let groupBy = state.initialGroupBy;
+      if (state.userSetGroupByKey !== null) {
+        groupBy = createGroupBy(
+          state.userSetGroupByKey,
+          state.colorGroupRegexString
+        );
+      }
+      const groups = groupRuns(
+        groupBy,
+        runsForAllExperiments,
+        state.runIdToExpId,
+        expNameByExpId
       );
-    }
-    const groups = groupRuns(
-      groupBy,
-      runsForAllExperiments,
-      state.runIdToExpId
-    );
 
-    Object.entries(groups.matches).forEach(([groupId, runs]) => {
-      const colorId = groupKeyToColorId.get(groupId) ?? groupKeyToColorId.size;
-      groupKeyToColorId.set(groupId, colorId);
+      Object.entries(groups.matches).forEach(([groupId, runs]) => {
+        const colorId =
+          groupKeyToColorId.get(groupId) ?? groupKeyToColorId.size;
+        groupKeyToColorId.set(groupId, colorId);
 
-      for (const run of runs) {
-        defaultRunColorIdForGroupBy.set(run.id, colorId);
+        for (const run of runs) {
+          defaultRunColorIdForGroupBy.set(run.id, colorId);
+        }
+      });
+
+      // unassign color for nonmatched runs to apply default unassigned style
+      for (const run of groups.nonMatches) {
+        defaultRunColorIdForGroupBy.set(run.id, -1);
       }
-    });
 
-    // unassign color for nonmatched runs to apply default unassigned style
-    for (const run of groups.nonMatches) {
-      defaultRunColorIdForGroupBy.set(run.id, -1);
+      return {
+        ...state,
+        defaultRunColorIdForGroupBy,
+        groupKeyToColorId,
+      };
     }
-
-    return {
-      ...state,
-      defaultRunColorIdForGroupBy,
-      groupKeyToColorId,
-    };
-  }),
+  ),
   on(
     runsActions.runGroupByChanged,
-    (state: RunsDataState, {experimentIds, groupBy}): RunsDataState => {
+    (
+      state: RunsDataState,
+      {experimentIds, groupBy, expNameByExpId}
+    ): RunsDataState => {
       // Reset the groupKeyToColorId
       const groupKeyToColorId = new Map<string, number>();
       const defaultRunColorIdForGroupBy = new Map(
@@ -255,7 +264,12 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
         .flatMap((experimentId) => state.runIds[experimentId])
         .map((runId) => state.runMetadata[runId]);
 
-      const groups = groupRuns(groupBy, allRuns, state.runIdToExpId);
+      const groups = groupRuns(
+        groupBy,
+        allRuns,
+        state.runIdToExpId,
+        expNameByExpId
+      );
 
       Object.entries(groups.matches).forEach(([groupId, runs]) => {
         const colorId =
@@ -273,7 +287,8 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
       }
 
       const updatedRegexString =
-        groupBy.key === GroupByKey.REGEX
+        groupBy.key === GroupByKey.REGEX ||
+        groupBy.key === GroupByKey.REGEX_BY_EXP
           ? groupBy.regexString
           : state.colorGroupRegexString;