Merge master into feature/auto-debug

Author: aws-toolkit-automation

package-lock.json

@@ -595,6 +595,7 @@
         "@aws-sdk/client-ec2": "<3.731.0",
         "@aws-sdk/client-ecr": "~3.693.0",
         "@aws-sdk/client-ecs": "~3.693.0",
+        "@aws-sdk/client-eks": "^3.583.0",
         "@aws-sdk/client-glue": "^3.852.0",
         "@aws-sdk/client-iam": "<3.731.0",
         "@aws-sdk/client-iot": "~3.693.0",
@@ -625,6 +626,7 @@
         "@aws/mynah-ui": "^4.35.4",
         "@gerhobbelt/gitignore-parser": "^0.2.0-9",
         "@iarna/toml": "^2.2.5",
+        "@kubernetes/client-node": "^0.20.0",
         "@smithy/fetch-http-handler": "^5.0.1",
         "@smithy/middleware-retry": "^4.0.3",
         "@smithy/node-http-handler": "^4.0.2",

packages/core/package.json

@@ -0,0 +1,86 @@
+#!/usr/bin/env bash
+
+# HyperPod Connection Script
+# 
+# This script establishes a connection to an AWS SageMaker HyperPod instance using AWS Systems Manager (SSM).
+# HyperPod is AWS's managed service for distributed machine learning training at scale.
+#
+# OVERVIEW:
+# The script acts as a wrapper around the AWS SSM CLI to create a secure session tunnel to a HyperPod 
+# compute instance. It validates required parameters, logs the connection attempt, and executes the 
+# SSM StartSession command with HyperPod-specific connection details.
+#
+# REQUIRED ENVIRONMENT VARIABLES:
+#   AWS_REGION          - AWS region where the HyperPod cluster is located (e.g., us-west-2)
+#   AWS_SSM_CLI         - Path to the AWS SSM CLI executable
+#   STREAM_URL          - WebSocket stream URL for the HyperPod session connection
+#   TOKEN               - Authentication token for the HyperPod session
+#   SESSION_ID          - Unique identifier for the HyperPod session
+#
+# OPTIONAL ENVIRONMENT VARIABLES:
+#   LOG_FILE_LOCATION   - Path to log file (default: /tmp/hyperpod_connect.log)
+#   DEBUG_LOG           - Enable debug logging (default: 0)
+#
+# USAGE:
+#   AWS_REGION=us-west-2 AWS_SSM_CLI=/usr/local/bin/session-manager-plugin \
+#   STREAM_URL=wss://... TOKEN=abc123... SESSION_ID=session-xyz \
+#   ./hyperpod_connect
+#
+# SECURITY NOTE:
+# This script handles sensitive authentication tokens. Ensure proper file permissions
+# and avoid logging sensitive values in production environments. 
+set -x
+set -e
+set -u
+
+_DATE_CMD=true
+
+if command > /dev/null 2>&1 -v date; then
+    _DATE_CMD=date
+elif command > /dev/null 2>&1 -v /bin/date; then
+    _DATE_CMD=/bin/date
+fi
+
+_log() {
+    echo "$("$_DATE_CMD" '+%Y/%m/%d %H:%M:%S')" "$@" >> "${LOG_FILE_LOCATION}" 2>&1
+}
+
+_require_nolog() {
+    if [ -z "${1:-}" ] || [ -z "${2:-}" ]; then
+        _log "error: missing required arg: $1"
+        exit 1
+    fi
+}
+
+_require() {
+    _require_nolog "$@"
+    _log "$1=$2"
+}
+
+_hyperpod() {
+    # Function inputs
+    local AWS_SSM_CLI=$1
+    local AWS_REGION=$2
+    local STREAM_URL=$3
+    local TOKEN=$4
+    local SESSION_ID=$5
+
+    exec "$AWS_SSM_CLI" "{\"streamUrl\":\"$STREAM_URL\",\"tokenValue\":\"$TOKEN\",\"sessionId\":\"$SESSION_ID\"}" "$AWS_REGION" "StartSession"
+}
+
+_main() {
+    # Set defaults for missing environment variables
+    DEBUG_LOG=${DEBUG_LOG:-0}
+    LOG_FILE_LOCATION=${LOG_FILE_LOCATION:-/tmp/hyperpod_connect.log}
+
+    _log "=============================================================================="
+    _require AWS_REGION "${AWS_REGION:-}"
+    _require AWS_SSM_CLI "${AWS_SSM_CLI:-}"
+    _require SESSION_ID "${SESSION_ID:-}"
+    _require_nolog STREAM_URL "${STREAM_URL:-}"
+    _require_nolog TOKEN "${TOKEN:-}"
+
+    _hyperpod "$AWS_SSM_CLI" "$AWS_REGION" "$STREAM_URL" "$TOKEN" "$SESSION_ID"
+}
+
+_main

packages/core/resources/hyperpod_connect

@@ -7,13 +7,20 @@ import * as path from 'path'
 import * as vscode from 'vscode'
 import { Commands } from '../../shared/vscode/commands2'
 import { SagemakerSpaceNode } from './explorer/sagemakerSpaceNode'
-import { SagemakerParentNode } from './explorer/sagemakerParentNode'
+import { SagemakerStudioNode } from './explorer/sagemakerStudioNode'
 import * as uriHandlers from './uriHandlers'
 import { openRemoteConnect, filterSpaceAppsByDomainUserProfiles, stopSpace } from './commands'
 import { updateIdleFile, startMonitoringTerminalActivity, ActivityCheckInterval } from './utils'
 import { ExtContext } from '../../shared/extensions'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { isSageMaker, UserActivity } from '../../shared/extensionUtilities'
+import { SagemakerDevSpaceNode } from './explorer/sagemakerDevSpaceNode'
+import {
+    filterDevSpacesByNamespaceCluster,
+    openHyperPodRemoteConnection,
+    stopHyperPodSpaceCommand,
+} from './hyperpodCommands'
+import { SagemakerHyperpodNode } from './explorer/sagemakerHyperpodNode'
 
 let terminalActivityInterval: NodeJS.Timeout | undefined
 
@@ -29,7 +36,7 @@ export async function activate(ctx: ExtContext): Promise<void> {
             })
         }),
 
-        Commands.register('aws.sagemaker.filterSpaceApps', async (node: SagemakerParentNode) => {
+        Commands.register('aws.sagemaker.filterSpaceApps', async (node: SagemakerStudioNode) => {
             await telemetry.sagemaker_filterSpaces.run(async () => {
                 await filterSpaceAppsByDomainUserProfiles(node)
             })
@@ -42,6 +49,30 @@ export async function activate(ctx: ExtContext): Promise<void> {
             await telemetry.sagemaker_stopSpace.run(async () => {
                 await stopSpace(node, ctx.extensionContext)
             })
+        }),
+
+        Commands.register('aws.hyperpod.filterDevSpaces', async (node: SagemakerHyperpodNode) => {
+            await telemetry.hyperpod_filterSpaces.run(async () => {
+                await filterDevSpacesByNamespaceCluster(node)
+            })
+        }),
+
+        Commands.register('aws.hyperpod.stopSpace', async (node: SagemakerDevSpaceNode) => {
+            if (!validateNode(node)) {
+                return
+            }
+            await telemetry.hyperpod_stopSpace.run(async () => {
+                await stopHyperPodSpaceCommand(node)
+            })
+        }),
+
+        Commands.register('aws.hyperpod.openRemoteConnection', async (node: SagemakerDevSpaceNode) => {
+            await telemetry.hyperpod_openRemoteConnection.run(async () => {
+                if (!validateNode(node)) {
+                    return
+                }
+                await openHyperPodRemoteConnection(node)
+            })
         })
     )
 

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

@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import * as nls from 'vscode-nls'
 import { SagemakerConstants } from './explorer/constants'
-import { SagemakerParentNode } from './explorer/sagemakerParentNode'
+import { SagemakerStudioNode } from './explorer/sagemakerStudioNode'
 import { DomainKeyDelimiter } from './utils'
 import { startVscodeRemote } from '../../shared/extensions/ssh'
 import { getLogger } from '../../shared/logger/logger'
@@ -30,15 +30,16 @@ import {
     SpaceStatus,
 } from './constants'
 import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
+import { node } from 'webpack'
 
 const localize = nls.loadMessageBundle()
 
-export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerParentNode): Promise<void> {
-    if (parentNode.domainUserProfiles.size === 0) {
-        // if parentNode has not been expanded, domainUserProfiles will be empty
+export async function filterSpaceAppsByDomainUserProfiles(studioNode: SagemakerStudioNode): Promise<void> {
+    if (studioNode.domainUserProfiles.size === 0) {
+        // if studioNode has not been expanded, domainUserProfiles will be empty
         // if so, this will attempt to populate domainUserProfiles
-        await parentNode.updateChildren()
-        if (parentNode.domainUserProfiles.size === 0) {
+        await studioNode.updateChildren()
+        if (studioNode.domainUserProfiles.size === 0) {
             getLogger().info(SagemakerConstants.NoSpaceToFilter)
             void vscode.window.showInformationMessage(SagemakerConstants.NoSpaceToFilter)
             return
@@ -47,7 +48,7 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
 
     // Sort by domain name and user profile
     const sortedDomainUserProfiles = new Map(
-        [...parentNode.domainUserProfiles].sort((a, b) => {
+        [...studioNode.domainUserProfiles].sort((a, b) => {
             const domainNameA = a[1].domain.DomainName || ''
             const domainNameB = b[1].domain.DomainName || ''
 
@@ -58,7 +59,7 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
         })
     )
 
-    const previousSelection = await parentNode.getSelectedDomainUsers()
+    const previousSelection = await studioNode.getSelectedDomainUsers()
     const items: (vscode.QuickPickItem & { key: string })[] = []
 
     for (const [key, userMetadata] of sortedDomainUserProfiles) {
@@ -84,8 +85,8 @@ export async function filterSpaceAppsByDomainUserProfiles(parentNode: SagemakerP
 
     const newSelection = result.map((r) => r.key)
     if (newSelection.length !== previousSelection.size || newSelection.some((key) => !previousSelection.has(key))) {
-        parentNode.saveSelectedDomainUsers(newSelection)
-        await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', parentNode)
+        studioNode.saveSelectedDomainUsers(newSelection)
+        await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', studioNode)
     }
 }
 
@@ -97,6 +98,9 @@ export async function deeplinkConnect(
     token: string,
     domain: string,
     appType?: string,
+    workspaceName?: string,
+    namespace?: string,
+    eksClusterArn?: string,
     isSMUS: boolean = false
 ) {
     getLogger().debug(
@@ -108,17 +112,34 @@ export async function deeplinkConnect(
         isSMUS
     )
 
+    getLogger().info(
+        `sm:deeplinkConnect: 
+        domain: ${domain}, 
+        appType: ${appType}, 
+        workspaceName: ${workspaceName}, 
+        namespace: ${namespace}, 
+        eksClusterArn: ${eksClusterArn}`
+    )
+
+    getLogger().info(
+        `sm:deeplinkConnect: domain: ${domain}, appType: ${appType}, workspaceName: ${workspaceName}, namespace: ${namespace}, eksClusterArn: ${eksClusterArn}`
+    )
+
     if (isRemoteWorkspace()) {
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
 
     try {
+        let connectionType = 'sm_dl'
+        if (domain === '') {
+            connectionType = 'sm_hp'
+        }
         const remoteEnv = await prepareDevEnvConnection(
             connectionIdentifier,
             ctx.extensionContext,
-            'sm_dl',
-            isSMUS,
+            connectionType,
+            isSMUS /* isSMUS */,
             undefined /* node */,
             session,
             wsUrl,
@@ -127,13 +148,20 @@ export async function deeplinkConnect(
             appType
         )
 
-        await startVscodeRemote(
-            remoteEnv.SessionProcess,
-            remoteEnv.hostname,
-            '/home/sagemaker-user',
-            remoteEnv.vscPath,
-            'sagemaker-user'
-        )
+        try {
+            await startVscodeRemote(
+                remoteEnv.SessionProcess,
+                remoteEnv.hostname,
+                '/home/sagemaker-user',
+                remoteEnv.vscPath,
+                'sagemaker-user'
+            )
+        } catch (remoteErr: any) {
+            throw new ToolkitError(
+                `Failed to establish remote connection: ${remoteErr.message}. Check Remote-SSH logs for details.`,
+                { cause: remoteErr, code: remoteErr.code || 'RemoteConnectionFailed' }
+            )
+        }
     } catch (err: any) {
         getLogger().error(
             'sm:OpenRemoteConnect: Unable to connect to target space with arn: %s error: %s isSMUS: %s',
@@ -143,6 +171,9 @@ export async function deeplinkConnect(
         )
 
         if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
+            void vscode.window.showErrorMessage(
+                `Remote connection failed: ${err.message || 'Unknown error'}. Check Output > Log (Window) for details.`
+            )
             throw err
         }
     }

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

@@ -4,6 +4,11 @@
  */
 
 export abstract class SagemakerConstants {
+    static readonly HyperPodPlaceHolderMessage = '[No HyperPod Spaces Found]'
+    static readonly NoDevSpaceToFilter = 'No dev spaces to filter'
+    static readonly SelectedClusterNamespacesState = 'aws.hyperpod.selectedClusterNamespaces'
+    static readonly FilterHyperpodPlaceholderKey = 'aws.filterHyperpodSpacesPlaceholder'
+    static readonly FilterHyperpodPlaceholderMessage = 'Filter dev spaces by name spaces or cluster (unselect to hide)'
     static readonly PlaceHolderMessage = '[No Sagemaker Spaces Found]'
     static readonly EnableIdentityFilteringSetting = 'aws.sagemaker.studio.spaces.enableIdentityFiltering'
     static readonly SelectedDomainUsersState = 'aws.sagemaker.selectedDomainUsers'