SageMaker HyperPod のログインノードに VS Code から Remote SSH してみる

SageMaker HyperPod のログインノードに VS Code から Remote SSH してみる

Clock Icon2025.01.10

こんにちは!クラウド事業本部コンサルティング部のたかくに(@takakuni_)です。

SageMaker HyperPod の Slurm WorkShop を見ていると、 Tips に VS Code からログインノードへ Remote SSH するシナリオがありました。

https://catalog.workshops.aws/sagemaker-hyperpod/en-US/05-advanced/05-vs-code

非常に便利だなと思ったため、実際に試してみます。

SageMaker クラスターの構築

まずは HyperPod クラスターを構築します。

今回はログインノードさえ用意できてれば OK のため、最小構成で作成します。

https://github.com/takakuni-classmethod/genai-blog/tree/main/sagemaker_hyperpod_101

EC2 の Remote SSH と比べて

EC2 の Remote SSH と比べて SageMaker HyperPod ならではの部分を解説します。

接続方法は EC2 と同じで、 Systems Manager の仕組みを利用して、Remote SSH を行います。

https://dev.classmethod.jp/articles/how-to-use-vscode-remote-ssh-with-aws-systems-manager/

ただし SageMaker HyperPod の場合、ノードにはキーペアを設定する項目がありません。

そのため、ローカルでキーペアを作成し、サーバーに公開鍵を登録する必要があります。

キーペアの作成

ということで、まずは、ローカルでキーペアを作成します。今回は hyperpod というファイル名にしました。

ssh-keygen -t rsa -q -f "$HOME/.ssh/hyperpod" -N ""

作成した公開鍵を確認します。

うまく作成できていますね。後ほど利用するため、メモしておきましょう。

takakuni@ sagemaker_hyperpod_101 % cat ~/.ssh/hyperpod.pub
ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCwTzpeozioaRNzfjml/iyHxfcYsi+u41LO4VZYVv07HJMf/HTRd3QxMxTu7bES1leQUw3zEItjFTrmqtZZnq292z/ZY/JJbDKv1GomBdd5sRw0POVkOuMOmSE0+3xwlxFG/UXlbUUdrtzLmMUpKCWz2uSOHaHQ0IZ2PGGYYgGTxzjv63xVvMV7aAfAEzU3MtWBUsXF0KKk9ZC0VUiLSIEm6BE3hMKQyUSOeE26+KOiS77KeOaXRlnWwrmsS1GBCPUIqIAz4QEiU82plJsInx1l3jnChgEUpeWBg9I+sLblm6SF4EG/EmT6SnuQt2hQ/XLjFgsTzUJYuMu4L5rU9738TjDn/kMRLeojpgOCksdV4d7TEUTFkeqfOB/xu8sfsfuGC+ozAnITrbQ5d9WC8ZBdDCy7AGxR60kEZqE5YlrBcAtE8iz9xyGSQu3uQqmMZAleH+F1CrizihWpSfg11mgTRfOtNbOdJube/G7bshnXDre2G1mXT4Lw1ysua6B3THM

続いて、サーバーに公開鍵を登録するため、Easy SSH を利用してノードにログインします。

適宜クラスター名やインスタンスグループ名を変更します。

curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
chmod +x easy-ssh.sh
./easy-ssh.sh -c controller-machine sagemaker-hyperpod

https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-run-jobs-access-nodes.html#sagemaker-hyperpod-run-jobs-access-nodes-tips

こちらも、うまくログインできていますね。

takakuni@ sagemaker_hyperpod_101 % ./easy-ssh.sh -c controller-machine sagemaker-hyperpod

=================================================

==== 🚀 HyperPod Cluster Easy SSH Script! 🚀 ====


=================================================
Cluster id: cf2jdza8s74w
Instance id: i-05c4ceea296d9faa4
Node Group: controller-machine
grep: /Users/takakuni/.ssh/config: No such file or directory
Would you like to add sagemaker-hyperpod to  ~/.ssh/config (yes/no)?
> no
❌ skipping adding ml-cluster to  ~/.ssh/config:
cat: /Users/takakuni.shinnosuke/.ssh/id_rsa.pub: No such file or directory
1. Detected SSH public key ~/.ssh/id_rsa.pub on the cluster. Skipping adding...

Now you can run:

$ ssh sagemaker-hyperpod

Starting session with SessionId: botocore-session-1736514901-ocrg4cseqnydr9io98qiu2pk88
#

Root からの直 SSH 接続は認められていないため、ユーザーを ubuntu に変更します。

# whoami
sudo su - ubuntu
whoamiroot
# ubuntu@ip-172-22-151-244:~$ whoami
ubuntu
ubuntu@ip-172-22-151-244:~$

最後に ~/.ssh/authorized_keys へ、先ほどの公開鍵を登録します。

このセットアップが面倒な場合は、ライフサイクルスクリプトの中で初期設定しても良いかもですね。

~/.ssh/authorized_keys
ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCwTzpeozioaRNzfjml/iyHxfcYsi+u41LO4VZYVv07HJMf/HTRd3QxMxTu7bES1leQUw3zEItjFTrmqtZZnq292z/ZY/JJbDKv1GomBdd5sRw0POVkOuMOmSE0+3xwlxFG/UXlbUUdrtzLmMUpKCWz2uSOHaHQ0IZ2PGGYYgGTxzjv63xVvMV7aAfAEzU3MtWBUsXF0KKk9ZC0VUiLSIEm6BE3hMKQyUSOeE26+KOiS77KeOaXRlnWwrmsS1GBCPUIqIAz4QEiU82plJsInx1l3jnChgEUpeWBg9I+sLblm6SF4EG/EmT6SnuQt2hQ/XLjFgsTzUJYuMu4L5rU9738TjDn/kMRLeojpgOCksdV4d7TEUTFkeqfOB/xu8sfsfuGC+ozAnITrbQ5d9WC8ZBdDCy7AGxR60kEZqE5YlrBcAtE8iz9xyGSQu3uQqmMZAleH+F1CrizihWpSfg11mgTRfOtNbOdJube/G7bshnXDre2G1mXT4Lw1ysua6B3THM=

セットアップ完了です。 exit で抜けちゃいます。

ubuntu@ip-172-22-151-244:~$ exit
logout
# exit

Exiting session with sessionId: botocore-session-1736514901-ocrg4cseqnydr9io98qiu2pk88.

takakuni@ sagemaker_hyperpod_101 %

ローカルの ~/.ssh/config を編集します。クラスター名やリージョン等は適宜変更します。

cat <<EOF >> ~/.ssh/config
Host sagemaker-hyperpod
  Hostname sagemaker-cluster:cf2jdza8s74w_controller-machine-i-05c4ceea296d9faa4
  User ubuntu
  Port 22
  IdentityFile ~/.ssh/hyperpod
  ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p' --profile takakuni-tf"
EOF

ローカルから ~/.ssh/config を利用して、接続テストを行います。

ssh sagemaker-hyperpod

うまく設定できていれば、以下のように接続できます。

takakuni@ sagemaker_hyperpod_101 % ssh sagemaker-hyperpod
========================================================================================
    ____              __  ___     __             __ __                  ___          __
   / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
  _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
 /___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
          /___/                                    /___/_/
                          HyperPod Instance AMI (Ubuntu 20.04)
========================================================================================

Welcome to Ubuntu 20.04.6 LTS (GNU/Linux 5.15.0-1072-aws x86_64v)


Utility libraries are installed in /usr/bin/python3.9.
To access them, use /usr/bin/python3.9.

AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
=============================================================================

 * Documentation:  https://help.ubuntu.com
 * Management:     https://landscape.canonical.com
 * Support:        https://ubuntu.com/pro

 System information as of Fri Jan 10 14:15:00 UTC 2025

  System load:  0.09               Processes:             192
  Usage of /:   52.3% of 96.73GB   Users logged in:       0
  Memory usage: 28%                IPv4 address for ens6: 172.22.151.244
  Swap usage:   0%

 * Ubuntu Pro delivers the most comprehensive open source security and
   compliance features.

   https://ubuntu.com/aws/pro

Expanded Security Maintenance for Applications is not enabled.

19 updates can be applied immediately.
11 of these updates are standard security updates.
To see these additional updates run: apt list --upgradable

41 additional security updates can be applied with ESM Apps.
Learn more about enabling ESM Apps service at https://ubuntu.com/esm


To replace an instance run:
   sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

To automatically resume jobs, please add the following in your job submission script:
   srun --auto-resume=1

Instance Type: t3.medium

=============================================================================
AMI Name: Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04)
Supported EC2 instances: G4dn, G5, G6, Gr6, G6e, P4d, P4de, P5, P5e, P5en, Trn1, Trn1n
NVIDIA driver version: 550.127.05
CUDA versions available: cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4
Default CUDA version is 12.1

Release notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
Developer Guide and Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/what-is-dlami.html
Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
=============================================================================

The programs included with the Ubuntu system are free software;
the exact distribution terms for each program are described in the
individual files in /usr/share/doc/*/copyright.

Ubuntu comes with ABSOLUTELY NO WARRANTY, to the extent permitted by
applicable law.

   ____              __  ___     __             __ __                  ___          __
  / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
 _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
/___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
         /___/                                    /___/_/

To replace an instance run:
   sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

To automatically resume jobs, please add the following in your job submission script:
   srun --auto-resume=1

You're on the controller
Instance Type: ml.t3.medium
ubuntu@ip-172-22-151-244:~$

それでは最後に、肝心の Remote SSH を実行します。

VS Code 左下の Remote SSH をクリックします。

2025-01-10 at 23.25.17-.envrc — genai-blog@2x.png

ホストに接続するから、先ほど構成したホストを選択しましょう。

2025-01-10 at 23.18.33-.envrc — genai-blog@2x.png

ここでうまく接続いかない場合は、プロファイルを付けてあげるなどしてみてください。

~/.ssh/config
Host sagemaker-hyperpod
  Hostname sagemaker-cluster:cf2jdza8s74w_controller-machine-i-05c4ceea296d9faa4
  User ubuntu
  Port 22
  IdentityFile ~/.ssh/hyperpod
+ ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p' --profile takakuni"
- ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"

Remote SSH で SSH 接続できました。

サーバー内部のファイルも認識してローカルから VS Code を介して触れるようになっています。非常に便利ですね。

2025-01-10 at 23.22.41-authorized_keys — ubuntu [SSH sagemaker-hyperpod]@2x.png

まとめ

以上、「SageMaker HyperPod のログインノードに Remote SSH してみた」でした。

普段の EC2 の Remote SSH 開通に比べて、公開鍵の登録がワンステップ必要ですが、それ以外は同じような操作感で繋げられることがわかりました。

非常に便利な機能ですので、ぜひ活用していきたいです。このブログがどなたかの参考になれば幸いです。

クラウド事業本部コンサルティング部のたかくに(@takakuni_)でした!

Share this article

facebook logohatena logotwitter logo

© Classmethod, Inc. All rights reserved.