SageMaker JupyterLab 環境に環境変数を設定したい
概要
カスタムイメージで構築した SageMaker Studio 管理下の Jupyter Lab 環境に、外部から環境変数を注入したいケースに遭遇した。
うまくいった方法に加えて、試行錯誤（うまくいかなかった方法）結果も記載する。
ユースケース
- SageMaker カスタムイメージを使用している
- Notebookの中で、クレデンシャルが必要な処理を実行したい
- APIキー、ライセンスキーなど
うまくいった方法
クレデンシャルは ssm パラメータストア（または Secrets Manager）に保存する。
シェルから、保存したパラメータを読み取って、環境変数に設定する。
.
├── Dockerfile
└── main.sh
FROM public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu
ARG NB_USER="sagemaker-user"
ARG NB_UID=1000
ARG NB_GID=100
ENV MAMBA_USER=$NB_USER
USER root
RUN apt-get update
RUN micromamba install sagemaker-inference --freeze-installed --yes --channel conda-forge --name base
USER $MAMBA_USER
COPY ./main.sh /usr/local/bin/
CMD ["main.sh"]
#!/bin/bash
# クレデンシャルを環境変数に設定
export CREDENTIAL = $(aws ssm get-parameter ...)
jupyter-lab --ServerApp.ip=0.0.0.0 --ServerApp.port=8888 --ServerApp.allow_origin=* --ServerApp.token='' --ServerApp.base_url=/jupyterlab/default
Dockerfile および jupyter 起動コマンド は、以下リンク先が元になっている。
うまくいかなかった方法
1. ライフサイクル設定で環境変数を設定
Sagemaker で用意されている「ライフサイクル設定」機能を利用。
#!/bin/bash
export CREDENTIAL = $(aws ssm get-parameter ...)
Jupyter Notebook および Jupyter 内のターミナルのいずれにも、環境変数は設定されていなかった。
恐らく、 Jupyter Notebook が動作するプロセスと、ライフサイクル設定を行うプロセスが異なるためだと思われる。
2. ~/.bashrc に、環境変数読み込みコードを追記
#!/bin/bash
export CREDENTIAL = $(aws ssm get-parameter ...)
bash -c "echo \"CREDENTIAL=${CREDENTIAL}\" >> ~/.bashrc"
Jupyter 内ターミナルには環境変数が設定されていたが、 Jupyter Notebook 内には未設定だった。
.bashrc がシェル起動時に実行されるスクリプトのため、当然と言えば当然か。
3. 環境変数読み込み処理とjupyter起動を別スクリプト化(うまくいった方法の前身)
.
├── Dockerfile
└── main.sh
FROM public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu
ARG NB_USER="sagemaker-user"
ARG NB_UID=1000
ARG NB_GID=100
ENV MAMBA_USER=$NB_USER
USER root
RUN apt-get update
RUN micromamba install sagemaker-inference --freeze-installed --yes --channel conda-forge --name base
USER $MAMBA_USER
COPY ./main.sh ./
CMD ["./main.sh"]
#!/bin/bash
# クレデンシャルを環境変数に設定
export CREDENTIAL = $(aws ssm get-parameter ...)
jupyter-lab --ServerApp.ip=0.0.0.0 --ServerApp.port=8888 --ServerApp.allow_origin=* --ServerApp.token='' --ServerApp.base_url=/jupyterlab/default
うまくいった方法と違い、起動用スクリプトをホームディレクトリに配置している。
このカスタムイメージを使用して Jupyter の起動を試みると、起動に失敗した。
ログを確認すると、以下のエラーメッセージが確認できた。
/home/sagemaker-user/main.sh: No such file or directory
ローカルでコンテナに入ったところ、 /home/sagemaker-user/main.sh は存在した。
なぜ main.sh が見つからないのか
Sagemaker 側で Jupyter 起動時にアタッチする EFS が原因だと思われる。
EFSを /home/sagemaker-user/ へアタッチするため、コンテナ内のホームディレクトリ配下に配置したファイルが全て消滅（上書き？）する。
main.sh を、ホームディレクトリである /home/sagemaker-user/ 配下 以外 に配置することで対処できる。
このとき、パスが通っている /usr/local/bin/ に配置すると、 どこからでも main.sh を実行できるようになって良い感じになる。
Discussion