Support for custom group colors

tensorboard/webapp/runs/actions/runs_actions.ts

@@ -78,6 +78,11 @@ export const runColorChanged = createAction(
   props<{runId: string; newColor: string}>()
 );
 
+export const groupColorChanged = createAction(
+  '[Runs] Group Color Changed',
+  props<{groupId: string; newColor: string}>()
+);
+
 export const runGroupByChanged = createAction(
   '[Runs] Run Group By Changed',
   props<{

tensorboard/webapp/runs/store/runs_reducers.ts

@@ -49,6 +49,7 @@ const {
 >(
   {
     runColorOverrideForGroupBy: new Map(),
+    groupColorOverride: new Map(),
     defaultRunColorIdForGroupBy: new Map(),
     groupKeyToColorId: new Map(),
     initialGroupBy: {key: GroupByKey.RUN},
@@ -271,13 +272,19 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
         expNameByExpId
       );
 
+      const colorOverrideFromGroup = state.groupColorOverride;
       Object.entries(groups.matches).forEach(([groupId, runs]) => {
         const colorId =
           groupKeyToColorId.get(groupId) ?? groupKeyToColorId.size;
         groupKeyToColorId.set(groupId, colorId);
 
+        const colorOverride = colorOverrideFromGroup.get(groupId);
         for (const run of runs) {
           defaultRunColorIdForGroupBy.set(run.id, colorId);
+
+          if (colorOverride) {
+            colorOverrideFromGroup.set(run.id, colorOverride);
+          }
         }
       });
 
@@ -298,8 +305,8 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
         userSetGroupByKey: groupBy.key,
         defaultRunColorIdForGroupBy,
         groupKeyToColorId,
-        // Resets the color override when the groupBy changes.
-        runColorOverrideForGroupBy: new Map(),
+        // Use group color override if set
+        runColorOverrideForGroupBy: colorOverrideFromGroup,
       };
     }
   ),
@@ -309,6 +316,12 @@ const dataReducer: ActionReducer<RunsDataState, Action> = createReducer(
 
     return {...state, runColorOverrideForGroupBy: nextRunColorOverride};
   }),
+  on(runsActions.groupColorChanged, (state, {groupId, newColor}) => {
+    const nextGroupColorOverride = new Map(state.groupColorOverride);
+    nextGroupColorOverride.set(groupId, newColor);
+
+    return {...state, groupColorOverride: nextGroupColorOverride};
+  }),
   on(runsActions.runSelectorRegexFilterChanged, (state, action) => {
     return {
       ...state,

tensorboard/webapp/runs/store/runs_selectors.ts

@@ -300,6 +300,13 @@ export const getRunColorOverride = createSelector(
   }
 );
 
+export const getGroupColorOverride = createSelector(
+  getDataState,
+  (state: RunsDataState): Map<string, string> => {
+    return state.groupColorOverride;
+  }
+);
+
 export const getDefaultRunColorIdMap = createSelector(
   getDataState,
   (state: RunsDataState): Map<string, number> => {

tensorboard/webapp/runs/store/runs_types.ts

@@ -51,6 +51,7 @@ export interface RunsDataNamespacedState {
   groupKeyToColorId: Map<string, number>;
   // Hex color string user has picked for a run.
   runColorOverrideForGroupBy: Map<RunId, string>;
+  groupColorOverride: Map<string, string>;
   initialGroupBy: GroupBy;
   userSetGroupByKey: GroupByKey | null;
   colorGroupRegexString: string;

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

@@ -65,19 +65,26 @@ <h4>Color group preview</h4>
       <div *ngIf="colorRunPairList.length; else empty" class="match-container">
         <ul
           class="group"
-          *ngFor="let colorRunsPair of colorRunPairList"
+          *ngFor="let colorRunsPair of colorRunPairList; trackBy: trackByGroups"
           [ngStyle]="{borderColor: colorRunsPair.color}"
         >
           <li>
-            <label
-              ><span
+            <label class="color-container">
+              <button
                 class="color-swatch"
-                [ngStyle]="{backgroundColor: colorRunsPair.color}"
-              ></span
-              ><code class="group-id" [title]="colorRunsPair.groupId"
-                >{{colorRunsPair.groupId}}</code
-              ></label
-            >
+                [style.background]="colorRunsPair.color"
+                [colorPicker]="colorRunsPair.color"
+                [cpDialogDisplay]="'popup'"
+                [cpPositionOffset]="-20"
+                [cpUseRootViewContainer]="true"
+                [cpOutputFormat]="'hex'"
+                (colorPickerChange)="groupColorChange.emit({groupId: colorRunsPair.groupId, newColor: $event})"
+              ></button>
+
+              <code class="group-id" [title]="colorRunsPair.groupId">
+                {{ colorRunsPair.groupId }}
+              </code>
+            </label>
             <ul>
               <ng-container *ngFor="let run of colorRunsPair.runs | slice:0:5;">
                 <li [title]="run.name">{{ run.name }}</li>

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

@@ -48,6 +48,10 @@ export class RegexEditDialogComponent {
   @Output() onSave = new EventEmitter();
   @Output() regexInputOnChange = new EventEmitter<string>();
   @Output() regexTypeOnChange = new EventEmitter<GroupByKey>();
+  @Output() groupColorChange = new EventEmitter<{
+    groupId: string;
+    newColor: string;
+  }>();
 
   // Constants
   REGEX_BY_RUN_STR = 'regex_by_run';
@@ -106,4 +110,12 @@ export class RegexEditDialogComponent {
         : GroupByKey.REGEX_BY_EXP
     );
   }
+
+  /**
+   * Same as in runs_data_table.ts
+   * We don't want to trigger rerenders when choosing a color and closing the colorPicker
+   */
+  trackByGroups(index: number, group: ColorGroup) {
+    return group.groupId;
+  }
 }

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

@@ -31,12 +31,13 @@ import {
   getEnableColorByExperiment,
 } from '../../../selectors';
 import {selectors as settingsSelectors} from '../../../settings/';
-import {runGroupByChanged} from '../../actions';
+import {groupColorChanged, runGroupByChanged} from '../../actions';
 import {
   getColorGroupRegexString,
   getRunIdsForExperiment,
   getRunGroupBy,
   getRuns,
+  getGroupColorOverride,
 } from '../../store/runs_selectors';
 import {getDashboardExperimentNames} from '../../../selectors';
 import {groupRuns} from '../../store/utils';
@@ -56,6 +57,7 @@ const INPUT_CHANGE_DEBOUNCE_INTERVAL_MS = 500;
     (onSave)="onSave()"
     (regexInputOnChange)="onRegexInputOnChange($event)"
     (regexTypeOnChange)="onRegexTypeOnChange($event)"
+    (groupColorChange)="onGroupColorChange($event)"
   ></regex-edit-dialog-component>`,
   styles: [
     `
@@ -111,7 +113,8 @@ export class RegexEditDialogContainer {
         this.runIdToEid$,
         this.expNameByExpId$,
         this.store.select(settingsSelectors.getColorPalette),
-        this.store.select(getDarkModeEnabled)
+        this.store.select(getDarkModeEnabled),
+        this.store.select(getGroupColorOverride)
       ),
       map(
         ([
@@ -122,6 +125,7 @@ export class RegexEditDialogContainer {
           expNameByExpId,
           colorPalette,
           darkModeEnabled,
+          groupColorOverride,
         ]) => {
           const groupBy = {
             key: regexType,
@@ -137,8 +141,12 @@ export class RegexEditDialogContainer {
           const colorRunPairList: ColorGroup[] = [];
 
           for (const [groupId, runs] of Object.entries(groups.matches)) {
-            let colorHex: string | undefined =
-              groupKeyToColorString.get(groupId);
+            let colorHex: string | undefined = groupColorOverride.get(groupId);
+
+            if (!colorHex) {
+              colorHex = groupKeyToColorString.get(groupId);
+            }
+
             if (!colorHex) {
               const color =
                 colorPalette.colors[
@@ -216,6 +224,10 @@ export class RegexEditDialogContainer {
     this.tentativeRegexType$.next(regexType);
   }
 
+  onGroupColorChange({groupId, newColor}: {groupId: string; newColor: string}) {
+    this.store.dispatch(groupColorChanged({groupId, newColor}));
+  }
+
   onSave(): void {
     combineLatest([
       this.groupByRegexString$,