Skip to content

Commit

Permalink
Display DAG in Markdown Preview editor (#4244)
Browse files Browse the repository at this point in the history
* wire up prototype for displaying DAG

* recommend installing mermaid support extension

* update package.json

* add unit test for extension recommendation

* add integration test for showing the DAG

* stub call to dag in extension test

* handle dag command throwing error
  • Loading branch information
mattseddon committed Jul 10, 2023
1 parent 25248b1 commit ee124ed
Show file tree
Hide file tree
Showing 18 changed files with 362 additions and 19 deletions.
23 changes: 12 additions & 11 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -143,17 +143,18 @@ These are the VS Code [settings] available for the Extension:

[settings]: https://code.visualstudio.com/docs/getstarted/settings

| **Option** | **Description** |
| -------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `dvc.dvcPath` | Path or shell command to the DVC binary. Required unless Microsoft's [Python extension] is installed and the `dvc` package found in its environment. |
| `dvc.pythonPath` | Path to the desired Python interpreter to use with DVC. Should only be utilized when using a virtual environment without Microsoft's [Python extension]. |
| `dvc.experimentsTableHeadMaxHeight` | Maximum height of experiment table head rows. |
| `dvc.focusedProjects` | A subset of paths to the workspace's available DVC projects. Using this option will override project auto-discovery. |
| `dvc.doNotInformMaxExperimentsPlotted` | Do not inform when plotting more experiments is blocked (maximum number selected). |
| `dvc.doNotShowSetupAfterInstall` | Do not prompt to show the setup page after installing. Useful for pre-configured development environments. |
| `dvc.doNotRecommendAddStudioToken` | Do not prompt to add a [studio.token] to the global DVC config, which enables automatic sharing of experiments to [Studio]. |
| `dvc.doNotRecommendRedHatExtension` | Do not prompt to install the Red Hat YAML extension, which helps with DVC YAML schema validation (`dvc.yaml` and `.dvc` files). |
| `dvc.doNotShowCliUnavailable` | Do not warn when the workspace contains a DVC project but the DVC binary is unavailable. |
| **Option** | **Description** |
| ------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `dvc.dvcPath` | Path or shell command to the DVC binary. Required unless Microsoft's [Python extension] is installed and the `dvc` package found in its environment. |
| `dvc.pythonPath` | Path to the desired Python interpreter to use with DVC. Should only be utilized when using a virtual environment without Microsoft's [Python extension]. |
| `dvc.experimentsTableHeadMaxHeight` | Maximum height of experiment table head rows. |
| `dvc.focusedProjects` | A subset of paths to the workspace's available DVC projects. Using this option will override project auto-discovery. |
| `dvc.doNotInformMaxExperimentsPlotted` | Do not inform when plotting more experiments is blocked (maximum number selected). |
| `dvc.doNotShowSetupAfterInstall` | Do not prompt to show the setup page after installing. Useful for pre-configured development environments. |
| `dvc.doNotRecommendAddStudioToken` | Do not prompt to add a [studio.token] to the global DVC config, which enables automatic sharing of experiments to [Studio]. |
| `dvc.doNotRecommendRedHatExtension` | Do not prompt to install the Red Hat YAML extension, which helps with DVC YAML schema validation (`dvc.yaml` and `.dvc` files). |
| `dvc.doNotRecommendMermaidSupportExtension` | Do not prompt to install the Markdown Preview Mermaid Support extension, which helps to visualize DVC pipeline DAGs. |
| `dvc.doNotShowCliUnavailable` | Do not warn when the workspace contains a DVC project but the DVC binary is unavailable. |

> **Note** that the `Setup The Workspace` command helps you set up the basic
> ones at the [Workspace level] (saved to `.vscode/setting.json`).
Expand Down
14 changes: 14 additions & 0 deletions extension/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -417,6 +417,11 @@
"command": "dvc.showOutput",
"category": "DVC"
},
{
"title": "Show Pipeline DAG",
"command": "dvc.showPipelineDAG",
"category": "DVC"
},
{
"title": "Show Plots",
"command": "dvc.showPlots",
Expand Down Expand Up @@ -606,6 +611,11 @@
"type": "boolean",
"default": null
},
"dvc.doNotRecommendMermaidSupportExtension": {
"description": "Do not prompt to install the Markdown Preview Mermaid Support extension, which helps to visualize DVC pipeline DAGs",
"type": "boolean",
"default": null
},
"dvc.doNotRecommendAddStudioToken": {
"description": "Do not prompt to add a studio.token to the global DVC config, which enables automatic sharing of experiments to Studio.",
"type": "boolean",
Expand Down Expand Up @@ -876,6 +886,10 @@
"command": "dvc.showCommands",
"when": "false"
},
{
"command": "dvc.showPipelineDAG",
"when": "dvc.commands.available && dvc.project.available"
},
{
"command": "dvc.showExperiments",
"when": "dvc.commands.available && dvc.project.available"
Expand Down
4 changes: 4 additions & 0 deletions extension/src/cli/dvc/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ import { join } from 'path'
export const UNEXPECTED_ERROR_CODE = 255
export const DOT_DVC = '.dvc'

export const TEMP_DAG_FILE = join(DOT_DVC, 'tmp', 'dag.md')

export const TEMP_PLOTS_DIR = join(DOT_DVC, 'tmp', 'plots')

const TEMP_EXP_DIR = join(DOT_DVC, 'tmp', 'exps')
Expand All @@ -22,6 +24,7 @@ export enum Command {
COMMIT = 'commit',
CONFIG = 'config',
DATA = 'data',
DAG = 'dag',
EXPERIMENT = 'exp',
INITIALIZE = 'init',
MOVE = 'move',
Expand Down Expand Up @@ -59,6 +62,7 @@ export enum Flag {
JSON = '--json',
KILL = '--kill',
LOCAL = '--local',
MD = '--md',
PROJECT = '--project',
NUM_COMMIT = '-n',
OUTPUT_PATH = '-o',
Expand Down
48 changes: 47 additions & 1 deletion extension/src/cli/dvc/reader.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ beforeEach(() => {
mockedGetProcessEnv.mockReturnValueOnce(mockedEnv)
})

describe('CliReader', () => {
describe('DvcReader', () => {
mockedDisposable.fn.mockReturnValueOnce({
track: function <T>(disposable: T): T {
return disposable
Expand Down Expand Up @@ -70,6 +70,52 @@ describe('CliReader', () => {
}
)

describe('dag', () => {
it('should match the expected output', async () => {
const cwd = __dirname
const dag = `\`\`\`mermaid
flowchart TD
node1["nested1/data/data.xml.dvc"]
node2["nested1/dvc.yaml:evaluate"]
node3["nested1/dvc.yaml:featurize"]
node4["nested1/dvc.yaml:prepare"]
node5["nested1/dvc.yaml:train"]
node1-->node4
node3-->node2
node3-->node5
node4-->node3
node5-->node2
node6["nested2/data/data.xml.dvc"]
\`\`\``
mockedCreateProcess.mockReturnValueOnce(getMockedProcess(dag))

const cliOutput = await dvcReader.dag(cwd)
expect(cliOutput).toStrictEqual(dag)
expect(mockedCreateProcess).toHaveBeenCalledWith({
args: ['dag', '--md'],
cwd,
env: mockedEnv,
executable: 'dvc'
})
})

it('should return the error if the cli returns any type of error', async () => {
const cwd = __dirname
const error = new Error('unexpected error - something something')
const unexpectedStderr = 'This is very unexpected'
;(error as MaybeConsoleError).exitCode = UNEXPECTED_ERROR_CODE
;(error as MaybeConsoleError).stderr = unexpectedStderr
mockedCreateProcess.mockImplementationOnce(() => {
throw error
})

const cliOutput = await dvcReader.dag(cwd)
expect(cliOutput).toStrictEqual(
'Error: unexpected error - something something'
)
})
})

describe('expShow', () => {
it('should match the expected output', async () => {
const cwd = __dirname
Expand Down
9 changes: 9 additions & 0 deletions extension/src/cli/dvc/reader.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ export const isDvcError = <
!!(Object.keys(dataOrError).length === 1 && (dataOrError as DvcError).error)

export const autoRegisteredCommands = {
DAG: 'dag',
DATA_STATUS: 'dataStatus',
EXP_SHOW: 'expShow',
GLOBAL_VERSION: 'globalVersion',
Expand All @@ -47,6 +48,14 @@ export class DvcReader extends DvcCli {
this
)

public async dag(cwd: string) {
try {
return await this.executeDvcProcess(cwd, Command.DAG, Flag.MD)
} catch (error: unknown) {
return (error as Error).toString()
}
}

public dataStatus(
cwd: string,
...args: Args
Expand Down
2 changes: 2 additions & 0 deletions extension/src/commands/external.ts
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ export enum RegisteredCommands {
EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopExperiment',
STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',

PIPELINE_SHOW_DAG = 'dvc.showPipelineDAG',

PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',
PLOTS_SHOW = 'dvc.showPlots',
PLOTS_SELECT = 'dvc.views.plotsPathsTree.selectPlots',
Expand Down
7 changes: 4 additions & 3 deletions extension/src/data/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@ export type ExperimentsOutput = {
}

export abstract class BaseData<
T extends { data: PlotsOutputOrError; revs: string[] } | ExperimentsOutput
T extends
| { data: PlotsOutputOrError; revs: string[] }
| ExperimentsOutput
| string
> extends DeferredDisposable {
public readonly onDidUpdate: Event<T>
public readonly onDidChangeDvcYaml: Event<void>
Expand Down Expand Up @@ -102,6 +105,4 @@ export abstract class BaseData<
}

abstract managedUpdate(path?: string): Promise<unknown>

protected abstract collectFiles(data: T): void
}
11 changes: 10 additions & 1 deletion extension/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,16 @@ import { registerSetupCommands } from './setup/commands/register'
import { Status } from './status'
import { registerPersistenceCommands } from './persistence/register'
import { showSetupOrExecuteCommand } from './commands/util'
import { WorkspacePipeline } from './pipeline/workspace'
import { registerPipelineCommands } from './pipeline/register'

class Extension extends Disposable {
protected readonly internalCommands: InternalCommands

private readonly resourceLocator: ResourceLocator
private readonly repositories: WorkspaceRepositories
private readonly experiments: WorkspaceExperiments
private readonly pipelines: WorkspacePipeline
private readonly plots: WorkspacePlots
private readonly setup: Setup
private readonly repositoriesTree: RepositoriesTree
Expand Down Expand Up @@ -117,6 +120,10 @@ class Extension extends Disposable {
new WorkspaceExperiments(this.internalCommands, context.workspaceState)
)

this.pipelines = this.dispose.track(
new WorkspacePipeline(this.internalCommands)
)

this.plots = this.dispose.track(
new WorkspacePlots(this.internalCommands, context.workspaceState)
)
Expand Down Expand Up @@ -184,6 +191,7 @@ class Extension extends Disposable {
this.internalCommands,
this.setup
)
registerPipelineCommands(this.pipelines, this.internalCommands)
registerPlotsCommands(this.plots, this.internalCommands, this.setup)
registerSetupCommands(this.setup, this.internalCommands)
this.internalCommands.registerExternalCommand(
Expand Down Expand Up @@ -272,7 +280,8 @@ class Extension extends Disposable {
await Promise.all([
this.repositories.create(this.getRoots()),
this.repositoriesTree.initialize(this.getRoots()),
this.experiments.create(this.getRoots(), this.resourceLocator)
this.experiments.create(this.getRoots(), this.resourceLocator),
this.pipelines.create(this.getRoots())
])
this.plots.create(this.getRoots(), this.resourceLocator, this.experiments)

Expand Down
27 changes: 27 additions & 0 deletions extension/src/pipeline/data.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
import { AvailableCommands, InternalCommands } from '../commands/internal'
import { BaseData } from '../data'

export class PipelineData extends BaseData<string> {
constructor(dvcRoot: string, internalCommands: InternalCommands) {
super(
dvcRoot,
internalCommands,
[{ name: 'update', process: () => this.update() }],
['dvc.yaml']
)

void this.managedUpdate()
}

public managedUpdate() {
return this.processManager.run('update')
}

public async update(): Promise<void> {
const dag = await this.internalCommands.executeCommand(
AvailableCommands.DAG,
this.dvcRoot
)
return this.notifyChanged(dag)
}
}
34 changes: 34 additions & 0 deletions extension/src/pipeline/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
import { join } from 'path'
import { appendFileSync, writeFileSync } from 'fs-extra'
import { PipelineData } from './data'
import { DeferredDisposable } from '../class/deferred'
import { InternalCommands } from '../commands/internal'
import { TEMP_DAG_FILE } from '../cli/dvc/constants'

export class Pipeline extends DeferredDisposable {
private readonly dvcRoot: string
private readonly data: PipelineData

constructor(dvcRoot: string, internalCommands: InternalCommands) {
super()
this.dvcRoot = dvcRoot
this.data = this.dispose.track(new PipelineData(dvcRoot, internalCommands))

void this.initialize()
}

public forceRerender() {
return appendFileSync(join(this.dvcRoot, TEMP_DAG_FILE), '\n')
}

private async initialize() {
this.dispose.track(
this.data.onDidUpdate(data =>
writeFileSync(join(this.dvcRoot, TEMP_DAG_FILE), data)
)
)

await this.data.isReady()
return this.deferred.resolve()
}
}
13 changes: 13 additions & 0 deletions extension/src/pipeline/register.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import { WorkspacePipeline } from './workspace'
import { RegisteredCommands } from '../commands/external'
import { InternalCommands } from '../commands/internal'

export const registerPipelineCommands = (
pipelines: WorkspacePipeline,
internalCommands: InternalCommands
): void => {
internalCommands.registerExternalCommand(
RegisteredCommands.PIPELINE_SHOW_DAG,
() => pipelines.showDag()
)
}
63 changes: 63 additions & 0 deletions extension/src/pipeline/workspace.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import { join } from 'path'
import { commands, Uri } from 'vscode'
import { Pipeline } from '.'
import { TEMP_DAG_FILE } from '../cli/dvc/constants'
import { BaseWorkspace } from '../workspace'
import {
MARKDOWN_MERMAID_EXTENSION_ID,
recommendMermaidSupportExtension
} from '../vscode/recommend'
import { InternalCommands } from '../commands/internal'
import { getOnDidChangeExtensions, isInstalled } from '../vscode/extensions'

export class WorkspacePipeline extends BaseWorkspace<Pipeline> {
private isMermaidSupportInstalled = isInstalled(MARKDOWN_MERMAID_EXTENSION_ID)

constructor(internalCommands: InternalCommands) {
super(internalCommands)

const onDidChangeExtensions = getOnDidChangeExtensions()
this.dispose.track(
onDidChangeExtensions(() => {
const wasMermaidInstalled = this.isMermaidSupportInstalled
this.isMermaidSupportInstalled = isInstalled(
MARKDOWN_MERMAID_EXTENSION_ID
)
if (!wasMermaidInstalled && this.isMermaidSupportInstalled) {
this.renderDagAsMermaid()
}
})
)
}

public createRepository(dvcRoot: string) {
const pipeline = this.dispose.track(
new Pipeline(dvcRoot, this.internalCommands)
)

this.setRepository(dvcRoot, pipeline)

return pipeline
}

public async showDag() {
const cwd = await this.getOnlyOrPickProject()

if (!cwd) {
return
}

void recommendMermaidSupportExtension()

return commands.executeCommand(
'markdown.showPreview',
Uri.file(join(cwd, TEMP_DAG_FILE))
)
}

private renderDagAsMermaid() {
for (const dvcRoot of this.getDvcRoots()) {
void this.getRepository(dvcRoot).forceRerender()
}
}
}
2 changes: 2 additions & 0 deletions extension/src/telemetry/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@ export interface IEventNamePropertyMapping {
[EventName.MODIFY_WORKSPACE_PARAMS_RESET_AND_RUN]: undefined
[EventName.STOP_EXPERIMENTS]: { stopped: boolean; wasRunning: boolean }

[EventName.PIPELINE_SHOW_DAG]: undefined

[EventName.PLOTS_PATH_TOGGLE]: undefined
[EventName.PLOTS_SHOW]: undefined
[EventName.PLOTS_SELECT]: undefined
Expand Down
Loading

0 comments on commit ee124ed

Please sign in to comment.