【SageMaker】Training Jobsでconda activate後に学習を実行する

本記事は、AI/ML on AWS Advent Calendar 2022の18日目の記事です。

以前の記事で、TabBERTモデル(IBM論文の付属コード)の環境構築と学習(事前学習、Fine-Tuning)をEC2上で行いました。
動作確認くらいであればこれで問題なかったのですが、いざ本番運用を考えてみると、以下のような問題がありそうでした。

- 学習中以外の時間にかかるEC2のコスト
- 推論環境を別途たてることになったときの環境再構築コスト
- 入出力データの管理の手間

https://qiita.com/suzuki0430/items/89dd028996f7dc774796

調べてみたところ、Amazon SageMakerというAWSの機械学習用マネジメントサービスを使えばこれらを解決できそうだったので、技術書展で購入した書籍を参考に、勉強がてらSageMakerを触ってみることにしました。

いざ触ってみるとつまづきポイントが多くて心が折れかけたのですが、1週間粘ってようやくJobを成功させることができました。

今回の記事はSageMaker Training Jobsの構築手順が主ですが、その中につまづいたポイントの記述も行っています。
同じ部分で悩んでいる方の参考になれば幸いです。

- Dockerfileのベースイメージ選択
- conda環境切替のためのDockerfile
- ノートブックインスタンスの作成
- コンテナとS3の入出力データのやりとり
- Training Job実行時のコンテナのパス(/opt/ml/~)
- Training Job実行前のconda環境切替(conda activate)
- 環境変数によるhyperparametersの引き渡し

# Docker環境構築
Dockerfile作成、イメージビルド、ECRプッシュまでの手順をまとめました。

## Dockerfile作成
SageMakerでは、Training Jobを実行する際にコンテナイメージから学習用のインスタンスを起動します。
以前、EC2で環境構築した際はDeep Learning AMIを使用し、ymlファイルからconda環境をつくりました。

コンテナで同様の環境をつくるにあたり、いろいろ試行錯誤しました。
先にDockerfile全文を記載します。

```Dockerfile
FROM 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker

COPY environment.yml .

RUN pip install --upgrade pip && \
    conda update -n base -c defaults conda && \
    conda env create -f environment.yml && \
    conda init bash && \
    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
    echo "conda activate tabformer-opt-sagemaker" >> ~/.bashrc

ENV CONDA_DEFAULT_ENV tabformer-opt-sagemaker && \
    PATH /opt/conda/envs/tabformer-opt-sagemaker/bin:$PATH

ENV LD_LIBRARY_PATH $LD_LIBRARY_PATH:/opt/conda/envs/tabformer-opt-sagemaker/lib
RUN echo $LD_LIBRARY_PATH
```

### ベースイメージ選定

```Dockerfile
FROM 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker
```

EC2では`Deep Learning AMI GPU PyTorch 1.13.0 (Ubuntu 20.04) 20221110`を使っていたので、それに相当するベースイメージを探しました。

調査に時間がかかってしまったのですが、AWSが公開しているdeep-learning-containersというものをみつけました。

https://github.com/aws/deep-learning-containers/blob/master/available_images.md

SageMaker Framework Containersから、以下を選択しました(PyTorch 1.13.0がなく、1.12.1を選択しました)。
![スクリーンショット 2022-12-17 10.41.51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/3a199d3a-e81f-5130-df44-880690c44a2d.png)

今回は以上のベースイメージを使用しましたが、Pyhonが動き、`build-essential`と`sagemaker-training`が入っている環境であれば、Training Jobは動作するとのことです。

### conda環境の構築

```Dockerfile
RUN pip install --upgrade pip && \
    conda update -n base -c defaults conda && \
    conda env create -f environment.yml && \
    conda init bash && \
    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
    echo "conda activate tabformer-opt-sagemaker" >> ~/.bashrc

ENV CONDA_DEFAULT_ENV tabformer-opt-sagemaker && \
    PATH /opt/conda/envs/tabformer-opt-sagemaker/bin:$PATH
```

以下の記事でも書いたのですが、元コードの`setup.yml`から動作するconda環境をつくれず、追加でライブラリを導入するなどいろいろいじりました。
そのため、最終的に動作したconda環境を新しいymlファイル`environment.yml`に書き出し、そのファイルを元にconda環境をつくるようにしました。

https://qiita.com/suzuki0430/items/0bda6a04019cef2eb319

また、コンテナを立ち上げたときにconda環境も立ち上がるよう、`.bashrc`に`conda activate ~`を書き込み、デフォルト環境`CONDA_DEFAULT_ENV`として設定しました。
※結局これだけではconda環境が立ち上がらず、後述のbashスクリプト内で`conda activate`しました。

https://note.yu9824.com/howto/2021/09/04/docker-conda-activate.html

#### 補足

```Dockerfile
ENV LD_LIBRARY_PATH $LD_LIBRARY_PATH:/opt/conda/envs/tabformer-opt-sagemaker/lib
RUN echo $LD_LIBRARY_PATH
```
以上のパスを通さないと学習が実行できなかったので追記しました。
私の環境だと必要でしたが、通常であればこちらの記述は不要です。

## イメージビルドとECRへのプッシュ
`deep-learning-containers`のイメージをベースとして使う場合、イメージをビルドする前にリージョンとECRレジストリを指定して、認証トークンの取得とDockerクライアントの認証を行う必要があります。

```bash
$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com
```
これでイメージビルドを行うことができます。

```bash
$ docker build -t <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/tabformer-opt:<タグ> .
```
また、イメージプッシュする際にも、自アカウントのECRレジストリへの認証を行います。

```bash
$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com
```
これでイメージプッシュを行うことができます。

```bash
$ docker push <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/tabformer-opt:<タグ>
```

# Training Jobの作成
Training Jobを作成するにあたり、ノートブックインスタンスに学習用のコードとJobを実行するコードを配置する必要があります。

## ノートブックインスタンスの作成
以下を付与したIAMロールを作成し、ノートブックインスタンス作成時にアタッチします。

- AmazonSageMakerFullAccess
- AmazonS3FullAccess
- AmazonEC2ContainerRegistryFullAccess

![スクリーンショット 2022-12-17 11.36.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/6ca46dee-74fd-c83c-3f98-82dbd00fdd30.png)

## フォルダ構成
作成したノートブックインスタンスからJupyterを開き、以下のようにフォルダやファイルを配置します(画像の`shell.sh`は不要です)。

- `tabformer-opt`:学習用のコード
- `pre-training_jobs.ipynb`：Jobの実行コード
- `pre-training.sh`：conda activate、Python実行用bashスクリプト

![スクリーンショット 2022-12-17 12.34.19.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/468570dc-7299-395c-fcd6-ca14c1f2576f.png)

## pre-training_jobs.ipynbの作成
`pre-training_jobs.ipynb`を作成し、前項で作成したコンテナイメージを読み込んでモデルの学習を実行できるようにします。
モデルの学習にはsagemaker.estimator.Estimatorを使用します。

https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html

こちらがコード全文です。

```py:pre-training_jobs.ipynb
import sagemaker
from sagemaker.estimator import Estimator

session = sagemaker.Session()
role = sagemaker.get_execution_role()

estimator = Estimator(
    image_uri=<作成したコンテナイメージのURL>,
    role=role,
    instance_type="ml.g4dn.2xlarge",
    instance_count=1,
    base_job_name="tabformer-opt-pre-training",
    output_path="s3://<バケット名>/sagemaker/output_data/pre_training",
    code_location="s3://<バケット名>/sagemaker/output_data/pre_training",
    sagemaker_session=session,
    entry_point="pre-training.sh",
    dependencies=["tabformer-opt"],
    hyperparameters={
        "mlm": True,
        "do_train": True,
        "field_hs": 64,
        "output_dir": "/opt/ml/model/",
        "data_root": "/opt/ml/input/data/input_data/",
        "data_fname": "<ファイル名>"
    }
)
estimator.fit({
    "input_data": "s3://<バケット名>/sagemaker/input_data/<ファイル名>.csv"
})
```
### コンテナパスとS3パスの対応
SageMakerではコンテナ起動時に`/opt/ml/~`のパスが使用されます。
S3とのデータ入出力を考えるとき、コンテナでのデータ配置とS3パスとの対応を理解しておく必要があります。

#### `/opt/ml/model`
こちらのパスにデータを保存すると、学習終了時にS3にデータが自動でアップロードされます。
S3の保存先は、Estimatorの引数`output_path`で指定します。

学習終了時ではなくリアルタイムでS3にデータをアップロードする場合は、`/opt/ml/checkpoints`にデータを保存します。
このとき、S3の保存先は`checkpoint_s3_uri`で指定します。

#### `/opt/ml/code`
`pre-training_jobs.ipynb`と同じディレクトリにあるフォルダとファイル(今回の場合`tabformer-opt`と`tabformer-opt.sh`)がこちらに配置されます。

#### `/opt/ml/input`
estimator.fitで入力データのあるS3パスを`<key>: <S3 URL>`で指定すると、コンテナでは`/opt/ml/input/data/<key>`以下に配置されます。

```py
estimator.fit({
    "input_data": "s3://<バケット名>/sagemaker/input_data/<ファイル名>.csv"
})
```

## pre-training.shの作成
こちらのbashスクリプトでは、conda環境の切替(`conda activate`)、フォルダ`tabformer-opt`の移動、学習を行います。

```sh
#!/bin/bash

"===conda activate==="
conda init bash
source /opt/conda/etc/profile.d/conda.sh
conda activate tabformer-opt-sagemaker

"===move directory==="
mv ./tabformer-opt/* .

"===pre-training==="
python main.py
```

### conda activate
Dockerfileにも`conda activate`を記述したのですが、`conda init bash`後のCommandNotFoundErrorを解消できませんでした。
そのため、bashスクリプト内で`source /opt/conda/etc/profile.d/conda.sh`を実行し、`conda activate`されるようにしました。

https://japan-challenge.net/conda-activate%E3%81%AE%E3%82%A8%E3%83%A9%E3%83%BC%E3%82%92conda-init%E3%81%A7%E8%A7%A3%E6%B1%BA%E3%81%99%E3%82%8B/

https://github.com/conda/conda/issues/8536

### move directory
学習のためには`tabformer-opt/main.py`を実行する必要があります。
`main.py`を実行するには`tabformer-opt`内の依存ファイルも読み込む必要があり、そのためにEstimatorで`dependencies=["tabformer-opt"]`を記述します。

あとは`entry_point: python tabformer-opt/main.py`とすればよさそうですが、これだとライブラリのカレントディレクトリが変わってしまい、importがうまくいかずJobが失敗します。

そのため、`mv ./tabformer-opt/* .`で`tabformer-opt`以下をカレントディレクトリに移動させています。

## 環境変数によるhyperparametersの引き渡し
`main.py`を実行するにあたり、今回のコードでは`--do_train` `--mlm`などをコマンドライン引数として渡す必要があります。
しかし、`tabformer-opt.sh`がentry_pointになったことにより、Estimatorのhyperparametersをコマンドライン引数として`main.py`に渡すことができません。

ここでSageMakerの環境変数`SM_HPS`を利用します。
`SM_HPS`にはhyperparametersのJSONがテキストに変換されたものが入っています。

```
SM_HPS={"data_fname":"summary","data_root":"/opt/ml/input/data/input_data/","do_train":"True","field_hs":64,"mlm":"True","output_dir":"/opt/ml/model/"}
```

これを`main.py`で以下のように読み込むと、SageMakerを使用しているとき(Training Jobsを実行しているとき)のみ、hyperparametersを利用するようになります。

```py:main.py
if __name__ == "__main__":

    parser = define_main_parser()
    opts = parser.parse_args()
    
    if "SM_HPS" in os.environ.keys():
        hps = json.loads(os.environ["SM_HPS"])
        for key, value in hps.items():
            if opts.__contains__(key):
                opts.__setattr__(key, value)

    main(opts)
```

# Training Jobの実行
`pre-training_jobs.ipynb`を実行すると、Training Jobが実行されます。
Jobが成功すると、ジョブのステータスが`Completed`となります。
![スクリーンショット 2022-12-17 14.54.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/4f6dd74c-1226-3f14-e0cc-16828648dfb3.png)

CloudWatchで実行ログを確認することができます。
![スクリーンショット 2022-12-17 14.55.48.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/af638c9d-44cd-f022-e598-f4a8732396d7.png)

Job実行後にS3を確認すると複数のオブジェクトがあり、`output`の`model.tar.gz`内にモデルが保存されています。
![スクリーンショット 2022-12-17 14.57.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/569054/67e013c9-118e-e07f-d4f5-0327351f545c.png)

# 参考資料

https://koshian2.booth.pm/items/4150087

https://amzn.asia/d/3vBnCSw