Amazon SageMaker Studio のノートブックでカスタムイメージを使う

## はじめに
Amazon SageMaker Studio になんで R（その他 Julia, TF最新バージョンなど）のカーネルがないんだ！と嘆いていらっしゃる方に朗報です。
この記事では、Amazon SageMaker Studio のノートブック環境をカスタムする方法をご紹介します。Amazon SageMaker は環境構築に Docker イメージを使っているので、自分で作成した Docker イメージを Amazon SageMaker Studio に設定してやれば、そのイメージをノートブックで使うことができるという寸法です。

Amazon SageMaker Studio のセットアップ方法は [こちら](https://qiita.com/clockpulse/items/d8688d2a73517a0d8e1c) の記事をご参照ください。

## Docker イメージの作成
それでは早速、Docker イメージを作ります。今回は [こちら](https://github.com/aws-samples/sagemaker-studio-custom-image-samples/tree/main/examples/tf23-image) のサンプルを拝借して、Tensorflow 2.3 を使うためのイメージを作ります。

### AWS Cloud9 の起動
今回は Docker イメージ作成環境として AWS Cloud9 を使用します。お手元の PC に Docker と AWS CLI 環境が構築されていればそちらを使うことももちろん可能です。AWS のコンソールから Cloud9 のコンソールに移動して [Create environment] をクリックし、以下の手順で環境を作成します。

* [Name] に任意の名前を設定して [Next step] をクリック
* 何も変更せず [Next step] をクリック
* 内容を確認し [Create environment] をクリック

ブラウザの新しいタブが作成され、1, 2 分するとそこに AWS Cloud9 の画面が表示されます。

以下の画像の赤枠で囲んだ部分でコマンドを実行していきます。はじめはエリアが狭いですが、マウスで枠をドラッグすればサイズを変えられます。
![c9.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/578dba60-580d-6e7e-aff0-3dea42852f51.png)


### Cloud9 のストレージのサイズを拡張
まずはじめに、Docker イメージを作る際に容量が足りなくなることがあるので、Cloud9 のストレージサイズを拡張します。以下のスクリプトを `resize.sh` などの名前で `/home/ec2-user/environment` に保存し、 `sh resize.sh 30` などのように実行します。引数の数字は、ストレージを何GB に拡張するかを指定するものです。スクリプトは [こちら](https://docs.aws.amazon.com/ja_jp/cloud9/latest/user-guide/move-environment.html) から拝借しました。


```bash:resize.sh
#!/bin/bash

# Specify the desired volume size in GiB as a command-line argument. If not specified, default to 20 GiB.
SIZE=${1:-20}

# Get the ID of the environment host Amazon EC2 instance.
INSTANCEID=$(curl http://169.254.169.254/latest/meta-data/instance-id)

# Get the ID of the Amazon EBS volume associated with the instance.
VOLUMEID=$(aws ec2 describe-instances \
  --instance-id $INSTANCEID \
  --query "Reservations[0].Instances[0].BlockDeviceMappings[0].Ebs.VolumeId" \
  --output text)

# Resize the EBS volume.
aws ec2 modify-volume --volume-id $VOLUMEID --size $SIZE

# Wait for the resize to finish.
while [ \
  "$(aws ec2 describe-volumes-modifications \
    --volume-id $VOLUMEID \
    --filters Name=modification-state,Values="optimizing","completed" \
    --query "length(VolumesModifications)"\
    --output text)" != "1" ]; do
sleep 1
done

#Check if we're on an NVMe filesystem
if [ $(readlink -f /dev/xvda) = "/dev/xvda" ]
then
  # Rewrite the partition table so that the partition takes up all the space that it can.
  sudo growpart /dev/xvda 1

  # Expand the size of the file system.
  # Check if we are on AL2
  STR=$(cat /etc/os-release)
  SUB="VERSION_ID=\"2\""
  if [[ "$STR" == *"$SUB"* ]]
  then
    sudo xfs_growfs -d /
  else
    sudo resize2fs /dev/xvda1
  fi

else
  # Rewrite the partition table so that the partition takes up all the space that it can.
  sudo growpart /dev/nvme0n1 1

  # Expand the size of the file system.
  # Check if we're on AL2
  STR=$(cat /etc/os-release)
  SUB="VERSION_ID=\"2\""
  if [[ "$STR" == *"$SUB"* ]]
  then
    sudo xfs_growfs -d /
  else
    sudo resize2fs /dev/nvme0n1p1
  fi
fi
```

### Amazon ECR のリポジトリ作成
ここでは、Docker イメージを push するための Amazon ECR のリポジトリを作成します。<font color="Crimson">`REGION` を使いたいリージョン名で書き換えてから</font>、以下のコマンドを Cloud9 で実行します。

```bash
REGION=<aws-region>
aws --region ${REGION} ecr create-repository \
    --repository-name smstudio-custom
```

### イメージのビルドとプッシュ
ここからは、イメージを作っていきます。まずサンプルコードを GitHub リポジトリからクローンします。Cloud9 で以下のコマンドを実行します。

```bash
git clone https://github.com/aws-samples/sagemaker-studio-custom-image-samples.git
```

今回は Tensorflow 2.3 を作るサンプルを使うので、以下のコマンドでカレントディレクトリを変更します。

```bash
cd sagemaker-studio-custom-image-samples/examples/tf23-image
```

次に、Docker イメージをビルドします。`ACCOUNT_ID` <font color="Crimson">をご自身の AWS アカウント ID で書き換えてから</font>、以下のコマンドを実行します。今回、Cloud9 のインスタンスタイプはデフォルトの t2.micro のため、大きなイメージをビルドする際は非常に時間がかかります。今回のイメージビルドには 1-2分程度かかりました。大きなイメージをビルドする場合は、よりスペックの高いインスタンスタイプを選択してください。

```bash
# Modify these as required. The Docker registry endpoint can be tuned based on your current region from https://docs.aws.amazon.com/general/latest/gr/ecr.html#ecr-docker-endpoints
ACCOUNT_ID=<account-id>

# Build the image
IMAGE_NAME=custom-tf23
aws --region ${REGION} ecr get-login-password | docker login --username AWS --password-stdin ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/smstudio-custom
docker build . -t ${IMAGE_NAME} -t ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/smstudio-custom:${IMAGE_NAME}
```

ビルドしたイメージをローカルでテストする方法は [こちら](https://github.com/aws-samples/sagemaker-studio-custom-image-samples/blob/main/DEVELOPMENT.md) をご参照ください。のちの設定で使用する UID や GIU の確認方法も記載されています。
さらに、以下のコマンドを実行して、ビルドしたイメージを Amazon ECR にプッシュします。

```bash
docker push ${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/smstudio-custom:${IMAGE_NAME}
```

## 作成したイメージを Amazon SageMaker に登録
Amazon ECR にプッシュしたイメージを Amazon SageMaker で使えるようにします。`ROLE_ARN` <font color="Crimson">に SageMakerFullAccess がついた SageMaker 用の IAM ロールの ARN を設定してから</font>以下のコマンドを実行します。Amazon SageMaker Studio のドメインを作成した際に作成した IAM ロールを使用しても OK です。新規で IAM ロールを作成する場合は [こちら](https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/id_roles_create_for-service.html#roles-creatingrole-service-console) をご参照ください。IAM ロール作成の際のユースケースの選択部分は SageMaker を選択してください。

```bash
# Role in your account to be used for the SageMaker Image
ROLE_ARN=<role-arn>

aws --region ${REGION} sagemaker create-image \
    --image-name ${IMAGE_NAME} \
    --role-arn ${ROLE_ARN}

aws --region ${REGION} sagemaker create-image-version \
    --image-name ${IMAGE_NAME} \
    --base-image "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com/smstudio-custom:${IMAGE_NAME}"

# Verify the image-version is created successfully. Do NOT proceed if image-version is in CREATE_FAILED state or in any other state apart from CREATED.
aws --region ${REGION} sagemaker describe-image-version --image-name ${IMAGE_NAME}
```

### 作成したイメージの AppImageConfig を作成
このコンフィグファイルでは、作成したイメージをノートブックで使用する際のカーネル名や、マウントディレクトリが定義されます。`app-image-config-input.json` で指定している `KernelSpecs/Name` には、作成したイメージをローカルで起動してから `jupyter-kernelspec list` を実行した際に表示されるカーネルから選択します。詳しい手順は [こちら](https://github.com/aws-samples/sagemaker-studio-custom-image-samples/blob/main/DEVELOPMENT.md) をご参照ください。

```bash
aws --region ${REGION} sagemaker create-app-image-config --cli-input-json file://app-image-config-input.json
```

### イメージをドメインにアタッチ
Amazon SageMaker に登録したイメージを AppImageConfig で指定して Amazon SageMaker ドメインにアタッチします。
<font color="Crimson">update-domain-input.json の中の DomainId にドメイン ID（AWS コンソールでは Studio ID と記載。SSO を使っている場合は Studio のアドレスと書かれた部分の d-xxxxx 部分）を設定してから</font>、以下のコマンドを実行します。

```bash
aws --region ${REGION} sagemaker update-domain --cli-input-json file://update-domain-input.json
```

これで、作成したイメージを Amazon SageMaker Studio で使用できるようになりました。Amazon SageMaker コンソールの左側のメニューから Amazon SageMaker Studio をクリックして SageMaker Studio コントロールパネルを表示すると、一番下の「ドメインにアタッチされたカスタムイメージ」のところにイメージが追加されていることが確認できます。
![スクリーンショット 2020-12-24 20.52.35.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/4c7eab87-d453-0191-5842-4a3234553354.png)

<font color="Crimson">新しいイメージをアタッチしたら、Amazon SageMaker Studio のメニューから File -> shutdown でいったん Jupyter server を停止し、再度 Amazon SageMaker Studio を起動してください。これを忘れると、カーネル一覧に新しいカーネルが表示されません。</font>

## 作成したイメージを Amazon SageMaker Studio ノートブックで使用
SageMaker Studio コントロールパネルから、Amazon SageMaker Studio を起動します。ノートブックを新規作成する場合は、Launcher の下の方に Notebooks and compute resources 
という項目があるので、そこの Select a SageMaker Image のプルダウンメニュー（画像の黄色い矢印）で作成したイメージを選択して Notebook ボタン（画像の赤い矢印）をクリックします。

![studio.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/fb7ce9d9-d97a-8fc2-5ed4-9678e78ef860.png)

すでに作成済みのノートブックでイメージを選択する場合は、ノートブック上部にあるカーネル名のエリア（画像の黄色い矢印）をクリックし、プルダウンメニューから所望のカーネルを選択します。
![nb.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/315319/4d626847-35f6-8ff5-7e5f-3b4234b1b7c1.png)

## まとめ
本記事では、Amazon SageMaker Studio で独自のカスタムイメージをノートブックのカーネルとして使用する方法をご紹介しました。この機能によって、自分の使いたい環境を簡単にセットアップできます。