color grouping: add group by regex under feature flag

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';