背景
最近、 DVC を使ってデータセットをバージョン管理している環境で、Amazon SageMaker (以下SageMaker) を使ったトレーニングを行う機会がありました。
DVCをSageMakerで使うやり方がいまいちよくわかっていなかったので、調べた結果を備忘録として載せておこうと思います。
前提
本記事では以下の環境を前提としています：
- DVCで管理しているデータがS3に保存されている
- Docker環境が利用可能
- Amazon ECR (以下ECR) リポジトリへのアクセス権限がある
- SageMakerの実行に必要なIAMロール、プロファイルが設定済み
- パッケージ管理にuvを使っている
DVCを使ったSageMakerのトレーニングフロー
DVCをSageMaker内で利用する方法はいくつかありますが、今回はSageMakerのトレーニングジョブを使ってトレーニングしていきます。
その際にECRへアップロードしたDockerイメージを用います。
データのダウンロードはDVCを使って、コンテナの実行時にS3からダウンロードします。
具体的には、以下の流れでトレーニングを実行します：
- DVCの設定情報、依存関係定義、トレーニングコードを含むDockerコンテナを作成
- Dockerイメージをビルドし、ECRにプッシュ
- SageMakerでECRのコンテナイメージを使用してトレーニングジョブを開始
-
DockerfileのENTRYPOINTに書かれた処理を実行- コンテナ起動時にDVCを使ってS3からデータをダウンロード
- トレーニングスクリプトの実行
ここからは実際のプロジェクト構成とその詳細な設定方法について説明していきます。
プロジェクト構成
本記事で使用するプロジェクトの構成は以下の通りです：
dvc-in-sagemaker/
├── Dockerfile # コンテナイメージのビルド設定
├── train_in_docker.sh # エントリーポイントスクリプト
├── pyproject.toml # Pythonプロジェクトの依存関係定義
├── .dvc/
│ └── config # DVCのリモートストレージ設定
├── data/
│ └── processed/
│ └── train.dvc # DVCで管理されるデータファイルの参照情報
└── src/
├── sagemaker/
│ ├── config.py # SageMaker実行設定
│ └── start_sagemaker_job.py # トレーニングジョブ開始スクリプト
└── train_src/
└── train.py # 機械学習トレーニングコード
各ファイルの役割は以下の通りです:
-
Dockerfile
- DVCとトレーニング環境を含むコンテナイメージを定義
- コンテナ起動時に
train_in_docker.shを実行する
-
train_in_docker.sh
-
DockerfileのENTRYPOINTで指定しており、コンテナ起動時にDVCでデータをダウンロードし、トレーニングを実行する
-
-
.dvc/config
- DVCで用いるS3などのリモートストレージへの接続設定
-
data/processed/train.dvc
- DVCで管理されているデータセットのメタデータ
-
src/train_src/train.py
- 実際の機械学習モデルのトレーニングロジック
Dockerイメージの作成
SageMakerで指定するDockerイメージでは、 ENTRYPOINT でインスタンス起動時に実行するコマンドを指定します。
今回は train_in_docker.sh を実行するようにし、その中でデータのダウンロードとトレーニングを行います。
FROM ghcr.io/astral-sh/uv:python3.11-bookworm-slim
WORKDIR /app
# トレーニングコードとUVのファイルをコピー、依存関係をインストール
COPY src/train_src/ src/train_src/
COPY .python-version .
COPY pyproject.toml .
COPY uv.lock .
RUN uv sync
# DVCのファイルをコピー
COPY .dvc/config .dvc/
COPY data/train_data.csv.dvc data/
# dvc コマンドをインストール
RUN pip install dvc dvc-s3
# git のインストール
RUN apt-get update && \
apt-get install -y git
RUN git init
# トレーニングジョブ開始時に実行されるスクリプトをコピー
COPY train_in_docker.sh /app/
RUN chmod +x /app/train_in_docker.sh
# CMD ではなく ENTRYPOINT を使用
ENTRYPOINT ["/app/train_in_docker.sh"]
#!/bin/bash
cd /app/
dvc pull
cd src/train_src
uv run train.py
Dockerイメージのプッシュ
作成した Dockerfile をビルドし、ECR上にプッシュします。
ECRのリポジトリページに表示される「プッシュコマンドを表示」から、プッシュに必要なコマンドを確認して実行できます。
ビルド時の注意点
私はM2 Macを使っているのですが、ECRの「プッシュコマンドを表示」のコマンドでビルドしたDockerイメージをSageMaker上で動かす際にフォーマットエラーが発生し、動作しませんでした。
docker build 時に --platform=linux/amd64 オプションをつけないとアーキテクチャの違いでトレーニングジョブが落ちてしまうので、注意しましょう。
SageMakerのトレーニングジョブの実行
あとは、SageMakerのGUIを操作してECRにプッシュしたコンテナでトレーニングジョブを設定すれば、トレーニングを実行できます。
ただ、毎回GUIを操作するのは手間なので、簡単にプログラムから実行できる方法を調べてみました。
以下のようなコードを用いてSageMakerでトレーニングジョブを実行することができます。
これにより、CI/CDでちょっと楽になるかもしれません。
import boto3
from sagemaker import Session
from sagemaker.estimator import Estimator
# 前提としてprofileは既に設定している必要があります
session = boto3.Session(profile_name='your-profile-name', region_name='ap-northeast-1')
sagemaker_session = Session(boto_session=session)
estimator = Estimator(
image_uri="your-ecr-image-uri",
role="your-sagemaker-execution-role-arn",
instance_count=1,
instance_type="ml.m5.large",
sagemaker_session=sagemaker_session,
# その他の設定...
)
estimator.fit(
wait=True,
)
認証とアクセス権限
DVCがS3からデータを取得するため、適切なIAMロールの設定が必要です。SageMakerの実行ロールには、DVCが使用するS3バケットへの読み取り権限を付与する必要があります。
メリットと課題
メリット
このアプローチの主なメリットは以下の通りです：
- S3にデータを保存しているため、SageMaker実行用のロールに最小権限だけ与えれば良い
- コンテナイメージにデータを含める必要がないため、ECRの保管コストが減る
課題
一方で、データをインスタンス実行中にダウンロードする必要があるため、ダウンロード時間分課金量も増えてしまうという課題もあります。
少量のデータなら大した時間はかかりませんが、データが増えてくると時間がかかります。
とはいえ、インスタンスを使用している時間の大部分はトレーニング時間になるため、データのダウンロード時間の影響は少ないかと思います。
終わりに
今回はS3とSageMakerがAWS内のサービス同士ということもあり、シンプルな構成になったかと思います。
データ管理は機械学習プロジェクトにおいては非常に重要なものになっています。
効率的な運用を目指して、引き続き調査を進めたいと思います。