stop making duplicate time series requests

Author: rileyajones

tensorboard/webapp/metrics/effects/index.ts

@@ -18,6 +18,7 @@ import {Action, createAction, createSelector, Store} from '@ngrx/store';
 import {forkJoin, merge, Observable, of} from 'rxjs';
 import {
   catchError,
+  combineLatestWith,
   filter,
   map,
   mergeMap,
@@ -67,6 +68,12 @@ const getCardFetchInfo = createSelector(
 
 const initAction = createAction('[Metrics Effects] Init');
 
+function parseRunIdFromSampledRunInfoName(eidRun: string): string {
+  if (!eidRun) return '';
+  const [, ...runIdChunks] = eidRun.split('/');
+  return runIdChunks.join('/');
+}
+
 @Injectable()
 export class MetricsEffects implements OnInitEffects {
   constructor(
@@ -75,6 +82,40 @@ export class MetricsEffects implements OnInitEffects {
     private readonly dataSource: MetricsDataSource
   ) {}
 
+  readonly tagToEid$: Observable<Record<string, Set<string>>> = this.store
+    .select(selectors.getMetricsTagMetadata)
+    .pipe(
+      combineLatestWith(this.store.select(selectors.getRunIdToExperimentId)),
+      map(([tagMetadata, runToEid]) => {
+        const imageTagToRuns = Object.fromEntries(
+          Object.entries(tagMetadata.images.tagRunSampledInfo).map(
+            ([tag, sampledRunInfo]) => {
+              const runIds = Object.keys(sampledRunInfo).map((runInfoKey) =>
+                parseRunIdFromSampledRunInfoName(runInfoKey)
+              );
+              return [tag, runIds];
+            }
+          )
+        );
+
+        const tagToEid: Record<string, Set<string>> = {};
+        function mapTagsToEid(tagToRun: Record<string, readonly string[]>) {
+          Object.entries(tagToRun).forEach(([tag, runIds]) => {
+            if (!tagToEid[tag]) {
+              tagToEid[tag] = new Set();
+            }
+            runIds.forEach((runId) => tagToEid[tag].add(runToEid[runId]));
+          });
+        }
+
+        mapTagsToEid(tagMetadata.scalars.tagToRuns);
+        mapTagsToEid(tagMetadata.histograms.tagToRuns);
+        mapTagsToEid(imageTagToRuns);
+
+        return tagToEid;
+      })
+    );
+
   /** @export */
   ngrxOnInitEffects(): Action {
     return initAction();
@@ -193,23 +234,31 @@ export class MetricsEffects implements OnInitEffects {
     fetchInfos: CardFetchInfo[],
     experimentIds: string[]
   ) {
-    /**
-     * TODO(psybuzz): if 2 cards require the same data, we should dedupe instead of
-     * making 2 identical requests.
-     */
-    const requests: TimeSeriesRequest[] = fetchInfos.map((fetchInfo) => {
-      const {plugin, tag, runId, sample} = fetchInfo;
-      const partialRequest: TimeSeriesRequest = isSingleRunPlugin(plugin)
-        ? {plugin, tag, runId: runId!}
-        : {plugin, tag, experimentIds};
-      if (sample !== undefined) {
-        partialRequest.sample = sample;
-      }
-      return partialRequest;
-    });
-
     // Fetch and handle responses.
-    return of(requests).pipe(
+    return this.tagToEid$.pipe(
+      map((tagToEid): TimeSeriesRequest[] => {
+        const requests = fetchInfos.map((fetchInfo) => {
+          const {plugin, tag, runId, sample} = fetchInfo;
+          const filteredEids = experimentIds.filter((eid) =>
+            tagToEid[tag]?.has(eid)
+          );
+
+          const partialRequest: TimeSeriesRequest = isSingleRunPlugin(plugin)
+            ? {plugin, tag, runId: runId!}
+            : {plugin, tag, experimentIds: filteredEids};
+          if (sample !== undefined) {
+            partialRequest.sample = sample;
+          }
+          return partialRequest;
+        });
+        const uniqueRequests = new Set(
+          requests.map((request) => JSON.stringify(request))
+        );
+
+        return Array.from(uniqueRequests).map(
+          (serialized) => JSON.parse(serialized) as TimeSeriesRequest
+        );
+      }),
       tap((requests) => {
         this.store.dispatch(actions.multipleTimeSeriesRequested({requests}));
       }),
@@ -300,4 +349,5 @@ export class MetricsEffects implements OnInitEffects {
 export const TEST_ONLY = {
   getCardFetchInfo,
   initAction,
+  parseRunIdFromSampledRunInfoName,
 };

tensorboard/webapp/metrics/effects/metrics_effects_test.ts

@@ -89,8 +89,52 @@ describe('metrics effects', () => {
       selectors.getMetricsTooltipSort,
       TooltipSort.ALPHABETICAL
     );
+
+    overrideTagMetadata();
+    overrideRunToEid();
   });
 
+  function overrideTagMetadata() {
+    store.overrideSelector(selectors.getMetricsTagMetadata, {
+      scalars: {
+        tagDescriptions: {} as any,
+        tagToRuns: {
+          tagA: ['run1'],
+          tagB: ['run2', 'run3'],
+          tagC: ['run4', 'run5'],
+          tagD: ['run6'],
+        },
+      },
+      histograms: {
+        tagDescriptions: {} as any,
+        tagToRuns: {
+          tagA: ['run1'],
+          tagB: ['run4'],
+        },
+      },
+      images: {
+        tagDescriptions: {},
+        tagRunSampledInfo: {
+          tagC: {
+            'defaultExperimentId/run1': {} as any,
+            'exp1/run3': {} as any,
+          },
+        },
+      },
+    });
+  }
+
+  function overrideRunToEid() {
+    store.overrideSelector(selectors.getRunIdToExperimentId, {
+      run1: 'exp1',
+      run2: 'exp1',
+      run3: 'exp2',
+      run4: 'defaultExperimentId',
+      run5: 'defaultExperimentId',
+      run6: 'defaultExperimentId',
+    });
+  }
+
   afterEach(() => {
     store?.resetSelectors();
   });
@@ -365,35 +409,25 @@ describe('metrics effects', () => {
           actions$.next(reloadAction());
 
           expect(fetchTagMetadataSpy).toHaveBeenCalled();
-          expect(fetchTimeSeriesSpy).toHaveBeenCalledTimes(2);
+          expect(fetchTimeSeriesSpy).toHaveBeenCalledTimes(1);
           expect(actualActions).toEqual([
             actions.metricsTagMetadataRequested(),
             actions.metricsTagMetadataLoaded({
               tagMetadata: buildDataSourceTagMetadata(),
             }),
 
-            // Currently we expect 2x the same requests if the cards are the same.
-            // Ideally we should dedupe requests for the same info.
             actions.multipleTimeSeriesRequested({
               requests: [
                 {
                   plugin: PluginType.SCALARS as MultiRunPluginType,
                   tag: 'tagA',
                   experimentIds: ['exp1'],
                 },
-                {
-                  plugin: PluginType.SCALARS as MultiRunPluginType,
-                  tag: 'tagA',
-                  experimentIds: ['exp1'],
-                },
               ],
             }),
             actions.fetchTimeSeriesLoaded({
               response: buildTimeSeriesResponse(),
             }),
-            actions.fetchTimeSeriesLoaded({
-              response: buildTimeSeriesResponse(),
-            }),
           ]);
         });
 
@@ -487,6 +521,8 @@ describe('metrics effects', () => {
       it('does not re-fetch time series, until a valid experiment id', () => {
         // Reset any `getExperimentIdsFromRoute` overrides above.
         store.resetSelectors();
+        overrideTagMetadata();
+        overrideRunToEid();
         store.overrideSelector(getActivePlugin, METRICS_PLUGIN_ID);
         store.overrideSelector(
           selectors.getVisibleCardIdSet,
@@ -510,6 +546,43 @@ describe('metrics effects', () => {
 
         expect(fetchTimeSeriesSpy).toHaveBeenCalledTimes(2);
       });
+
+      it('does not send requests to experiments lacking a cards tag', () => {
+        store.overrideSelector(getActivePlugin, METRICS_PLUGIN_ID);
+        store.overrideSelector(selectors.getExperimentIdsFromRoute, [
+          'exp1',
+          'exp2',
+        ]);
+        store.overrideSelector(
+          selectors.getVisibleCardIdSet,
+          new Set(['card1', 'card2'])
+        );
+        provideCardFetchInfo([
+          {id: 'card1', tag: 'tagA'},
+          {id: 'card2', tag: 'tagB'},
+        ]);
+        store.refreshState();
+
+        const effectFetchTimeSeriesSpy = spyOn(
+          effects as any,
+          'fetchTimeSeries'
+        ).and.stub();
+
+        actions$.next(coreActions.manualReload());
+
+        expect(effectFetchTimeSeriesSpy).toHaveBeenCalledTimes(2);
+        expect(effectFetchTimeSeriesSpy).toHaveBeenCalledWith({
+          plugin: 'scalars',
+          tag: 'tagA',
+          experimentIds: ['exp1'],
+        });
+
+        expect(effectFetchTimeSeriesSpy).toHaveBeenCalledWith({
+          plugin: 'scalars',
+          tag: 'tagB',
+          experimentIds: ['exp1', 'exp2'],
+        });
+      });
     });
 
     describe('loadTimeSeriesForVisibleCardsWithoutData', () => {
@@ -778,4 +851,24 @@ describe('metrics effects', () => {
       }
     });
   });
+
+  describe('#utilities', () => {
+    describe('parseRunIdFromSampledRunInfoName', () => {
+      it('removes prefixed experiment id', () => {
+        expect(
+          TEST_ONLY.parseRunIdFromSampledRunInfoName('experimentId/someRun')
+        ).toEqual('someRun');
+      });
+
+      it('preserves "/" characters in run names', () => {
+        expect(
+          TEST_ONLY.parseRunIdFromSampledRunInfoName('experimentId/some/run')
+        ).toEqual('some/run');
+      });
+
+      it('returns an empty string when an empty string is provided', () => {
+        expect(TEST_ONLY.parseRunIdFromSampledRunInfoName('')).toEqual('');
+      });
+    });
+  });
 });