Skip to content

Commit 3c1c282

Browse files
NewtonDerNewton Der
NewtonDer
and
Newton Der
authored
feat(sagemaker): Show notification if instanceType has insufficient memory (#2157)
## Problem When the user has not started the space before (instanceType not defined), the extension would fall back to using `ml.t3.medium`, which has insufficient memory. Or if the user has attempted to restart the space with an instanceType with insufficient memory, the request to `StartSpace` would fail. Instance types with insufficient memory: - `ml.t3.medium` - `ml.c7i.large` - `ml.c6i.large` - `ml.c6id.large` - `ml.c5.large` ## Solution Show an error notification if user is attempting to start a space with insufficient memory. Suggest the user to use an upgraded instance type with more memory depending on the one they attempted to use. If the user confirms, then the call to `StartSpace` will continue with the recommended instanceType. - `ml.t3.medium` --> `ml.t3.large` - `ml.c7i.large` --> `ml.c7i.xlarge` - `ml.c6i.large` --> `ml.c6i.xlarge` - `ml.c6id.large` --> `ml.c6id.xlarge` - `ml.c5.large` --> `ml.c5.xlarge` --- - 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: Newton Der <[email protected]>
1 parent 6d9764b commit 3c1c282

File tree

4 files changed

+129
-26
lines changed

4 files changed

+129
-26
lines changed

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

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import { ExtContext } from '../../shared/extensions'
1818
import { SagemakerClient } from '../../shared/clients/sagemaker'
1919
import { ToolkitError } from '../../shared/errors'
2020
import { showConfirmationMessage } from '../../shared/utilities/messages'
21+
import { InstanceTypeError } from './constants'
2122

2223
const localize = nls.loadMessageBundle()
2324

@@ -158,14 +159,22 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
158159
export async function openRemoteConnect(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
159160
if (node.getStatus() === 'Stopped') {
160161
const client = new SagemakerClient(node.regionCode)
161-
await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
162-
await tryRefreshNode(node)
163-
const appType = node.spaceApp.SpaceSettingsSummary?.AppType
164-
if (!appType) {
165-
throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.')
162+
163+
try {
164+
await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
165+
await tryRefreshNode(node)
166+
const appType = node.spaceApp.SpaceSettingsSummary?.AppType
167+
if (!appType) {
168+
throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.')
169+
}
170+
await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
171+
await tryRemoteConnection(node, ctx)
172+
} catch (err: any) {
173+
// Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
174+
if (err.code !== InstanceTypeError) {
175+
throw err
176+
}
166177
}
167-
await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
168-
await tryRemoteConnection(node, ctx)
169178
} else if (node.getStatus() === 'Running') {
170179
await tryRemoteConnection(node, ctx)
171180
}

‎packages/core/src/awsService/sagemaker/constants.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
export const InstanceTypeError = 'InstanceTypeError'
7+
8+
export const InstanceTypeMinimum = 'ml.t3.large'
9+
10+
export const InstanceTypeInsufficientMemory: Record<string, string> = {
11+
'ml.t3.medium': 'ml.t3.large',
12+
'ml.c7i.large': 'ml.c7i.xlarge',
13+
'ml.c6i.large': 'ml.c6i.xlarge',
14+
'ml.c6id.large': 'ml.c6id.xlarge',
15+
'ml.c5.large': 'ml.c5.xlarge',
16+
}
17+
18+
export const InstanceTypeInsufficientMemoryMessage = (
19+
spaceName: string,
20+
chosenInstanceType: string,
21+
recommendedInstanceType: string
22+
) => {
23+
return `Unable to create app for [${spaceName}] because instanceType [${chosenInstanceType}] is not supported for remote access enabled spaces. Use instanceType with at least 8 GiB memory. Would you like to start your space with instanceType [${recommendedInstanceType}]?`
24+
}
25+
26+
export const InstanceTypeNotSelectedMessage = (spaceName: string) => {
27+
return `No instanceType specified for [${spaceName}]. ${InstanceTypeMinimum} is the default instance type, which meets minimum 8 GiB memory requirements for remote access. Continuing will start your space with instanceType [${InstanceTypeMinimum}] and remotely connect.`
28+
}

‎packages/core/src/shared/clients/sagemaker.ts‎

Lines changed: 67 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,17 @@ import { isEmpty } from 'lodash'
3737
import { sleep } from '../utilities/timeoutUtils'
3838
import { ClientWrapper } from './clientWrapper'
3939
import { AsyncCollection } from '../utilities/asyncCollection'
40+
import {
41+
InstanceTypeError,
42+
InstanceTypeMinimum,
43+
InstanceTypeInsufficientMemory,
44+
InstanceTypeInsufficientMemoryMessage,
45+
InstanceTypeNotSelectedMessage,
46+
} from '../../awsService/sagemaker/constants'
4047
import { getDomainSpaceKey } from '../../awsService/sagemaker/utils'
4148
import { getLogger } from '../logger/logger'
4249
import { ToolkitError } from '../errors'
50+
import { yes, no, continueText, cancel } from '../localizedText'
4351

4452
export interface SagemakerSpaceApp extends SpaceDetails {
4553
App?: AppDetails
@@ -85,7 +93,9 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
8593
}
8694

8795
public async startSpace(spaceName: string, domainId: string) {
88-
let spaceDetails
96+
let spaceDetails: DescribeSpaceCommandOutput
97+
98+
// Get existing space details
8999
try {
90100
spaceDetails = await this.describeSpace({
91101
DomainId: domainId,
@@ -95,6 +105,54 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
95105
throw this.handleStartSpaceError(err)
96106
}
97107

108+
// Get app type
109+
const appType = spaceDetails.SpaceSettings?.AppType
110+
if (appType !== 'JupyterLab' && appType !== 'CodeEditor') {
111+
throw new ToolkitError(`Unsupported AppType "${appType}" for space "${spaceName}"`)
112+
}
113+
114+
// Get app resource spec
115+
const requestedResourceSpec =
116+
appType === 'JupyterLab'
117+
? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec
118+
: spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec
119+
120+
let instanceType = requestedResourceSpec?.InstanceType
121+
122+
// Is InstanceType defined and has enough memory?
123+
if (instanceType && instanceType in InstanceTypeInsufficientMemory) {
124+
// Prompt user to select one with sufficient memory (1 level up from their chosen one)
125+
const response = await vscode.window.showErrorMessage(
126+
InstanceTypeInsufficientMemoryMessage(
127+
spaceDetails.SpaceName || '',
128+
instanceType,
129+
InstanceTypeInsufficientMemory[instanceType]
130+
),
131+
yes,
132+
no
133+
)
134+
135+
if (response === no) {
136+
throw new ToolkitError('InstanceType has insufficient memory.', { code: InstanceTypeError })
137+
}
138+
139+
instanceType = InstanceTypeInsufficientMemory[instanceType]
140+
} else if (!instanceType) {
141+
// Prompt user to select the minimum supported instance type
142+
const response = await vscode.window.showErrorMessage(
143+
InstanceTypeNotSelectedMessage(spaceDetails.SpaceName || ''),
144+
continueText,
145+
cancel
146+
)
147+
148+
if (response === cancel) {
149+
throw new ToolkitError('InstanceType not defined.', { code: InstanceTypeError })
150+
}
151+
152+
instanceType = InstanceTypeMinimum
153+
}
154+
155+
// Get remote access flag
98156
if (!spaceDetails.SpaceSettings?.RemoteAccess || spaceDetails.SpaceSettings?.RemoteAccess === 'DISABLED') {
99157
try {
100158
await this.updateSpace({
@@ -110,23 +168,17 @@ export class SagemakerClient extends ClientWrapper<SageMakerClient> {
110168
}
111169
}
112170

113-
const appType = spaceDetails.SpaceSettings?.AppType
114-
if (appType !== 'JupyterLab' && appType !== 'CodeEditor') {
115-
throw new ToolkitError(`Unsupported AppType "${appType}" for space "${spaceName}"`)
116-
}
117-
118-
const requestedResourceSpec =
119-
appType === 'JupyterLab'
120-
? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec
121-
: spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec
122-
123-
const fallbackResourceSpec: ResourceSpec = {
124-
InstanceType: 'ml.t3.medium',
171+
const resourceSpec: ResourceSpec = {
172+
// Default values
125173
SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu',
126174
SageMakerImageVersionAlias: '3.2.0',
127-
}
128175

129-
const resourceSpec = requestedResourceSpec?.InstanceType ? requestedResourceSpec : fallbackResourceSpec
176+
// The existing resource spec
177+
...requestedResourceSpec,
178+
179+
// The instance type user has chosen
180+
InstanceType: instanceType,
181+
}
130182

131183
const cleanedResourceSpec =
132184
resourceSpec && 'EnvironmentArn' in resourceSpec

‎packages/core/src/test/shared/clients/sagemakerClient.test.ts‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
131131
AppType: 'CodeEditor',
132132
CodeEditorAppSettings: {
133133
DefaultResourceSpec: {
134-
InstanceType: 'ml.t3.medium',
134+
InstanceType: 'ml.t3.large',
135135
SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
136136
SageMakerImageVersionAlias: '1.0.0',
137137
},
@@ -155,7 +155,13 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
155155
SpaceSettings: {
156156
RemoteAccess: 'ENABLED',
157157
AppType: 'CodeEditor',
158-
CodeEditorAppSettings: {},
158+
CodeEditorAppSettings: {
159+
DefaultResourceSpec: {
160+
InstanceType: 'ml.t3.large',
161+
SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img',
162+
SageMakerImageVersionAlias: '1.0.0',
163+
},
164+
},
159165
},
160166
})
161167

@@ -184,7 +190,11 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
184190
SpaceSettings: {
185191
RemoteAccess: 'ENABLED',
186192
AppType: 'JupyterLab',
187-
JupyterLabAppSettings: {},
193+
JupyterLabAppSettings: {
194+
DefaultResourceSpec: {
195+
InstanceType: 'ml.t3.large',
196+
},
197+
},
188198
},
189199
})
190200

@@ -194,7 +204,11 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () {
194204

195205
sinon.assert.calledOnceWithExactly(
196206
createAppStub,
197-
sinon.match.hasNested('ResourceSpec.InstanceType', 'ml.t3.medium')
207+
sinon.match.hasNested('ResourceSpec', {
208+
InstanceType: 'ml.t3.large',
209+
SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu',
210+
SageMakerImageVersionAlias: '3.2.0',
211+
})
198212
)
199213
})
200214

0 commit comments

Comments
 (0)