SageMaker NotebookでGitリポジトリにSSHでpush/pullできるようにする

2019-03-04 machinelearning aws python

Sagemaker NotebookはAWSの機械学習のワークフローを提供するSageMakerの一部であるマネージドなJupyter Notebooksで、可視化などはもちろん、ここから複数インタンスでの学習ジョブを実行したりすることができる。

Git統合によってノートブック作成時にGitHubなどのリポジトリを指定すると前もって持ってきてくれるようになったが、今のところHTTPSエンドポイントにしか対応していないようで、ユーザー名・パスワードまたはトークンといった個人に紐づく認証情報が必要になる。今回はこの機能を使わずに、ライフサイクル設定でssh鍵を置き、これでpush/pullできるようにする。

パスフレーズなしの鍵を作って公開鍵を対象リポジトリのDeployKeyに登録してread/writeできるようにする。

$ mkdir sagemaker-sshkey
$ cd sagemaker-sshkey
$ ssh-keygen -t rsa -b 4096 -f id_rsa -N ""
$ pbcopy < id_rsa.pub

秘密鍵はSSMのParameter Storeに登録する。

AWS Systems Manager (SSM)のParameter Storeに認証情報を置き参照する - sambaiz-net

$ aws ssm put-parameter --name "sagemaker-sshkey" --value "`cat id_rsa`" --type String --overwrite
$ aws ssm get-parameters --names "sagemaker-sshkey"

ライフサイクル設定でノートブック開始時に次のスクリプトが実行されるようにする。このスクリプトはrootで実行される。Parameter Storeが読める権限をNotebookのIAMに付けておく。

#!/bin/bash

set -e

su - ec2-user <<EOF
cd /home/ec2-user
aws ssm get-parameters --names "sagemaker-sshkey" | jq -r ".Parameters[0].Value" > .ssh/id_rsa
echo -e "Host github.com\n\tStrictHostKeyChecking no" > .ssh/config
chmod 600 .ssh/*
git config --global user.name "Sagemaker"
cd SageMaker
if [ ! -e "sagemaker-pytorch-example" ]; then
  git clone [email protected]:sambaiz/sagemaker-pytorch-example.git
fi
EOF

あとはノートブック作成時にこのライフサイクル設定を選べば、最初からリポジトリがpullされ、左のGitタブからpush/pullできる状態で立ち上がる。

Git