diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts index ce0b7edd7db..717d1608ff1 100644 --- a/packages/core/src/awsService/sagemaker/commands.ts +++ b/packages/core/src/awsService/sagemaker/commands.ts @@ -191,13 +191,17 @@ export async function openRemoteConnect( void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage) return } + + const spaceName = node.spaceApp.SpaceName! await tryRefreshNode(node) + + // for Stopped SM spaces - check instance type before showing progress if (node.getStatus() === 'Stopped') { // In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client. const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode) try { - await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!) + await client.startSpace(spaceName, node.spaceApp.DomainId!) await tryRefreshNode(node) const appType = node.spaceApp.SpaceSettingsSummary?.AppType if (!appType) { @@ -205,18 +209,42 @@ export async function openRemoteConnect( code: 'undefinedAppType', }) } - await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType) - await tryRemoteConnection(node, ctx) + + // Only start showing progress after instance type validation + return await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + cancellable: false, + title: `Connecting to ${spaceName}`, + }, + async (progress) => { + progress.report({ message: 'Starting the space.' }) + await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType) + await tryRemoteConnection(node, ctx, progress) + } + ) } catch (err: any) { // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory - if (err.code !== InstanceTypeError) { - throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, { - cause: err as Error, - code: err.code, - }) + // just return without showing progress + if (err.code === InstanceTypeError) { + return } + throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, { + cause: err as Error, + code: err.code, + }) } } else if (node.getStatus() === 'Running') { - await tryRemoteConnection(node, ctx) + // For running spaces, show progress + return await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + cancellable: false, + title: `Connecting to ${spaceName}`, + }, + async (progress) => { + await tryRemoteConnection(node, ctx, progress) + } + ) } } diff --git a/packages/core/src/awsService/sagemaker/model.ts b/packages/core/src/awsService/sagemaker/model.ts index 3117fc3e22e..e25e8791d4f 100644 --- a/packages/core/src/awsService/sagemaker/model.ts +++ b/packages/core/src/awsService/sagemaker/model.ts @@ -27,12 +27,14 @@ const logger = getLogger('sagemaker') export async function tryRemoteConnection( node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode, - ctx: vscode.ExtensionContext + ctx: vscode.ExtensionContext, + progress: vscode.Progress<{ message?: string; increment?: number }> ) { const spaceArn = (await node.getSpaceArn()) as string const isSMUS = node instanceof SagemakerUnifiedStudioSpaceNode const remoteEnv = await prepareDevEnvConnection(spaceArn, ctx, 'sm_lc', isSMUS, node) try { + progress.report({ message: 'Opening remote session' }) await startVscodeRemote( remoteEnv.SessionProcess, remoteEnv.hostname,