Merge remote-tracking branch 'origin/main'

ivan-khvostishkov · ivan-khvostishkov · commit ec0c80bb5933 · 2025-02-26T14:23:04.000+01:00
diff --git a/IAM_SSM_Setup.md b/IAM_SSM_Setup.md
@@ -12,15 +12,15 @@ SageMaker SSH Helper relies on the AWS Systems Manager service to create SSH tun
 
 ### Automated setup with CDK and Cloud9
 
-a. From AWS Console, pop up [CloudShell](https://aws.amazon.com/cloudshell/) environment. Alternatively, you can the commands run in your local terminal. In this case, make sure you've installed Node.js and CDK and fulfilled [all other CDK prerequisites](https://docs.aws.amazon.com/cdk/v2/guide/getting_started.html#getting_started_prerequisites). In both cases you need to have an admin role.
+a. From AWS Console, pop up [CloudShell](https://aws.amazon.com/cloudshell/) environment (use the button located at the bottom left corner of the browser window). Alternatively, you can the commands run in your local terminal. In this case, make sure you've installed Node.js and CDK and fulfilled [all other CDK prerequisites](https://docs.aws.amazon.com/cdk/v2/guide/getting_started.html#getting_started_prerequisites). In both cases you need to have an admin role.
 
 b. Define your SageMaker role, local user role, AWS account ID and AWS Region as variables by executing the following commands in the terminal line by line:
 
 ```shell
 SAGEMAKER_ROLE_ARN=...
 USER_ROLE_ARN=...
-ACCOUNT_ID=
-REGION=
+ACCOUNT_ID=...
+REGION=...
 ```
 
 Note that if you connect to AWS from your local CLI as an IAM user, you will need to assume a `USER_ROLE_ARN` when connecting to SageMaker. 
@@ -31,18 +31,19 @@ b. Execute the following commands (you can copy-paste them as a whole script):
 
 ```shell
 pip install 'sagemaker-ssh-helper[cdk]'
+npm install -g aws-cdk
 
 cdk bootstrap aws://"$ACCOUNT_ID"/"$REGION"
 
-APP="python -m sagemaker_ssh_helper.cdk.iam_ssm_app"
+IAM_APP="python -m sagemaker_ssh_helper.cdk.iam_ssm_app"
 
-AWS_REGION="$REGION" cdk -a "$APP" deploy SSH-IAM-SSM-Stack \
+AWS_REGION="$REGION" cdk -a "$IAM_APP" deploy SSH-IAM-SSM-Stack \
   -c sagemaker_role="$SAGEMAKER_ROLE_ARN" \
   -c user_role="$USER_ROLE_ARN"
 
-APP="python -m sagemaker_ssh_helper.cdk.advanced_tier_app"
+SSM_APP="python -m sagemaker_ssh_helper.cdk.advanced_tier_app"
 
-AWS_REGION="$REGION" cdk -a "$APP" deploy SSM-Advanced-Tier-Stack
+AWS_REGION="$REGION" cdk -a "$SSM_APP" deploy SSM-Advanced-Tier-Stack
 ```
 
 In the above code we define local variable `APP` to execute CDK apps, and export `AWS_REGION` environment variable upon execution that is set to `REGION` local variable defined earlier.
@@ -52,16 +53,13 @@ Local variables `SAGEMAKER_ROLE_ARN` and `USER_ROLE_ARN` are passed as parameter
 c. To enable SageMaker SSH Helper in additional AWS Regions, run these commands per region (adjust `REGION` variable each time):
 
 ```shell
-ACCOUNT_ID=
-REGION=
+REGION=...
 ```
 
 ```shell
 cdk bootstrap aws://"$ACCOUNT_ID"/"$REGION"
 
-APP="python -m sagemaker_ssh_helper.cdk.advanced_tier_app"
-
-AWS_REGION="$REGION" cdk -a "$APP" deploy SSM-Advanced-Tier-Stack
+AWS_REGION="$REGION" cdk -a "$SSM_APP" deploy SSM-Advanced-Tier-Stack
 ```
 
 *Note:* If you will run the jobs from SageMaker Studio instead of your local machine, specify `USER_ROLE_ARN` the same as `SAGEMAKER_ROLE_ARN`.
diff --git a/README.md b/README.md
@@ -90,6 +90,7 @@ Install the latest stable version of library from the [PyPI repository](https://
 ```shell
 pip install sagemaker-ssh-helper
 ```
+
 **Caution:** It's always recommended to install the library into a Python venv, not into the system env. If you want to use later the SSH plugins of your IDE that will use the system env and system Python, you should add the venv into the system PATH, as described in the section [Remote code execution with PyCharm / VSCode over SSH](#remote-interpreter).
 
 If you're working on Windows, see [FAQ](FAQ.md#is-windows-supported).
@@ -583,6 +584,7 @@ Follow the steps in the next section for the IDE configuration, to prepare the `
 1. On the local machine, make sure that you installed the latest [AWS CLI v2](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) and the [AWS Session Manager CLI plugin](https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html). To do so, perform the automated installation with the [sm-local-configure](sagemaker_ssh_helper/sm-local-configure) script:
 
 ```shell
+# pip install sagemaker-ssh-helper
 sm-local-configure
 ```
 
@@ -752,21 +754,25 @@ If everything is set up correctly, PyCharm will stop at your breakpoint, highlig
 
 ## <a name="studio"></a>Local IDE integration with SageMaker Studio over SSH for PyCharm / VSCode
 
-[![Download Demo (.mov)](https://user-images.githubusercontent.com/87804596/205895890-e5e87f8b-1ca6-4ce6-bac1-5cb6e6f61dde.png)](https://aws-blogs-artifacts-public.s3.amazonaws.com/artifacts/ML-4988/SSH_Helper-Remote-IDE.mov)
-[Download Demo (.mov)](https://aws-blogs-artifacts-public.s3.amazonaws.com/artifacts/ML-4988/SSH_Helper-Remote-IDE.mov)
-
-> **Note**: This demo is recorded with a previous version of SSH Helper and may be not up-to-date with the recent features. Check the documentation for the most up-to-date steps.
+![SSH_Helper-Remote-IDE.png](https://user-images.githubusercontent.com/87804596/205895890-e5e87f8b-1ca6-4ce6-bac1-5cb6e6f61dde.png)
 
 For your local IDE integration with SageMaker Studio, follow the same steps as for configuring the IDE for [Remote code execution](#remote-interpreter), but instead of submitting the training / processing / inference code to SageMaker with Python SDK, execute the Jupyter notebook, as described in the next steps.
 
+> **Important:** Make sure you read the "Getting started" section and didn't skip the steps from [Setting up your AWS account with IAM and SSM configuration](IAM_SSM_Setup.md).
+
+
 1. Copy [SageMaker_SSH_IDE.ipynb](SageMaker_SSH_IDE.ipynb) into SageMaker Studio and run it. 
 
 Note that the `main` branch of this repo can contain changes that are not compatible with the version of `sagemaker-ssh-helper` that you installed from pip.
 
 To be completely sure that you're using the version of the notebook that corresponds to the installed library, take a copy of the notebook from your filesystem after you install SSH Helper package, e.g.:
 
 ```bash
-cp /opt/conda/sm_ssh/SageMaker_SSH_IDE.ipynb /root/
+pip install sagemaker-ssh-helper
+SM_SSH_PREFIX=`python -c 'import sys; print(sys.prefix);'`
+echo $SM_SSH_PREFIX
+# check the output, e.g. /opt/conda
+cp $SM_SSH_PREFIX/sm_ssh/SageMaker_SSH_IDE.ipynb ~/user-default-efs/
 ```
 
 You can also check the version with `pip freeze | grep sagemaker-ssh-helper` and take the notebook from [the corresponding release tag](https://github.com/aws-samples/sagemaker-ssh-helper/tags).
@@ -780,7 +786,7 @@ You might want to change the `LOCAL_USER_ID` variable upon the first run, to pre
 
 2. Configure remote interpreter in PyCharm / VS Code to connect to SageMaker Studio
 
-Use `app_name.user_profile_name.domain_id.studio.sagemaker` or `app_name.studio.sagemaker` as the `fqdn` to connect.
+Use `app_name.app_space_name.domain_id.studio.sagemaker` or `app_name.studio.sagemaker` as the `fqdn` to connect.
 
 To see available apps to connect to, you may run the `list` command:
 
@@ -870,15 +876,15 @@ Instead of your local user ID put the SageMaker Studio user ID (you can get it b
 3. On the System (!) terminal (not image terminal), run:
 
 ```shell
-sm-ssh connect app_name.user_profile_name.domain_id.studio.sagemaker
+sm-ssh connect app_name.app_space_name.domain_id.studio.sagemaker
 ```
 
 Alternatively, use SSH command to forward the VNC port and add more ports to the command, e.g., `-L localhost:8787:localhost:8787` to forward the Dask dashboard that is running inside the kernel gateway:
 
 ```shell
 ssh -L localhost:5901:localhost:5901 \
   -L localhost:8787:localhost:8787 \
-  app_name.user_profile_name.domain_id.studio.sagemaker
+  app_name.app_space_name.domain_id.studio.sagemaker
 ```
 
 4. Navigate to `https://d-egm0dexample.studio.eu-west-1.sagemaker.aws/jupyter/default/proxy/6080/vnc.html?host=d-egm0dexample.studio.eu-west-1.sagemaker.aws&port=443&path=jupyter/default/proxy/6080/websockify`
diff --git a/SageMaker_SSH_IDE.ipynb b/SageMaker_SSH_IDE.ipynb
@@ -73,8 +73,18 @@
    "outputs": [],
    "source": [
     "%%sh\n",
-    "sm-ssh-ide configure\n",
-    "#sm-ssh-ide configure --ssh-only"
+    "cd /opt/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%sh\n",
+    "#sm-ssh-ide configure\n",
+    "sm-ssh-ide configure --ssh-only"
    ]
   },
   {
@@ -110,7 +120,10 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "collapsed": false
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    }
    },
    "outputs": [],
    "source": [
@@ -213,8 +226,8 @@
    "outputs": [],
    "source": [
     "%%sh\n",
-    "sm-ssh-ide stop\n",
-    "sm-ssh-ide start"
+    "sudo sm-ssh-ide stop\n",
+    "sudo sm-ssh-ide start"
    ]
   },
   {
@@ -231,7 +244,7 @@
    "outputs": [],
    "source": [
     "%%sh\n",
-    "sm-ssh-ide ssm-agent"
+    "sudo sm-ssh-ide ssm-agent"
    ]
   },
   {
@@ -851,9 +864,9 @@
   ],
   "instance_type": "ml.m5.large",
   "kernelspec": {
-   "display_name": "Python 3 (Data Science 2.0)",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:eu-west-1:470317259841:image/sagemaker-data-science-38"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -865,7 +878,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.11.11"
   }
  },
  "nbformat": 4,
diff --git a/sagemaker_ssh_helper/sm-helper-functions b/sagemaker_ssh_helper/sm-helper-functions
@@ -16,6 +16,7 @@ function _silent_install() {
   install_path="/usr/local/bin"
 
   dir=$(dirname "$0")
+  dir=$(cd "$dir" && pwd)
   script=$1
   mod=$2
 
@@ -27,6 +28,7 @@ function _silent_install() {
 
 function _install_helper_scripts() {
   dir=$(dirname "$0")
+  dir=$(cd "$dir" && pwd)
 
   # Scripts may not be available in PATH during bootstrap, so we manually copy scripts into /usr/local/bin/.
   _silent_install sm-helper-functions
@@ -186,17 +188,22 @@ function _print_sm_app_name() {
   echo -n "$sm_resource_metadata_json" | sed -e 's/^.*"ResourceName":\"\([^"]*\)\".*$/\1/'
 }
 
+# shellcheck disable=SC2001
+function _print_sm_execution_role() {
+  sm_resource_metadata_json=$(tr -d "\n" < /opt/ml/metadata/resource-metadata.json)
+  echo -n "$sm_resource_metadata_json" | sed -e 's/^.*"ExecutionRoleArn":\"\([^"]*\)\".*$/\1/'
+}
+
 # shellcheck disable=SC2001
 function _print_sm_domain_id() {
   sm_resource_metadata_json=$(tr -d "\n" < /opt/ml/metadata/resource-metadata.json)
   echo -n "$sm_resource_metadata_json" | sed -e 's/^.*"DomainId":\"\([^"]*\)\".*$/\1/'
 }
 
 # shellcheck disable=SC2001
-function _print_sm_user_profile_name() {
-  # FIXME: Check for "SpaceName" - spaces are not supported yet
+function _print_sm_app_space_name() {
   sm_resource_metadata_json=$(tr -d "\n" < /opt/ml/metadata/resource-metadata.json)
-  echo -n "$sm_resource_metadata_json" | sed -e 's/^.*"UserProfileName":\"\([^"]*\)\".*$/\1/'
+  echo -n "$sm_resource_metadata_json" | sed -e 's/^.*"SpaceName":\"\([^"]*\)\".*$/\1/'
 }
 
 function _print_sm_studio_python() {
diff --git a/sagemaker_ssh_helper/sm-init-ssm b/sagemaker_ssh_helper/sm-init-ssm
@@ -59,4 +59,9 @@ response=$(aws ssm create-activation \
 acode=$(echo $response | jq --raw-output '.ActivationCode')
 aid=$(echo $response | jq --raw-output '.ActivationId')
 
-echo Yes | amazon-ssm-agent -register -id "$aid" -code "$acode" -region "$CURRENT_REGION"
+if [[ "$1" == "--sudo" ]]; then
+    echo Yes | sudo amazon-ssm-agent -register -id "$aid" -code "$acode" -region "$CURRENT_REGION"
+else
+    # Should be already root
+    echo Yes | amazon-ssm-agent -register -id "$aid" -code "$acode" -region "$CURRENT_REGION"
+fi
diff --git a/sagemaker_ssh_helper/sm-setup-ssh b/sagemaker_ssh_helper/sm-setup-ssh
@@ -36,7 +36,7 @@ if [[ "$1" == "configure" ]]; then
       exit 0
   fi
 
-  chmod 1777 /tmp
+  chmod 1777 /tmp 2>/dev/null || echo "Cannot fix permissions of /tmp. Probably, already fixed and all is OK"
   mkdir -p ~/.ssh
   [ -d /usr/share/man/man1 ] || mkdir /usr/share/man/man1
 
diff --git a/sagemaker_ssh_helper/sm-ssh-ide b/sagemaker_ssh_helper/sm-ssh-ide
@@ -14,34 +14,41 @@ set -e
 set -o pipefail
 
 dir=$(dirname "$0")
+dir=$(cd "$dir" && pwd)
 source "$dir"/sm-helper-functions
 
-if [[ "$1" == "configure" ]]; then
-    # Everything that requires root and Internet goes to 'configure'
-    if [[ -f /opt/sagemaker-ssh-helper/.ssh-ide-configured ]]; then
-        echo "sm-ssh-ide: Already configured. Remove /opt/sagemaker-ssh-helper/.ssh-ide-configured to force reconfiguration."
-        exit 0
-    fi
-
-    echo "sm-ssh-ide: Configuring services..."
-
-    OPTIONS=$2
-    mkdir -p /opt/sagemaker-ssh-helper/
-    echo "$OPTIONS" > /opt/sagemaker-ssh-helper/.sm-ssh-ide-options
-
+if [[ "$1" == "install-helper-scripts" ]]; then
+    echo "sm-ssh-ide: Installing helper scripts to /usr/local/bin/"
     _install_helper_scripts
+elif [[ "$1" == "configure" ]]; then
+    # Everything that requires Internet goes to 'configure'
 
-    cat >/etc/profile.d/sm-ssh-ide.sh <<EOF
+    sudo tee /etc/profile.d/sm-ssh-ide.sh >/dev/null <<EOF
+echo "Configuring SageMaker SSH Helper from /etc/profile.d/sm-ssh-ide.sh" >/dev/null
 export XAUTHORITY="/tmp/.Xauthority-\$USER"
 export ICEAUTHORITY="/tmp/.ICEauthority-\$USER"
 touch "/tmp/.Xauthority-\$USER"
 touch "/tmp/.ICEauthority-\$USER"
 chmod 600 "/tmp/.Xauthority-\$USER"
 chmod 600 "/tmp/.ICEauthority-\$USER"
 EOF
+
     source /etc/profile.d/sm-ssh-ide.sh
 
-    sm-setup-ssh configure
+    if [[ -f /opt/sagemaker-ssh-helper/.ssh-ide-configured ]]; then
+        echo "sm-ssh-ide: Already configured. Remove /opt/sagemaker-ssh-helper/.ssh-ide-configured to force reconfiguration."
+        exit 0
+    fi
+
+    echo "sm-ssh-ide: Configuring services..."
+
+    OPTIONS=$2
+    sudo mkdir -p /opt/sagemaker-ssh-helper/
+    echo "$OPTIONS" | sudo tee /opt/sagemaker-ssh-helper/.sm-ssh-ide-options >/dev/null
+
+    sudo "$dir"/sm-ssh-ide install-helper-scripts
+
+    sudo sm-setup-ssh configure
 
     if [[ "$OPTIONS" != "--ssh-only" ]]; then
 
@@ -75,13 +82,13 @@ EOF
       echo "sm-ssh-ide: Skipping VNC and GUI apps install"
     fi
 
-    touch /opt/sagemaker-ssh-helper/.ssh-ide-configured
+    sudo touch /opt/sagemaker-ssh-helper/.ssh-ide-configured
 
     echo "sm-ssh-ide: Finished configuration. Done"
 
-elif [[ "$1" == "get-user-profile-name" ]]; then
+elif [[ "$1" == "get-app-space-name" ]]; then
   _assert_is_ssh_ide_inside_studio
-  USER_PROFILE=$(_print_sm_user_profile_name)
+  USER_PROFILE=$(_print_sm_app_space_name)
   echo "$USER_PROFILE"
 
 elif [[ "$1" == "get-domain-id" ]]; then
@@ -93,14 +100,14 @@ elif [[ "$1" == "get-metadata" ]]; then
   _assert_is_ssh_ide_inside_studio
 
   echo "App name: $(_print_sm_app_name)"
-  echo "User profile name: $(_print_sm_user_profile_name)"
+  echo "App space name: $(_print_sm_app_space_name)"
   echo "Domain: $(_print_sm_domain_id)"
   echo ""
   echo "Connect from local machine: "
   echo "sm-ssh connect $(_print_sm_app_name).studio.sagemaker"
-  echo "sm-ssh connect $(_print_sm_app_name).$(_print_sm_user_profile_name).$(_print_sm_domain_id).studio.sagemaker"
+  echo "sm-ssh connect $(_print_sm_app_name).$(_print_sm_app_space_name).$(_print_sm_domain_id).studio.sagemaker"
   echo "ssh -A $(_print_sm_app_name).studio.sagemaker"
-  echo "ssh -A $(_print_sm_app_name).$(_print_sm_user_profile_name).$(_print_sm_domain_id).studio.sagemaker"
+  echo "ssh -A $(_print_sm_app_name).$(_print_sm_app_space_name).$(_print_sm_domain_id).studio.sagemaker"
 
 elif [[ "$1" == "set-jb-license-server" ]]; then
 
@@ -119,7 +126,7 @@ elif [[ "$1" == "set-jb-license-server" ]]; then
         echo "sm-ssh-ide: Skipping the update of /etc/hosts with PyCharm license server (already there)"
     else
         echo "sm-ssh-ide: Updating /etc/hosts with PyCharm license server from ~/.sm-jb-license-server"
-        echo "127.0.0.1  $JB_LICENSE_SERVER_HOST" >> /etc/hosts
+        echo "127.0.0.1  $JB_LICENSE_SERVER_HOST" | sudo tee -a /etc/hosts
     fi
 
 elif [[ "$1" == "set-vnc-password" ]]; then
@@ -159,13 +166,8 @@ elif [[ "$1" == "init-ssm" ]]; then
     LOCAL_USER_ID="$(cat ~/.sm-ssh-owner)"
     echo "sm-ssh-ide: Will use local user ID: $LOCAL_USER_ID"
 
-    user_profile_json=$(aws sagemaker describe-user-profile \
-      --domain-id "$(_print_sm_domain_id)" \
-      --user-profile-name "$(_print_sm_user_profile_name)" \
-      --output json \
-      | tr -d "\n")
-    execution_role=$(echo "$user_profile_json" |  { grep "ExecutionRole" || true; } \
-      | sed -e 's/^.*"ExecutionRole": \"\([^"]*\)\".*$/\1/')
+    execution_role=$(_print_sm_execution_role)
+    echo "sm-ssh-ide: Detected SageMaker execution role ARN: $execution_role"
 
     SSH_SSM_ROLE=$(echo "$execution_role" | sed -e 's/^.*:role\/\(.*\)/\1/')
 
@@ -188,7 +190,7 @@ elif [[ "$1" == "init-ssm" ]]; then
     export SSH_SSM_ROLE
     export SSH_OWNER_TAG=$LOCAL_USER_ID
 
-    sm-init-ssm
+    sm-init-ssm --sudo
 
     echo "sm-ssh-ide: SSM initialized. Done"
 
