Show Experiment Remote Status (#4324)

* prototype solution

* add remote exp refs test fixture

* review and refactor

* address review comments

extension/src/cli/git/reader.ts

@@ -1,6 +1,6 @@
 import { resolve } from 'path'
 import { GitCli } from '.'
-import { Command, Flag } from './constants'
+import { Command, DEFAULT_REMOTE, Flag } from './constants'
 import { getOptions } from './options'
 import { typeCheckCommands } from '..'
 import { trimAndSplit } from '../../util/stdout'
@@ -10,6 +10,7 @@ export const autoRegisteredCommands = {
   GIT_GET_BRANCHES: 'getBranches',
   GIT_GET_COMMIT_MESSAGES: 'getCommitMessages',
   GIT_GET_NUM_COMMITS: 'getNumCommits',
+  GIT_GET_REMOTE_EXPERIMENT_REFS: 'getRemoteExperimentRefs',
   GIT_GET_REMOTE_URL: 'getRemoteUrl',
   GIT_GET_REPOSITORY_ROOT: 'getGitRepositoryRoot',
   GIT_HAS_CHANGES: 'hasChanges',
@@ -68,6 +69,18 @@ export class GitReader extends GitCli {
     }
   }
 
+  public async getRemoteExperimentRefs(cwd: string): Promise<string> {
+    const options = getOptions({
+      args: [Command.LS_REMOTE, DEFAULT_REMOTE, 'refs/exps/*'],
+      cwd
+    })
+    try {
+      return await this.executeProcess(options)
+    } catch {
+      return ''
+    }
+  }
+
   public async getRemoteUrl(cwd: string): Promise<string> {
     const options = getOptions({ args: [Command.LS_REMOTE, Flag.GET_URL], cwd })
     try {

extension/src/data/index.ts

@@ -13,6 +13,7 @@ export type ExperimentsOutput = {
   availableNbCommits: { [branch: string]: number }
   expShow: ExpShowOutput
   gitLog: string
+  remoteExpRefs: string
   rowOrder: { branch: string; sha: string }[]
 }
 

extension/src/experiments/commands/index.ts

@@ -43,7 +43,11 @@ const convertUrlTextToLink = (stdout: string) => {
 }
 
 export const getPushExperimentCommand =
-  (internalCommands: InternalCommands, setup: Setup) =>
+  (
+    experiments: WorkspaceExperiments,
+    internalCommands: InternalCommands,
+    setup: Setup
+  ) =>
   ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) => {
     const studioAccessToken = setup.getStudioAccessToken()
     if (
@@ -56,23 +60,38 @@ export const getPushExperimentCommand =
     }
 
     return Toast.showProgress('exp push', async progress => {
+      const repository = experiments.getRepository(dvcRoot)
+
+      const updateOnCompletion = () => {
+        return repository.unsetPushing(ids)
+      }
+
       progress.report({ increment: 0 })
 
       progress.report({ increment: 25, message: `Pushing ${ids.join(' ')}...` })
 
       const remainingProgress = 75
 
-      const stdout = await internalCommands.executeCommand(
-        AvailableCommands.EXP_PUSH,
-        dvcRoot,
-        ...ids
-      )
+      repository.setPushing(ids)
+
+      try {
+        const stdout = await internalCommands.executeCommand(
+          AvailableCommands.EXP_PUSH,
+          dvcRoot,
+          ...ids
+        )
+
+        void updateOnCompletion()
 
-      progress.report({
-        increment: remainingProgress,
-        message: convertUrlTextToLink(stdout)
-      })
+        progress.report({
+          increment: remainingProgress,
+          message: convertUrlTextToLink(stdout)
+        })
 
-      return Toast.delayProgressClosing(15000)
+        return Toast.delayProgressClosing(15000)
+      } catch (error: unknown) {
+        void updateOnCompletion()
+        throw error
+      }
     })
   }

extension/src/experiments/commands/register.ts

@@ -277,7 +277,7 @@ export const registerExperimentCommands = (
 
   internalCommands.registerExternalCliCommand(
     RegisteredCliCommands.EXPERIMENT_VIEW_PUSH,
-    getPushExperimentCommand(internalCommands, setup)
+    getPushExperimentCommand(experiments, internalCommands, setup)
   )
 
   internalCommands.registerExternalCliCommand(

extension/src/experiments/data/index.ts

@@ -42,6 +42,19 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   }
 
   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
+    })
+  }
+
+  private async updateExpShow() {
     await this.updateBranches()
     const branches = this.experiments.getBranchesToShow()
     let gitLog = ''
@@ -66,12 +79,7 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     )
 
     this.collectFiles({ expShow })
-
-    return this.notifyChanged({ availableNbCommits, expShow, gitLog, rowOrder })
-  }
-
-  protected collectFiles({ expShow }: { expShow: ExpShowOutput }) {
-    this.collectedFiles = collectFiles(expShow, this.collectedFiles)
+    return { availableNbCommits, expShow, gitLog, rowOrder }
   }
 
   private async collectGitLogAndOrder(
@@ -119,6 +127,10 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
     this.experiments.setBranches(allBranches)
   }
 
+  private collectFiles({ expShow }: { expShow: ExpShowOutput }) {
+    this.collectedFiles = collectFiles(expShow, this.collectedFiles)
+  }
+
   private async watchExpGitRefs(): Promise<void> {
     const gitRoot = await this.internalCommands.executeCommand(
       AvailableCommands.GIT_GET_REPOSITORY_ROOT,
@@ -149,4 +161,11 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
       }
     )
   }
+
+  private updateRemoteExpRefs() {
+    return this.internalCommands.executeCommand(
+      AvailableCommands.GIT_GET_REMOTE_EXPERIMENT_REFS,
+      this.dvcRoot
+    )
+  }
 }

extension/src/experiments/index.ts

@@ -177,7 +177,8 @@ export class Experiments extends BaseRepository<TableData> {
     availableNbCommits,
     expShow,
     gitLog,
-    rowOrder
+    rowOrder,
+    remoteExpRefs
   }: ExperimentsOutput) {
     const hadCheckpoints = this.hasCheckpoints()
     const dvcLiveOnly = await this.checkSignalFile()
@@ -188,7 +189,8 @@ export class Experiments extends BaseRepository<TableData> {
         gitLog,
         dvcLiveOnly,
         rowOrder,
-        availableNbCommits
+        availableNbCommits,
+        remoteExpRefs
       )
     ])
 
@@ -199,6 +201,16 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  public setPushing(ids: string[]) {
+    this.experiments.setPushing(ids)
+    return this.notifyChanged()
+  }
+
+  public unsetPushing(ids: string[]) {
+    this.experiments.unsetPushing(ids)
+    return this.update()
+  }
+
   public hasCheckpoints() {
     return this.experiments.hasCheckpoints()
   }

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] = ['', false]
+const DEFAULT_DATA: [string, boolean, string] = ['', 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] = [
+  const expShowWithTwoCommits: [ExpShowOutput, string, boolean, string] = [
     generateTestExpShowOutput(
       {},
       {

extension/src/experiments/model/collect.ts

@@ -9,6 +9,7 @@ import { extractColumns } from '../columns/extract'
 import {
   CommitData,
   Experiment,
+  GitRemoteStatus,
   RunningExperiment,
   isQueued,
   isRunning
@@ -29,6 +30,7 @@ import { RegisteredCommands } from '../../commands/external'
 import { Resource } from '../../resourceLocator'
 import { shortenForLabel } from '../../util/string'
 import { COMMITS_SEPARATOR } from '../../cli/git/constants'
+import { trimAndSplit } from '../../util/stdout'
 
 export type ExperimentItem = {
   command?: {
@@ -233,6 +235,21 @@ 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
@@ -249,7 +266,8 @@ const collectRunningExperiment = (
 const collectExpRange = (
   acc: ExperimentsAccumulator,
   baseline: Experiment,
-  expRange: ExpRange
+  expRange: ExpRange,
+  remoteExpShas: Set<string>
 ): void => {
   const { revs, executor } = expRange
   if (!revs?.length) {
@@ -289,6 +307,7 @@ const collectExpRange = (
   }
 
   collectExecutorInfo(experiment, executor)
+  collectRemoteStatus(experiment, remoteExpShas)
   collectRunningExperiment(acc, experiment)
 
   addToMapArray(acc.experimentsByCommit, baselineId, experiment)
@@ -340,10 +359,20 @@ 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
+  dvcLiveOnly: boolean,
+  remoteExpRefs: string
 ): ExperimentsAccumulator => {
   const acc: ExperimentsAccumulator = {
     cliError: undefined,
@@ -358,6 +387,7 @@ export const collectExperiments = (
   }
 
   const commitData = collectCommitsData(gitLog)
+  const remoteExpShas = collectRemoteExpShas(remoteExpRefs)
 
   for (const expState of expShow) {
     const baseline = collectExpState(acc, expState, commitData)
@@ -368,7 +398,7 @@ export const collectExperiments = (
     }
 
     for (const expRange of experiments) {
-      collectExpRange(acc, baseline, expRange)
+      collectExpRange(acc, baseline, expRange, remoteExpShas)
     }
   }