feat(sagemaker): Merge sagemaker to master (#7681)

## Notes
- feat(sagemaker): Merging Feature/sagemaker-connect-phase-2 to master
- Reference PR: #7677
---

- 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: aws-toolkit-automation <[email protected]>
Co-authored-by: Roger Zhang <[email protected]>
Co-authored-by: Reed Hamilton <[email protected]>
Co-authored-by: Jacob Chung <[email protected]>
Co-authored-by: aws-ides-bot <[email protected]>
Co-authored-by: aws-asolidu <[email protected]>
Co-authored-by: Newton Der <[email protected]>
Co-authored-by: Newton Der <[email protected]>

Author: 8 people

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

@@ -3,18 +3,27 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+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 * 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'
+
+let terminalActivityInterval: NodeJS.Timeout | undefined
 
 export async function activate(ctx: ExtContext): Promise<void> {
     ctx.extensionContext.subscriptions.push(
         uriHandlers.register(ctx),
         Commands.register('aws.sagemaker.openRemoteConnection', async (node: SagemakerSpaceNode) => {
+            if (!validateNode(node)) {
+                return
+            }
             await telemetry.sagemaker_openRemoteConnection.run(async () => {
                 await openRemoteConnect(node, ctx.extensionContext)
             })
@@ -27,9 +36,47 @@ export async function activate(ctx: ExtContext): Promise<void> {
         }),
 
         Commands.register('aws.sagemaker.stopSpace', async (node: SagemakerSpaceNode) => {
+            if (!validateNode(node)) {
+                return
+            }
             await telemetry.sagemaker_stopSpace.run(async () => {
                 await stopSpace(node, ctx.extensionContext)
             })
         })
     )
+
+    // If running in SageMaker AI Space, track user activity for autoshutdown feature
+    if (isSageMaker('SMAI')) {
+        // Use /tmp/ directory so the file is cleared on each reboot to prevent stale timestamps.
+        const tmpDirectory = '/tmp/'
+        const idleFilePath = path.join(tmpDirectory, '.sagemaker-last-active-timestamp')
+
+        const userActivity = new UserActivity(ActivityCheckInterval)
+        userActivity.onUserActivity(() => updateIdleFile(idleFilePath))
+
+        terminalActivityInterval = startMonitoringTerminalActivity(idleFilePath)
+
+        // Write initial timestamp
+        await updateIdleFile(idleFilePath)
+
+        ctx.extensionContext.subscriptions.push(userActivity, {
+            dispose: () => {
+                if (terminalActivityInterval) {
+                    clearInterval(terminalActivityInterval)
+                    terminalActivityInterval = undefined
+                }
+            },
+        })
+    }
+}
+
+/**
+ * Checks if a node  is undefined and shows a warning message if so.
+ */
+function validateNode(node: unknown): boolean {
+    if (!node) {
+        void vscode.window.showWarningMessage('Space information is being refreshed. Please try again shortly.')
+        return false
+    }
+    return true
 }

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

@@ -18,6 +18,8 @@ import { ExtContext } from '../../shared/extensions'
 import { SagemakerClient } from '../../shared/clients/sagemaker'
 import { ToolkitError } from '../../shared/errors'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
+import { RemoteSessionError } from '../../shared/remoteSession'
+import { ConnectFromRemoteWorkspaceMessage, InstanceTypeError } from './constants'
 
 const localize = nls.loadMessageBundle()
 
@@ -90,34 +92,36 @@ export async function deeplinkConnect(
     )
 
     if (isRemoteWorkspace()) {
-        void vscode.window.showErrorMessage(
-            'You are in a remote workspace, skipping deeplink connect. Please open from a local workspace.'
-        )
+        void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
 
-    const remoteEnv = await prepareDevEnvConnection(
-        connectionIdentifier,
-        ctx.extensionContext,
-        'sm_dl',
-        session,
-        wsUrl,
-        token,
-        domain
-    )
-
     try {
+        const remoteEnv = await prepareDevEnvConnection(
+            connectionIdentifier,
+            ctx.extensionContext,
+            'sm_dl',
+            session,
+            wsUrl,
+            token,
+            domain
+        )
+
         await startVscodeRemote(
             remoteEnv.SessionProcess,
             remoteEnv.hostname,
             '/home/sagemaker-user',
             remoteEnv.vscPath,
             'sagemaker-user'
         )
-    } catch (err) {
+    } catch (err: any) {
         getLogger().error(
             `sm:OpenRemoteConnect: Unable to connect to target space with arn: ${connectionIdentifier} error: ${err}`
         )
+
+        if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
+            throw err
+        }
     }
 }
 
@@ -156,16 +160,29 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
 }
 
 export async function openRemoteConnect(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
+    if (isRemoteWorkspace()) {
+        void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
+        return
+    }
+
     if (node.getStatus() === 'Stopped') {
         const client = new SagemakerClient(node.regionCode)
-        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.')
+
+        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.')
+            }
+            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) {
+                throw err
+            }
         }
-        await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
-        await tryRemoteConnection(node, ctx)
     } else if (node.getStatus() === 'Running') {
         await tryRemoteConnection(node, ctx)
     }

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

@@ -0,0 +1,31 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const ConnectFromRemoteWorkspaceMessage =
+    'Unable to establish new remote connection. Your last active VS Code window is connected to a remote workspace. To open a new SageMaker Studio connection, select your local VS Code window and try again.'
+
+export const InstanceTypeError = 'InstanceTypeError'
+
+export const InstanceTypeMinimum = 'ml.t3.large'
+
+export const InstanceTypeInsufficientMemory: Record<string, string> = {
+    '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',
+}
+
+export const InstanceTypeInsufficientMemoryMessage = (
+    spaceName: string,
+    chosenInstanceType: string,
+    recommendedInstanceType: string
+) => {
+    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}]?`
+}
+
+export const InstanceTypeNotSelectedMessage = (spaceName: string) => {
+    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.`
+}

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

@@ -10,9 +10,9 @@ import globals from '../../shared/extensionGlobals'
 import { ToolkitError } from '../../shared/errors'
 import { DevSettings } from '../../shared/settings'
 import { Auth } from '../../auth/auth'
-import { parseRegionFromArn } from './utils'
 import { SpaceMappings, SsmConnectionInfo } from './types'
 import { getLogger } from '../../shared/logger/logger'
+import { parseArn } from './detached-server/utils'
 
 const mappingFileName = '.sagemaker-space-profiles'
 const mappingFilePath = path.join(os.homedir(), '.aws', mappingFileName)
@@ -81,9 +81,14 @@ export async function persistSSMConnection(
     wsUrl?: string,
     token?: string
 ): Promise<void> {
-    const region = parseRegionFromArn(appArn)
+    const { region } = parseArn(appArn)
     const endpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] ?? ''
 
+    // TODO: Hardcoded to 'jupyterlab' due to a bug in Studio that only supports refreshing
+    // the token for both CodeEditor and JupyterLab Apps in the jupyterlab subdomain.
+    // This will be fixed shortly after NYSummit launch to support refresh URL in CodeEditor subdomain.
+    const appSubDomain = 'jupyterlab'
+
     let envSubdomain: string
 
     if (endpoint.includes('beta')) {
@@ -101,8 +106,7 @@ export async function persistSSMConnection(
             ? `studio.${region}.sagemaker.aws`
             : `${envSubdomain}.studio.${region}.asfiovnxocqpcry.com`
 
-    const refreshUrl = `https://studio-${domain}.${baseDomain}/api/remoteaccess/token`
-
+    const refreshUrl = `https://studio-${domain}.${baseDomain}/${appSubDomain}`
     await setSpaceCredentials(appArn, refreshUrl, {
         sessionId: session ?? '-',
         url: wsUrl ?? '-',

packages/core/src/awsService/sagemaker/detached-server/routes/getSessionAsync.ts

@@ -8,6 +8,7 @@
 import { IncomingMessage, ServerResponse } from 'http'
 import url from 'url'
 import { SessionStore } from '../sessionStore'
+import { open, parseArn, readServerInfo } from '../utils'
 
 export async function handleGetSessionAsync(req: IncomingMessage, res: ServerResponse): Promise<void> {
     const parsedUrl = url.parse(req.url || '', true)
@@ -37,38 +38,30 @@ export async function handleGetSessionAsync(req: IncomingMessage, res: ServerRes
                 })
             )
             return
-        } else {
-            res.writeHead(200, { 'Content-Type': 'text/plain' })
-            res.end(
-                `No session found for connection identifier: ${connectionIdentifier}. Reconnecting for deeplink is not supported yet.`
-            )
-            return
         }
 
-        // Temporarily disabling reconnect logic for the 7/3 Phase 1 launch.
-        // Will re-enable in the next release around 7/14.
-
-        // const status = await store.getStatus(connectionIdentifier, requestId)
-        // if (status === 'pending') {
-        //     res.writeHead(204)
-        //     res.end()
-        //     return
-        // } else if (status === 'not-started') {
-        //     const serverInfo = await readServerInfo()
-        //     const refreshUrl = await store.getRefreshUrl(connectionIdentifier)
+        const status = await store.getStatus(connectionIdentifier, requestId)
+        if (status === 'pending') {
+            res.writeHead(204)
+            res.end()
+            return
+        } else if (status === 'not-started') {
+            const serverInfo = await readServerInfo()
+            const refreshUrl = await store.getRefreshUrl(connectionIdentifier)
+            const { spaceName } = parseArn(connectionIdentifier)
 
-        //     const url = `${refreshUrl}?connection_identifier=${encodeURIComponent(
-        //         connectionIdentifier
-        //     )}&request_id=${encodeURIComponent(requestId)}&call_back_url=${encodeURIComponent(
-        //         `http://localhost:${serverInfo.port}/refresh_token`
-        //     )}`
+            const url = `${refreshUrl}/${encodeURIComponent(spaceName)}?reconnect_identifier=${encodeURIComponent(
+                connectionIdentifier
+            )}&reconnect_request_id=${encodeURIComponent(requestId)}&reconnect_callback_url=${encodeURIComponent(
+                `http://localhost:${serverInfo.port}/refresh_token`
+            )}`
 
-        //     await open(url)
-        //     res.writeHead(202, { 'Content-Type': 'text/plain' })
-        //     res.end('Session is not ready yet. Please retry in a few seconds.')
-        //     await store.markPending(connectionIdentifier, requestId)
-        //     return
-        // }
+            await open(url)
+            res.writeHead(202, { 'Content-Type': 'text/plain' })
+            res.end('Session is not ready yet. Please retry in a few seconds.')
+            await store.markPending(connectionIdentifier, requestId)
+            return
+        }
     } catch (err) {
         console.error('Error handling session async request:', err)
         res.writeHead(500, { 'Content-Type': 'text/plain' })

packages/core/src/awsService/sagemaker/detached-server/sessionStore.ts

@@ -49,7 +49,8 @@ export class SessionStore {
 
         const asyncEntry = requests[requestId]
         if (asyncEntry?.status === 'fresh') {
-            await this.markConsumed(connectionId, requestId)
+            delete requests[requestId]
+            await writeMapping(mapping)
             return asyncEntry
         }
 

packages/core/src/awsService/sagemaker/detached-server/utils.ts

@@ -44,18 +44,18 @@ export async function readServerInfo(): Promise<ServerInfo> {
 }
 
 /**
- * Parses a SageMaker ARN to extract region and account ID.
+ * Parses a SageMaker ARN to extract region, account ID, and space name.
  * Supports formats like:
- *   arn:aws:sagemaker:<region>:<account_id>:space/<space_name>
+ *   arn:aws:sagemaker:<region>:<account_id>:space/<domain>/<space_name>
  *   or sm_lc_arn:aws:sagemaker:<region>:<account_id>:space__d-xxxx__<name>
  *
  * If the input is prefixed with an identifier (e.g. "sagemaker-user@"), the function will strip it.
  *
  * @param arn - The full SageMaker ARN string
- * @returns An object containing the region and accountId
+ * @returns An object containing the region, accountId, and spaceName
  * @throws If the ARN format is invalid
  */
-export function parseArn(arn: string): { region: string; accountId: string } {
+export function parseArn(arn: string): { region: string; accountId: string; spaceName: string } {
     const cleanedArn = arn.includes('@') ? arn.split('@')[1] : arn
     const regex = /^arn:aws:sagemaker:(?<region>[^:]+):(?<account_id>\d+):space[/:].+$/i
     const match = cleanedArn.match(regex)
@@ -64,9 +64,16 @@ export function parseArn(arn: string): { region: string; accountId: string } {
         throw new Error(`Invalid SageMaker ARN format: "${arn}"`)
     }
 
+    // Extract space name from the end of the ARN (after the last forward slash)
+    const spaceName = cleanedArn.split('/').pop()
+    if (!spaceName) {
+        throw new Error(`Could not extract space name from ARN: "${arn}"`)
+    }
+
     return {
         region: match.groups.region,
         accountId: match.groups.account_id,
+        spaceName: spaceName,
     }
 }