From 4ed2c0b267a0611a2ab8402d60445f4d94518e15 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Mon, 1 Apr 2024 06:45:45 +0000
Subject: [PATCH 1/7] defined savedPinsDataSource and added related tests

---
 tensorboard/webapp/BUILD                      |   1 +
 tensorboard/webapp/metrics/data_source/BUILD  |  26 ++++
 .../webapp/metrics/data_source/index.ts       |   2 +
 .../data_source/saved_pins_data_source.ts     |  45 +++++++
 .../saved_pins_data_source_module.ts          |  21 ++++
 .../saved_pins_data_source_test.ts            | 117 ++++++++++++++++++
 .../webapp/metrics/data_source/types.ts       |   2 +
 7 files changed, 214 insertions(+)
 create mode 100644 tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
 create mode 100644 tensorboard/webapp/metrics/data_source/saved_pins_data_source_module.ts
 create mode 100644 tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts

diff --git a/tensorboard/webapp/BUILD b/tensorboard/webapp/BUILD
index cbd29b5968..749240747e 100644
--- a/tensorboard/webapp/BUILD
+++ b/tensorboard/webapp/BUILD
@@ -275,6 +275,7 @@ tf_ng_web_test_suite(
         "//tensorboard/webapp/metrics:test_lib",
         "//tensorboard/webapp/metrics:utils_test",
         "//tensorboard/webapp/metrics/data_source:metrics_data_source_test",
+        "//tensorboard/webapp/metrics/data_source:saved_pins_data_source_test",
         "//tensorboard/webapp/metrics/effects:effects_test",
         "//tensorboard/webapp/metrics/store:store_test",
         "//tensorboard/webapp/metrics/views:views_test",
diff --git a/tensorboard/webapp/metrics/data_source/BUILD b/tensorboard/webapp/metrics/data_source/BUILD
index 441ee5a39b..9fefec787e 100644
--- a/tensorboard/webapp/metrics/data_source/BUILD
+++ b/tensorboard/webapp/metrics/data_source/BUILD
@@ -14,6 +14,7 @@ tf_ng_module(
     deps = [
         ":backend_types",
         ":types",
+        ":saved_pins_data_source",
         "//tensorboard/webapp/feature_flag",
         "//tensorboard/webapp/feature_flag/store",
         "//tensorboard/webapp/feature_flag/store:types",
@@ -37,6 +38,18 @@ tf_ng_module(
     ],
 )
 
+tf_ng_module(
+    name = "saved_pins_data_source",
+    srcs = [
+        "saved_pins_data_source.ts",
+        "saved_pins_data_source_module.ts",
+    ],
+    deps = [
+        ":types",
+        "@npm//@angular/core",
+    ],
+)
+
 tf_ts_library(
     name = "types",
     srcs = [
@@ -96,3 +109,16 @@ tf_ts_library(
         "@npm//@types/jasmine",
     ],
 )
+
+tf_ts_library(
+    name = "saved_pins_data_source_test",
+    testonly = True,
+    srcs = [
+        "saved_pins_data_source_test.ts",
+    ],
+    deps = [
+        ":saved_pins_data_source",
+        "//tensorboard/webapp/angular:expect_angular_core_testing",
+        "@npm//@types/jasmine",
+    ],
+)
diff --git a/tensorboard/webapp/metrics/data_source/index.ts b/tensorboard/webapp/metrics/data_source/index.ts
index b09db8f686..405c3d1c72 100644
--- a/tensorboard/webapp/metrics/data_source/index.ts
+++ b/tensorboard/webapp/metrics/data_source/index.ts
@@ -13,5 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 export * from './metrics_data_source';
+export * from './saved_pins_data_source';
 export * from './metrics_data_source_module';
+export * from './saved_pins_data_source_module';
 export * from './types';
diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
new file mode 100644
index 0000000000..44ddfec27f
--- /dev/null
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
@@ -0,0 +1,45 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+import {Injectable} from '@angular/core';
+import {Tag} from './types';
+
+const SAVED_SCALAR_PINS_KEY = 'tb-saved-scalar-pins';
+
+@Injectable()
+export class SavedPinsDataSource {
+  saveScalarPin(tag: Tag): void {
+    const existingPins = this.getSavedScalarPins();
+    if (!existingPins.includes(tag)) {
+      existingPins.push(tag);
+    }
+    localStorage.setItem(SAVED_SCALAR_PINS_KEY, JSON.stringify(existingPins));
+  }
+
+  removeScalarPin(tag: Tag): void {
+    const existingPins = this.getSavedScalarPins();
+    localStorage.setItem(
+      SAVED_SCALAR_PINS_KEY,
+      JSON.stringify(existingPins.filter((pin) => pin !== tag))
+    );
+  }
+
+  getSavedScalarPins(): Tag[] {
+    const savedPins = localStorage.getItem(SAVED_SCALAR_PINS_KEY);
+    if (savedPins) {
+      return JSON.parse(savedPins) as Tag[];
+    }
+    return [];
+  }
+}
diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_module.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_module.ts
new file mode 100644
index 0000000000..9bde26b6a6
--- /dev/null
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_module.ts
@@ -0,0 +1,21 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+import {NgModule} from '@angular/core';
+import {SavedPinsDataSource} from './saved_pins_data_source';
+
+@NgModule({
+  providers: [SavedPinsDataSource],
+})
+export class SavedPinsDataSourceModule {}
diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
new file mode 100644
index 0000000000..31dd6b358c
--- /dev/null
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
@@ -0,0 +1,117 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+import {TestBed} from '@angular/core/testing';
+import {SavedPinsDataSource} from './saved_pins_data_source';
+
+const SAVED_SCALAR_PINS_KEY = 'tb-saved-scalar-pins';
+
+describe('SavedPinsDataSource Test', () => {
+  let mockStorage: Record<string, string>;
+  let dataSource: SavedPinsDataSource;
+
+  beforeEach(async () => {
+    await TestBed.configureTestingModule({
+      providers: [SavedPinsDataSource],
+    });
+
+    dataSource = TestBed.inject(SavedPinsDataSource);
+
+    mockStorage = {};
+    spyOn(window.localStorage, 'setItem').and.callFake(
+      (key: string, value: string) => {
+        if (key !== SAVED_SCALAR_PINS_KEY) {
+          throw new Error('incorrect key used');
+        }
+
+        mockStorage[key] = value;
+      }
+    );
+
+    spyOn(window.localStorage, 'getItem').and.callFake((key: string) => {
+      if (key !== SAVED_SCALAR_PINS_KEY) {
+        throw new Error('incorrect key used');
+      }
+
+      return mockStorage[key];
+    });
+  });
+
+  describe('getSavedScalarPins', () => {
+    it('gets the saved scalar pins', () => {
+      window.localStorage.setItem(
+        SAVED_SCALAR_PINS_KEY,
+        JSON.stringify(['new_tag'])
+      );
+
+      const result = dataSource.getSavedScalarPins();
+
+      expect(result).toEqual(['new_tag']);
+    });
+
+    it('returns empty list if there is no saved pins', () => {
+      const result = dataSource.getSavedScalarPins();
+
+      expect(result).toEqual([]);
+    });
+  });
+
+  describe('saveScalarPin', () => {
+    it('stores the provided tag in the local storage', () => {
+      dataSource.saveScalarPin('tag1');
+
+      expect(dataSource.getSavedScalarPins().length).toEqual(1);
+      expect(dataSource.getSavedScalarPins()).toEqual(['tag1']);
+    });
+
+    it('adds the provided tag to the existing list', () => {
+      localStorage.setItem(SAVED_SCALAR_PINS_KEY, JSON.stringify(['tag1']));
+
+      dataSource.saveScalarPin('tag2');
+
+      expect(dataSource.getSavedScalarPins().length).toEqual(2);
+      expect(dataSource.getSavedScalarPins()).toEqual(['tag1', 'tag2']);
+    });
+
+    it('does not addd the provided tag if it already exists', () => {
+      localStorage.setItem(
+        SAVED_SCALAR_PINS_KEY,
+        JSON.stringify(['tag1', 'tag2'])
+      );
+
+      dataSource.saveScalarPin('tag2');
+
+      expect(dataSource.getSavedScalarPins().length).toEqual(2);
+      expect(dataSource.getSavedScalarPins()).toEqual(['tag1', 'tag2']);
+    });
+  });
+
+  describe('removeScalarPin', () => {
+    it('removes the given tag if it exists', () => {
+      dataSource.saveScalarPin('tag3');
+
+      dataSource.removeScalarPin('tag3');
+
+      expect(dataSource.getSavedScalarPins().length).toEqual(0);
+    });
+
+    it('does not remove anything if the given tag does not exist', () => {
+      dataSource.saveScalarPin('tag1');
+
+      dataSource.removeScalarPin('tag3');
+
+      expect(dataSource.getSavedScalarPins().length).toEqual(1);
+    });
+  });
+});
diff --git a/tensorboard/webapp/metrics/data_source/types.ts b/tensorboard/webapp/metrics/data_source/types.ts
index b85b6f89ea..9ca725dd81 100644
--- a/tensorboard/webapp/metrics/data_source/types.ts
+++ b/tensorboard/webapp/metrics/data_source/types.ts
@@ -183,3 +183,5 @@ export function isFailedTimeSeriesResponse(
 ): response is TimeSeriesFailedResponse {
   return response.hasOwnProperty('error');
 }
+
+export type Tag = string;

From 586e503f756fee3ef69b13b369655d8020bd1711 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Mon, 1 Apr 2024 06:47:46 +0000
Subject: [PATCH 2/7] add saving pin logic in the metrics_effect

---
 tensorboard/webapp/metrics/effects/index.ts   |  34 ++++-
 .../metrics/effects/metrics_effects_test.ts   | 128 ++++++++++++++++--
 tensorboard/webapp/metrics/metrics_module.ts  |   7 +-
 tensorboard/webapp/metrics/testing.ts         |  20 +++
 4 files changed, 173 insertions(+), 16 deletions(-)

diff --git a/tensorboard/webapp/metrics/effects/index.ts b/tensorboard/webapp/metrics/effects/index.ts
index d2f5684ace..65a97e5e69 100644
--- a/tensorboard/webapp/metrics/effects/index.ts
+++ b/tensorboard/webapp/metrics/effects/index.ts
@@ -42,13 +42,14 @@ import {
   TagMetadata,
   TimeSeriesRequest,
   TimeSeriesResponse,
+  SavedPinsDataSource,
 } from '../data_source/index';
 import {
   getCardLoadState,
   getCardMetadata,
   getMetricsTagMetadataLoadState,
 } from '../store';
-import {CardId, CardMetadata} from '../types';
+import {CardId, CardMetadata, PluginType} from '../types';
 
 export type CardFetchInfo = CardMetadata & {
   id: CardId;
@@ -73,7 +74,8 @@ export class MetricsEffects implements OnInitEffects {
   constructor(
     private readonly actions$: Actions,
     private readonly store: Store<State>,
-    private readonly dataSource: MetricsDataSource
+    private readonly metricsDataSource: MetricsDataSource,
+    private readonly savedPinsDataSource: SavedPinsDataSource
   ) {}
 
   /** @export */
@@ -141,7 +143,7 @@ export class MetricsEffects implements OnInitEffects {
       this.store.dispatch(actions.metricsTagMetadataRequested());
     }),
     switchMap(([, , experimentIds]) => {
-      return this.dataSource.fetchTagMetadata(experimentIds!).pipe(
+      return this.metricsDataSource.fetchTagMetadata(experimentIds!).pipe(
         tap((tagMetadata: TagMetadata) => {
           this.store.dispatch(actions.metricsTagMetadataLoaded({tagMetadata}));
         }),
@@ -174,7 +176,7 @@ export class MetricsEffects implements OnInitEffects {
   }
 
   private fetchTimeSeries(request: TimeSeriesRequest) {
-    return this.dataSource.fetchTimeSeries([request]).pipe(
+    return this.metricsDataSource.fetchTimeSeries([request]).pipe(
       tap((responses: TimeSeriesResponse[]) => {
         const errors = responses.filter(isFailedTimeSeriesResponse);
         if (errors.length) {
@@ -261,6 +263,23 @@ export class MetricsEffects implements OnInitEffects {
     })
   );
 
+  private readonly loadSavedPins$ = this.actions$.pipe(
+    ofType(actions.cardPinStateToggled),
+    withLatestFrom(this.getVisibleCardFetchInfos()),
+    map(([{cardId, canCreateNewPins, wasPinned}, fetchInfos]) => {
+      const card = fetchInfos.find((value) => value.id === cardId);
+      // Saving only scalar pinned cards.
+      if (!card || card.plugin !== PluginType.SCALARS) {
+        return;
+      }
+      if (wasPinned) {
+        this.savedPinsDataSource.removeScalarPin(card.tag);
+      } else if (canCreateNewPins) {
+        this.savedPinsDataSource.saveScalarPin(card.tag);
+      }
+    })
+  );
+
   /**
    * In general, this effect dispatch the following actions:
    *
@@ -292,7 +311,12 @@ export class MetricsEffects implements OnInitEffects {
         /**
          * Subscribes to: card visibility, reloads.
          */
-        this.loadTimeSeries$
+        this.loadTimeSeries$,
+
+        /**
+         * Subscribes to: cardPinStateToggled.
+         */
+        this.loadSavedPins$
       );
     },
     {dispatch: false}
diff --git a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
index 2851c19a7d..e9f4d2ab35 100644
--- a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
+++ b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
@@ -38,6 +38,7 @@ import {
   TagMetadata,
   TimeSeriesRequest,
   TimeSeriesResponse,
+  SavedPinsDataSource,
 } from '../data_source';
 import {getMetricsTagMetadataLoadState} from '../store';
 import {
@@ -46,12 +47,14 @@ import {
   buildMetricsState,
   createScalarStepData,
   provideTestingMetricsDataSource,
+  provideTestingSavedPinsDataSource,
 } from '../testing';
 import {CardId, TooltipSort} from '../types';
 import {CardFetchInfo, MetricsEffects, TEST_ONLY} from './index';
 
 describe('metrics effects', () => {
-  let dataSource: MetricsDataSource;
+  let metricsDataSource: MetricsDataSource;
+  let savedPinsDataSource: SavedPinsDataSource;
   let effects: MetricsEffects;
   let store: MockStore<State>;
   let actions$: Subject<Action>;
@@ -66,6 +69,7 @@ describe('metrics effects', () => {
       providers: [
         provideMockActions(actions$),
         provideTestingMetricsDataSource(),
+        provideTestingSavedPinsDataSource(),
         MetricsEffects,
         provideMockStore({
           initialState: {
@@ -81,7 +85,8 @@ describe('metrics effects', () => {
       actualActions.push(action);
     });
     effects = TestBed.inject(MetricsEffects);
-    dataSource = TestBed.inject(MetricsDataSource);
+    metricsDataSource = TestBed.inject(MetricsDataSource);
+    savedPinsDataSource = TestBed.inject(SavedPinsDataSource);
     store.overrideSelector(selectors.getExperimentIdsFromRoute, null);
     store.overrideSelector(selectors.getMetricsIgnoreOutliers, false);
     store.overrideSelector(selectors.getMetricsScalarSmoothing, 0.3);
@@ -107,7 +112,7 @@ describe('metrics effects', () => {
       beforeEach(() => {
         fetchTagMetadataSubject = new Subject();
         fetchTagMetadataSpy = spyOn(
-          dataSource,
+          metricsDataSource,
           'fetchTagMetadata'
         ).and.returnValue(fetchTagMetadataSubject);
       });
@@ -307,10 +312,10 @@ describe('metrics effects', () => {
 
       beforeEach(() => {
         fetchTagMetadataSpy = spyOn(
-          dataSource,
+          metricsDataSource,
           'fetchTagMetadata'
         ).and.returnValue(of(buildDataSourceTagMetadata()));
-        fetchTimeSeriesSpy = spyOn(dataSource, 'fetchTimeSeries');
+        fetchTimeSeriesSpy = spyOn(metricsDataSource, 'fetchTimeSeries');
         selectSpy = spyOn(store, 'select').and.callThrough();
       });
 
@@ -530,7 +535,7 @@ describe('metrics effects', () => {
 
       it('does not fetch when nothing is visible', () => {
         fetchTimeSeriesSpy = spyOn(
-          dataSource,
+          metricsDataSource,
           'fetchTimeSeries'
         ).and.returnValue(of(sampleBackendResponses));
         store.overrideSelector(selectors.getExperimentIdsFromRoute, ['exp1']);
@@ -553,7 +558,7 @@ describe('metrics effects', () => {
 
       it('fetches only once when hiding then showing a card', () => {
         fetchTimeSeriesSpy = spyOn(
-          dataSource,
+          metricsDataSource,
           'fetchTimeSeries'
         ).and.returnValue(of(sampleBackendResponses));
         store.overrideSelector(selectors.getExperimentIdsFromRoute, ['exp1']);
@@ -608,7 +613,7 @@ describe('metrics effects', () => {
 
       it('does not fetch when a loaded card exits and re-enters', () => {
         fetchTimeSeriesSpy = spyOn(
-          dataSource,
+          metricsDataSource,
           'fetchTimeSeries'
         ).and.returnValue(of(sampleBackendResponses));
         store.overrideSelector(selectors.getExperimentIdsFromRoute, ['exp1']);
@@ -704,7 +709,7 @@ describe('metrics effects', () => {
             sample: 5,
           },
         ];
-        fetchTimeSeriesSpy = spyOn(dataSource, 'fetchTimeSeries');
+        fetchTimeSeriesSpy = spyOn(metricsDataSource, 'fetchTimeSeries');
         fetchTimeSeriesSpy
           .withArgs([expectedRequests[0]])
           .and.returnValue(of([sampleBackendResponses[0]]));
@@ -758,7 +763,7 @@ describe('metrics effects', () => {
                 loadState,
               })
             );
-          fetchTimeSeriesSpy = spyOn(dataSource, 'fetchTimeSeries');
+          fetchTimeSeriesSpy = spyOn(metricsDataSource, 'fetchTimeSeries');
 
           store.overrideSelector(
             selectors.getVisibleCardIdSet,
@@ -777,5 +782,108 @@ describe('metrics effects', () => {
         });
       }
     });
+
+    describe('loadSavedPins', () => {
+      let saveScalarPinSpy: jasmine.Spy;
+      let removeScalarPinSpy: jasmine.Spy;
+
+      beforeEach(() => {
+        saveScalarPinSpy = spyOn(savedPinsDataSource, 'saveScalarPin');
+        removeScalarPinSpy = spyOn(savedPinsDataSource, 'removeScalarPin');
+
+        store.overrideSelector(TEST_ONLY.getCardFetchInfo, {
+          id: 'card1',
+          plugin: PluginType.SCALARS,
+          tag: 'tagA',
+          runId: null,
+          loadState: DataLoadState.LOADED,
+        });
+        store.overrideSelector(
+          selectors.getVisibleCardIdSet,
+          new Set(['card1'])
+        );
+        store.refreshState();
+      });
+
+      it('removes scalar pin if the given card was pinned', () => {
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card1',
+            wasPinned: true,
+            canCreateNewPins: true,
+          })
+        );
+
+        expect(removeScalarPinSpy).toHaveBeenCalled();
+        expect(removeScalarPinSpy).toHaveBeenCalledWith('tagA');
+        expect(saveScalarPinSpy).not.toHaveBeenCalled();
+      });
+
+      it('pins the card if the given card was not pinned and canCreateNewPins is true', () => {
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card1',
+            wasPinned: false,
+            canCreateNewPins: true,
+          })
+        );
+
+        expect(saveScalarPinSpy).toHaveBeenCalled();
+        expect(saveScalarPinSpy).toHaveBeenCalledWith('tagA');
+        expect(removeScalarPinSpy).not.toHaveBeenCalled();
+      });
+
+      it('does not pin the card if the given card was not pinned and canCreateNewPins is false', () => {
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card1',
+            wasPinned: false,
+            canCreateNewPins: false,
+          })
+        );
+
+        expect(saveScalarPinSpy).not.toHaveBeenCalled();
+        expect(removeScalarPinSpy).not.toHaveBeenCalled();
+      });
+
+      it('does not pin the card if the plugin type is not a scalar', () => {
+        store.overrideSelector(TEST_ONLY.getCardFetchInfo, {
+          id: 'card2',
+          plugin: PluginType.HISTOGRAMS,
+          tag: 'tagA',
+          runId: null,
+          loadState: DataLoadState.LOADED,
+        });
+        store.overrideSelector(
+          selectors.getVisibleCardIdSet,
+          new Set(['card2'])
+        );
+        store.refreshState();
+
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card2',
+            wasPinned: false,
+            canCreateNewPins: true,
+          })
+        );
+
+        expect(saveScalarPinSpy).not.toHaveBeenCalled();
+        expect(removeScalarPinSpy).not.toHaveBeenCalled();
+      });
+
+      it('does not pin the card if there is no matching card', () => {
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card3',
+            wasPinned: false,
+            canCreateNewPins: true,
+          })
+        );
+
+        expect(saveScalarPinSpy).not.toHaveBeenCalled();
+        expect(removeScalarPinSpy).not.toHaveBeenCalled();
+      });
+    });
   });
 });
diff --git a/tensorboard/webapp/metrics/metrics_module.ts b/tensorboard/webapp/metrics/metrics_module.ts
index 2e9cedc6aa..8e85e15490 100644
--- a/tensorboard/webapp/metrics/metrics_module.ts
+++ b/tensorboard/webapp/metrics/metrics_module.ts
@@ -22,7 +22,11 @@ import {CoreModule} from '../core/core_module';
 import {PersistentSettingsConfigModule} from '../persistent_settings/persistent_settings_config_module';
 import {PluginRegistryModule} from '../plugins/plugin_registry_module';
 import * as actions from './actions';
-import {MetricsDataSourceModule, METRICS_PLUGIN_ID} from './data_source';
+import {
+  MetricsDataSourceModule,
+  METRICS_PLUGIN_ID,
+  SavedPinsDataSourceModule,
+} from './data_source';
 import {MetricsEffects} from './effects';
 import {
   getMetricsCardMinWidth,
@@ -145,6 +149,7 @@ export function getRangeSelectionHeadersFactory() {
       reducers,
       METRICS_STORE_CONFIG_TOKEN
     ),
+    SavedPinsDataSourceModule,
     EffectsModule.forFeature([MetricsEffects]),
     AlertActionModule.registerAlertActions(alertActionProvider),
     PersistentSettingsConfigModule.defineGlobalSetting(
diff --git a/tensorboard/webapp/metrics/testing.ts b/tensorboard/webapp/metrics/testing.ts
index 4dc9757d89..a941997693 100644
--- a/tensorboard/webapp/metrics/testing.ts
+++ b/tensorboard/webapp/metrics/testing.ts
@@ -27,6 +27,8 @@ import {
   TagMetadata as DataSourceTagMetadata,
   TimeSeriesRequest,
   TimeSeriesResponse,
+  SavedPinsDataSource,
+  Tag,
 } from './data_source';
 import * as selectors from './store/metrics_selectors';
 import {
@@ -393,3 +395,21 @@ export function buildStepIndexMetadata(
     ...override,
   };
 }
+
+@Injectable()
+export class TestingSavedPinsDataSource {
+  saveScalarPin(tag: Tag) {}
+
+  removeScalarPin(tag: Tag) {}
+
+  getSavedScalarPins() {
+    return [];
+  }
+}
+
+export function provideTestingSavedPinsDataSource() {
+  return [
+    TestingSavedPinsDataSource,
+    {provide: SavedPinsDataSource, useExisting: TestingSavedPinsDataSource},
+  ];
+}

From c29f40da516fb909e072c63c3c1372a6bf3d1a74 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Mon, 1 Apr 2024 07:58:46 +0000
Subject: [PATCH 3/7] fix BUILD lint error

---
 tensorboard/webapp/metrics/data_source/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorboard/webapp/metrics/data_source/BUILD b/tensorboard/webapp/metrics/data_source/BUILD
index 9fefec787e..56dd7b164e 100644
--- a/tensorboard/webapp/metrics/data_source/BUILD
+++ b/tensorboard/webapp/metrics/data_source/BUILD
@@ -13,8 +13,8 @@ tf_ng_module(
     ],
     deps = [
         ":backend_types",
-        ":types",
         ":saved_pins_data_source",
+        ":types",
         "//tensorboard/webapp/feature_flag",
         "//tensorboard/webapp/feature_flag/store",
         "//tensorboard/webapp/feature_flag/store:types",

From 5602eb92aa7941c3f08957e0dd60931ff9b61554 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Tue, 2 Apr 2024 06:03:27 +0000
Subject: [PATCH 4/7] rename effect to addOrRemovePin$ and remove some
 redundant test codes

---
 .../metrics/data_source/saved_pins_data_source_test.ts       | 5 +----
 tensorboard/webapp/metrics/effects/index.ts                  | 4 ++--
 tensorboard/webapp/metrics/effects/metrics_effects_test.ts   | 4 +---
 3 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
index 31dd6b358c..0bea1b1c48 100644
--- a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
@@ -71,7 +71,6 @@ describe('SavedPinsDataSource Test', () => {
     it('stores the provided tag in the local storage', () => {
       dataSource.saveScalarPin('tag1');
 
-      expect(dataSource.getSavedScalarPins().length).toEqual(1);
       expect(dataSource.getSavedScalarPins()).toEqual(['tag1']);
     });
 
@@ -80,7 +79,6 @@ describe('SavedPinsDataSource Test', () => {
 
       dataSource.saveScalarPin('tag2');
 
-      expect(dataSource.getSavedScalarPins().length).toEqual(2);
       expect(dataSource.getSavedScalarPins()).toEqual(['tag1', 'tag2']);
     });
 
@@ -92,7 +90,6 @@ describe('SavedPinsDataSource Test', () => {
 
       dataSource.saveScalarPin('tag2');
 
-      expect(dataSource.getSavedScalarPins().length).toEqual(2);
       expect(dataSource.getSavedScalarPins()).toEqual(['tag1', 'tag2']);
     });
   });
@@ -111,7 +108,7 @@ describe('SavedPinsDataSource Test', () => {
 
       dataSource.removeScalarPin('tag3');
 
-      expect(dataSource.getSavedScalarPins().length).toEqual(1);
+      expect(dataSource.getSavedScalarPins()).toEqual(['tag1']);
     });
   });
 });
diff --git a/tensorboard/webapp/metrics/effects/index.ts b/tensorboard/webapp/metrics/effects/index.ts
index 65a97e5e69..e4c5b3ca2d 100644
--- a/tensorboard/webapp/metrics/effects/index.ts
+++ b/tensorboard/webapp/metrics/effects/index.ts
@@ -263,7 +263,7 @@ export class MetricsEffects implements OnInitEffects {
     })
   );
 
-  private readonly loadSavedPins$ = this.actions$.pipe(
+  private readonly addOrRemovePin$ = this.actions$.pipe(
     ofType(actions.cardPinStateToggled),
     withLatestFrom(this.getVisibleCardFetchInfos()),
     map(([{cardId, canCreateNewPins, wasPinned}, fetchInfos]) => {
@@ -316,7 +316,7 @@ export class MetricsEffects implements OnInitEffects {
         /**
          * Subscribes to: cardPinStateToggled.
          */
-        this.loadSavedPins$
+        this.addOrRemovePin$
       );
     },
     {dispatch: false}
diff --git a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
index e9f4d2ab35..76db439f78 100644
--- a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
+++ b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
@@ -783,7 +783,7 @@ describe('metrics effects', () => {
       }
     });
 
-    describe('loadSavedPins', () => {
+    describe('addOrRemovePin', () => {
       let saveScalarPinSpy: jasmine.Spy;
       let removeScalarPinSpy: jasmine.Spy;
 
@@ -814,7 +814,6 @@ describe('metrics effects', () => {
           })
         );
 
-        expect(removeScalarPinSpy).toHaveBeenCalled();
         expect(removeScalarPinSpy).toHaveBeenCalledWith('tagA');
         expect(saveScalarPinSpy).not.toHaveBeenCalled();
       });
@@ -828,7 +827,6 @@ describe('metrics effects', () => {
           })
         );
 
-        expect(saveScalarPinSpy).toHaveBeenCalled();
         expect(saveScalarPinSpy).toHaveBeenCalledWith('tagA');
         expect(removeScalarPinSpy).not.toHaveBeenCalled();
       });

From 6fdd93e03ab0c15b722f3993c629a433a44f89cc Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Tue, 2 Apr 2024 15:37:17 +0000
Subject: [PATCH 5/7] change localStorage to window.localStorage

---
 .../webapp/metrics/data_source/saved_pins_data_source.ts | 9 ++++++---
 .../metrics/data_source/saved_pins_data_source_test.ts   | 7 +++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
index 44ddfec27f..da02bf4b11 100644
--- a/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source.ts
@@ -24,19 +24,22 @@ export class SavedPinsDataSource {
     if (!existingPins.includes(tag)) {
       existingPins.push(tag);
     }
-    localStorage.setItem(SAVED_SCALAR_PINS_KEY, JSON.stringify(existingPins));
+    window.localStorage.setItem(
+      SAVED_SCALAR_PINS_KEY,
+      JSON.stringify(existingPins)
+    );
   }
 
   removeScalarPin(tag: Tag): void {
     const existingPins = this.getSavedScalarPins();
-    localStorage.setItem(
+    window.localStorage.setItem(
       SAVED_SCALAR_PINS_KEY,
       JSON.stringify(existingPins.filter((pin) => pin !== tag))
     );
   }
 
   getSavedScalarPins(): Tag[] {
-    const savedPins = localStorage.getItem(SAVED_SCALAR_PINS_KEY);
+    const savedPins = window.localStorage.getItem(SAVED_SCALAR_PINS_KEY);
     if (savedPins) {
       return JSON.parse(savedPins) as Tag[];
     }
diff --git a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
index 0bea1b1c48..325acc32d2 100644
--- a/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
+++ b/tensorboard/webapp/metrics/data_source/saved_pins_data_source_test.ts
@@ -75,7 +75,10 @@ describe('SavedPinsDataSource Test', () => {
     });
 
     it('adds the provided tag to the existing list', () => {
-      localStorage.setItem(SAVED_SCALAR_PINS_KEY, JSON.stringify(['tag1']));
+      window.localStorage.setItem(
+        SAVED_SCALAR_PINS_KEY,
+        JSON.stringify(['tag1'])
+      );
 
       dataSource.saveScalarPin('tag2');
 
@@ -83,7 +86,7 @@ describe('SavedPinsDataSource Test', () => {
     });
 
     it('does not addd the provided tag if it already exists', () => {
-      localStorage.setItem(
+      window.localStorage.setItem(
         SAVED_SCALAR_PINS_KEY,
         JSON.stringify(['tag1', 'tag2'])
       );

From 4806818cff6dbbd38aa02b6cbd9a3822d5367e01 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Tue, 2 Apr 2024 16:20:41 +0000
Subject: [PATCH 6/7] define enableGlobalPins feature flag

---
 .../webapp/feature_flag/store/feature_flag_metadata.ts     | 5 +++++
 .../webapp/feature_flag/store/feature_flag_selectors.ts    | 7 +++++++
 tensorboard/webapp/feature_flag/types.ts                   | 2 ++
 3 files changed, 14 insertions(+)

diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_metadata.ts b/tensorboard/webapp/feature_flag/store/feature_flag_metadata.ts
index a348320a10..63643582bb 100644
--- a/tensorboard/webapp/feature_flag/store/feature_flag_metadata.ts
+++ b/tensorboard/webapp/feature_flag/store/feature_flag_metadata.ts
@@ -120,6 +120,11 @@ export const FeatureFlagMetadataMap: FeatureFlagMetadataMapType<FeatureFlags> =
       queryParamOverride: 'enableSuggestedCards',
       parseValue: parseBoolean,
     },
+    enableGlobalPins: {
+      defaultValue: false,
+      queryParamOverride: 'enableGlobalPins',
+      parseValue: parseBoolean,
+    },
   };
 
 /**
diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
index c772d691ab..c594091dca 100644
--- a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
+++ b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
@@ -153,3 +153,10 @@ export const getIsScalarColumnContextMenusEnabled = createSelector(
     return flags.enableScalarColumnContextMenus;
   }
 );
+
+export const getEnableGlobalPins = createSelector(
+  getFeatureFlags,
+  (flags: FeatureFlags): boolean => {
+    return flags.enableGlobalPins;
+  }
+);
\ No newline at end of file
diff --git a/tensorboard/webapp/feature_flag/types.ts b/tensorboard/webapp/feature_flag/types.ts
index 6304e2b988..74d24e1534 100644
--- a/tensorboard/webapp/feature_flag/types.ts
+++ b/tensorboard/webapp/feature_flag/types.ts
@@ -50,4 +50,6 @@ export interface FeatureFlags {
   // Adds a new section at the top of the time series metrics view
   // containing suggested cards based on the users previous interactions.
   enableSuggestedCards: boolean;
+  // Persists pinned scalar cards across multiple experiments.
+  enableGlobalPins: boolean;
 }

From 0bc3a45328c9d404c9a9b0361f3f4a64b212ca25 Mon Sep 17 00:00:00 2001
From: Seayeon Lee <seayeon@google.com>
Date: Tue, 2 Apr 2024 16:26:11 +0000
Subject: [PATCH 7/7] protect writing localstorage with feature flag

---
 .../store/feature_flag_selectors.ts           |  2 +-
 tensorboard/webapp/metrics/effects/BUILD      |  1 +
 tensorboard/webapp/metrics/effects/index.ts   | 36 ++++++++++++-------
 .../metrics/effects/metrics_effects_test.ts   | 24 +++++++++++--
 4 files changed, 48 insertions(+), 15 deletions(-)

diff --git a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
index c594091dca..1816bed3aa 100644
--- a/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
+++ b/tensorboard/webapp/feature_flag/store/feature_flag_selectors.ts
@@ -159,4 +159,4 @@ export const getEnableGlobalPins = createSelector(
   (flags: FeatureFlags): boolean => {
     return flags.enableGlobalPins;
   }
-);
\ No newline at end of file
+);
diff --git a/tensorboard/webapp/metrics/effects/BUILD b/tensorboard/webapp/metrics/effects/BUILD
index b0979fc828..1fe50ad839 100644
--- a/tensorboard/webapp/metrics/effects/BUILD
+++ b/tensorboard/webapp/metrics/effects/BUILD
@@ -45,6 +45,7 @@ tf_ts_library(
         "//tensorboard/webapp/metrics/actions",
         "//tensorboard/webapp/metrics/data_source",
         "//tensorboard/webapp/metrics/store",
+        "//tensorboard/webapp/testing:utils",
         "//tensorboard/webapp/types",
         "//tensorboard/webapp/util:dom",
         "//tensorboard/webapp/webapp_data_source:http_client_testing",
diff --git a/tensorboard/webapp/metrics/effects/index.ts b/tensorboard/webapp/metrics/effects/index.ts
index e4c5b3ca2d..317cc96d79 100644
--- a/tensorboard/webapp/metrics/effects/index.ts
+++ b/tensorboard/webapp/metrics/effects/index.ts
@@ -265,19 +265,31 @@ export class MetricsEffects implements OnInitEffects {
 
   private readonly addOrRemovePin$ = this.actions$.pipe(
     ofType(actions.cardPinStateToggled),
-    withLatestFrom(this.getVisibleCardFetchInfos()),
-    map(([{cardId, canCreateNewPins, wasPinned}, fetchInfos]) => {
-      const card = fetchInfos.find((value) => value.id === cardId);
-      // Saving only scalar pinned cards.
-      if (!card || card.plugin !== PluginType.SCALARS) {
-        return;
-      }
-      if (wasPinned) {
-        this.savedPinsDataSource.removeScalarPin(card.tag);
-      } else if (canCreateNewPins) {
-        this.savedPinsDataSource.saveScalarPin(card.tag);
+    withLatestFrom(
+      this.getVisibleCardFetchInfos(),
+      this.store.select(selectors.getEnableGlobalPins)
+    ),
+    map(
+      ([
+        {cardId, canCreateNewPins, wasPinned},
+        fetchInfos,
+        enableGlobalPins,
+      ]) => {
+        if (!enableGlobalPins) {
+          return;
+        }
+        const card = fetchInfos.find((value) => value.id === cardId);
+        // Saving only scalar pinned cards.
+        if (!card || card.plugin !== PluginType.SCALARS) {
+          return;
+        }
+        if (wasPinned) {
+          this.savedPinsDataSource.removeScalarPin(card.tag);
+        } else if (canCreateNewPins) {
+          this.savedPinsDataSource.saveScalarPin(card.tag);
+        }
       }
-    })
+    )
   );
 
   /**
diff --git a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
index 76db439f78..85204ed278 100644
--- a/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
+++ b/tensorboard/webapp/metrics/effects/metrics_effects_test.ts
@@ -51,6 +51,7 @@ import {
 } from '../testing';
 import {CardId, TooltipSort} from '../types';
 import {CardFetchInfo, MetricsEffects, TEST_ONLY} from './index';
+import {buildMockState} from '../../testing/utils';
 
 describe('metrics effects', () => {
   let metricsDataSource: MetricsDataSource;
@@ -73,8 +74,10 @@ describe('metrics effects', () => {
         MetricsEffects,
         provideMockStore({
           initialState: {
-            ...appStateFromMetricsState(buildMetricsState()),
-            ...coreTesting.createState(coreTesting.createCoreState()),
+            ...buildMockState({
+              ...appStateFromMetricsState(buildMetricsState()),
+              ...coreTesting.createState(coreTesting.createCoreState()),
+            }),
           },
         }),
       ],
@@ -791,6 +794,7 @@ describe('metrics effects', () => {
         saveScalarPinSpy = spyOn(savedPinsDataSource, 'saveScalarPin');
         removeScalarPinSpy = spyOn(savedPinsDataSource, 'removeScalarPin');
 
+        store.overrideSelector(selectors.getEnableGlobalPins, true);
         store.overrideSelector(TEST_ONLY.getCardFetchInfo, {
           id: 'card1',
           plugin: PluginType.SCALARS,
@@ -882,6 +886,22 @@ describe('metrics effects', () => {
         expect(saveScalarPinSpy).not.toHaveBeenCalled();
         expect(removeScalarPinSpy).not.toHaveBeenCalled();
       });
+
+      it('does not pin the card if getEnableGlobalPins is false', () => {
+        store.overrideSelector(selectors.getEnableGlobalPins, false);
+        store.refreshState();
+
+        actions$.next(
+          actions.cardPinStateToggled({
+            cardId: 'card1',
+            wasPinned: false,
+            canCreateNewPins: true,
+          })
+        );
+
+        expect(saveScalarPinSpy).not.toHaveBeenCalled();
+        expect(removeScalarPinSpy).not.toHaveBeenCalled();
+      });
     });
   });
 });