SageMaker TrainingにおけるDockerイメージとConda環境の最適化

# はじめに
以前、以下の記事でDockerコンテナ内にConda環境を設定し、Amazon SageMakerでのトレーニングを実行する手順を書きました。

https://qiita.com/suzuki0430/items/47228a8af26ccb4a88a2

また、以下の記事では、デバッグ環境(EC2/Deep Learning AMI)とトレーニング環境(SageMaker)で環境を一致させるために、condaの設定ファイル`environment.yml`をデバッグ環境から書き出して、それを使ってトレーニング環境のDockerイメージを構築しました。

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

DockerのベースイメージにSageMaker Training用のイメージを選択した結果、ベースイメージと`environment.yml`のライブラリの重複が原因で、11 GBくらいの巨大なイメージが生成されてしまいました。

当時はとりあえず早く学習を回したい！という気持ちだったので、そのコンテナイメージを使い続けていました。
ただ、やはりイメージの読み込みに時間がかかってしまったり、主要ライブラリ(`transformers`や`torch`など)の将来的なバージョンアップで苦労するのが目にみえていたので、このたび`environment.yml`と`Dockerfile`を整理しました。

デバック環境で構築したConda環境をSageMakerで使いたいけれど、Dockerイメージをどうやって構築したらいいかいまいちわからない方に役立つ内容になっていると思います。

# 実装
以下が最終的な`Dockerfile`と`environment.yml`です。

```Dockerfile:Dockerfile
# Amazon SageMakerのPyTorch GPUトレーニング用ベースイメージを使用
FROM 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker

# Condaのパスを設定（ベースイメージにCondaが含まれているため、インストールは不要）
ENV PATH /opt/conda/bin:$PATH

# Conda環境を初期化
RUN conda init bash

# environment.ymlファイルをコンテナにコピー
COPY environment.yml /root/environment.yml

# environment環境を作成
RUN conda env create -f /root/environment.yml

# Conda環境を自動的にアクティブにする
RUN echo "conda activate environment" >> ~/.bashrc

# コンテナ起動時にenvironment環境でbashを起動
CMD [ "/bin/bash" ]
```

```yml:enrironment.yml
name: environment
channels:
  - anaconda
  - pytorch
  - huggingface
  - conda-forge
dependencies:
  - python=3.8.13
  - pandas=1.4.3
  - scikit-learn=1.1.1
  - transformers=4.26.0
  - numpy=1.23.1
  - libgcc
  - imbalanced-learn=0.10.1
  - pip:
      - wandb
      - python-dotenv
```

SageMakerのPyTorch GPUトレーニング用のベースイメージを使用し、Conda環境での重複インストールを避けるために、`environment.yml`から`pytorch`や`cudatoolkit`の記載を削除しました。

これによってイメージサイズは約8 GBとなり、以前のイメージから3 GBほど削減することができました。

# 補足1: イメージビルドとECRへのプッシュ
作成した`Dockerfile`を元にビルドして、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/<イメージ名>:<タグ名> .
```
※ビルドを行うマシンがARM系のMac(M1~3)の場合は以下のコマンドでビルドを実行します。

```bash
$ docker buildx build --platform linux/amd64 -f Dockerfile -t <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/<イメージ名>:<タグ名> .
```

また、イメージプッシュする際にも、自アカウントの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/<イメージ名>:<タグ名>
```

# 補足2: SageMaker側のエントリーポイントとなるシェルスクリプトの設定
SageMaker Trainingジョブを設定する際には、以下のようなPythonスクリプトを使用しますが、この設定だけでは自動的にConda環境に切り替わらないため、追加のステップが必要です。

```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="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"],
    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"
})
```

そのため、`entry_point`で指定しているシェルスクリプトでConda環境の切り替えを行う必要があります。

```sh:pre-training.sh
#!/bin/bash

# conda activate
source /opt/conda/etc/profile.d/conda.sh
conda activate environment

# pre-training
python main.py
```

# 補足3: 失敗例
さらなるイメージサイズ軽量化を目指してつくった`Dockerfile`と`environment.yml`の例も紹介します。
この例では、デバッグ環境に通常のUbuntu AMIを使用したEC2インスタンスを想定しています。

まず、`Dockerfile`の方はベースイメージに`ubuntu`を使用しています。
このベースイメージには`conda`が含まれていないため、wgetを使用してMinicondaをインストールしました(Conda環境とあわせてpy38を指定しています)。
また、SageMaker Trainingで必要な`build-essential`と`sagemaker-training`を別途インストールしています。

※`sagemaker-training`がないと`[FATAL tini (7)] exec train failed: No such file or directory`というエラーが出ます。このパッケージは、SageMaker Trainingで必要となるため、それが欠如しているとエラーが発生します。

https://github.com/aws/sagemaker-python-sdk/issues/4113

```Dockerfile:Dockerfile
# Ubuntu 22.04 LTSをベースイメージとして使用
FROM ubuntu:22.04

# 必要なパッケージをインストール
RUN apt-get update && \
    apt-get install -y wget bzip2 ca-certificates build-essential && \
    rm -rf /var/lib/apt/lists/*

# Minicondaをダウンロードしてインストール
RUN wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-py38_23.11.0-2-Linux-x86_64.sh -O ~/miniconda.sh && \
    /bin/bash ~/miniconda.sh -b -p /opt/conda && \
    rm ~/miniconda.sh && \
    /opt/conda/bin/conda clean --all --yes

# Condaのパスを設定
ENV PATH /opt/conda/bin:$PATH

# Conda環境を初期化
RUN conda init bash

# environment.ymlファイルをコンテナにコピー
COPY environment.yml /root/environment.yml

# environment環境を作成
RUN conda env create -f /root/environment.yml

# Conda環境を自動的にアクティブにする
RUN echo "conda activate environment" >> ~/.bashrc

# sagemaker-training パッケージをインストール
RUN pip install sagemaker-training

# コンテナ起動時にtabformer-opt-v2環境でbashを起動
CMD [ "/bin/bash" ]
```
一方、`environment.yml`の方では、構築先に含まれない各種ライブラリの依存関係をしています。
そのため、採用した`environment.yml`と比べると、指定する依存関係が多いです。

```yml:environment.yml
name: environment
channels:
  - anaconda
  - pytorch
  - huggingface
  - conda-forge
dependencies:
  - python=3.8.13
  - pip>=21.0
  - pytorch=1.12.1
  - torchaudio=0.12.1
  - torchvision=0.13.1
  - pandas=1.4.3
  - scikit-learn=1.1.1
  - transformers=4.26.0
  - numpy
  - libgcc
  - cudatoolkit=11.6
  - imbalanced-learn=0.10.1
  - pip:
      - transformers==4.26.0
      - wandb
      - python-dotenv
```

こちらの例では、なんとイメージサイズを約3 GBまで減らすことができ、デバッグ環境やSageMaker環境でのトレーニング自体は開始できたのですが、なぜかGPUではなくCPUが使われてしまいました。

そのため、残念ながら実採用には至りませんでした。

この経験から、SageMakerのPyTorch GPUトレーニング用ベースイメージに含まれるライブラリを基に`environment.yml`を構築することの重要性がわかりました。