Ensure the correct fetched status is applied when overriding plot rev…

…isions (#3557)
iterative · Apr 2, 2023 · 7c5a9e4 · 7c5a9e4
1 parent c0443af
commit 7c5a9e4
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 16 deletions.
diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts
@@ -236,6 +236,7 @@ describe('collectOverrideRevisionDetails', () => {
           }
         ] as SelectedExperimentWithColor[],
         new Set(['a', 'c', 'd', 'e']),
+        new Set(['a', 'c', 'd', 'e']),
         {},
         (id: string) =>
           ({
@@ -348,6 +349,7 @@ describe('collectOverrideRevisionDetails', () => {
             status: ExperimentStatus.SUCCESS
           }
         ] as SelectedExperimentWithColor[],
+        new Set([]),
         new Set(['a', 'c', 'd', 'e']),
         { [runningId]: EXPERIMENT_WORKSPACE_ID },
         (id: string) =>
@@ -382,31 +384,31 @@ describe('collectOverrideRevisionDetails', () => {
     expect(overrideRevisions).toStrictEqual([
       {
         displayColor: '#4299e1',
-        fetched: true,
+        fetched: false,
         firstThreeColumns: [],
         group: 'a',
         id: 'a',
         revision: 'a'
       },
       {
         displayColor: '#13adc7',
-        fetched: true,
+        fetched: false,
         firstThreeColumns: [],
         group: undefined,
         id: EXPERIMENT_WORKSPACE_ID,
         revision: EXPERIMENT_WORKSPACE_ID
       },
       {
         displayColor: '#48bb78',
-        fetched: true,
+        fetched: false,
         firstThreeColumns: [],
         group: 'c',
         id: 'c',
         revision: 'c'
       },
       {
         displayColor: '#f56565',
-        fetched: true,
+        fetched: false,
         firstThreeColumns: [],
         group: 'd',
         id: 'd',
@@ -460,6 +462,7 @@ describe('collectOverrideRevisionDetails', () => {
             status: ExperimentStatus.SUCCESS
           }
         ] as SelectedExperimentWithColor[],
+        new Set([]),
         new Set(['a', 'c', 'd', 'e']),
         {},
         (id: string) =>
@@ -499,7 +502,7 @@ describe('collectOverrideRevisionDetails', () => {
     ])
   })
 
-  it('should override the revision details for finished but unfetched checkpoint tips', () => {
+  it('should override the revision details for finished but unfetched checkpoint tips (based on available data)', () => {
     const justFinishedRunningId = 'exp-was-running'
     const { overrideComparison, overrideRevisions } =
       collectOverrideRevisionDetails(
@@ -517,6 +520,7 @@ describe('collectOverrideRevisionDetails', () => {
           { label: 'c' },
           { label: 'd' }
         ] as SelectedExperimentWithColor[],
+        new Set([]),
         new Set(['a', 'c', 'd', 'e']),
         { [justFinishedRunningId]: justFinishedRunningId },
         (id: string) =>

diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts
@@ -551,16 +551,17 @@ export const collectCommitRevisionDetails = (
   return commitRevisions
 }
 
-const getRevision = (
+const getOverrideRevision = (
   displayColor: Color,
   experiment: Experiment,
-  firstThreeColumns: string[]
+  firstThreeColumns: string[],
+  fetchedRevs: Set<string>
 ): Revision => {
   const { commit, displayNameOrParent, logicalGroupName, id, label } =
     experiment
   const revision: Revision = {
     displayColor,
-    fetched: true,
+    fetched: fetchedRevs.has(label),
     firstThreeColumns: getRevisionFirstThreeColumns(
       firstThreeColumns,
       experiment
@@ -580,18 +581,20 @@ const overrideWithWorkspace = (
   selectedWithOverrides: Revision[],
   displayColor: Color,
   label: string,
-  firstThreeColumns: string[]
+  firstThreeColumns: string[],
+  fetchedRevs: Set<string>
 ): void => {
   orderMapping[label] = EXPERIMENT_WORKSPACE_ID
   selectedWithOverrides.push(
-    getRevision(
+    getOverrideRevision(
       displayColor,
       {
         id: EXPERIMENT_WORKSPACE_ID,
         label: EXPERIMENT_WORKSPACE_ID,
         logicalGroupName: undefined
       },
-      firstThreeColumns
+      firstThreeColumns,
+      fetchedRevs
     )
   )
 }
@@ -610,15 +613,18 @@ const isExperimentThatWillDisappearAtEnd = (
 const getMostRecentFetchedCheckpointRevision = (
   selectedRevision: SelectedExperimentWithColor,
   fetchedRevs: Set<string>,
+  revisionsWithData: Set<string>,
   checkpoints: Experiment[] | undefined,
   firstThreeColumns: string[]
 ): Revision => {
   const mostRecent =
-    checkpoints?.find(({ label }) => fetchedRevs.has(label)) || selectedRevision
-  return getRevision(
+    checkpoints?.find(({ label }) => revisionsWithData.has(label)) ||
+    selectedRevision
+  return getOverrideRevision(
     selectedRevision.displayColor,
     mostRecent,
-    firstThreeColumns
+    firstThreeColumns,
+    fetchedRevs
   )
 }
 
@@ -627,6 +633,7 @@ const overrideRevisionDetail = (
   selectedWithOverrides: Revision[],
   selectedRevision: SelectedExperimentWithColor,
   fetchedRevs: Set<string>,
+  revisionsWithData: Set<string>,
   checkpoints: Experiment[] | undefined,
   firstThreeColumns: string[]
 ) => {
@@ -635,6 +642,7 @@ const overrideRevisionDetail = (
   const mostRecent = getMostRecentFetchedCheckpointRevision(
     selectedRevision,
     fetchedRevs,
+    revisionsWithData,
     checkpoints,
     firstThreeColumns
   )
@@ -647,6 +655,7 @@ const collectRevisionDetail = (
   selectedWithOverrides: Revision[],
   selectedRevision: SelectedExperimentWithColor,
   fetchedRevs: Set<string>,
+  revisionsWithData: Set<string>,
   unfinishedRunningExperiments: { [id: string]: string },
   getCheckpoints: (id: string) => Experiment[] | undefined,
   firstThreeColumns: string[]
@@ -662,7 +671,8 @@ const collectRevisionDetail = (
       selectedWithOverrides,
       displayColor,
       label,
-      firstThreeColumns
+      firstThreeColumns,
+      fetchedRevs
     )
   }
 
@@ -679,21 +689,28 @@ const collectRevisionDetail = (
       selectedWithOverrides,
       selectedRevision,
       fetchedRevs,
+      revisionsWithData,
       getCheckpoints(id),
       firstThreeColumns
     )
   }
 
   orderMapping[label] = label
   selectedWithOverrides.push(
-    getRevision(displayColor, selectedRevision, firstThreeColumns)
+    getOverrideRevision(
+      displayColor,
+      selectedRevision,
+      firstThreeColumns,
+      fetchedRevs
+    )
   )
 }
 
 export const collectOverrideRevisionDetails = (
   comparisonOrder: string[],
   selectedRevisions: SelectedExperimentWithColor[],
   fetchedRevs: Set<string>,
+  revisionsWithData: Set<string>,
   unfinishedRunningExperiments: { [id: string]: string },
   getCheckpoints: (id: string) => Experiment[] | undefined,
   firstThreeColumns: string[]
@@ -710,6 +727,7 @@ export const collectOverrideRevisionDetails = (
       selectedWithOverrides,
       selectedRevision,
       fetchedRevs,
+      revisionsWithData,
       unfinishedRunningExperiments,
       getCheckpoints,
       firstThreeColumns

diff --git a/extension/src/plots/model/index.ts b/extension/src/plots/model/index.ts
@@ -232,6 +232,10 @@ export class PlotsModel extends ModelWithPersistence {
         this.comparisonOrder,
         this.experiments.getSelectedRevisions(),
         this.fetchedRevs,
+        new Set([
+          ...Object.keys(this.comparisonData),
+          ...Object.keys(this.revisionData)
+        ]),
         finishedExperiments,
         id => this.experiments.getCheckpoints(id),
         this.experiments.getFirstThreeColumnOrder()