SageMakerのノートブックインスタンス上に独自のカーネルを永続的に追加する

はじめに

SageMakerノートブックインスタンスにはデフォルトでいくつかのカーネルが用意されていますが、独自のカーネルを追加して使いたい、というケースが有るかと思います。
例えばPython3.8を使いたい、となると2021年2月現在Python3.6のカーネルしか提供されていません。
その場合、Python3.8のconda環境を自分でインストールして、それをノートブックインスタンスのカーネルに追加するということが必要になります。
なのでこの記事では、その方法について紹介したいと思います。

とりあえずやってみる

まずは何も考えずにノートブックインスタンス上にPython3.8のconda環境をインストールしてみます。

conda環境のインストール

ノートブックインスタンス上でターミナルを開き、以下のコマンドを打ってみます。
※2つ目のコマンドがないとノートブックのカーネル一覧に表示されません

conda create --name myenv python=3.8
python -m ipykernel install --user --name myenv --display-name "myenv"

結果

更新ボタンを押して新しくノートブックを作成しようとすると、新しく作成した環境がカーネル一覧に追加されていることが分かります。

ノートブックインスタンスを再起動してみる

これでめでたしめでたし、とならないのがSageMakerのやっかいなところ。
ノートブックインスタンスを一旦停止してもう一度起動すると、作成したconda環境が消えてしまいます。

原因

SageMakerのノートブックインスタンスはEC2インスタンスなので、再起動するごとにEBSボリューム(/home/ec2-user/SageMaker/)に保存されたデータ以外はデフォルトの状態にリセットされてしまいます。
conda環境のデフォルトのインストール先は/home/ec2-user/anaconda3/なので、再起動するとリセットされてしまうというわけです。
なので、永続的に独自のconda環境を利用しようと考えると、conda環境をEBSボリュームにインストールする必要があります。

対応策

公式のドキュメントを見てみると、ノートブックのライフサイクル設定を追加する方法が紹介されています。
その方法は、ノートブックの開始時に独自のconda環境をEBSボリューム(/home/ec2-user/SageMaker/)にインストールして、Jupyter のカーネルとして使用できるよう追加するようなスクリプトを実行する、という感じです。
しかし、このライフサイクル設定には5分以内に処理を完了させないとノートブックインスタンスの起動が失敗するという制限があるので、容量の大きいconda環境をインストールしようとすると、オーバーしてしまう可能性があります。
また、conda環境は一度EBSボリュームに保存してしまえば起動時にインストール処理を行う必要はありません。
なので、conda環境のインストールは手で前もって実行しておいて、ライフサイクル設定ではEBSボリュームに追加したconda環境を認識させる処理のみを記述したいと思います。

やってみる

ということで実際に実行してみます。

EBSボリュームへのconda環境のインストール

まずはEBSボリューム(/home/ec2-user/SageMaker/)に独自のconda環境をインストールします。
ノートブックインスタンスでターミナルを開き次のコマンドを実行します。

conda create --prefix /home/ec2-user/SageMaker/kernels/myenv python=3.8 ipykernel
source activate /home/ec2-user/SageMaker/kernels/myenv
python -m ipykernel install --user --name myenv --display-name myenv
source deactivate

これで更新ボタンを押すとカーネル一覧に新しくEBSボリュームに追加したconda環境が追加されています。

ライフサイクルの設定

しかし、2行目から4行目のconda環境のカーネルへの追加はノートブックインスタンスが再起動するごとにリセットされてしまいます。
なので、ライフサイクル設定を用いてノートブックインスタンスが起動するごとにこれらのコマンドを打つようにします。
SageMakerのマネジメントコンソールからライフサイクル設定を開き、設定の作成を開きます。

すると、次のような画面が表示されますので、ここに起動時に実行するシェルスクリプトを記述します。

スクリプトはこんな感じです。

#!/bin/bash

sudo -u ec2-user -i <<'EOF'

if [ -d "/home/ec2-user/SageMaker/kernels" ]; then
  for env in /home/ec2-user/SageMaker/kernels/*; do
    source activate $env
    python -m ipykernel install --user --name $(basename "$env") --display-name "$(basename "$env")"
    source deactivate
done
fi

EOF

次に、ノートブックインスタンスを選択して設定の更新を開きます。
ライフサイクル設定で、さきほど作成したライフサイクルを選択します。

結果

これで一通り完了したので、もう一度ノートブックインスタンスを再起動してみます。
カーネルの一覧を確認すると、無事新しく作成したconda環境があります。これで完了です。

注意事項

独自のconda環境のインストール先(/home/ec2-user/SageMaker/kernels)にconda環境以外のデータやファイルを置かないでください（起動時にエラーが出ます）。

#参考
https://aws.amazon.com/jp/premiumsupport/knowledge-center/sagemaker-lifecycle-script-timeout/
https://medium.com/decathlondevelopers/making-jupyter-kernels-remanent-in-aws-sagemaker-a130bc47eab7