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

はじめに

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

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

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

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

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

実装

以下が最終的なDockerfileとenvironment.ymlです。

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" ]

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クライアントの認証を行う必要があります。

$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com

これでイメージビルドを行うことができます。

$ docker build -t <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/<イメージ名>:<タグ名> .

※ビルドを行うマシンがARM系のMac(M1~3)の場合は以下のコマンドでビルドを実行します。

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

また、イメージプッシュする際にも、自アカウントのECRレジストリへの認証を行います。

$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com

これでイメージプッシュを行うことができます。

$ docker push <アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/<イメージ名>:<タグ名>

補足2: SageMaker側のエントリーポイントとなるシェルスクリプトの設定

SageMaker Trainingジョブを設定する際には、以下のようなPythonスクリプトを使用しますが、この設定だけでは自動的にConda環境に切り替わらないため、追加のステップが必要です。

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環境の切り替えを行う必要があります。

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で必要となるため、それが欠如しているとエラーが発生します。

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と比べると、指定する依存関係が多いです。

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を構築することの重要性がわかりました。