Bug Fix: Stop making duplicate time series requests

tensorboard/webapp/metrics/effects/BUILD

@@ -17,6 +17,7 @@ tf_ng_module(
         "//tensorboard/webapp/metrics/data_source",
         "//tensorboard/webapp/metrics/store",
         "//tensorboard/webapp/types",
+        "//tensorboard/webapp/util:types",
         "@npm//@angular/core",
         "@npm//@ngrx/effects",
         "@npm//@ngrx/store",

tensorboard/webapp/metrics/effects/index.ts

@@ -19,13 +19,15 @@ import {forkJoin, merge, Observable, of} from 'rxjs';
 import {
   catchError,
   throttleTime,
+  combineLatestWith,
   filter,
   map,
   mergeMap,
   switchMap,
   take,
   tap,
   withLatestFrom,
+  shareReplay,
 } from 'rxjs/operators';
 import * as routingActions from '../../app_routing/actions';
 import {State} from '../../app_state';
@@ -42,13 +44,16 @@ import {
   TagMetadata,
   TimeSeriesRequest,
   TimeSeriesResponse,
+  NonSampledPluginType,
 } from '../data_source/index';
 import {
   getCardLoadState,
   getCardMetadata,
   getMetricsTagMetadataLoadState,
+  TagMetadata as StoreTagMetadata,
 } from '../store';
-import {CardId, CardMetadata} from '../types';
+import {CardId, CardMetadata, PluginType} from '../types';
+import {DeepReadonly} from '../../util/types';
 
 export type CardFetchInfo = CardMetadata & {
   id: CardId;
@@ -68,6 +73,31 @@ const getCardFetchInfo = createSelector(
 
 const initAction = createAction('[Metrics Effects] Init');
 
+function generateMultiRunTagsToEidMapping(
+  tagMetadata: DeepReadonly<StoreTagMetadata>,
+  runToEid: Record<string, string>
+): Record<string, Set<string>> {
+  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]));
+    });
+  }
+
+  for (const pluginType in tagMetadata) {
+    if (isSingleRunPlugin(pluginType as PluginType)) {
+      continue;
+    }
+
+    mapTagsToEid(tagMetadata[pluginType as NonSampledPluginType].tagToRuns);
+  }
+
+  return tagToEid;
+}
+
 @Injectable()
 export class MetricsEffects implements OnInitEffects {
   constructor(
@@ -76,6 +106,22 @@ export class MetricsEffects implements OnInitEffects {
     private readonly dataSource: MetricsDataSource
   ) {}
 
+  /**
+   * Computes a record of tag to the experiments it appears in.
+   *
+   * The computation is done by translating Plugin -> Tag -> Run -> ExpId
+   *
+   * Sampled plugins are ignored because they are associated with runs, not experiments.
+   */
+  readonly multiRunTagsToEid$: Observable<Record<string, Set<string>>> =
+    this.store.select(selectors.getMetricsTagMetadata).pipe(
+      combineLatestWith(this.store.select(selectors.getRunIdToExperimentId)),
+      map(([tagMetadata, runToEid]) => {
+        return generateMultiRunTagsToEidMapping(tagMetadata, runToEid);
+      }),
+      shareReplay(1)
+    );
+
   /** @export */
   ngrxOnInitEffects(): Action {
     return initAction();
@@ -195,23 +241,44 @@ 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.multiRunTagsToEid$.pipe(
+      take(1),
+      map((tagToEid): TimeSeriesRequest[] => {
+        const requests = fetchInfos
+          .map((fetchInfo) => {
+            const {plugin, tag, runId, sample} = fetchInfo;
+
+            if (isSingleRunPlugin(plugin)) {
+              if (!runId) {
+                return;
+              }
+              return {
+                plugin,
+                tag,
+                runId,
+                sample,
+              };
+            }
+
+            const filteredEids = experimentIds.filter((eid) =>
+              tagToEid[tag]?.has(eid)
+            );
+            if (!filteredEids.length) {
+              return;
+            }
+
+            return {plugin, tag, experimentIds: filteredEids};
+          })
+          .filter(Boolean);
+        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}));
       }),
@@ -302,4 +369,5 @@ export class MetricsEffects implements OnInitEffects {
 export const TEST_ONLY = {
   getCardFetchInfo,
   initAction,
+  generateMultiRunTagsToEidMapping,
 };