AI/ML on AWS Advent Calendar 2022 の 4日目です。
完全に出し遅れました。大変申し訳ないです ><
さて、本記事では、SageMaker の学習（Training）, Processing Job, 推論エンドポイント等のコンテナにSSHできる sagemaker-ssh-helper を試してみたという内容でお送りします。
sagemaker-ssh-helperとは
これ
SageMaker SSH helper uses AWS Systems Manager (SSM) Session Manager, to register the SageMaker container in SSM, followed by creating an SSM session between your client machine and the SageMaker container. You can then create an SSH connection on top of the SSM session, that allows opening a Linux shell, and/or configuring bidirectional SSH port forwarding to enable applications like remote development/debugging/desktop, and others.
とのことで、SageMaker によって動くコンテナに Systems Manager(SSM) の Session Manager によって接続する仕組み。
あくまで、コンテナの中に入るのであって、コンテナの下周りのインスタンスに入れるわけではないのでご注意を。
使い道
使い方と仕組み
基本的に README がめっちゃ親切なので、そちら見るのが一番。
が、仕組みとか知るためにちょっと見てみる。
モジュールのインストール
$ pip install sagemaker-ssh-helper
モジュールインポート
このモジュールを学習/処理/推論コード内にインポートする方法が以下2種類ある模様。
1. Dependencyから注入
dependencies を記述できる場合には、そこからコンテナにモジュールのパスを dependency_dir() を使って注入する。
※dependency_dir() でSDK実行環境状の sagemaker-ssh-helper のパスを取得する。
https://github.com/aws-samples/sagemaker-ssh-helper/blob/main/sagemaker_ssh_helper/wrapper.py#L47-L48
学習時の場合
estimator = PyTorch(entry_point='train.py', source_dir='code', dependencies=[SSHEstimatorWrapper.dependency_dir()], # <--NEW-- role=role, framework_version='1.9.1', py_version='py38', instance_count=1, instance_type='ml.m5.xlarge')
加えて、train.py 内でインポート。
import sagemaker_ssh_helper
sagemaker_ssh_helper.setup_and_start_ssh()
推論時の場合
model = estimator.create_model(entry_point='inference.py', source_dir='source_dir/', dependencies=[SSHModelWrapper.dependency_dir()]) # <--NEW--
加えて、infrenece.py 内でlibをパスに追加しインポート。
※推論エンドポイントでは、code/lib に dependencies が追加される。
import os import sys sys.path.append(os.path.join(os.path.dirname(__file__), "lib")) import sagemaker_ssh_helper sagemaker_ssh_helper.setup_and_start_ssh()
Freamework processorの場合
torch_processor.run(
source_dir="source_dir/",
dependencies=[SSHProcessorWrapper.dependency_dir()], # <--NEW--
加えて、process.py内でインポート。
import sagemaker_ssh_helper
sagemaker_ssh_helper.setup_and_start_ssh()
2. inputsから注入
dependencies を記述できない場合には、以下のように呼び出し時の inputs に augmented_input() を指定することで、
augmented_input() では内部で、dependency_dir を呼び出して、Processing JobのProcessingInputとしてライブラリを渡す。
Script Processorsの場合
spark_processor.run(
submit_app="source_dir/process.py",
inputs=[ssh_wrapper.augmented_input()] # <--NEW--
加えて、process.py 内にて input をパスに追加してインポート。
import sys sys.path.append("/opt/ml/processing/input/") import sagemaker_ssh_helper sagemaker_ssh_helper.setup_and_start_ssh()
セッションの開始
以下のコードを学習(train.py)/処理(process.py)/推論(inference.py)内で実行。
sagemaker_ssh_helper.setup_and_start_ssh()
内部的にはサブプロセスで各種シェルを実行させている。
- SSHサーバーの実行
- Sytems ManagerへのActivation
- SSM Agentの実行
- 接続待ちループ
https://github.com/aws-samples/sagemaker-ssh-helper/blob/main/sagemaker_ssh_helper/__init__.py#L4
使ってみた
ということで、こちらのサンプルを使って入ってみる。
aws-ml-jp/sagemaker-pipelines-sample.ipynb at main · aws-samples/aws-ml-jp · GitHub
SageMaker Processing で SSH
preprocess.pyに以下追加
import sys sys.path.append("/opt/ml/processing/input/") import sagemaker_ssh_helper sagemaker_ssh_helper.setup_and_start_ssh()
sklearnの実行時に inputs を渡す。
ssh_wrapper = SSHProcessorWrapper.create(sklearn_processor, connection_wait_time_seconds=600) # <--NEW-- sklearn_processor.run( wait='False', code='./preprocess/preprocess.py', # ProcessingInput は指定したものを全て S3 から processing インスタンスにコピーされる。 Destination でコピー先を指定できる。 inputs=[ ssh_wrapper.augmented_input(), # <--NEW--
これで実行をすると以下のように sm-wait の表示が出る。
これは接続を待つためのwaitループ。
connection_wait_time_secondsで指定した分だけ待機する。
2022-12-05 15:42:47 INFO [amazon-ssm-agent] [LongRunningWorkerContainer] [WorkerProvider] Worker ssm-agent-worker (pid:1345) started 2022-12-05 15:42:47 INFO [amazon-ssm-agent] [LongRunningWorkerContainer] Monitor long running worker health every 60 seconds sm-wait: Still waiting (31)... sm-wait: Still waiting (32)... sm-wait: Still waiting (33)... sm-wait: Still waiting (34)... sm-wait: Still waiting (35)... sm-wait: Still waiting (36)... sm-wait: Still waiting (37)... sm-wait: Still waiting (38)... sm-wait: Still waiting (39)...
SSMコンソールのフリートマネージャーで見てると mi- で始まるノードが登録されていることがわかる。
ノードアクションからSession Managerによる接続が可能。
接続で来たので、色々打ってみる。
root@ip-10-0-128-155:/opt/ml# pwd /opt/ml root@ip-10-0-128-155:/opt/ml# ll total 28 drwxr-xr-x 1 root root 4096 Dec 5 15:42 ./ drwxr-xr-x 1 root root 4096 Sep 21 16:40 ../ drw-r--r-- 2 root root 4096 Dec 5 15:39 config/ drwxr-xr-x 2 root root 4096 Sep 21 16:40 models/ drwxr-xr-x 4 root root 4096 Dec 5 15:42 output/ drwxr-xr-x 4 root root 4096 Dec 5 15:42 processing/ root@ip-10-0-128-155:/opt/ml# cd processing/ root@ip-10-0-128-155:/opt/ml/processing# ll total 16 drwxr-xr-x 4 root root 4096 Dec 5 15:42 ./ drwxr-xr-x 1 root root 4096 Dec 5 15:42 ../ drwxr-xr-x 5 root root 4096 Dec 5 15:42 input/ drwxr-xr-x 5 root root 4096 Dec 5 15:42 output/ root@ip-10-0-128-155:/opt/ml/processing# ll input/ total 20 drwxr-xr-x 5 root root 4096 Dec 5 15:42 ./ drwxr-xr-x 4 root root 4096 Dec 5 15:42 ../ drwxr-xr-x 2 root root 4096 Dec 5 15:42 code/ drwxr-xr-x 2 root root 4096 Dec 5 15:42 raw_data/ drwxr-xr-x 3 root root 4096 Dec 5 15:42 sagemaker_ssh_helper/
connection_wait_time_secondsが経過するとKillされる。
sm-wait: Still waiting (600)... sm-wait: Waiting complete. No 'pkill' command received. Do you need to increase timeout?
学習とか推論でもやってみようと思っていたんだけど、完全に力尽きた。
READMEが充実してるんでそれを見てください、、
注意点
上述のように、コンテナ起動時にSSMにActivationして、Session Managerで接続するという仕組みとなるため、SSMのアドバンスドインスタンス層を有効にする必要がある。
アドバンストインスタンス層が必要なシナリオ EC2 以外のノードに接続するために Session Manager を使用したい。
アドバンスドインスタンスはインスタンスごとに時間料金がかかるため注意が必要。
https://aws.amazon.com/jp/systems-manager/pricing/#:~:text=%E3%81%82%E3%82%8A%E3%81%BE%E3%81%99%E3%80%82-,%E3%82%AA%E3%83%B3%E3%83%97%E3%83%AC%E3%83%9F%E3%82%B9%E3%82%A4%E3%83%B3%E3%82%B9%E3%82%BF%E3%83%B3%E3%82%B9%E7%AE%A1%E7%90%86,-%E3%82%A2%E3%82%AB%E3%82%A6%E3%83%B3%E3%83%88%E3%83%AC%E3%83%99%E3%83%AB%E3%81%8A%E3%82%88%E3%81%B3
アドバンスドオンプレミスインスタンスごとに時間あたり 0.00695 USD
登録解除については以下にあるように明示的に解除が必要。
https://docs.aws.amazon.com/ja_jp/systems-manager/latest/userguide/sysman-managed-instance-activation.html
以前に登録したすべてのオンプレミスサーバーおよび VM は、明示的に登録を解除するまで、Systems Manager マネージドインスタンスとして登録されたままになります。Systems Manager コンソールの [Fleet Manager] ページと [Managed Instances (マネージドインスタンス)] タブで、マネージドインスタンスの登録を解除するには、AWS CLI コマンド deregister-managed-instance を使用するか、API コール DeregisterManagedInstance を使用します。
しかし現状、sagemaker-ssh-helperのビルトインの登録解除の仕組みは無い模様。
https://github.com/aws-samples/sagemaker-ssh-helper/blob/1617abc4b097e8ed7ea7256ec9130e80e50d3a25/FAQ.md#how-can-i-clean-up-system-manager-after-receiving-error-registration-failed-due-to-error-registering-the-instance-with-aws-ssm-registrationlimitexceeded-registration-limit-of-20000-reached-for-ssm-on-prem-managed-instances
Currently, there's no built-in machinsm to deregister them when a job is completed. ...
解除の方法として以下紹介されている。
manually via the UI, or using deregister_old_instances_from_ssm.py.
あるいは、EC2以外のマネージドインスタンスへのSession Managerが使えなくなるが、スタンダード層に設定変更をすると課金が掛からなくなるので、必要な時のみアドバンスドインスタンスにしてその後はスタンダードに戻すでも良さそう。
deregister_old_instances_from_ssm.py を使ってみる
deregister_old_instances_from_ssm.py が登録解除のスクリプトとして公開されているので使ってみる。
from sagemaker_ssh_helper.deregister_old_instances_from_ssm import main as deregister_instances_main deregister_instances_main()
これだとsagemaker-ssh-helper以外から登録されたマネージドインスタンスもすべて解除されるので要注意。
This utility will deregister from SSM all SageMaker SSH Helper related managed instances.
WARNING: you should be careful NOT deregister managed instances that are not related to SM SSH Helper.
Getting SSM managed instances using pagination
Appended 5 instances
Found 5 managed instances in total in SSM
Filtering to SageMaker SSH Helper related instances only
Filtered out 0 instances missing a regex match for re.match(".*sagemaker.*","IamRole")
Filtered out 0 instances missing a regex match for re.match("ConnectionLost","PingStatus")
Will filter through 5 instances, verifying they are tagged with "SSHOwner".
Filtered out 0 instances missing the tag "SSHOwner"
Found 5 managed instances related to SSH Helper
Do you want to deregister these 5 instances? (y/n) y
0: Deregistered SSM instance mi-020c30c0b18c11121
1: Deregistered SSM instance mi-08b65321aec40a632
2: Deregistered SSM instance mi-009d2770c3791428c
3: Deregistered SSM instance mi-09ab6e0f30f9258a2
4: Deregistered SSM instance mi-0d1ee0f339d6e5f97
Successfully deregistered 5 out of 5 instances to deregister.
Done.
ちゃんと消えた。
