Skip to content

Commit 2e1a219

Browse files
feat(sagemaker): add progress indicator during space remote connection process (#8247)
## Problem Users had no idea what was happening during connecting to remote space, which can take few seconds to connect. User clicks "Connect" → Nothing visible happens for few seconds → Either success or failure ## Solution Adds progress tracking for SageMaker space connections which shows connection status with space name in progress dialog Test cases are not being added because it is just showing the user about progress during Space operations for good user experience. ## Testing Tested for both SM-AI and SMUS Spaces. 1. https://github.com/user-attachments/assets/20b02db3-7b6b-44d3-8e83-66d4e468da2c 2. https://github.com/user-attachments/assets/e2196483-b935-45d3-ae00-86e4163d298e --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Laxman Reddy <[email protected]>
1 parent 7457cc6 commit 2e1a219

File tree

2 files changed

+40
-10
lines changed

2 files changed

+40
-10
lines changed

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

Lines changed: 37 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -191,32 +191,60 @@ export async function openRemoteConnect(
191191
void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
192192
return
193193
}
194+
195+
const spaceName = node.spaceApp.SpaceName!
194196
await tryRefreshNode(node)
197+
198+
// for Stopped SM spaces - check instance type before showing progress
195199
if (node.getStatus() === 'Stopped') {
196200
// In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
197201
const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
198202

199203
try {
200-
await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
204+
await client.startSpace(spaceName, node.spaceApp.DomainId!)
201205
await tryRefreshNode(node)
202206
const appType = node.spaceApp.SpaceSettingsSummary?.AppType
203207
if (!appType) {
204208
throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
205209
code: 'undefinedAppType',
206210
})
207211
}
208-
await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
209-
await tryRemoteConnection(node, ctx)
212+
213+
// Only start showing progress after instance type validation
214+
return await vscode.window.withProgress(
215+
{
216+
location: vscode.ProgressLocation.Notification,
217+
cancellable: false,
218+
title: `Connecting to ${spaceName}`,
219+
},
220+
async (progress) => {
221+
progress.report({ message: 'Starting the space.' })
222+
await client.waitForAppInService(node.spaceApp.DomainId!, spaceName, appType)
223+
await tryRemoteConnection(node, ctx, progress)
224+
}
225+
)
210226
} catch (err: any) {
211227
// Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
212-
if (err.code !== InstanceTypeError) {
213-
throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
214-
cause: err as Error,
215-
code: err.code,
216-
})
228+
// just return without showing progress
229+
if (err.code === InstanceTypeError) {
230+
return
217231
}
232+
throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
233+
cause: err as Error,
234+
code: err.code,
235+
})
218236
}
219237
} else if (node.getStatus() === 'Running') {
220-
await tryRemoteConnection(node, ctx)
238+
// For running spaces, show progress
239+
return await vscode.window.withProgress(
240+
{
241+
location: vscode.ProgressLocation.Notification,
242+
cancellable: false,
243+
title: `Connecting to ${spaceName}`,
244+
},
245+
async (progress) => {
246+
await tryRemoteConnection(node, ctx, progress)
247+
}
248+
)
221249
}
222250
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,14 @@ const logger = getLogger('sagemaker')
2727

2828
export async function tryRemoteConnection(
2929
node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
30-
ctx: vscode.ExtensionContext
30+
ctx: vscode.ExtensionContext,
31+
progress: vscode.Progress<{ message?: string; increment?: number }>
3132
) {
3233
const spaceArn = (await node.getSpaceArn()) as string
3334
const isSMUS = node instanceof SagemakerUnifiedStudioSpaceNode
3435
const remoteEnv = await prepareDevEnvConnection(spaceArn, ctx, 'sm_lc', isSMUS, node)
3536
try {
37+
progress.report({ message: 'Opening remote session' })
3638
await startVscodeRemote(
3739
remoteEnv.SessionProcess,
3840
remoteEnv.hostname,

0 commit comments

Comments
 (0)