Skip to content

Commit 7cceda0

Browse files
laileni-awsaws-ajanggarkaprava08roger-zhangg
authored
feat(sagemaker): Releasing feature sagemaker parker (#8340)
## Problem ## Solution --- - 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: Austin Jang <[email protected]> Co-authored-by: Arkaprava De <[email protected]> Co-authored-by: Roger Zhang <[email protected]>
1 parent 0e5c394 commit 7cceda0

25 files changed

+8012
-4085
lines changed

‎package-lock.json‎

Lines changed: 6474 additions & 3848 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/core/package.json‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,7 @@
595595
"@aws-sdk/client-ec2": "<3.731.0",
596596
"@aws-sdk/client-ecr": "~3.693.0",
597597
"@aws-sdk/client-ecs": "~3.693.0",
598+
"@aws-sdk/client-eks": "^3.583.0",
598599
"@aws-sdk/client-glue": "^3.852.0",
599600
"@aws-sdk/client-iam": "<3.731.0",
600601
"@aws-sdk/client-iot": "~3.693.0",
@@ -625,6 +626,7 @@
625626
"@aws/mynah-ui": "^4.35.4",
626627
"@gerhobbelt/gitignore-parser": "^0.2.0-9",
627628
"@iarna/toml": "^2.2.5",
629+
"@kubernetes/client-node": "^0.20.0",
628630
"@smithy/fetch-http-handler": "^5.0.1",
629631
"@smithy/middleware-retry": "^4.0.3",
630632
"@smithy/node-http-handler": "^4.0.2",

‎packages/core/resources/hyperpod_connect‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
#!/usr/bin/env bash
2+
3+
# HyperPod Connection Script
4+
#
5+
# This script establishes a connection to an AWS SageMaker HyperPod instance using AWS Systems Manager (SSM).
6+
# HyperPod is AWS's managed service for distributed machine learning training at scale.
7+
#
8+
# OVERVIEW:
9+
# The script acts as a wrapper around the AWS SSM CLI to create a secure session tunnel to a HyperPod
10+
# compute instance. It validates required parameters, logs the connection attempt, and executes the
11+
# SSM StartSession command with HyperPod-specific connection details.
12+
#
13+
# REQUIRED ENVIRONMENT VARIABLES:
14+
# AWS_REGION - AWS region where the HyperPod cluster is located (e.g., us-west-2)
15+
# AWS_SSM_CLI - Path to the AWS SSM CLI executable
16+
# STREAM_URL - WebSocket stream URL for the HyperPod session connection
17+
# TOKEN - Authentication token for the HyperPod session
18+
# SESSION_ID - Unique identifier for the HyperPod session
19+
#
20+
# OPTIONAL ENVIRONMENT VARIABLES:
21+
# LOG_FILE_LOCATION - Path to log file (default: /tmp/hyperpod_connect.log)
22+
# DEBUG_LOG - Enable debug logging (default: 0)
23+
#
24+
# USAGE:
25+
# AWS_REGION=us-west-2 AWS_SSM_CLI=/usr/local/bin/session-manager-plugin \
26+
# STREAM_URL=wss://... TOKEN=abc123... SESSION_ID=session-xyz \
27+
# ./hyperpod_connect
28+
#
29+
# SECURITY NOTE:
30+
# This script handles sensitive authentication tokens. Ensure proper file permissions
31+
# and avoid logging sensitive values in production environments.
32+
set -x
33+
set -e
34+
set -u
35+
36+
_DATE_CMD=true
37+
38+
if command > /dev/null 2>&1 -v date; then
39+
_DATE_CMD=date
40+
elif command > /dev/null 2>&1 -v /bin/date; then
41+
_DATE_CMD=/bin/date
42+
fi
43+
44+
_log() {
45+
echo "$("$_DATE_CMD" '+%Y/%m/%d %H:%M:%S')" "$@" >> "${LOG_FILE_LOCATION}" 2>&1
46+
}
47+
48+
_require_nolog() {
49+
if [ -z "${1:-}" ] || [ -z "${2:-}" ]; then
50+
_log "error: missing required arg: $1"
51+
exit 1
52+
fi
53+
}
54+
55+
_require() {
56+
_require_nolog "$@"
57+
_log "$1=$2"
58+
}
59+
60+
_hyperpod() {
61+
# Function inputs
62+
local AWS_SSM_CLI=$1
63+
local AWS_REGION=$2
64+
local STREAM_URL=$3
65+
local TOKEN=$4
66+
local SESSION_ID=$5
67+
68+
exec "$AWS_SSM_CLI" "{\"streamUrl\":\"$STREAM_URL\",\"tokenValue\":\"$TOKEN\",\"sessionId\":\"$SESSION_ID\"}" "$AWS_REGION" "StartSession"
69+
}
70+
71+
_main() {
72+
# Set defaults for missing environment variables
73+
DEBUG_LOG=${DEBUG_LOG:-0}
74+
LOG_FILE_LOCATION=${LOG_FILE_LOCATION:-/tmp/hyperpod_connect.log}
75+
76+
_log "=============================================================================="
77+
_require AWS_REGION "${AWS_REGION:-}"
78+
_require AWS_SSM_CLI "${AWS_SSM_CLI:-}"
79+
_require SESSION_ID "${SESSION_ID:-}"
80+
_require_nolog STREAM_URL "${STREAM_URL:-}"
81+
_require_nolog TOKEN "${TOKEN:-}"
82+
83+
_hyperpod "$AWS_SSM_CLI" "$AWS_REGION" "$STREAM_URL" "$TOKEN" "$SESSION_ID"
84+
}
85+
86+
_main

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

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,20 @@ import * as path from 'path'
77
import * as vscode from 'vscode'
88
import { Commands } from '../../shared/vscode/commands2'
99
import { SagemakerSpaceNode } from './explorer/sagemakerSpaceNode'
10-
import { SagemakerParentNode } from './explorer/sagemakerParentNode'
10+
import { SagemakerStudioNode } from './explorer/sagemakerStudioNode'
1111
import * as uriHandlers from './uriHandlers'
1212
import { openRemoteConnect, filterSpaceAppsByDomainUserProfiles, stopSpace } from './commands'
1313
import { updateIdleFile, startMonitoringTerminalActivity, ActivityCheckInterval } from './utils'
1414
import { ExtContext } from '../../shared/extensions'
1515
import { telemetry } from '../../shared/telemetry/telemetry'
1616
import { isSageMaker, UserActivity } from '../../shared/extensionUtilities'
17+
import { SagemakerDevSpaceNode } from './explorer/sagemakerDevSpaceNode'
18+
import {
19+
filterDevSpacesByNamespaceCluster,
20+
openHyperPodRemoteConnection,
21+
stopHyperPodSpaceCommand,
22+
} from './hyperpodCommands'
23+
import { SagemakerHyperpodNode } from './explorer/sagemakerHyperpodNode'
1724

1825
let terminalActivityInterval: NodeJS.Timeout | undefined
1926

@@ -29,7 +36,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
2936
})
3037
}),
3138

32-
Commands.register('aws.sagemaker.filterSpaceApps', async (node: SagemakerParentNode) => {
39+
Commands.register('aws.sagemaker.filterSpaceApps', async (node: SagemakerStudioNode) => {
3340
await telemetry.sagemaker_filterSpaces.run(async () => {
3441
await filterSpaceAppsByDomainUserProfiles(node)
3542
})
@@ -42,6 +49,30 @@ export async function activate(ctx: ExtContext): Promise<void> {
4249
await telemetry.sagemaker_stopSpace.run(async () => {
4350
await stopSpace(node, ctx.extensionContext)
4451
})
52+
}),
53+
54+
Commands.register('aws.hyperpod.filterDevSpaces', async (node: SagemakerHyperpodNode) => {
55+
await telemetry.hyperpod_filterSpaces.run(async () => {
56+
await filterDevSpacesByNamespaceCluster(node)
57+
})
58+
}),
59+
60+
Commands.register('aws.hyperpod.stopSpace', async (node: SagemakerDevSpaceNode) => {
61+
if (!validateNode(node)) {
62+
return
63+
}
64+
await telemetry.hyperpod_stopSpace.run(async () => {
65+
await stopHyperPodSpaceCommand(node)
66+
})
67+
}),
68+
69+
Commands.register('aws.hyperpod.openRemoteConnection', async (node: SagemakerDevSpaceNode) => {
70+
await telemetry.hyperpod_openRemoteConnection.run(async () => {
71+
if (!validateNode(node)) {
72+
return
73+
}
74+
await openHyperPodRemoteConnection(node)
75+
})
4576
})
4677
)
4778

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

Lines changed: 50 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import * as vscode from 'vscode'
77
import * as nls from 'vscode-nls'
88
import { SagemakerConstants } from './explorer/constants'
9-
import { SagemakerParentNode } from './explorer/sagemakerParentNode'
9+
import { SagemakerStudioNode } from './explorer/sagemakerStudioNode'
1010
import { DomainKeyDelimiter } from './utils'
1111
import { startVscodeRemote } from '../../shared/extensions/ssh'
1212
import { getLogger } from '../../shared/logger/logger'
@@ -30,15 +30,16 @@ import {
3030
SpaceStatus,
3131
} from './constants'
3232
import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
33+
import { node } from 'webpack'
3334

3435
const localize = nls.loadMessageBundle()
3536

36-
export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerParentNode): Promise<void> {
37-
if (parentNode.domainUserProfiles.size === 0) {
38-
// if parentNode has not been expanded, domainUserProfiles will be empty
37+
export async function filterSpaceAppsByDomainUserProfiles(studioNode: SagemakerStudioNode): Promise<void> {
38+
if (studioNode.domainUserProfiles.size === 0) {
39+
// if studioNode has not been expanded, domainUserProfiles will be empty
3940
// if so, this will attempt to populate domainUserProfiles
40-
await parentNode.updateChildren()
41-
if (parentNode.domainUserProfiles.size === 0) {
41+
await studioNode.updateChildren()
42+
if (studioNode.domainUserProfiles.size === 0) {
4243
getLogger().info(SagemakerConstants.NoSpaceToFilter)
4344
void vscode.window.showInformationMessage(SagemakerConstants.NoSpaceToFilter)
4445
return
@@ -47,7 +48,7 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
4748

4849
// Sort by domain name and user profile
4950
const sortedDomainUserProfiles = new Map(
50-
[...parentNode.domainUserProfiles].sort((a, b) => {
51+
[...studioNode.domainUserProfiles].sort((a, b) => {
5152
const domainNameA = a[1].domain.DomainName || ''
5253
const domainNameB = b[1].domain.DomainName || ''
5354

@@ -58,7 +59,7 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
5859
})
5960
)
6061

61-
const previousSelection = await parentNode.getSelectedDomainUsers()
62+
const previousSelection = await studioNode.getSelectedDomainUsers()
6263
const items: (vscode.QuickPickItem & { key: string })[] = []
6364

6465
for (const [key, userMetadata] of sortedDomainUserProfiles) {
@@ -84,8 +85,8 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
8485

8586
const newSelection = result.map((r) => r.key)
8687
if (newSelection.length !== previousSelection.size || newSelection.some((key) => !previousSelection.has(key))) {
87-
parentNode.saveSelectedDomainUsers(newSelection)
88-
await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', parentNode)
88+
studioNode.saveSelectedDomainUsers(newSelection)
89+
await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', studioNode)
8990
}
9091
}
9192

@@ -97,6 +98,9 @@ export async function deeplinkConnect(
9798
token: string,
9899
domain: string,
99100
appType?: string,
101+
workspaceName?: string,
102+
namespace?: string,
103+
eksClusterArn?: string,
100104
isSMUS: boolean = false
101105
) {
102106
getLogger().debug(
@@ -108,17 +112,34 @@ export async function deeplinkConnect(
108112
isSMUS
109113
)
110114

115+
getLogger().info(
116+
`sm:deeplinkConnect:
117+
domain: ${domain},
118+
appType: ${appType},
119+
workspaceName: ${workspaceName},
120+
namespace: ${namespace},
121+
eksClusterArn: ${eksClusterArn}`
122+
)
123+
124+
getLogger().info(
125+
`sm:deeplinkConnect: domain: ${domain}, appType: ${appType}, workspaceName: ${workspaceName}, namespace: ${namespace}, eksClusterArn: ${eksClusterArn}`
126+
)
127+
111128
if (isRemoteWorkspace()) {
112129
void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
113130
return
114131
}
115132

116133
try {
134+
let connectionType = 'sm_dl'
135+
if (domain === '') {
136+
connectionType = 'sm_hp'
137+
}
117138
const remoteEnv = await prepareDevEnvConnection(
118139
connectionIdentifier,
119140
ctx.extensionContext,
120-
'sm_dl',
121-
isSMUS,
141+
connectionType,
142+
isSMUS /* isSMUS */,
122143
undefined /* node */,
123144
session,
124145
wsUrl,
@@ -127,13 +148,20 @@ export async function deeplinkConnect(
127148
appType
128149
)
129150

130-
await startVscodeRemote(
131-
remoteEnv.SessionProcess,
132-
remoteEnv.hostname,
133-
'/home/sagemaker-user',
134-
remoteEnv.vscPath,
135-
'sagemaker-user'
136-
)
151+
try {
152+
await startVscodeRemote(
153+
remoteEnv.SessionProcess,
154+
remoteEnv.hostname,
155+
'/home/sagemaker-user',
156+
remoteEnv.vscPath,
157+
'sagemaker-user'
158+
)
159+
} catch (remoteErr: any) {
160+
throw new ToolkitError(
161+
`Failed to establish remote connection: ${remoteErr.message}. Check Remote-SSH logs for details.`,
162+
{ cause: remoteErr, code: remoteErr.code || 'RemoteConnectionFailed' }
163+
)
164+
}
137165
} catch (err: any) {
138166
getLogger().error(
139167
'sm:OpenRemoteConnect: Unable to connect to target space with arn: %s error: %s isSMUS: %s',
@@ -143,6 +171,9 @@ export async function deeplinkConnect(
143171
)
144172

145173
if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
174+
void vscode.window.showErrorMessage(
175+
`Remote connection failed: ${err.message || 'Unknown error'}. Check Output > Log (Window) for details.`
176+
)
146177
throw err
147178
}
148179
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,11 @@
44
*/
55

66
export abstract class SagemakerConstants {
7+
static readonly HyperPodPlaceHolderMessage = '[No HyperPod Spaces Found]'
8+
static readonly NoDevSpaceToFilter = 'No dev spaces to filter'
9+
static readonly SelectedClusterNamespacesState = 'aws.hyperpod.selectedClusterNamespaces'
10+
static readonly FilterHyperpodPlaceholderKey = 'aws.filterHyperpodSpacesPlaceholder'
11+
static readonly FilterHyperpodPlaceholderMessage = 'Filter dev spaces by name spaces or cluster (unselect to hide)'
712
static readonly PlaceHolderMessage = '[No Sagemaker Spaces Found]'
813
static readonly EnableIdentityFilteringSetting = 'aws.sagemaker.studio.spaces.enableIdentityFiltering'
914
static readonly SelectedDomainUsersState = 'aws.sagemaker.selectedDomainUsers'

0 commit comments

Comments
 (0)