diff --git a/packages/core/src/awsService/sagemaker/activation.ts b/packages/core/src/awsService/sagemaker/activation.ts index 49a0244c48e..da8392ebad4 100644 --- a/packages/core/src/awsService/sagemaker/activation.ts +++ b/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 { 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 { }), 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 } diff --git a/packages/core/src/awsService/sagemaker/commands.ts b/packages/core/src/awsService/sagemaker/commands.ts index 22a00a25219..0075d7e5dff 100644 --- a/packages/core/src/awsService/sagemaker/commands.ts +++ b/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,23 +92,21 @@ 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, @@ -114,10 +114,14 @@ export async function deeplinkConnect( 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) } diff --git a/packages/core/src/awsService/sagemaker/constants.ts b/packages/core/src/awsService/sagemaker/constants.ts new file mode 100644 index 00000000000..1fc51a1d20d --- /dev/null +++ b/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 = { + '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.` +} diff --git a/packages/core/src/awsService/sagemaker/credentialMapping.ts b/packages/core/src/awsService/sagemaker/credentialMapping.ts index 205fc5fbad4..60d4e94260e 100644 --- a/packages/core/src/awsService/sagemaker/credentialMapping.ts +++ b/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 { - 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 ?? '-', diff --git a/packages/core/src/awsService/sagemaker/detached-server/routes/getSessionAsync.ts b/packages/core/src/awsService/sagemaker/detached-server/routes/getSessionAsync.ts index 32c7c876945..e59b1b9dd10 100644 --- a/packages/core/src/awsService/sagemaker/detached-server/routes/getSessionAsync.ts +++ b/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 { 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' }) diff --git a/packages/core/src/awsService/sagemaker/detached-server/sessionStore.ts b/packages/core/src/awsService/sagemaker/detached-server/sessionStore.ts index 312765de263..04098f68c89 100644 --- a/packages/core/src/awsService/sagemaker/detached-server/sessionStore.ts +++ b/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 } diff --git a/packages/core/src/awsService/sagemaker/detached-server/utils.ts b/packages/core/src/awsService/sagemaker/detached-server/utils.ts index 50e80e536f3..9ac6b6b0303 100644 --- a/packages/core/src/awsService/sagemaker/detached-server/utils.ts +++ b/packages/core/src/awsService/sagemaker/detached-server/utils.ts @@ -44,18 +44,18 @@ export async function readServerInfo(): Promise { } /** - * 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:::space/ + * arn:aws:sagemaker:::space// * or sm_lc_arn:aws:sagemaker:::space__d-xxxx__ * * 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:(?[^:]+):(?\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, } } diff --git a/packages/core/src/awsService/sagemaker/explorer/sagemakerParentNode.ts b/packages/core/src/awsService/sagemaker/explorer/sagemakerParentNode.ts index 193a11cf972..dd445f344fb 100644 --- a/packages/core/src/awsService/sagemaker/explorer/sagemakerParentNode.ts +++ b/packages/core/src/awsService/sagemaker/explorer/sagemakerParentNode.ts @@ -23,6 +23,7 @@ import { getRemoteAppMetadata } from '../remoteUtils' export const parentContextValue = 'awsSagemakerParentNode' export type SelectedDomainUsers = [string, string[]][] +export type SelectedDomainUsersByRegion = [string, SelectedDomainUsers][] export interface UserProfileMetadata { domain: DescribeDomainResponse @@ -133,12 +134,12 @@ export class SagemakerParentNode extends AWSTreeNodeBase { } public async getSelectedDomainUsers(): Promise> { - const selectedDomainUsersMap = new Map( - globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) + const selectedDomainUsersByRegionMap = new Map( + globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) ) + const selectedDomainUsersMap = new Map(selectedDomainUsersByRegionMap.get(this.regionCode)) const defaultSelectedDomainUsers = await this.getDefaultSelectedDomainUsers() - const cachedDomainUsers = selectedDomainUsersMap.get(this.callerIdentity.Arn || '') if (cachedDomainUsers && cachedDomainUsers.length > 0) { @@ -149,13 +150,19 @@ export class SagemakerParentNode extends AWSTreeNodeBase { } public saveSelectedDomainUsers(selectedDomainUsers: string[]) { - const selectedDomainUsersMap = new Map( - globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) + const selectedDomainUsersByRegionMap = new Map( + globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) ) + const selectedDomainUsersMap = new Map(selectedDomainUsersByRegionMap.get(this.regionCode)) + if (this.callerIdentity.Arn) { selectedDomainUsersMap?.set(this.callerIdentity.Arn, selectedDomainUsers) - globals.globalState.tryUpdate(SagemakerConstants.SelectedDomainUsersState, [...selectedDomainUsersMap]) + selectedDomainUsersByRegionMap?.set(this.regionCode, [...selectedDomainUsersMap]) + + globals.globalState.tryUpdate(SagemakerConstants.SelectedDomainUsersState, [ + ...selectedDomainUsersByRegionMap, + ]) } } diff --git a/packages/core/src/awsService/sagemaker/explorer/sagemakerSpaceNode.ts b/packages/core/src/awsService/sagemaker/explorer/sagemakerSpaceNode.ts index 16fd00d95cb..6151224a510 100644 --- a/packages/core/src/awsService/sagemaker/explorer/sagemakerSpaceNode.ts +++ b/packages/core/src/awsService/sagemaker/explorer/sagemakerSpaceNode.ts @@ -162,15 +162,17 @@ export class SagemakerSpaceNode extends AWSTreeNodeBase implements AWSResourceNo public async refreshNode(): Promise { await this.updateSpaceAppStatus() - await tryRefreshNode(this) + await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', this) } } export async function tryRefreshNode(node?: SagemakerSpaceNode) { if (node) { - const n = node instanceof SagemakerSpaceNode ? node.parent : node try { - await n.refreshNode() + // For SageMaker spaces, refresh just the individual space node to avoid expensive + // operation of refreshing all spaces in the domain + await node.updateSpaceAppStatus() + await vscode.commands.executeCommand('aws.refreshAwsExplorerNode', node) } catch (e) { getLogger().error('refreshNode failed: %s', (e as Error).message) } diff --git a/packages/core/src/awsService/sagemaker/model.ts b/packages/core/src/awsService/sagemaker/model.ts index 9acf481f2f0..20a667a0bfa 100644 --- a/packages/core/src/awsService/sagemaker/model.ts +++ b/packages/core/src/awsService/sagemaker/model.ts @@ -121,7 +121,7 @@ export async function startLocalServer(ctx: vscode.ExtensionContext) { const errLog = path.join(storagePath, 'sagemaker-local-server.err.log') const infoFilePath = path.join(storagePath, 'sagemaker-local-server-info.json') - logger.info(`local server logs at ${storagePath}/sagemaker-local-server.*.log`) + logger.info(`sagemaker-local-server.*.log at ${storagePath}`) const customEndpoint = DevSettings.instance.get('endpoints', {})['sagemaker'] diff --git a/packages/core/src/awsService/sagemaker/remoteUtils.ts b/packages/core/src/awsService/sagemaker/remoteUtils.ts index ffd7210eea1..9ff8d8ca177 100644 --- a/packages/core/src/awsService/sagemaker/remoteUtils.ts +++ b/packages/core/src/awsService/sagemaker/remoteUtils.ts @@ -5,8 +5,9 @@ import { fs } from '../../shared/fs/fs' import { SagemakerClient } from '../../shared/clients/sagemaker' -import { parseRegionFromArn, RemoteAppMetadata } from './utils' +import { RemoteAppMetadata } from './utils' import { getLogger } from '../../shared/logger/logger' +import { parseArn } from './detached-server/utils' export async function getRemoteAppMetadata(): Promise { try { @@ -21,7 +22,7 @@ export async function getRemoteAppMetadata(): Promise { throw new Error('DomainId or SpaceName not found in metadata file') } - const region = parseRegionFromArn(metadata.ResourceArn) + const { region } = parseArn(metadata.ResourceArn) const client = new SagemakerClient(region) const spaceDetails = await client.describeSpace({ DomainId: domainId, SpaceName: spaceName }) diff --git a/packages/core/src/awsService/sagemaker/uriHandlers.ts b/packages/core/src/awsService/sagemaker/uriHandlers.ts index 8ee91c03d88..17c3c512272 100644 --- a/packages/core/src/awsService/sagemaker/uriHandlers.ts +++ b/packages/core/src/awsService/sagemaker/uriHandlers.ts @@ -7,17 +7,20 @@ import * as vscode from 'vscode' import { SearchParams } from '../../shared/vscode/uriHandler' import { deeplinkConnect } from './commands' import { ExtContext } from '../../shared/extensions' +import { telemetry } from '../../shared/telemetry/telemetry' export function register(ctx: ExtContext) { async function connectHandler(params: ReturnType) { - await deeplinkConnect( - ctx, - params.connection_identifier, - params.session, - `${params.ws_url}&cell-number=${params['cell-number']}`, - params.token, - params.domain - ) + await telemetry.sagemaker_deeplinkConnect.run(async () => { + await deeplinkConnect( + ctx, + params.connection_identifier, + params.session, + `${params.ws_url}&cell-number=${params['cell-number']}`, + params.token, + params.domain + ) + }) } return vscode.Disposable.from(ctx.uriHandler.onPath('/connect/sagemaker', connectHandler, parseConnectParams)) diff --git a/packages/core/src/awsService/sagemaker/utils.ts b/packages/core/src/awsService/sagemaker/utils.ts index 602cb17f6ed..33cc5880bee 100644 --- a/packages/core/src/awsService/sagemaker/utils.ts +++ b/packages/core/src/awsService/sagemaker/utils.ts @@ -4,9 +4,12 @@ */ import * as cp from 'child_process' // eslint-disable-line no-restricted-imports +import * as path from 'path' import { AppStatus, SpaceStatus } from '@aws-sdk/client-sagemaker' import { SagemakerSpaceApp } from '../../shared/clients/sagemaker' import { sshLogFileLocation } from '../../shared/sshConfig' +import { fs } from '../../shared/fs/fs' +import { getLogger } from '../../shared/logger/logger' export const DomainKeyDelimiter = '__' @@ -94,11 +97,60 @@ export function spawnDetachedServer(...args: Parameters) { return cp.spawn(...args) } -export function parseRegionFromArn(arn: string): string { - const parts = arn.split(':') - if (parts.length < 4) { - throw new Error(`Invalid ARN: "${arn}"`) +export const ActivityCheckInterval = 60000 + +/** + * Updates the idle file with the current timestamp + */ +export async function updateIdleFile(idleFilePath: string): Promise { + try { + const timestamp = new Date().toISOString() + await fs.writeFile(idleFilePath, timestamp) + } catch (error) { + getLogger().error(`Failed to update SMAI idle file: ${error}`) + } +} + +/** + * Checks for terminal activity by reading the /dev/pts directory and comparing modification times of the files. + * + * The /dev/pts directory is used in Unix-like operating systems to represent pseudo-terminal (PTY) devices. + * Each active terminal session is assigned a PTY device. These devices are represented as files within the /dev/pts directory. + * When a terminal session has activity, such as when a user inputs commands or output is written to the terminal, + * the modification time (mtime) of the corresponding PTY device file is updated. By monitoring the modification + * times of the files in the /dev/pts directory, we can detect terminal activity. + * + * If activity is detected (i.e., if any PTY device file was modified within the CHECK_INTERVAL), this function + * updates the last activity timestamp. + */ +export async function checkTerminalActivity(idleFilePath: string): Promise { + try { + const files = await fs.readdir('/dev/pts') + const now = Date.now() + + for (const [fileName] of files) { + const filePath = path.join('/dev/pts', fileName) + try { + const stats = await fs.stat(filePath) + const mtime = new Date(stats.mtime).getTime() + if (now - mtime < ActivityCheckInterval) { + await updateIdleFile(idleFilePath) + return + } + } catch (err) { + getLogger().error(`Error reading file stats:`, err) + } + } + } catch (err) { + getLogger().error(`Error reading /dev/pts directory:`, err) } +} - return parts[3] // region is the 4th part +/** + * Starts monitoring terminal activity by setting an interval to check for activity in the /dev/pts directory. + */ +export function startMonitoringTerminalActivity(idleFilePath: string): NodeJS.Timeout { + return setInterval(async () => { + await checkTerminalActivity(idleFilePath) + }, ActivityCheckInterval) } diff --git a/packages/core/src/shared/clients/sagemaker.ts b/packages/core/src/shared/clients/sagemaker.ts index d24a0f74869..8a8e138dd85 100644 --- a/packages/core/src/shared/clients/sagemaker.ts +++ b/packages/core/src/shared/clients/sagemaker.ts @@ -37,9 +37,17 @@ import { isEmpty } from 'lodash' import { sleep } from '../utilities/timeoutUtils' import { ClientWrapper } from './clientWrapper' import { AsyncCollection } from '../utilities/asyncCollection' +import { + InstanceTypeError, + InstanceTypeMinimum, + InstanceTypeInsufficientMemory, + InstanceTypeInsufficientMemoryMessage, + InstanceTypeNotSelectedMessage, +} from '../../awsService/sagemaker/constants' import { getDomainSpaceKey } from '../../awsService/sagemaker/utils' import { getLogger } from '../logger/logger' import { ToolkitError } from '../errors' +import { yes, no, continueText, cancel } from '../localizedText' export interface SagemakerSpaceApp extends SpaceDetails { App?: AppDetails @@ -85,7 +93,9 @@ export class SagemakerClient extends ClientWrapper { } public async startSpace(spaceName: string, domainId: string) { - let spaceDetails + let spaceDetails: DescribeSpaceCommandOutput + + // Get existing space details try { spaceDetails = await this.describeSpace({ DomainId: domainId, @@ -95,6 +105,54 @@ export class SagemakerClient extends ClientWrapper { throw this.handleStartSpaceError(err) } + // Get app type + const appType = spaceDetails.SpaceSettings?.AppType + if (appType !== 'JupyterLab' && appType !== 'CodeEditor') { + throw new ToolkitError(`Unsupported AppType "${appType}" for space "${spaceName}"`) + } + + // Get app resource spec + const requestedResourceSpec = + appType === 'JupyterLab' + ? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec + : spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec + + let instanceType = requestedResourceSpec?.InstanceType + + // Is InstanceType defined and has enough memory? + if (instanceType && instanceType in InstanceTypeInsufficientMemory) { + // Prompt user to select one with sufficient memory (1 level up from their chosen one) + const response = await vscode.window.showErrorMessage( + InstanceTypeInsufficientMemoryMessage( + spaceDetails.SpaceName || '', + instanceType, + InstanceTypeInsufficientMemory[instanceType] + ), + yes, + no + ) + + if (response === no) { + throw new ToolkitError('InstanceType has insufficient memory.', { code: InstanceTypeError }) + } + + instanceType = InstanceTypeInsufficientMemory[instanceType] + } else if (!instanceType) { + // Prompt user to select the minimum supported instance type + const response = await vscode.window.showErrorMessage( + InstanceTypeNotSelectedMessage(spaceDetails.SpaceName || ''), + continueText, + cancel + ) + + if (response === cancel) { + throw new ToolkitError('InstanceType not defined.', { code: InstanceTypeError }) + } + + instanceType = InstanceTypeMinimum + } + + // Get remote access flag if (!spaceDetails.SpaceSettings?.RemoteAccess || spaceDetails.SpaceSettings?.RemoteAccess === 'DISABLED') { try { await this.updateSpace({ @@ -110,23 +168,17 @@ export class SagemakerClient extends ClientWrapper { } } - const appType = spaceDetails.SpaceSettings?.AppType - if (appType !== 'JupyterLab' && appType !== 'CodeEditor') { - throw new ToolkitError(`Unsupported AppType "${appType}" for space "${spaceName}"`) - } - - const requestedResourceSpec = - appType === 'JupyterLab' - ? spaceDetails.SpaceSettings?.JupyterLabAppSettings?.DefaultResourceSpec - : spaceDetails.SpaceSettings?.CodeEditorAppSettings?.DefaultResourceSpec - - const fallbackResourceSpec: ResourceSpec = { - InstanceType: 'ml.t3.medium', + const resourceSpec: ResourceSpec = { + // Default values SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu', SageMakerImageVersionAlias: '3.2.0', - } - const resourceSpec = requestedResourceSpec?.InstanceType ? requestedResourceSpec : fallbackResourceSpec + // The existing resource spec + ...requestedResourceSpec, + + // The instance type user has chosen + InstanceType: instanceType, + } const cleanedResourceSpec = resourceSpec && 'EnvironmentArn' in resourceSpec diff --git a/packages/core/src/shared/remoteSession.ts b/packages/core/src/shared/remoteSession.ts index 282b629df51..b45bdb3ca9c 100644 --- a/packages/core/src/shared/remoteSession.ts +++ b/packages/core/src/shared/remoteSession.ts @@ -25,6 +25,11 @@ import { EvaluationResult } from '@aws-sdk/client-iam' const policyAttachDelay = 5000 +export enum RemoteSessionError { + ExtensionVersionTooLow = 'ExtensionVersionTooLow', + MissingExtension = 'MissingExtension', +} + export interface MissingTool { readonly name: 'code' | 'ssm' | 'ssh' readonly reason?: string @@ -114,13 +119,13 @@ export async function ensureRemoteSshInstalled(): Promise { if (isExtensionInstalled(VSCODE_EXTENSION_ID.remotessh)) { throw new ToolkitError('Remote SSH extension version is too low', { cancelled: true, - code: 'ExtensionVersionTooLow', + code: RemoteSessionError.ExtensionVersionTooLow, details: { expected: vscodeExtensionMinVersion.remotessh }, }) } else { throw new ToolkitError('Remote SSH extension not installed', { cancelled: true, - code: 'MissingExtension', + code: RemoteSessionError.MissingExtension, }) } } diff --git a/packages/core/src/shared/sshConfig.ts b/packages/core/src/shared/sshConfig.ts index db20c173393..bba23b9a4d8 100644 --- a/packages/core/src/shared/sshConfig.ts +++ b/packages/core/src/shared/sshConfig.ts @@ -40,6 +40,9 @@ export class SshConfig { } protected async getProxyCommand(command: string): Promise> { + // Use %n for SageMaker to preserve original hostname case (avoids SSH canonicalization lowercasing and DNS lookup) + const hostnameToken = this.scriptPrefix === 'sagemaker_connect' ? '%n' : '%h' + if (this.isWin()) { // Some older versions of OpenSSH (7.8 and below) have a bug where attempting to use powershell.exe directly will fail without an absolute path const proc = new ChildProcess('powershell.exe', ['-Command', '(get-command powershell.exe).Path']) @@ -47,9 +50,9 @@ export class SshConfig { if (r.exitCode !== 0) { return Result.err(new ToolkitError('Failed to get absolute path for powershell', { cause: r.error })) } - return Result.ok(`"${r.stdout}" -ExecutionPolicy RemoteSigned -File "${command}" %h`) + return Result.ok(`"${r.stdout}" -ExecutionPolicy RemoteSigned -File "${command}" ${hostnameToken}`) } else { - return Result.ok(`'${command}' '%h'`) + return Result.ok(`'${command}' '${hostnameToken}'`) } } diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json index 55f4f8934cf..3bc103d81e3 100644 --- a/packages/core/src/shared/telemetry/vscodeTelemetry.json +++ b/packages/core/src/shared/telemetry/vscodeTelemetry.json @@ -268,6 +268,15 @@ } ] }, + { + "name": "sagemaker_deeplinkConnect", + "description": "Connect to SageMake Space via a deeplink", + "metadata": [ + { + "type": "result" + } + ] + }, { "name": "amazonq_didSelectProfile", "description": "Emitted after the user's Q Profile has been set, whether the user was prompted with a dialog, or a profile was automatically assigned after signing in.", diff --git a/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts b/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts index 1d17651a042..06f19a5e890 100644 --- a/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts +++ b/packages/core/src/test/awsService/sagemaker/credentialMapping.test.ts @@ -12,7 +12,7 @@ import globals from '../../../shared/extensionGlobals' describe('credentialMapping', () => { describe('persistLocalCredentials', () => { - const appArn = 'arn:aws:sagemaker:us-west-2:123456789012:app/domain/space' + const appArn = 'arn:aws:sagemaker:us-west-2:123456789012:space/d-f0lwireyzpjp/test-space' let sandbox: sinon.SinonSandbox @@ -78,8 +78,8 @@ describe('credentialMapping', () => { }) describe('persistSSMConnection', () => { - const appArn = 'arn:aws:sagemaker:us-west-2:123456789012:app/domain/space' - const domain = 'my-domain' + const appArn = 'arn:aws:sagemaker:us-west-2:123456789012:space/d-f0lwireyzpjp/test-space' + const domain = 'd-f0lwireyzpjp' let sandbox: sinon.SinonSandbox beforeEach(() => { @@ -102,6 +102,16 @@ describe('credentialMapping', () => { sandbox.stub(fs, 'existsFile').resolves(false) const writeStub = sandbox.stub(fs, 'writeFile').resolves() + // Stub the AWS API call + const mockDescribeSpace = sandbox.stub().resolves({ + SpaceSettings: { + AppType: 'JupyterLab', + }, + }) + sandbox.stub(require('../../../shared/clients/sagemaker'), 'SagemakerClient').returns({ + describeSpace: mockDescribeSpace, + }) + await persistSSMConnection(appArn, domain) const raw = writeStub.firstCall.args[1] @@ -121,6 +131,16 @@ describe('credentialMapping', () => { sandbox.stub(fs, 'existsFile').resolves(false) const writeStub = sandbox.stub(fs, 'writeFile').resolves() + // Stub the AWS API call + const mockDescribeSpace = sandbox.stub().resolves({ + SpaceSettings: { + AppType: 'JupyterLab', + }, + }) + sandbox.stub(require('../../../shared/clients/sagemaker'), 'SagemakerClient').returns({ + describeSpace: mockDescribeSpace, + }) + await persistSSMConnection(appArn, domain, 'sess', 'wss://ws', 'token') const raw = writeStub.firstCall.args[1] @@ -140,6 +160,16 @@ describe('credentialMapping', () => { sandbox.stub(fs, 'existsFile').resolves(false) const writeStub = sandbox.stub(fs, 'writeFile').resolves() + // Stub the AWS API call + const mockDescribeSpace = sandbox.stub().resolves({ + SpaceSettings: { + AppType: 'JupyterLab', + }, + }) + sandbox.stub(require('../../../shared/clients/sagemaker'), 'SagemakerClient').returns({ + describeSpace: mockDescribeSpace, + }) + await persistSSMConnection(appArn, domain) const raw = writeStub.firstCall.args[1] @@ -150,5 +180,31 @@ describe('credentialMapping', () => { 'loadtest.studio.us-west-2.asfiovnxocqpcry.com' ) }) + + // TODO: Skipped due to hardcoded appSubDomain. Currently 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. Additionally, appType will be determined by + // the deeplink URI rather than the describeSpace call from the toolkit. + it.skip('throws error when app type is unsupported', async () => { + sandbox.stub(DevSettings.instance, 'get').returns({}) + sandbox.stub(fs, 'existsFile').resolves(false) + + // Stub the AWS API call to return an unsupported app type + const mockDescribeSpace = sandbox.stub().resolves({ + SpaceSettings: { + AppType: 'UnsupportedApp', + }, + }) + sandbox.stub(require('../../../shared/clients/sagemaker'), 'SagemakerClient').returns({ + describeSpace: mockDescribeSpace, + }) + + await assert.rejects(() => persistSSMConnection(appArn, domain), { + name: 'Error', + message: + 'Unsupported or missing app type for space. Expected JupyterLab or CodeEditor, got: UnsupportedApp', + }) + }) }) }) diff --git a/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSession.test.ts b/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSession.test.ts index 1e09fdbc8da..5b3f176a29f 100644 --- a/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSession.test.ts +++ b/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSession.test.ts @@ -42,6 +42,7 @@ describe('handleGetSession', () => { sinon.stub(utils, 'parseArn').returns({ region: 'us-west-2', accountId: '123456789012', + spaceName: 'space-name', }) await handleGetSession(req as http.IncomingMessage, res as http.ServerResponse) @@ -56,6 +57,7 @@ describe('handleGetSession', () => { sinon.stub(utils, 'parseArn').returns({ region: 'us-west-2', accountId: '123456789012', + spaceName: 'space-name', }) sinon.stub(utils, 'startSagemakerSession').rejects(new Error('session error')) @@ -71,6 +73,7 @@ describe('handleGetSession', () => { sinon.stub(utils, 'parseArn').returns({ region: 'us-west-2', accountId: '123456789012', + spaceName: 'space-name', }) sinon.stub(utils, 'startSagemakerSession').resolves({ SessionId: 'abc123', diff --git a/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSessionAsync.test.ts b/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSessionAsync.test.ts index f8d76912b2b..8d3ab8563ee 100644 --- a/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSessionAsync.test.ts +++ b/packages/core/src/test/awsService/sagemaker/detached-server/routes/getSessionAsync.test.ts @@ -8,6 +8,7 @@ import * as sinon from 'sinon' import assert from 'assert' import { SessionStore } from '../../../../../awsService/sagemaker/detached-server/sessionStore' import { handleGetSessionAsync } from '../../../../../awsService/sagemaker/detached-server/routes/getSessionAsync' +import * as utils from '../../../../../awsService/sagemaker/detached-server/utils' describe('handleGetSessionAsync', () => { let req: Partial @@ -51,46 +52,46 @@ describe('handleGetSessionAsync', () => { }) }) - // Temporarily disabling reconnect logic for the 7/3 Phase 1 launch. - // Will re-enable in the next release around 7/14. - - // it('responds with 204 if session is pending', async () => { - // req = { url: '/session_async?connection_identifier=abc&request_id=req123' } - // storeStub.getFreshEntry.returns(Promise.resolve(undefined)) - // storeStub.getStatus.returns(Promise.resolve('pending')) + it('responds with 204 if session is pending', async () => { + req = { url: '/session_async?connection_identifier=abc&request_id=req123' } + storeStub.getFreshEntry.returns(Promise.resolve(undefined)) + storeStub.getStatus.returns(Promise.resolve('pending')) - // await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) + await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) - // assert(resWriteHead.calledWith(204)) - // assert(resEnd.calledOnce) - // }) + assert(resWriteHead.calledWith(204)) + assert(resEnd.calledOnce) + }) - // it('responds with 202 if status is not-started and opens browser', async () => { - // req = { url: '/session_async?connection_identifier=abc&request_id=req123' } + it('responds with 202 if status is not-started and opens browser', async () => { + req = { url: '/session_async?connection_identifier=abc&request_id=req123' } - // storeStub.getFreshEntry.returns(Promise.resolve(undefined)) - // storeStub.getStatus.returns(Promise.resolve('not-started')) - // storeStub.getRefreshUrl.returns(Promise.resolve('https://example.com/refresh')) - // storeStub.markPending.returns(Promise.resolve()) + storeStub.getFreshEntry.returns(Promise.resolve(undefined)) + storeStub.getStatus.returns(Promise.resolve('not-started')) + storeStub.getRefreshUrl.returns(Promise.resolve('https://example.com/refresh')) + storeStub.markPending.returns(Promise.resolve()) - // sinon.stub(utils, 'readServerInfo').resolves({ pid: 1234, port: 4567 }) - // sinon.stub(utils, 'open').resolves() - // await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) + sinon.stub(utils, 'readServerInfo').resolves({ pid: 1234, port: 4567 }) + sinon + .stub(utils, 'parseArn') + .returns({ region: 'us-east-1', accountId: '123456789012', spaceName: 'test-space' }) + sinon.stub(utils, 'open').resolves() + await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) - // assert(resWriteHead.calledWith(202)) - // assert(resEnd.calledWithMatch(/Session is not ready yet/)) - // assert(storeStub.markPending.calledWith('abc', 'req123')) - // }) + assert(resWriteHead.calledWith(202)) + assert(resEnd.calledWithMatch(/Session is not ready yet/)) + assert(storeStub.markPending.calledWith('abc', 'req123')) + }) - // it('responds with 500 if unexpected error occurs', async () => { - // req = { url: '/session_async?connection_identifier=abc&request_id=req123' } - // storeStub.getFreshEntry.throws(new Error('fail')) + it('responds with 500 if unexpected error occurs', async () => { + req = { url: '/session_async?connection_identifier=abc&request_id=req123' } + storeStub.getFreshEntry.throws(new Error('fail')) - // await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) + await handleGetSessionAsync(req as http.IncomingMessage, res as http.ServerResponse) - // assert(resWriteHead.calledWith(500)) - // assert(resEnd.calledWith('Unexpected error')) - // }) + assert(resWriteHead.calledWith(500)) + assert(resEnd.calledWith('Unexpected error')) + }) afterEach(() => { sinon.restore() diff --git a/packages/core/src/test/awsService/sagemaker/detached-server/sessionStore.test.ts b/packages/core/src/test/awsService/sagemaker/detached-server/sessionStore.test.ts index 0bb46b7d24b..2a7828a4951 100644 --- a/packages/core/src/test/awsService/sagemaker/detached-server/sessionStore.test.ts +++ b/packages/core/src/test/awsService/sagemaker/detached-server/sessionStore.test.ts @@ -59,7 +59,7 @@ describe('SessionStore', () => { assert(writeMappingStub.calledOnce) }) - it('returns async fresh entry and marks consumed', async () => { + it('returns async fresh entry and deletes it', async () => { const store = new SessionStore() // Disable initial-connection freshness readMappingStub.returns({ @@ -77,6 +77,10 @@ describe('SessionStore', () => { assert.ok(result, 'Expected result to be defined') assert.strictEqual(result.sessionId, 'a') assert(writeMappingStub.calledOnce) + + // Verify the entry was deleted from the mapping + const updated = writeMappingStub.firstCall.args[0] + assert.strictEqual(updated.deepLink[connectionId].requests[requestId], undefined) }) it('returns undefined if no fresh entries exist', async () => { diff --git a/packages/core/src/test/awsService/sagemaker/detached-server/utils.test.ts b/packages/core/src/test/awsService/sagemaker/detached-server/utils.test.ts index 66a47747bf9..1eeb5708d11 100644 --- a/packages/core/src/test/awsService/sagemaker/detached-server/utils.test.ts +++ b/packages/core/src/test/awsService/sagemaker/detached-server/utils.test.ts @@ -8,29 +8,22 @@ import { parseArn } from '../../../../awsService/sagemaker/detached-server/utils describe('parseArn', () => { it('parses a standard SageMaker ARN with forward slash', () => { - const arn = 'arn:aws:sagemaker:us-west-2:123456789012:space/my-space-name' + const arn = 'arn:aws:sagemaker:us-west-2:123456789012:space/domain-name/my-space-name' const result = parseArn(arn) assert.deepStrictEqual(result, { region: 'us-west-2', accountId: '123456789012', - }) - }) - - it('parses a standard SageMaker ARN with colon', () => { - const arn = 'arn:aws:sagemaker:eu-central-1:123456789012:space:space-name' - const result = parseArn(arn) - assert.deepStrictEqual(result, { - region: 'eu-central-1', - accountId: '123456789012', + spaceName: 'my-space-name', }) }) it('parses an ARN prefixed with sagemaker-user@', () => { - const arn = 'sagemaker-user@arn:aws:sagemaker:ap-southeast-1:123456789012:space/foo' + const arn = 'sagemaker-user@arn:aws:sagemaker:ap-southeast-1:123456789012:space/foo/my-space-name' const result = parseArn(arn) assert.deepStrictEqual(result, { region: 'ap-southeast-1', accountId: '123456789012', + spaceName: 'my-space-name', }) }) diff --git a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts index a0a0f807b73..8fccfe4bfd9 100644 --- a/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts +++ b/packages/core/src/test/awsService/sagemaker/explorer/sagemakerParentNode.test.ts @@ -8,7 +8,13 @@ import * as vscode from 'vscode' import { DescribeDomainResponse } from '@amzn/sagemaker-client' import { GetCallerIdentityResponse } from 'aws-sdk/clients/sts' import { SagemakerClient, SagemakerSpaceApp } from '../../../../shared/clients/sagemaker' -import { SagemakerParentNode } from '../../../../awsService/sagemaker/explorer/sagemakerParentNode' +import { SagemakerConstants } from '../../../../awsService/sagemaker/explorer/constants' +import { + SagemakerParentNode, + SelectedDomainUsers, + SelectedDomainUsersByRegion, +} from '../../../../awsService/sagemaker/explorer/sagemakerParentNode' +import { globals } from '../../../../shared' import { DefaultStsClient } from '../../../../shared/clients/stsClient' import { assertNodeListOnlyHasPlaceholderNode } from '../../../utilities/explorerNodeAssertions' import assert from 'assert' @@ -26,6 +32,71 @@ describe('sagemakerParentNode', function () { ['domain1', { DomainId: 'domain1', DomainName: 'domainName1' }], ['domain2', { DomainId: 'domain2', DomainName: 'domainName2' }], ]) + const spaceAppsMap: Map = new Map([ + [ + 'domain1__name1', + { + SpaceName: 'name1', + DomainId: 'domain1', + OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, + Status: 'InService', + DomainSpaceKey: 'domain1__name1', + }, + ], + [ + 'domain2__name2', + { + SpaceName: 'name2', + DomainId: 'domain2', + OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, + Status: 'InService', + DomainSpaceKey: 'domain2__name2', + }, + ], + ]) + const spaceAppsMapPending: Map = new Map([ + [ + 'domain1__name3', + { + SpaceName: 'name3', + DomainId: 'domain1', + OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, + Status: 'InService', + DomainSpaceKey: 'domain1__name3', + App: { + Status: 'InService', + }, + }, + ], + [ + 'domain2__name4', + { + SpaceName: 'name4', + DomainId: 'domain2', + OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, + Status: 'InService', + DomainSpaceKey: 'domain2__name4', + App: { + Status: 'Pending', + }, + }, + ], + ]) + const iamUser = { + UserId: 'test-userId', + Account: '123456789012', + Arn: 'arn:aws:iam::123456789012:user/user2', + } + const assumedRoleUser = { + UserId: 'test-userId', + Account: '123456789012', + Arn: 'arn:aws:sts::123456789012:assumed-role/UserRole/user2', + } + const ssoUser = { + UserId: 'test-userId', + Account: '123456789012', + Arn: 'arn:aws:sts::123456789012:assumed-role/AWSReservedSSO_MyPermissionSet_abcd1234/user2', + } const getConfigTrue = { get: () => true, } @@ -55,50 +126,15 @@ describe('sagemakerParentNode', function () { fetchSpaceAppsAndDomainsStub.returns( Promise.resolve([new Map(), new Map()]) ) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:iam::123456789012:user/test-user', - }) - ) + getCallerIdentityStub.returns(Promise.resolve(iamUser)) const childNodes = await testNode.getChildren() assertNodeListOnlyHasPlaceholderNode(childNodes) }) it('has child nodes', async function () { - const spaceAppsMap: Map = new Map([ - [ - 'domain1__name1', - { - SpaceName: 'name1', - DomainId: 'domain1', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, - Status: 'InService', - DomainSpaceKey: 'domain1__name1', - }, - ], - [ - 'domain2__name2', - { - SpaceName: 'name2', - DomainId: 'domain2', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, - Status: 'InService', - DomainSpaceKey: 'domain2__name2', - }, - ], - ]) - fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMap, domainsMap])) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:iam::123456789012:user/test-user', - }) - ) + getCallerIdentityStub.returns(Promise.resolve(iamUser)) sinon .stub(vscode.workspace, 'getConfiguration') .returns(getConfigFalse as unknown as vscode.WorkspaceConfiguration) @@ -110,43 +146,8 @@ describe('sagemakerParentNode', function () { }) it('adds pending nodes to polling nodes set', async function () { - const spaceAppsMap: Map = new Map([ - [ - 'domain1__name3', - { - SpaceName: 'name3', - DomainId: 'domain1', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, - Status: 'InService', - DomainSpaceKey: 'domain1__name3', - App: { - Status: 'InService', - }, - }, - ], - [ - 'domain2__name4', - { - SpaceName: 'name4', - DomainId: 'domain2', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, - Status: 'InService', - DomainSpaceKey: 'domain2__name4', - App: { - Status: 'Pending', - }, - }, - ], - ]) - - fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMap, domainsMap])) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:iam::123456789012:user/test-user', - }) - ) + fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMapPending, domainsMap])) + getCallerIdentityStub.returns(Promise.resolve(iamUser)) await testNode.updateChildren() assert.strictEqual(testNode.pollingSet.size, 1) @@ -154,37 +155,8 @@ describe('sagemakerParentNode', function () { }) it('filters spaces owned by user profiles that match the IAM user', async function () { - const spaceAppsMap: Map = new Map([ - [ - 'domain1__name1', - { - SpaceName: 'name1', - DomainId: 'domain1', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, - Status: 'InService', - DomainSpaceKey: 'domain1__name1', - }, - ], - [ - 'domain2__name2', - { - SpaceName: 'name2', - DomainId: 'domain2', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, - Status: 'InService', - DomainSpaceKey: 'domain2__name2', - }, - ], - ]) - fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMap, domainsMap])) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:iam::123456789012:user/user2', - }) - ) + getCallerIdentityStub.returns(Promise.resolve(iamUser)) sinon .stub(vscode.workspace, 'getConfiguration') .returns(getConfigTrue as unknown as vscode.WorkspaceConfiguration) @@ -195,37 +167,8 @@ describe('sagemakerParentNode', function () { }) it('filters spaces owned by user profiles that match the IAM assumed-role session name', async function () { - const spaceAppsMap: Map = new Map([ - [ - 'domain1__name1', - { - SpaceName: 'name1', - DomainId: 'domain1', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, - Status: 'InService', - DomainSpaceKey: 'domain1__name1', - }, - ], - [ - 'domain2__name2', - { - SpaceName: 'name2', - DomainId: 'domain2', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, - Status: 'InService', - DomainSpaceKey: 'domain2__name2', - }, - ], - ]) - fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMap, domainsMap])) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:sts::123456789012:assumed-role/UserRole/user2', - }) - ) + getCallerIdentityStub.returns(Promise.resolve(assumedRoleUser)) sinon .stub(vscode.workspace, 'getConfiguration') .returns(getConfigTrue as unknown as vscode.WorkspaceConfiguration) @@ -236,37 +179,8 @@ describe('sagemakerParentNode', function () { }) it('filters spaces owned by user profiles that match the Identity Center user', async function () { - const spaceAppsMap: Map = new Map([ - [ - 'domain1__name1', - { - SpaceName: 'name1', - DomainId: 'domain1', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user1-abcd' }, - Status: 'InService', - DomainSpaceKey: 'domain1__name1', - }, - ], - [ - 'domain2__name2', - { - SpaceName: 'name2', - DomainId: 'domain2', - OwnershipSettingsSummary: { OwnerUserProfileName: 'user2-efgh' }, - Status: 'InService', - DomainSpaceKey: 'domain2__name2', - }, - ], - ]) - fetchSpaceAppsAndDomainsStub.returns(Promise.resolve([spaceAppsMap, domainsMap])) - getCallerIdentityStub.returns( - Promise.resolve({ - UserId: 'test-userId', - Account: '123456789012', - Arn: 'arn:aws:sts::123456789012:assumed-role/AWSReservedSSO_MyPermissionSet_abcd1234/user2', - }) - ) + getCallerIdentityStub.returns(Promise.resolve(ssoUser)) sinon .stub(vscode.workspace, 'getConfiguration') .returns(getConfigFalse as unknown as vscode.WorkspaceConfiguration) @@ -276,6 +190,81 @@ describe('sagemakerParentNode', function () { assert.strictEqual(childNodes[0].label, 'name2 (Stopped)', 'Unexpected node label') }) + describe('getSelectedDomainUsers', function () { + let originalState: Map + + beforeEach(async function () { + testNode = new SagemakerParentNode(testRegion, client) + originalState = new Map( + globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) + ) + }) + + afterEach(async function () { + await globals.globalState.update(SagemakerConstants.SelectedDomainUsersState, [...originalState]) + }) + + it('gets cached selectedDomainUsers for a given region', async function () { + await globals.globalState.update(SagemakerConstants.SelectedDomainUsersState, [ + [testRegion, [['arn:aws:iam::123456789012:user/user2', ['domain2__user-cached']]]], + ]) + testNode.callerIdentity = iamUser + sinon + .stub(vscode.workspace, 'getConfiguration') + .returns(getConfigTrue as unknown as vscode.WorkspaceConfiguration) + + const result = await testNode.getSelectedDomainUsers() + assert.deepStrictEqual( + [...result], + ['domain2__user-cached'], + 'Should match only cached selected domain user' + ) + }) + + it('gets default selectedDomainUsers', async function () { + await globals.globalState.update(SagemakerConstants.SelectedDomainUsersState, []) + testNode.spaceApps = spaceAppsMap + testNode.callerIdentity = iamUser + sinon + .stub(vscode.workspace, 'getConfiguration') + .returns(getConfigTrue as unknown as vscode.WorkspaceConfiguration) + + const result = await testNode.getSelectedDomainUsers() + assert.deepStrictEqual( + [...result], + ['domain2__user2-efgh'], + 'Should match only default selected domain user' + ) + }) + }) + + describe('saveSelectedDomainUsers', function () { + let originalState: Map + + beforeEach(async function () { + testNode = new SagemakerParentNode(testRegion, client) + originalState = new Map( + globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) + ) + }) + + afterEach(async function () { + await globals.globalState.update(SagemakerConstants.SelectedDomainUsersState, [...originalState]) + }) + + it('saves selectedDomainUsers for a given region', async function () { + testNode.callerIdentity = iamUser + testNode.saveSelectedDomainUsers(['domain1__user-1', 'domain2__user-2']) + + const selectedDomainUsersByRegionMap = new Map( + globals.globalState.get(SagemakerConstants.SelectedDomainUsersState, []) + ) + const selectedDomainUsers = new Map(selectedDomainUsersByRegionMap.get(testRegion)) + + assert.deepStrictEqual(selectedDomainUsers.get(iamUser.Arn), ['domain1__user-1', 'domain2__user-2']) + }) + }) + describe('getLocalSelectedDomainUsers', function () { const createSpaceApp = (ownerName: string): SagemakerSpaceApp => ({ SpaceName: 'space1', diff --git a/packages/core/src/test/awsService/sagemaker/remoteUtils.test.ts b/packages/core/src/test/awsService/sagemaker/remoteUtils.test.ts index b2e1071e0db..4168dfdeee4 100644 --- a/packages/core/src/test/awsService/sagemaker/remoteUtils.test.ts +++ b/packages/core/src/test/awsService/sagemaker/remoteUtils.test.ts @@ -38,8 +38,10 @@ describe('getRemoteAppMetadata', function () { beforeEach(() => { sandbox = sinon.createSandbox() fsStub = sandbox.stub(fs, 'readFileText') - parseRegionStub = sandbox.stub().returns('us-west-2') - sandbox.replace(require('../../../awsService/sagemaker/utils'), 'parseRegionFromArn', parseRegionStub) + parseRegionStub = sandbox + .stub() + .returns({ region: 'us-west-2', accountId: '123456789012', spaceName: 'test-space' }) + sandbox.replace(require('../../../awsService/sagemaker/detached-server/utils'), 'parseArn', parseRegionStub) describeSpaceStub = sandbox.stub().resolves(mockSpaceDetails) sandbox.stub(SagemakerClient.prototype, 'describeSpace').callsFake(describeSpaceStub) diff --git a/packages/core/src/test/awsService/sagemaker/utils.test.ts b/packages/core/src/test/awsService/sagemaker/utils.test.ts index b7376790106..c73fd6968fc 100644 --- a/packages/core/src/test/awsService/sagemaker/utils.test.ts +++ b/packages/core/src/test/awsService/sagemaker/utils.test.ts @@ -4,8 +4,11 @@ */ import { AppStatus, SpaceStatus } from '@aws-sdk/client-sagemaker' -import { generateSpaceStatus } from '../../../awsService/sagemaker/utils' +import { generateSpaceStatus, ActivityCheckInterval } from '../../../awsService/sagemaker/utils' import * as assert from 'assert' +import * as sinon from 'sinon' +import { fs } from '../../../shared/fs/fs' +import * as utils from '../../../awsService/sagemaker/utils' describe('generateSpaceStatus', function () { it('returns Failed if space status is Failed', function () { @@ -64,3 +67,84 @@ describe('generateSpaceStatus', function () { ) }) }) + +describe('checkTerminalActivity', function () { + let sandbox: sinon.SinonSandbox + let fsReaddirStub: sinon.SinonStub + let fsStatStub: sinon.SinonStub + let fsWriteFileStub: sinon.SinonStub + + beforeEach(function () { + sandbox = sinon.createSandbox() + fsReaddirStub = sandbox.stub(fs, 'readdir') + fsStatStub = sandbox.stub(fs, 'stat') + fsWriteFileStub = sandbox.stub(fs, 'writeFile') + }) + + afterEach(function () { + sandbox.restore() + }) + + it('should write to idle file when recent terminal activity is detected', async function () { + const idleFilePath = '/tmp/test-idle-file' + const recentTime = Date.now() - ActivityCheckInterval / 2 // Recent activity + + fsReaddirStub.resolves([ + ['pts1', 1], + ['pts2', 1], + ]) // Mock file entries + fsStatStub.onFirstCall().resolves({ mtime: new Date(recentTime) }) + fsWriteFileStub.resolves() + + await utils.checkTerminalActivity(idleFilePath) + + // Verify that fs.writeFile was called (which means updateIdleFile was called) + assert.strictEqual(fsWriteFileStub.callCount, 1) + assert.strictEqual(fsWriteFileStub.firstCall.args[0], idleFilePath) + + // Verify the timestamp is a valid ISO string + const timestamp = fsWriteFileStub.firstCall.args[1] + assert.strictEqual(typeof timestamp, 'string') + assert.ok(!isNaN(Date.parse(timestamp))) + }) + + it('should stop checking once activity is detected', async function () { + const idleFilePath = '/tmp/test-idle-file' + const recentTime = Date.now() - ActivityCheckInterval / 2 + + fsReaddirStub.resolves([ + ['pts1', 1], + ['pts2', 1], + ['pts3', 1], + ]) + fsStatStub.onFirstCall().resolves({ mtime: new Date(recentTime) }) // First file has activity + fsWriteFileStub.resolves() + + await utils.checkTerminalActivity(idleFilePath) + + // Should only call stat once since activity was detected on first file + assert.strictEqual(fsStatStub.callCount, 1) + // Should write to file once + assert.strictEqual(fsWriteFileStub.callCount, 1) + }) + + it('should handle stat error gracefully and continue checking other files', async function () { + const idleFilePath = '/tmp/test-idle-file' + const recentTime = Date.now() - ActivityCheckInterval / 2 + const statError = new Error('File not found') + + fsReaddirStub.resolves([ + ['pts1', 1], + ['pts2', 1], + ]) + fsStatStub.onFirstCall().rejects(statError) // First file fails + fsStatStub.onSecondCall().resolves({ mtime: new Date(recentTime) }) // Second file succeeds + fsWriteFileStub.resolves() + + await utils.checkTerminalActivity(idleFilePath) + + // Should continue and find activity on second file + assert.strictEqual(fsStatStub.callCount, 2) + assert.strictEqual(fsWriteFileStub.callCount, 1) + }) +}) diff --git a/packages/core/src/test/shared/clients/sagemakerClient.test.ts b/packages/core/src/test/shared/clients/sagemakerClient.test.ts index 94a07dd32eb..888d2222692 100644 --- a/packages/core/src/test/shared/clients/sagemakerClient.test.ts +++ b/packages/core/src/test/shared/clients/sagemakerClient.test.ts @@ -131,7 +131,7 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () { AppType: 'CodeEditor', CodeEditorAppSettings: { DefaultResourceSpec: { - InstanceType: 'ml.t3.medium', + InstanceType: 'ml.t3.large', SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img', SageMakerImageVersionAlias: '1.0.0', }, @@ -155,7 +155,13 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () { SpaceSettings: { RemoteAccess: 'ENABLED', AppType: 'CodeEditor', - CodeEditorAppSettings: {}, + CodeEditorAppSettings: { + DefaultResourceSpec: { + InstanceType: 'ml.t3.large', + SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:img', + SageMakerImageVersionAlias: '1.0.0', + }, + }, }, }) @@ -184,7 +190,11 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () { SpaceSettings: { RemoteAccess: 'ENABLED', AppType: 'JupyterLab', - JupyterLabAppSettings: {}, + JupyterLabAppSettings: { + DefaultResourceSpec: { + InstanceType: 'ml.t3.large', + }, + }, }, }) @@ -194,7 +204,11 @@ describe('SagemakerClient.fetchSpaceAppsAndDomains', function () { sinon.assert.calledOnceWithExactly( createAppStub, - sinon.match.hasNested('ResourceSpec.InstanceType', 'ml.t3.medium') + sinon.match.hasNested('ResourceSpec', { + InstanceType: 'ml.t3.large', + SageMakerImageArn: 'arn:aws:sagemaker:us-west-2:542918446943:image/sagemaker-distribution-cpu', + SageMakerImageVersionAlias: '3.2.0', + }) ) }) diff --git a/packages/core/src/test/shared/sshConfig.test.ts b/packages/core/src/test/shared/sshConfig.test.ts index 96ca450ae14..03841644e24 100644 --- a/packages/core/src/test/shared/sshConfig.test.ts +++ b/packages/core/src/test/shared/sshConfig.test.ts @@ -82,6 +82,16 @@ describe('VscodeRemoteSshConfig', async function () { const command = result.unwrap() assert.strictEqual(command, testProxyCommand) }) + + it('uses %n token for sagemaker_connect to preserve hostname case', async function () { + const sagemakerConfig = new MockSshConfig('sshPath', 'testHostNamePrefix', 'sagemaker_connect') + sagemakerConfig.testIsWin = false + + const result = await sagemakerConfig.getProxyCommandWrapper('sagemaker_connect') + assert.ok(result.isOk()) + const command = result.unwrap() + assert.strictEqual(command, `'sagemaker_connect' '%n'`) + }) }) describe('matchSshSection', async function () { diff --git a/packages/toolkit/.changes/next-release/Bug Fix-25f95c85-8b96-4cbd-bc3e-da833340be06.json b/packages/toolkit/.changes/next-release/Bug Fix-25f95c85-8b96-4cbd-bc3e-da833340be06.json new file mode 100644 index 00000000000..aa8be24a11c --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-25f95c85-8b96-4cbd-bc3e-da833340be06.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "SageMaker: Enable per-region manual filtering of Spaces" +} diff --git a/packages/toolkit/.changes/next-release/Bug Fix-2e0b8ccd-08eb-46e6-947d-27ef5701aca8.json b/packages/toolkit/.changes/next-release/Bug Fix-2e0b8ccd-08eb-46e6-947d-27ef5701aca8.json new file mode 100644 index 00000000000..14364c43d14 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-2e0b8ccd-08eb-46e6-947d-27ef5701aca8.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "SageMaker: Show error message when connecting remotely from a remote workspace" +} diff --git a/packages/toolkit/.changes/next-release/Bug Fix-3d681d73-4171-44f3-a5ee-50f192a398ca.json b/packages/toolkit/.changes/next-release/Bug Fix-3d681d73-4171-44f3-a5ee-50f192a398ca.json new file mode 100644 index 00000000000..b1c50349e5d --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-3d681d73-4171-44f3-a5ee-50f192a398ca.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "SageMaker: Prompt user to use upgraded instance type if the chosen one has insufficient memory" +} diff --git a/packages/toolkit/.changes/next-release/Bug Fix-da7c8255-3962-49eb-b4e6-6091eb788bca.json b/packages/toolkit/.changes/next-release/Bug Fix-da7c8255-3962-49eb-b4e6-6091eb788bca.json new file mode 100644 index 00000000000..3b45f594ad6 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-da7c8255-3962-49eb-b4e6-6091eb788bca.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "SageMaker: Resolve connection issues to SageMaker Spaces with capital letters in the name" +} diff --git a/packages/toolkit/.changes/next-release/Feature-3fa7c727-cb0a-439c-9bfe-34a2a1fa99de.json b/packages/toolkit/.changes/next-release/Feature-3fa7c727-cb0a-439c-9bfe-34a2a1fa99de.json new file mode 100644 index 00000000000..db431135a2c --- /dev/null +++ b/packages/toolkit/.changes/next-release/Feature-3fa7c727-cb0a-439c-9bfe-34a2a1fa99de.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "SageMaker: Add support for deep-linked Space reconnection" +} diff --git a/packages/toolkit/.changes/next-release/Feature-d97769de-7dd4-4fe6-a3ba-c951994e6231.json b/packages/toolkit/.changes/next-release/Feature-d97769de-7dd4-4fe6-a3ba-c951994e6231.json new file mode 100644 index 00000000000..8117953da02 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Feature-d97769de-7dd4-4fe6-a3ba-c951994e6231.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "SageMaker: Enable auto-shutdown support for Spaces" +}