Add:progress indicator during SageMaker space connection process

Author: bhavya2109sharma

packages/core/src/awsService/sagemaker/commands.ts

@@ -191,32 +191,69 @@ export async function openRemoteConnect(
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
-    await tryRefreshNode(node)
-    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 tryRefreshNode(node)
-            const appType = node.spaceApp.SpaceSettingsSummary?.AppType
-            if (!appType) {
-                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
-                    code: 'undefinedAppType',
-                })
-            }
-            await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
-            await tryRemoteConnection(node, ctx)
-        } catch (err: any) {
-            // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
-            if (err.code !== InstanceTypeError) {
+
+    const spaceName = node.spaceApp.SpaceName!
+
+    try {
+        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(spaceName, node.spaceApp.DomainId!)
+                await tryRefreshNode(node)
+                const appType = node.spaceApp.SpaceSettingsSummary?.AppType
+                if (!appType) {
+                    throw new ToolkitError(
+                        'AppType is undefined for the selected space. Cannot start remote connection.',
+                        {
+                            code: 'undefinedAppType',
+                        }
+                    )
+                }
+
+                // 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. This can take around 2 minutes.' })
+                        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
+                // 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') {
+            // 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)
+                }
+            )
         }
-    } else if (node.getStatus() === 'Running') {
-        await tryRemoteConnection(node, ctx)
+    } catch (err: any) {
+        getLogger().error(`sm:openRemoteConnect: ${err}`)
+        throw err
     }
 }

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,