Run remote experiments data update in the background (#4342)

Author: mattseddon

extension/src/data/index.ts

@@ -9,14 +9,22 @@ import { uniqueValues } from '../util/array'
 import { DeferredDisposable } from '../class/deferred'
 import { isPathInSubProject, isSameOrChild } from '../fileSystem'
 
-export type ExperimentsOutput = {
+type LocalExperimentsOutput = {
   availableNbCommits: { [branch: string]: number }
   expShow: ExpShowOutput
   gitLog: string
-  remoteExpRefs: string
   rowOrder: { branch: string; sha: string }[]
 }
 
+type RemoteExperimentsOutput = { remoteExpRefs: string }
+
+export type ExperimentsOutput = LocalExperimentsOutput | RemoteExperimentsOutput
+
+export const isRemoteExperimentsOutput = (
+  data: ExperimentsOutput
+): data is RemoteExperimentsOutput =>
+  (data as RemoteExperimentsOutput).remoteExpRefs !== undefined
+
 export abstract class BaseData<
   T extends
     | { data: PlotsOutputOrError; revs: string[] }
@@ -58,15 +66,13 @@ export abstract class BaseData<
     this.staticFiles = staticFiles
 
     this.watchFiles()
-
-    this.waitForInitialData()
   }
 
   protected notifyChanged(data: T) {
     this.updated.fire(data)
   }
 
-  private waitForInitialData() {
+  protected waitForInitialData() {
     const waitForInitialData = this.dispose.track(
       this.onDidUpdate(() => {
         this.dispose.untrack(waitForInitialData)

extension/src/experiments/data/index.ts

@@ -8,7 +8,11 @@ import { getRelativePattern } from '../../fileSystem/relativePattern'
 import { createFileSystemWatcher } from '../../fileSystem/watcher'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
 import { ExpShowOutput } from '../../cli/dvc/contract'
-import { BaseData, ExperimentsOutput } from '../../data'
+import {
+  BaseData,
+  ExperimentsOutput,
+  isRemoteExperimentsOutput
+} from '../../data'
 import { Args, DOT_DVC, ExperimentFlag } from '../../cli/dvc/constants'
 import { COMMITS_SEPARATOR, gitPath } from '../../cli/git/constants'
 import { getGitPath } from '../../fileSystem'
@@ -35,23 +39,18 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
 
     void this.watchExpGitRefs()
     void this.managedUpdate()
+    this.waitForInitialLocalData()
   }
 
   public managedUpdate() {
     return this.processManager.run('update')
   }
 
   public async update(): Promise<void> {
-    const [{ availableNbCommits, expShow, gitLog, rowOrder }, remoteExpRefs] =
-      await Promise.all([this.updateExpShow(), this.updateRemoteExpRefs()])
-
-    return this.notifyChanged({
-      availableNbCommits,
-      expShow,
-      gitLog,
-      remoteExpRefs,
-      rowOrder
-    })
+    await Promise.all([
+      this.notifyChanged(await this.updateExpShow()),
+      this.notifyChanged(await this.updateRemoteExpRefs())
+    ])
   }
 
   private async updateExpShow() {
@@ -162,10 +161,27 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     )
   }
 
-  private updateRemoteExpRefs() {
-    return this.internalCommands.executeCommand(
-      AvailableCommands.GIT_GET_REMOTE_EXPERIMENT_REFS,
-      this.dvcRoot
+  private async updateRemoteExpRefs() {
+    const [remoteExpRefs] = await Promise.all([
+      this.internalCommands.executeCommand(
+        AvailableCommands.GIT_GET_REMOTE_EXPERIMENT_REFS,
+        this.dvcRoot
+      ),
+      this.isReady()
+    ])
+    return { remoteExpRefs }
+  }
+
+  private waitForInitialLocalData() {
+    const waitForInitialData = this.dispose.track(
+      this.onDidUpdate(data => {
+        if (isRemoteExperimentsOutput(data)) {
+          return
+        }
+        this.dispose.untrack(waitForInitialData)
+        waitForInitialData.dispose()
+        this.deferred.resolve()
+      })
     )
   }
 }

extension/src/experiments/index.ts

@@ -33,7 +33,7 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { ExperimentsOutput } from '../data'
+import { ExperimentsOutput, isRemoteExperimentsOutput } from '../data'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -173,24 +173,25 @@ export class Experiments extends BaseRepository<TableData> {
     return this.data.managedUpdate()
   }
 
-  public async setState({
-    availableNbCommits,
-    expShow,
-    gitLog,
-    rowOrder,
-    remoteExpRefs
-  }: ExperimentsOutput) {
+  public async setState(data: ExperimentsOutput) {
+    if (isRemoteExperimentsOutput(data)) {
+      const { remoteExpRefs } = data
+      this.experiments.transformAndSetRemote(remoteExpRefs)
+      return this.webviewMessages.sendWebviewMessage()
+    }
+
+    const { expShow, gitLog, rowOrder, availableNbCommits } = data
+
     const hadCheckpoints = this.hasCheckpoints()
     const dvcLiveOnly = await this.checkSignalFile()
     await Promise.all([
       this.columns.transformAndSet(expShow),
-      this.experiments.transformAndSet(
+      this.experiments.transformAndSetLocal(
         expShow,
         gitLog,
         dvcLiveOnly,
         rowOrder,
-        availableNbCommits,
-        remoteExpRefs
+        availableNbCommits
       )
     ])
 
@@ -206,9 +207,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
-  public unsetPushing(ids: string[]) {
+  public async unsetPushing(ids: string[]) {
+    await this.update()
     this.experiments.unsetPushing(ids)
-    return this.update()
+    return this.notifyChanged()
   }
 
   public hasCheckpoints() {
@@ -579,7 +581,7 @@ export class Experiments extends BaseRepository<TableData> {
   private setupInitialData() {
     const waitForInitialData = this.dispose.track(
       this.onDidChangeExperiments(async () => {
-        await this.pipeline.isReady()
+        await Promise.all([this.data.isReady(), this.pipeline.isReady()])
         this.deferred.resolve()
         this.dispose.untrack(waitForInitialData)
         waitForInitialData.dispose()

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

@@ -2,7 +2,7 @@ import { collectExperiments } from './collect'
 import { generateTestExpShowOutput } from '../../test/util/experiments'
 import { ExpShowOutput } from '../../cli/dvc/contract'
 
-const DEFAULT_DATA: [string, boolean, string] = ['', false, '']
+const DEFAULT_DATA: [string, boolean] = ['', false]
 
 describe('collectExperiments', () => {
   it('should return an empty array if no commits are present', () => {
@@ -13,7 +13,7 @@ describe('collectExperiments', () => {
     expect(commits).toStrictEqual([])
   })
 
-  const expShowWithTwoCommits: [ExpShowOutput, string, boolean, string] = [
+  const expShowWithTwoCommits: [ExpShowOutput, string, boolean] = [
     generateTestExpShowOutput(
       {},
       {

extension/src/experiments/model/collect.ts

@@ -9,7 +9,6 @@ import { extractColumns } from '../columns/extract'
 import {
   CommitData,
   Experiment,
-  GitRemoteStatus,
   RunningExperiment,
   isQueued,
   isRunning
@@ -235,21 +234,6 @@ const collectExecutorInfo = (
   }
 }
 
-const collectRemoteStatus = (
-  experiment: Experiment,
-  remoteExpShas: Set<string>
-): void => {
-  if (
-    !experiment.sha ||
-    ![undefined, ExperimentStatus.SUCCESS].includes(experiment.status)
-  ) {
-    return
-  }
-  experiment.gitRemoteStatus = remoteExpShas.has(experiment.sha)
-    ? GitRemoteStatus.ON_REMOTE
-    : GitRemoteStatus.NOT_ON_REMOTE
-}
-
 const collectRunningExperiment = (
   acc: ExperimentsAccumulator,
   experiment: Experiment
@@ -266,8 +250,7 @@ const collectRunningExperiment = (
 const collectExpRange = (
   acc: ExperimentsAccumulator,
   baseline: Experiment,
-  expRange: ExpRange,
-  remoteExpShas: Set<string>
+  expRange: ExpRange
 ): void => {
   const { revs, executor } = expRange
   if (!revs?.length) {
@@ -307,7 +290,6 @@ const collectExpRange = (
   }
 
   collectExecutorInfo(experiment, executor)
-  collectRemoteStatus(experiment, remoteExpShas)
   collectRunningExperiment(acc, experiment)
 
   addToMapArray(acc.experimentsByCommit, baselineId, experiment)
@@ -359,20 +341,10 @@ const collectCliError = (
   }
 }
 
-const collectRemoteExpShas = (remoteExpRefs: string): Set<string> => {
-  const remoteExpShas = new Set<string>()
-  for (const ref of trimAndSplit(remoteExpRefs)) {
-    const [sha] = ref.split(/\s/)
-    remoteExpShas.add(sha)
-  }
-  return remoteExpShas
-}
-
 export const collectExperiments = (
   expShow: ExpShowOutput,
   gitLog: string,
-  dvcLiveOnly: boolean,
-  remoteExpRefs: string
+  dvcLiveOnly: boolean
 ): ExperimentsAccumulator => {
   const acc: ExperimentsAccumulator = {
     cliError: undefined,
@@ -387,7 +359,6 @@ export const collectExperiments = (
   }
 
   const commitData = collectCommitsData(gitLog)
-  const remoteExpShas = collectRemoteExpShas(remoteExpRefs)
 
   for (const expState of expShow) {
     const baseline = collectExpState(acc, expState, commitData)
@@ -398,7 +369,7 @@ export const collectExperiments = (
     }
 
     for (const expRange of experiments) {
-      collectExpRange(acc, baseline, expRange, remoteExpShas)
+      collectExpRange(acc, baseline, expRange)
     }
   }
 
@@ -534,3 +505,12 @@ export const collectBranches = (
 
   return { branches, currentBranch }
 }
+
+export const collectRemoteExpShas = (remoteExpRefs: string): Set<string> => {
+  const remoteExpShas = new Set<string>()
+  for (const ref of trimAndSplit(remoteExpRefs)) {
+    const [sha] = ref.split(/\s/)
+    remoteExpShas.add(sha)
+  }
+  return remoteExpShas
+}