color grouping: add group by regex under feature flag (#5117)

We want to bring up color group by regex gradually by adding a feature flag `enableColorGroupByRegex`
The feature flag is added in global (also shown in OSS) and we want to flip the flag in .corp only later (no concrete plan to bring color grouping in OSS yet)

Author: japie1235813

tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts

@@ -91,6 +91,13 @@ export const getEnabledColorGroup = createSelector(getFeatureFlags, (flags) => {
   return flags.enabledColorGroup;
 });
 
+export const getEnabledColorGroupByRegex = createSelector(
+  getFeatureFlags,
+  (flags) => {
+    return flags.enabledColorGroupByRegex;
+  }
+);
+
 export const getIsMetricsImageSupportEnabled = createSelector(
   getFeatureFlags,
   (flags) => {

tensorboard/webapp/feature_flag/store/feature_flag_selectors_test.ts

@@ -216,6 +216,31 @@ describe('feature_flag_selectors', () => {
     });
   });
 
+  describe('#getEnabledColorGroupByRegex', () => {
+    it('returns the proper value', () => {
+      let state = buildState(
+        buildFeatureFlagState({
+          defaultFlags: buildFeatureFlag({
+            enabledColorGroupByRegex: false,
+          }),
+        })
+      );
+      expect(selectors.getEnabledColorGroupByRegex(state)).toEqual(false);
+
+      state = buildState(
+        buildFeatureFlagState({
+          defaultFlags: buildFeatureFlag({
+            enabledColorGroupByRegex: false,
+          }),
+          flagOverrides: {
+            enabledColorGroupByRegex: true,
+          },
+        })
+      );
+      expect(selectors.getEnabledColorGroupByRegex(state)).toEqual(true);
+    });
+  });
+
   describe('#getIsMetricsImageSupportEnabled', () => {
     it('returns the proper value', () => {
       let state = buildState(

tensorboard/webapp/feature_flag/store/feature_flag_store_config_provider.ts

@@ -24,6 +24,7 @@ export const initialState: FeatureFlagState = {
     defaultEnableDarkMode: false,
     enableDarkModeOverride: null,
     enabledColorGroup: false,
+    enabledColorGroupByRegex: false,
     enabledExperimentalPlugins: [],
     inColab: false,
     scalarsBatchSize: undefined,

tensorboard/webapp/feature_flag/testing.ts

@@ -23,6 +23,7 @@ export function buildFeatureFlag(
     defaultEnableDarkMode: false,
     enableDarkModeOverride: null,
     enabledColorGroup: false,
+    enabledColorGroupByRegex: false,
     enabledExperimentalPlugins: [],
     inColab: false,
     scalarsBatchSize: undefined,

tensorboard/webapp/feature_flag/types.ts

@@ -27,6 +27,8 @@ export interface FeatureFlags {
   isAutoDarkModeAllowed: boolean;
   // Whether to enable experimental semantic color grouping feature.
   enabledColorGroup: boolean;
+  // Whether to enable color grouping by regex.
+  enabledColorGroupByRegex: boolean;
   // experimental plugins to manually enable.
   enabledExperimentalPlugins: string[];
   // Whether the TensorBoard is being run inside Colab output cell.

tensorboard/webapp/routes/dashboard_deeplink_provider.ts

@@ -30,6 +30,7 @@ import {GroupBy, GroupByKey} from '../runs/types';
 import * as selectors from '../selectors';
 import {
   ENABLE_COLOR_GROUP_QUERY_PARAM_KEY,
+  ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY,
   EXPERIMENTAL_PLUGIN_QUERY_PARAM_KEY,
 } from '../webapp_data_source/tb_feature_flag_data_source_types';
 import {
@@ -98,6 +99,14 @@ export class DashboardDeepLinkProvider extends DeepLinkProvider {
             value: String(overriddenFeatureFlags.enabledColorGroup),
           });
         }
+        if (
+          typeof overriddenFeatureFlags.enabledColorGroupByRegex === 'boolean'
+        ) {
+          queryParams.push({
+            key: ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY,
+            value: String(overriddenFeatureFlags.enabledColorGroupByRegex),
+          });
+        }
         return queryParams;
       })
     );

tensorboard/webapp/routes/dashboard_deeplink_provider_test.ts

@@ -347,6 +347,33 @@ describe('core deeplink provider', () => {
         []
       );
     });
+
+    it('serializes enableColorGroupByRegex state', () => {
+      store.overrideSelector(selectors.getOverriddenFeatureFlags, {
+        enabledColorGroupByRegex: true,
+      });
+      store.refreshState();
+
+      expect(queryParamsSerialized[queryParamsSerialized.length - 1]).toEqual([
+        {key: 'enableColorGroupByRegex', value: 'true'},
+      ]);
+
+      store.overrideSelector(selectors.getOverriddenFeatureFlags, {
+        enabledColorGroupByRegex: false,
+      });
+      store.refreshState();
+
+      expect(queryParamsSerialized[queryParamsSerialized.length - 1]).toEqual([
+        {key: 'enableColorGroupByRegex', value: 'false'},
+      ]);
+
+      store.overrideSelector(selectors.getOverriddenFeatureFlags, {});
+      store.refreshState();
+
+      expect(queryParamsSerialized[queryParamsSerialized.length - 1]).toEqual(
+        []
+      );
+    });
   });
 
   describe('runs', () => {

tensorboard/webapp/runs/views/runs_table/runs_group_menu_button_component.ng.html

@@ -51,10 +51,10 @@
     </span>
     <label>Run</label>
   </button>
-  <!--
   <button
     mat-menu-item
     role="menuitemradio"
+    *ngIf="showGroupByRegex"
     [attr.aria-checked]="selectedGroupBy.key === GroupByKey.REGEX"
     (click)="onGroupByChange.emit({key: GroupByKey.REGEX, regexString: ''})"
   >
@@ -66,5 +66,4 @@
     </span>
     <label>Regex</label>
   </button>
-  -->
 </mat-menu>

tensorboard/webapp/runs/views/runs_table/runs_group_menu_button_component.ts

@@ -31,8 +31,8 @@ import {GroupBy, GroupByKey} from '../../types';
 export class RunsGroupMenuButtonComponent {
   readonly GroupByKey = GroupByKey;
 
-  @Input()
-  selectedGroupBy!: GroupBy;
+  @Input() selectedGroupBy!: GroupBy;
+  @Input() showGroupByRegex!: boolean;
 
   @Output()
   onGroupByChange = new EventEmitter<GroupBy>();

tensorboard/webapp/runs/views/runs_table/runs_group_menu_button_container.ts

@@ -17,6 +17,7 @@ import {Store} from '@ngrx/store';
 import {Observable} from 'rxjs';
 
 import {State} from '../../../app_state';
+import {getEnabledColorGroupByRegex} from '../../../selectors';
 import {runGroupByChanged} from '../../actions';
 import {getRunGroupBy} from '../../store/runs_selectors';
 import {GroupBy} from '../../types';
@@ -29,12 +30,17 @@ import {GroupBy} from '../../types';
   template: `
     <runs-group-menu-button-component
       [selectedGroupBy]="selectedGroupBy$ | async"
+      [showGroupByRegex]="showGroupByRegex$ | async"
       (onGroupByChange)="onGroupByChange($event)"
     ></runs-group-menu-button-component>
   `,
   changeDetection: ChangeDetectionStrategy.OnPush,
 })
 export class RunsGroupMenuButtonContainer {
+  showGroupByRegex$: Observable<boolean> = this.store.select(
+    getEnabledColorGroupByRegex
+  );
+
   @Input() experimentIds!: string[];
 
   constructor(private readonly store: Store<State>) {}

tensorboard/webapp/runs/views/runs_table/runs_table_test.ts

@@ -54,6 +54,7 @@ import {DiscreteFilter, IntervalFilter} from '../../../hparams/types';
 import {
   getCurrentRouteRunSelection,
   getEnabledColorGroup,
+  getEnabledColorGroupByRegex,
   getExperiment,
   getExperimentIdToAliasMap,
   getRouteId,
@@ -249,6 +250,7 @@ describe('runs_table', () => {
     );
     store.overrideSelector(getRouteId, '123');
     store.overrideSelector(getEnabledColorGroup, false);
+    store.overrideSelector(getEnabledColorGroupByRegex, false);
     store.overrideSelector(getRunGroupBy, {key: GroupByKey.RUN});
     dispatchSpy = spyOn(store, 'dispatch').and.callFake((action: Action) => {
       actualActions.push(action);
@@ -559,9 +561,9 @@ describe('runs_table', () => {
         expect(menuButton).toBeTruthy();
       });
 
-      /* TODO(japie1235813): Brings back group by regex. */
-      it('renders "Experiment" and "Run"', () => {
+      it('renders "Experiment", "Run", and "Regex"', () => {
         store.overrideSelector(getEnabledColorGroup, true);
+        store.overrideSelector(getEnabledColorGroupByRegex, true);
         const fixture = createComponent(
           ['book'],
           [RunsTableColumn.RUN_NAME, RunsTableColumn.RUN_COLOR]
@@ -577,14 +579,15 @@ describe('runs_table', () => {
 
         expect(
           items.map((element) => element.querySelector('label')!.textContent)
-        ).toEqual(['Experiment', 'Run']);
+        ).toEqual(['Experiment', 'Run', 'Regex']);
       });
 
       it(
         'renders a check icon and aria-checked for the current groupBy menu ' +
           'item',
         () => {
           store.overrideSelector(getEnabledColorGroup, true);
+          store.overrideSelector(getEnabledColorGroupByRegex, true);
           store.overrideSelector(getRunGroupBy, {key: GroupByKey.EXPERIMENT});
           const fixture = createComponent(
             ['book'],
@@ -599,15 +602,13 @@ describe('runs_table', () => {
 
           const items = getOverlayMenuItems();
 
-          /* TODO(japie1235813): Brings back group by regex. */
           expect(
             items.map((element) => element.getAttribute('aria-checked'))
-          ).toEqual(['true', 'false']);
+          ).toEqual(['true', 'false', 'false']);
           expect(
             items.map((element) => Boolean(element.querySelector('mat-icon')))
-          ).toEqual([true, false]);
+          ).toEqual([true, false, false]);
 
-          /* TODO(japie1235813): Brings back group by regex.
           store.overrideSelector(getRunGroupBy, {
             key: GroupByKey.REGEX,
             regexString: 'hello',
@@ -621,12 +622,12 @@ describe('runs_table', () => {
           expect(
             items.map((element) => Boolean(element.querySelector('mat-icon')))
           ).toEqual([false, false, true]);
-          */
         }
       );
 
       it('dispatches `runGroupByChanged` when a menu item is clicked', () => {
         store.overrideSelector(getEnabledColorGroup, true);
+        store.overrideSelector(getEnabledColorGroupByRegex, true);
         store.overrideSelector(getRunGroupBy, {key: GroupByKey.EXPERIMENT});
         const fixture = createComponent(
           ['book'],
@@ -641,7 +642,7 @@ describe('runs_table', () => {
 
         const items = getOverlayMenuItems();
 
-        const [experiments, runs] = items as HTMLElement[];
+        const [experiments, runs, regex] = items as HTMLElement[];
         experiments.click();
 
         expect(dispatchSpy).toHaveBeenCalledWith(
@@ -659,17 +660,15 @@ describe('runs_table', () => {
           })
         );
 
-        /* TODO(japie1235813): Brings back group by regex.
         regex.click();
         expect(dispatchSpy).toHaveBeenCalledWith(
           runGroupByChanged({
             experimentIds: ['book'],
-            // regexString is hardcoded to '' for now; should be fixed when
-            // regex support is properly implemented.
+            // TODO(japie1235813): regexString is hardcoded to '' for now;
+            // should be fixed when regex support is properly implemented.
             groupBy: {key: GroupByKey.REGEX, regexString: ''},
           })
         );
-        */
       });
 
       it(

tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source.ts

@@ -16,6 +16,7 @@ import {Injectable} from '@angular/core';
 
 import {
   ENABLE_COLOR_GROUP_QUERY_PARAM_KEY,
+  ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY,
   ENABLE_DARK_MODE_QUERY_PARAM_KEY,
   EXPERIMENTAL_PLUGIN_QUERY_PARAM_KEY,
   SCALARS_BATCH_SIZE_PARAM_KEY,
@@ -68,6 +69,11 @@ export class QueryParamsFeatureFlagDataSource extends TBFeatureFlagDataSource {
         params.get(ENABLE_COLOR_GROUP_QUERY_PARAM_KEY) !== 'false';
     }
 
+    if (params.has(ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY)) {
+      featureFlags.enabledColorGroupByRegex =
+        params.get(ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY) !== 'false';
+    }
+
     if (params.has(ENABLE_DARK_MODE_QUERY_PARAM_KEY)) {
       featureFlags.defaultEnableDarkMode =
         params.get(ENABLE_DARK_MODE_QUERY_PARAM_KEY) !== 'false';

tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_test.ts

@@ -95,6 +95,27 @@ describe('tb_feature_flag_data_source', () => {
         expect(dataSource.getFeatures()).toEqual({enabledColorGroup: true});
       });
 
+      it('returns enableColorGroupByRegex from the query params', () => {
+        spyOn(TEST_ONLY.util, 'getParams').and.returnValues(
+          new URLSearchParams('enableColorGroupByRegex=false'),
+          new URLSearchParams('enableColorGroupByRegex='),
+          new URLSearchParams('enableColorGroupByRegex=true'),
+          new URLSearchParams('enableColorGroupByRegex=foo')
+        );
+        expect(dataSource.getFeatures()).toEqual({
+          enabledColorGroupByRegex: false,
+        });
+        expect(dataSource.getFeatures()).toEqual({
+          enabledColorGroupByRegex: true,
+        });
+        expect(dataSource.getFeatures()).toEqual({
+          enabledColorGroupByRegex: true,
+        });
+        expect(dataSource.getFeatures()).toEqual({
+          enabledColorGroupByRegex: true,
+        });
+      });
+
       it('returns all flag values when they are all set', () => {
         spyOn(TEST_ONLY.util, 'getParams').and.returnValue(
           new URLSearchParams(

tensorboard/webapp/webapp_data_source/tb_feature_flag_data_source_types.ts

@@ -37,4 +37,7 @@ export const SCALARS_BATCH_SIZE_PARAM_KEY = 'scalarsBatchSize';
 
 export const ENABLE_COLOR_GROUP_QUERY_PARAM_KEY = 'enableColorGroup';
 
+export const ENABLE_COLOR_GROUP_BY_REGEX_QUERY_PARAM_KEY =
+  'enableColorGroupByRegex';
+
 export const ENABLE_DARK_MODE_QUERY_PARAM_KEY = 'darkMode';