Set up local Visual Studio Code - Amazon SageMaker AI
Set up your local environmentConnect to your local VS Code

Set up local Visual Studio Code

After administrators complete the instructions in Connect your local Visual Studio Code to SageMaker spaces, you can connect your local Visual Studio Code to your remote SageMaker spaces.

Topics

Set up your local environment

Install Visual Studio Code on your local machine. For information on the requirements, see Connect your local Visual Studio Code to SageMaker spaces.

Connect to your local VS Code

Before you can establish a connection from your local Visual Studio Code to your remote SageMaker spaces, your administrator must Set up remote access. Your administrator sets up a specific method for you to establish a connection. Choose the method that was set up for you.

Use the following procedure to establish a connection using deep link.

  1. Launch Amazon SageMaker Studio.

  2. In the Studio UI, navigate to your space.

  3. Choose Open space with.

  4. Choose VS Code. When you do so you may be prompted to Open Visual Studio Code. When you choose to do so, your local VS Code opens with another pop-up to confirm. Once completed, the remote connection established.

Method 2: AWS Toolkit for Visual Studio Code

Use the following procedure to establish a connection using the AWS Toolkit for Visual Studio Code.

  1. Open VS Code.

  2. Open the AWS Toolkit extension.

  3. Connect to AWS.

  4. In the AWS Explorer, expand SageMaker AI.

  5. Find your Studio space.

  6. Choose the Connect icon next to your space to start it.

    Note

    • Stop and restart the space in the Toolkit for Visual Studio to enable remote access, if not already connected.

    • If the space is not using a supported instance size, you will be asked to change the instance.

Method 3: Connect from the terminal via SSH CLI

Use the following procedure to establish a connection using the SSH CLI.

For this setup, ensure that you have the AWS Command Line Interface (AWS CLI) and the Session Manager plugin for the AWS CLI installed.

Create a shell script (for example, /home/user/sagemaker_connect.sh):

#!/bin/bash
# Disable the -x option if printing each command is not needed.
set -exuo pipefail

SPACE_ARN="$1"
AWS_PROFILE="${2:-}"

# Validate ARN and extract region
if [[ "$SPACE_ARN" =~ ^arn:aws[-a-z]*:sagemaker:([a-z0-9-]+):[0-9]{12}:space\/[^\/]+\/[^\/]+$ ]]; then
    AWS_REGION="${BASH_REMATCH[1]}"
else
    echo "Error: Invalid SageMaker Studio Space ARN format."
    exit 1
fi

# Optional profile flag
PROFILE_ARG=()
if [[ -n "$AWS_PROFILE" ]]; then
    PROFILE_ARG=(--profile "$AWS_PROFILE")
fi

# Start session
START_SESSION_JSON=$(aws sagemaker start-session \
    --resource-identifier "$SPACE_ARN" \
    --region "${AWS_REGION}" \
    "${PROFILE_ARG[@]}")

# Extract fields using grep and sed
SESSION_ID=$(echo "$START_SESSION_JSON" | grep -o '"SessionId": "[^"]*"' | sed 's/.*: "//;s/"$//')
STREAM_URL=$(echo "$START_SESSION_JSON" | grep -o '"StreamUrl": "[^"]*"' | sed 's/.*: "//;s/"$//')
TOKEN=$(echo "$START_SESSION_JSON" | grep -o '"TokenValue": "[^"]*"' | sed 's/.*: "//;s/"$//')

# Validate extracted values
if [[ -z "$SESSION_ID" || -z "$STREAM_URL" || -z "$TOKEN" ]]; then
    echo "Error: Failed to extract session information from sagemaker start session response."
    exit 1
fi

# Call session-manager-plugin
session-manager-plugin \
    "{\"streamUrl\":\"$STREAM_URL\",\"tokenValue\":\"$TOKEN\",\"sessionId\":\"$SESSION_ID\"}" \
    "$AWS_REGION" "StartSession"

  1. Make the script executable:

    chmod +x /home/user/sagemaker_connect.sh

  2. Configure $HOME/.ssh/config to add the following entry:

Host my_space_name_abc
  HostName 'arn:PARTITION:sagemaker:REGION:ACCOUNT:space/DOMAIN_ID/space_name_abc'
  ProxyCommand '/home/user/sagemaker_connect.sh' '%h'
  ForwardAgent yes
  AddKeysToAgent yes
  StrictHostKeyChecking accept-new

For example, the PARTITION can be aws.

If you need to use a named AWS credential profile, change the proxy command as follows:

  ProxyCommand '/home/user/sagemaker_connect.sh' '%h' YOUR_CREDENTIAL_PROFILE_NAME

  • Connect via SSH or run SCP command:

ssh my_space_name_abc
scp file_abc my_space_name_abc:/tmp/