Add progress ring to plots ribbon while data is loading (#2841)

* add progress ring to ribbon while plots are loading

* override revisions when there is a selected running checkpoint experiment

* prevent workspace infinite loop

* style progress ring

* revert stale changes

* restyle progress ring in ribbon

* fetch all available revisions from experiments for id to label mapping

* refactor code

* only push workspace into revs if it is not already included

extension/src/plots/data/index.ts

@@ -59,6 +59,10 @@ export class PlotsData extends BaseData<{
       ...args
     )
 
+    if (!revs.includes('workspace') && args.length < 2) {
+      revs.push('workspace')
+    }
+
     this.notifyChanged({ data, revs })
 
     const files = this.collectFiles({ data })

extension/src/plots/index.ts

@@ -140,8 +140,7 @@ export class Plots extends BaseRepository<TPlotsData> {
       this.paths.hasPaths() &&
       definedAndNonEmpty(this.plots.getUnfetchedRevisions())
     ) {
-      this.webviewMessages.sendCheckpointPlotsMessage()
-      return this.data.managedUpdate()
+      this.data.managedUpdate()
     }
 
     return this.webviewMessages.sendWebviewMessage()

extension/src/plots/model/collect.test.ts

@@ -9,20 +9,23 @@ import {
   collectTemplates,
   collectMetricOrder,
   collectWorkspaceRunningCheckpoint,
-  collectWorkspaceRaceConditionData
+  collectWorkspaceRaceConditionData,
+  collectOverrideRevisionDetails
 } from './collect'
 import plotsDiffFixture from '../../test/fixtures/plotsDiff/output'
 import expShowFixture from '../../test/fixtures/expShow/base/output'
 import modifiedFixture from '../../test/fixtures/expShow/modified/output'
 import checkpointPlotsFixture from '../../test/fixtures/expShow/base/checkpointPlots'
-import { ExperimentsOutput } from '../../cli/dvc/contract'
+import { ExperimentsOutput, ExperimentStatus } from '../../cli/dvc/contract'
 import {
   definedAndNonEmpty,
   sameContents,
   uniqueValues
 } from '../../util/array'
 import { TemplatePlot } from '../webview/contract'
 import { getCLIBranchId } from '../../test/fixtures/plotsDiff/util'
+import { SelectedExperimentWithColor } from '../../experiments/model'
+import { Experiment } from '../../experiments/webview/contract'
 
 const logsLossPath = join('logs', 'loss.tsv')
 
@@ -358,3 +361,112 @@ describe('collectWorkspaceRaceConditionData', () => {
     )
   })
 })
+
+describe('collectOverrideRevisionDetails', () => {
+  it('should override the revision details for running checkpoint tips', () => {
+    const runningId = 'b'
+
+    const { overrideOrder, overrideRevisions, unfinishedRunningExperiments } =
+      collectOverrideRevisionDetails(
+        ['a', 'b', 'c', 'd'],
+        [
+          { label: 'a' },
+          {
+            checkpoint_tip: 'b',
+            displayColor: '#13adc7',
+            id: runningId,
+            label: 'b',
+            sha: 'b',
+            status: ExperimentStatus.RUNNING
+          },
+          { label: 'c' },
+          { label: 'd' }
+        ] as SelectedExperimentWithColor[],
+        new Set(['a', 'c', 'd', 'e']),
+        new Set(),
+        (id: string) => ({ [runningId]: [{ label: 'e' }] as Experiment[] }[id])
+      )
+    expect(overrideOrder).toStrictEqual(['a', 'e', 'c', 'd'])
+    expect(overrideRevisions).toStrictEqual([
+      { label: 'a' },
+      {
+        displayColor: '#13adc7',
+        label: 'e'
+      },
+      { label: 'c' },
+      { label: 'd' }
+    ])
+    expect(unfinishedRunningExperiments).toStrictEqual(new Set([runningId]))
+  })
+
+  it('should override the revision details for finished but unfetched checkpoint tips', () => {
+    const justFinishedRunningId = 'exp-was-running'
+    const { overrideOrder, overrideRevisions, unfinishedRunningExperiments } =
+      collectOverrideRevisionDetails(
+        ['a', 'b', 'c', 'd'],
+        [
+          { label: 'a' },
+          {
+            checkpoint_tip: 'b',
+            displayColor: '#13adc7',
+            id: justFinishedRunningId,
+            label: 'b',
+            sha: 'b',
+            status: ExperimentStatus.SUCCESS
+          },
+          { label: 'c' },
+          { label: 'd' }
+        ] as SelectedExperimentWithColor[],
+        new Set(['a', 'c', 'd', 'e']),
+        new Set([justFinishedRunningId]),
+        (id: string) =>
+          ({ [justFinishedRunningId]: [{ label: 'e' }] as Experiment[] }[id])
+      )
+    expect(overrideOrder).toStrictEqual(['a', 'e', 'c', 'd'])
+    expect(overrideRevisions).toStrictEqual([
+      { label: 'a' },
+      {
+        displayColor: '#13adc7',
+        label: 'e'
+      },
+      { label: 'c' },
+      { label: 'd' }
+    ])
+    expect(unfinishedRunningExperiments).toStrictEqual(
+      new Set([justFinishedRunningId])
+    )
+  })
+
+  it('should remove the id from the unfinishedRunningExperiments set once the revision has been fetched', () => {
+    const justFinishedRunningId = 'exp-was-running'
+    const { overrideOrder, overrideRevisions, unfinishedRunningExperiments } =
+      collectOverrideRevisionDetails(
+        ['a', 'b', 'c', 'd'],
+        [
+          { label: 'a' },
+          {
+            checkpoint_tip: 'b',
+            displayColor: '#13adc7',
+            id: justFinishedRunningId,
+            label: 'b',
+            sha: 'b',
+            status: ExperimentStatus.SUCCESS
+          },
+          { label: 'c' },
+          { label: 'd' }
+        ] as SelectedExperimentWithColor[],
+        new Set(['a', 'b', 'c', 'd', 'e']),
+        new Set([justFinishedRunningId]),
+        (id: string) =>
+          ({ [justFinishedRunningId]: [{ label: 'e' }] as Experiment[] }[id])
+      )
+    expect(overrideOrder).toStrictEqual(['a', 'b', 'c', 'd'])
+    expect(overrideRevisions).toStrictEqual([
+      { label: 'a' },
+      expect.objectContaining({ label: 'b' }),
+      { label: 'c' },
+      { label: 'd' }
+    ])
+    expect(unfinishedRunningExperiments).toStrictEqual(new Set([]))
+  })
+})

extension/src/plots/model/collect.ts

@@ -28,7 +28,11 @@ import {
   decodeColumn,
   appendColumnToPath
 } from '../../experiments/columns/paths'
-import { MetricOrParamColumns } from '../../experiments/webview/contract'
+import {
+  Experiment,
+  isRunning,
+  MetricOrParamColumns
+} from '../../experiments/webview/contract'
 import { addToMapArray } from '../../util/map'
 import { TemplateOrder } from '../paths/collect'
 import { extendVegaSpec, isMultiViewPlot } from '../vega/util'
@@ -42,6 +46,7 @@ import {
   unmergeConcatenatedFields
 } from '../multiSource/collect'
 import { StrokeDashEncoding } from '../multiSource/constants'
+import { SelectedExperimentWithColor } from '../../experiments/model'
 
 type CheckpointPlotAccumulator = {
   iterations: Record<string, number>
@@ -711,3 +716,112 @@ export const collectBranchRevisionDetails = (
   }
   return branchRevisions
 }
+
+const getMostRecentFetchedCheckpointRevision = (
+  selectedRevision: SelectedExperimentWithColor,
+  fetchedRevs: Set<string>,
+  checkpoints: Experiment[] | undefined
+): SelectedExperimentWithColor => {
+  const mostRecent =
+    checkpoints?.find(({ label }) => fetchedRevs.has(label)) || selectedRevision
+  return {
+    ...mostRecent,
+    displayColor: selectedRevision.displayColor
+  } as SelectedExperimentWithColor
+}
+
+const overrideRevisionDetail = (
+  orderMapping: { [label: string]: string },
+  selectedWithFetchedRunningCheckpointRevs: SelectedExperimentWithColor[],
+  selectedRevision: SelectedExperimentWithColor,
+  fetchedRevs: Set<string>,
+  unfinishedRunningExperiments: Set<string>,
+  checkpoints: Experiment[] | undefined
+) => {
+  const { id, status, label } = selectedRevision
+
+  if (isRunning(status)) {
+    unfinishedRunningExperiments.add(id)
+  }
+
+  const mostRecent = getMostRecentFetchedCheckpointRevision(
+    selectedRevision,
+    fetchedRevs,
+    checkpoints
+  )
+  orderMapping[label] = mostRecent.label
+  selectedWithFetchedRunningCheckpointRevs.push(mostRecent)
+}
+
+const collectRevisionDetail = (
+  orderMapping: { [label: string]: string },
+  selectedWithFetchedRunningCheckpointRevs: SelectedExperimentWithColor[],
+  selectedRevision: SelectedExperimentWithColor,
+  fetchedRevs: Set<string>,
+  unfinishedRunningExperiments: Set<string>,
+  getCheckpoints: (id: string) => Experiment[] | undefined
+  // eslint-disable-next-line sonarjs/cognitive-complexity
+) => {
+  const { label, status, id, checkpoint_tip, sha } = selectedRevision
+
+  const isCheckpointTip = sha === checkpoint_tip
+  const running = isRunning(status)
+  const preventRevisionsDisappearingAtEnd =
+    isCheckpointTip && unfinishedRunningExperiments.has(id)
+
+  if (
+    !fetchedRevs.has(label) &&
+    (running || preventRevisionsDisappearingAtEnd)
+  ) {
+    return overrideRevisionDetail(
+      orderMapping,
+      selectedWithFetchedRunningCheckpointRevs,
+      selectedRevision,
+      fetchedRevs,
+      unfinishedRunningExperiments,
+      getCheckpoints(id)
+    )
+  }
+
+  if (!running && isCheckpointTip) {
+    unfinishedRunningExperiments.delete(id)
+  }
+
+  orderMapping[label] = label
+  selectedWithFetchedRunningCheckpointRevs.push(selectedRevision)
+}
+
+export const collectOverrideRevisionDetails = (
+  comparisonOrder: string[],
+  selectedRevisions: SelectedExperimentWithColor[],
+  fetchedRevs: Set<string>,
+  unfinishedRunningExperiments: Set<string>,
+  getCheckpoints: (id: string) => Experiment[] | undefined
+): {
+  overrideOrder: string[]
+  overrideRevisions: SelectedExperimentWithColor[]
+  unfinishedRunningExperiments: Set<string>
+} => {
+  const orderMapping: { [label: string]: string } = {}
+  const selectedWithFetchedRunningCheckpointRevs: SelectedExperimentWithColor[] =
+    []
+
+  for (const selectedRevision of selectedRevisions) {
+    collectRevisionDetail(
+      orderMapping,
+      selectedWithFetchedRunningCheckpointRevs,
+      selectedRevision,
+      fetchedRevs,
+      unfinishedRunningExperiments,
+      getCheckpoints
+    )
+  }
+
+  return {
+    overrideOrder: comparisonOrder
+      .map(label => orderMapping[label])
+      .filter(Boolean),
+    overrideRevisions: selectedWithFetchedRunningCheckpointRevs,
+    unfinishedRunningExperiments
+  }
+}