diff --git a/packages/databricks-vscode/DATABRICKS.quickstart.md b/packages/databricks-vscode/DATABRICKS.quickstart.md index 688416493..00ae9b7fd 100644 --- a/packages/databricks-vscode/DATABRICKS.quickstart.md +++ b/packages/databricks-vscode/DATABRICKS.quickstart.md @@ -48,12 +48,12 @@ The Databricks extension for Visual Studio Code enables you to connect to your r If your folder has multiple [Declarative Automation Bundles](#dabs), you can select which one to use by clicking "Open Existing Databricks project" button and selecting the desired project. -## Select a cluster +## Select a compute -The extension uses an interactive cluster to run code. To select an interactive cluster: +The extension uses an interactive compute to run code. To select an interactive compute: 1. Open the Databricks panel by clicking on the Databricks icon on the left -2. Click on the "Select Cluster" button. +2. Click on the "Select Compute" button. - If you wish to change the selected cluster, click on the "Configure Cluster" gear icon, next to the name of the selected cluster. ## Run Python code diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 10a0ffa86..27465bcd6 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -52,6 +52,13 @@ "types": "out/extension.d.ts", "contributes": { "commands": [ + { + "command": "databricks.ssh.startTunnel", + "title": "Start SSH Tunnel", + "category": "Databricks", + "icon": "$(remote)", + "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode" + }, { "command": "databricks.connection.logout", "title": "Logout", @@ -73,20 +80,20 @@ }, { "command": "databricks.connection.attachCluster", - "title": "Attach cluster", + "title": "Attach compute", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(plug)" }, { "command": "databricks.connection.attachClusterQuickPick", - "title": "Configure cluster", + "title": "Configure compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(gear)" }, { "command": "databricks.connection.detachCluster", - "title": "Detach cluster", + "title": "Detach compute", "category": "Databricks", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "icon": "$(debug-disconnect)" @@ -151,14 +158,14 @@ }, { "command": "databricks.cluster.start", - "title": "Start Cluster", + "title": "Start Compute", "icon": "$(debug-start)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" }, { "command": "databricks.cluster.stop", - "title": "Stop Cluster", + "title": "Stop Compute", "icon": "$(stop-circle)", "enablement": "databricks.context.activated && databricks.context.loggedIn && !databricks.context.remoteMode", "category": "Databricks" @@ -542,7 +549,7 @@ { "id": "clusterView", "when": "databricks.feature.views.cluster && !databricks.context.remoteMode", - "name": "Clusters" + "name": "Computes" }, { "id": "dabsResourceExplorerView", @@ -1101,6 +1108,11 @@ "command": "databricks.sync.stop", "when": "view == configurationView && viewItem =~ /^databricks.*sync.*is-running.*$/ && databricks.context.bundle.isDevTarget", "group": "navigation@0" + }, + { + "command": "databricks.ssh.startTunnel", + "when": "view == configurationView && viewItem == databricks.configuration.sshTunnel", + "group": "navigation@0" } ], "editor/title": [ @@ -1233,7 +1245,7 @@ "submenus": [ { "id": "databricks.cluster.filter", - "label": "Filter clusters ...", + "label": "Filter compute ...", "icon": "$(filter)" }, { @@ -1439,7 +1451,7 @@ "databricks.clusters.onlyShowAccessibleClusters": { "type": "boolean", "default": false, - "description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)" + "description": "Enable/disable filtering for only accessible computes (computes on which the current user can run code)" }, "databricks.overrideDatabricksConfigFile": { "type": "string", @@ -1455,7 +1467,7 @@ "views.workspace" ], "enumDescriptions": [ - "Show cluster view in the explorer.", + "Show compute view in the explorer.", "Show workspace browser in the explorer." ], "type": "string" diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index bd440c1d2..fc92c5486 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -128,6 +128,48 @@ describe(__filename, function () { assert.equal([command, ...args].join(" "), syncCommand); }); + it("should create ssh connect commands", () => { + const logFilePath = getTempLogFilePath(); + const cli = createCliWrapper(logFilePath); + + // Logging is configured via env vars, not CLI flags, so no --log-* + // args appear on the ssh connect command line. + + // Serverless: no --cluster / --auto-start-cluster. + let {args} = cli.getSshConnectCommand({compute: {type: "serverless"}}); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + ]); + + // Serverless GPU: --accelerator, no --cluster / --auto-start-cluster. + ({args} = cli.getSshConnectCommand({ + compute: {type: "serverless", accelerator: "GPU_1xA10"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + "--accelerator=GPU_1xA10", + ]); + + // Dedicated cluster: --cluster and --auto-start-cluster. + ({args} = cli.getSshConnectCommand({ + compute: {type: "cluster", clusterId: "1234-clusterid"}, + })); + assert.deepStrictEqual(args, [ + "ssh", + "connect", + "--ide=vscode", + "--auto-approve", + "--cluster=1234-clusterid", + "--auto-start-cluster", + ]); + }); + it("should list profiles when no config file exists", async () => { const logFilePath = getTempLogFilePath(); const cli = createCliWrapper(logFilePath); diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 17ca9170b..f7a3f89f8 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -10,6 +10,7 @@ import { Uri, commands, CancellationToken, + env, } from "vscode"; import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs"; import {promisify} from "node:util"; @@ -574,6 +575,53 @@ export class CliWrapper { }); } + /** + * Env vars for interactive CLI commands run in a terminal (e.g. `ssh + * connect`). Auth is forwarded via env vars, matching the bundle init flow. + */ + getSshConnectEnvVars(authProvider: AuthProvider) { + return removeUndefinedKeys({ + ...EnvVarGenerators.getEnvVarsForCli( + this.extensionContext, + workspaceConfigs.databrickscfgLocation + ), + ...EnvVarGenerators.getProxyEnvVars(), + ...this.getLogginEnvVars(), + ...authProvider.toEnv(), + // eslint-disable-next-line @typescript-eslint/naming-convention + DATABRICKS_OUTPUT_FORMAT: "text", + }); + } + + /** + * Constructs the `databricks ssh connect` command args for opening a remote + * IDE window. Serverless is the default when no cluster is given. + * + * The --ide flag matches the host editor so the CLI opens the right remote + * window: Cursor identifies itself via env.uriScheme === "cursor", + * everything else (VS Code, Insiders) uses vscode. + * + * Logging is configured out of band via the DATABRICKS_LOG_* env vars (see + * getSshConnectEnvVars), so we do not pass --log-* flags here. + */ + getSshConnectCommand(opts: { + compute: + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + }): {args: string[]} { + const ide = env.uriScheme === "cursor" ? "cursor" : "vscode"; + const args = ["ssh", "connect", `--ide=${ide}`, "--auto-approve"]; + if (opts.compute.type === "cluster") { + // Start a stopped single-user cluster when connecting. + args.push(`--cluster=${opts.compute.clusterId}`); + args.push("--auto-start-cluster"); + } else if (opts.compute.accelerator) { + // Serverless GPU: request a specific accelerator type. + args.push(`--accelerator=${opts.compute.accelerator}`); + } + return {args}; + } + async bundleInit( templateDirPath: string, outputDirPath: string, diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 1f4e63e1d..b9ac7005a 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -133,7 +133,7 @@ export class ConnectionCommands implements Disposable { ClusterItem | QuickPickItem >(); quickPick.title = - typeof title === "string" ? title : "Select Cluster"; + typeof title === "string" ? title : "Select Compute"; quickPick.keepScrollPosition = true; quickPick.busy = true; quickPick.canSelectMany = false; diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index fd7e799ec..050455f83 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -49,6 +49,7 @@ import {Events, Metadata} from "./telemetry/constants"; import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller"; import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits"; import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton"; +import {SshCommands} from "./ssh/SshCommands"; import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager"; import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue"; import { @@ -825,6 +826,17 @@ export async function activate( ) ); + // SSH tunnel (remote development) group + const sshCommands = new SshCommands(connectionManager, clusterModel, cli); + context.subscriptions.push( + sshCommands, + telemetry.registerCommand( + "databricks.ssh.startTunnel", + sshCommands.startTunnelCommand, + sshCommands + ) + ); + // Cluster group const clusterTreeDataProvider = new ClusterListDataProvider(clusterModel); const clusterCommands = new ClusterCommands( diff --git a/packages/databricks-vscode/src/ssh/SshCommands.ts b/packages/databricks-vscode/src/ssh/SshCommands.ts new file mode 100644 index 000000000..5d91186af --- /dev/null +++ b/packages/databricks-vscode/src/ssh/SshCommands.ts @@ -0,0 +1,236 @@ +import { + Disposable, + QuickPick, + QuickPickItem, + QuickPickItemKind, + ThemeIcon, + window, +} from "vscode"; +import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; +import {ClusterModel} from "../cluster/ClusterModel"; +import {ConnectionManager} from "../configuration/ConnectionManager"; +import { + ClusterItem, + formatQuickPickClusterDetails, +} from "../configuration/ConnectionCommands"; +import {CliWrapper} from "../cli/CliWrapper"; +import {onError} from "../utils/onErrorDecorator"; + +const SERVERLESS_LABEL = "$(cloud) Serverless"; + +type Compute = + | {type: "serverless"; accelerator?: string} + | {type: "cluster"; clusterId: string}; + +/** + * A serverless QuickPick item, optionally carrying a GPU accelerator type + * forwarded to `databricks ssh connect --accelerator`. Plain serverless has no + * accelerator. + */ +interface ServerlessItem extends QuickPickItem { + accelerator?: string; +} + +// Serverless compute options shown at the top of the picker: plain serverless +// plus the serverless GPU accelerator types supported by the CLI. +const SERVERLESS_ITEMS: ServerlessItem[] = [ + { + label: SERVERLESS_LABEL, + alwaysShow: true, + }, + { + label: "$(cloud) Serverless GPU 1xA10", + alwaysShow: true, + accelerator: "GPU_1xA10", + }, + { + label: "$(cloud) Serverless GPU 8xH100", + alwaysShow: true, + accelerator: "GPU_8xH100", + }, +]; + +export class SshCommands implements Disposable { + private disposables: Disposable[] = []; + + constructor( + private readonly connectionManager: ConnectionManager, + private readonly clusterModel: ClusterModel, + private readonly cli: CliWrapper + ) {} + + @onError({popup: {prefix: "Error starting SSH tunnel."}}) + async startTunnelCommand() { + const workspaceClient = this.connectionManager.workspaceClient; + const me = this.connectionManager.databricksWorkspace?.userName; + if (!workspaceClient || !me) { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return; + } + + const compute = await this.pickCompute(me); + if (compute === undefined) { + return; + } + await this.launchSshTunnel(compute); + } + + private pickCompute(me: string): Promise { + return new Promise((resolve) => { + const quickPick = window.createQuickPick< + ClusterItem | ServerlessItem + >(); + quickPick.title = "Select compute for SSH tunnel"; + quickPick.keepScrollPosition = true; + quickPick.busy = true; + quickPick.canSelectMany = false; + + const staticItems: ServerlessItem[] = [ + ...SERVERLESS_ITEMS, + { + label: "", + kind: QuickPickItemKind.Separator, + }, + ]; + quickPick.items = staticItems; + + const refreshItems = () => { + // Only dedicated single-user clusters owned by the current user + // can be used for an SSH tunnel. + const clusters = (this.clusterModel.roots ?? []).filter((c) => + c.isValidSingleUser(me) + ); + quickPick.items = staticItems.concat( + clusters.map((c) => { + const treeItem = + ClusterListDataProvider.clusterNodeToTreeItem(c); + return { + label: `$(${ + (treeItem.iconPath as ThemeIcon).id + }) ${c.name!} (${c.id})`, + detail: formatQuickPickClusterDetails(c), + cluster: c, + }; + }) + ); + // Clear the spinner only once clusters have actually loaded. + // On a cold first open the loader is still fetching, so we keep + // spinning and let onDidChange repaint when clusters arrive. + if (clusters.length > 0) { + quickPick.busy = false; + } + this.preselect(quickPick); + }; + + // Fallback so the spinner can't hang forever for a user with no + // eligible clusters (onDidChange may never add any). + const spinnerTimeout = setTimeout(() => { + quickPick.busy = false; + }, 10_000); + + // Register the change listener before triggering refresh() so no + // onDidChange fired by the (re)started loader can be missed. + const disposables: Disposable[] = [ + this.clusterModel.onDidChange(refreshItems), + quickPick, + {dispose: () => clearTimeout(spinnerTimeout)}, + ]; + + // Paint whatever is already cached first (fast path on reopen), then + // trigger a reload; fresh results stream in via onDidChange. + refreshItems(); + this.clusterModel.refresh(); + quickPick.show(); + + quickPick.onDidAccept(() => { + const selectedItem = quickPick.selectedItems[0]; + disposables.forEach((d) => d.dispose()); + if (selectedItem === undefined) { + resolve(undefined); + } else if ("cluster" in selectedItem) { + resolve({ + type: "cluster", + clusterId: selectedItem.cluster.id, + }); + } else { + resolve({ + type: "serverless", + accelerator: selectedItem.accelerator, + }); + } + }); + + quickPick.onDidHide(() => { + disposables.forEach((d) => d.dispose()); + // resolve(undefined); + }); + }); + } + + /** + * Pre-selects the compute the user already has configured locally: the + * attached single-user cluster if any, otherwise serverless. + */ + private preselect(quickPick: QuickPick) { + const currentCluster = this.connectionManager.cluster; + if (currentCluster?.isSingleUser()) { + const match = quickPick.items.find( + (i): i is ClusterItem => + "cluster" in i && i.cluster.id === currentCluster.id + ); + if (match) { + quickPick.activeItems = [match]; + return; + } + } + if (this.connectionManager.serverless) { + const serverlessItem = quickPick.items.find( + (i) => i.label === SERVERLESS_LABEL + ); + if (serverlessItem) { + quickPick.activeItems = [serverlessItem]; + } + } + } + + private async launchSshTunnel(compute: Compute) { + const authProvider = + this.connectionManager.databricksWorkspace?.authProvider; + const userName = this.connectionManager.databricksWorkspace?.userName; + if (!authProvider || !userName) { + window.showErrorMessage( + "Please connect to a Databricks workspace before starting an SSH tunnel." + ); + return; + } + + const {args} = this.cli.getSshConnectCommand({compute}); + + const env: Record = { + ...this.cli.getSshConnectEnvVars(authProvider), + // The remote window opens at the user's home folder. Forward the + // file the user is currently editing so the remote extension can + // auto-open it (see the remote-mode branch in extension.ts). The + // CLI has no folder/file flag, so we pass it out of band. + /* eslint-disable @typescript-eslint/naming-convention */ + DATABRICKS_REMOTE_HOME_FOLDER: `/Users/${userName}`, + /* eslint-enable @typescript-eslint/naming-convention */ + }; + + const terminal = window.createTerminal({ + name: "Databricks SSH Tunnel", + isTransient: true, + env, + strictEnv: false, + }); + this.disposables.push(terminal); + terminal.show(); + terminal.sendText(`${this.cli.escapedCliPath} ${args.join(" ")}`); + } + + dispose() { + this.disposables.forEach((d) => d.dispose()); + } +} diff --git a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts index 6a0814f74..6303e132c 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ClusterComponent.ts @@ -114,7 +114,7 @@ export class ClusterComponent extends BaseComponent { // We are logged in -> Select cluster prompt return [ { - label: LabelUtils.highlightedLabel("Select a cluster"), + label: LabelUtils.highlightedLabel("Select a compute"), collapsibleState: TreeItemCollapsibleState.None, contextValue: getContextValue("none"), iconPath: new ThemeIcon( @@ -123,7 +123,7 @@ export class ClusterComponent extends BaseComponent { ), id: TREE_ICON_ID, command: { - title: "Select a cluster", + title: "Select a compute", command: "databricks.connection.attachClusterQuickPick", }, }, diff --git a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts index 2fd63c23c..b977f496b 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/ConfigurationDataProvider.ts @@ -20,6 +20,7 @@ import {logging} from "@databricks/sdk-experimental"; import {Loggers} from "../../logger"; import {FeatureManager} from "../../feature-manager/FeatureManager"; import {EnvironmentComponent} from "./EnvironmentComponent"; +import {SshTunnelComponent} from "./SshTunnelComponent"; import {WorkspaceFolderComponent} from "./WorkspaceFolderComponent"; import {WorkspaceFolderManager} from "../../vscode-objs/WorkspaceFolderManager"; import {CodeSynchronizer} from "../../sync"; @@ -53,6 +54,7 @@ export class ConfigurationDataProvider private readonly workspaceFolderManager: WorkspaceFolderManager ) { this.components = [ + new SshTunnelComponent(this.connectionManager), new WorkspaceFolderComponent(this.workspaceFolderManager), new BundleTargetComponent(this.configModel), new AuthTypeComponent( diff --git a/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts new file mode 100644 index 000000000..6a697a398 --- /dev/null +++ b/packages/databricks-vscode/src/ui/configuration-view/SshTunnelComponent.ts @@ -0,0 +1,45 @@ +import {ThemeIcon, TreeItemCollapsibleState} from "vscode"; +import {BaseComponent} from "./BaseComponent"; +import {ConfigurationTreeItem} from "./types"; +import {ConnectionManager} from "../../configuration/ConnectionManager"; + +export class SshTunnelComponent extends BaseComponent { + constructor(private readonly connectionManager: ConnectionManager) { + super(); + this.disposables.push( + this.connectionManager.onDidChangeState(() => { + this.onDidChangeEmitter.fire(); + }) + ); + } + + private async getRoot(): Promise { + if (this.connectionManager.state !== "CONNECTED") { + return []; + } + + return [ + { + label: "Start SSH Tunnel", + iconPath: new ThemeIcon("remote"), + contextValue: "databricks.configuration.sshTunnel", + collapsibleState: TreeItemCollapsibleState.None, + tooltip: "Start Databricks remote development via SSH", + command: { + title: "Start SSH Tunnel", + command: "databricks.ssh.startTunnel", + }, + }, + ]; + } + + public async getChildren( + parent?: ConfigurationTreeItem + ): Promise { + if (parent === undefined) { + return this.getRoot(); + } + + return []; + } +}