GitHub Actionsを使ったAmazon SageMaker Pipelinesの自動化

2024/06/21に公開

TL;DR

Amazon SageMaker Pipelinesを自動で再実行するために、GitHub Actionsの定期実行機能を使用しました。パイプライン内で最新のデータを格納するステップを含んでいるため、定期実行によって常に最新のデータで再学習できます。他の再学習方法もありますが、この方法が簡単かつ効果的だったため選択しました。

対象の読者

  • SageMakerでの機械学習モデルの定期的な再学習を自動化したいエンジニア
  • 過去の足跡を辿りたくなった私

イントロダクション

私はこれまで、ローカルからSageMaker Pipelinesを実行していましたが、これを定期的に再実行できるようにしたいと考えました。そこで、GitHub Actionsを使ってCron化することにしました。最初は、ローカルで使っているDockerコンテナをGitHub Actions上でも使用し、SageMaker Pipelinesを起動するコマンドを叩けるようにすることを考えました。しかし、必要なパッケージはsagemakerくらいだったため、直接インストールする方法に変更しました。

フォルダ構造とDockerコンテナの設定

以下のフォルダ構造と設定ファイルを使用して、DockerコンテナをGitHub Actions上で使用しようと試みました。しかし、モデルを動かすわけではないので、必要なパッケージはsagemakerくらいで限定的でありこの案はやめました。実行環境ごとに必要なパッケージを切り掛けられるのも、SageMaker Pipelinesのいいところです。

一応、DockerコンテナをGitHub Actions上で起動する際の設定を載せておきます。

フォルダ構造

.
├── .github
│   └── workflows
│       └── main.yml
└── docker
  ├── Dockerfile
  ├── requirements.txt
  └── action.yml

action.yml

name: 'Custom Docker Action'
description: 'Runs a custom Docker container'
runs:
  using: 'docker'
  image: 'Dockerfile'

ワークフロー (main.yml)

name: Use Custom Docker Action

on: push

jobs:
  example-job:
    runs-on: ubuntu-latest
    
    steps:
      - name: Checkout repository
        uses: actions/checkout@v2
        
      - name: Build and run Docker container
        uses: ./docker
          
      - name: Run command inside Docker container
        run: docker run -v ${{ github.workspace }}:/code -w /code my-docker-image python your_script.py

なぜ定期実行を選択したのか?

再学習は機械学習モデルの精度を維持・向上させるために行われます。以下に代表的な再学習の方法をいくつか紹介します。

1. 定期再学習(Periodic Retraining)

一定の時間間隔でモデルを再学習させる方法です。例えば、毎日、毎週、毎月などのスケジュールで再学習を行います。

  • メリット: システムがシンプルで、スケジュールに従って確実に再学習が行われる。
  • デメリット: 新しいデータが到着していない場合でも再学習が行われるため、計算資源の無駄が発生する可能性がある。

2. トリガー再学習(Triggered Retraining)

新しいデータが到着した際にモデルを再学習させる方法です。例えば、新しいデータバッチが一定のサイズに達した時や特定のイベントが発生した時に再学習を行います。

  • メリット: 新しいデータが利用可能になった時点で再学習が行われるため、最新のデータを即座に反映できる。
  • デメリット: 再学習のタイミングが不規則になりやすく、システムの複雑度が増す。

3. インクリメンタル再学習(Incremental Learning)

モデルを部分的に再学習させる方法です。新しいデータが到着した際に、そのデータのみを用いて既存のモデルを更新します。

  • メリット: 全データセットを使用する再学習に比べて計算コストが低く、迅速にモデルを更新できる。
  • デメリット: モデルの品質が新しいデータに大きく依存するため、データの偏りやノイズの影響を受けやすい。

4. オンライン再学習(Online Learning)

データが到着するたびに、即座にモデルを更新する方法です。ストリーミングデータに適しており、リアルタイムにモデルを更新します。

  • メリット: データが到着次第即座にモデルが更新されるため、最も最新の状態を常に維持できる。
  • デメリット: 高頻度の更新が必要なため、システムの負荷が高くなる。

5. バッチ再学習(Batch Learning)

定期的に新しいデータをまとめて収集し、そのバッチを用いてモデルを再学習させる方法です。

  • メリット: 大量のデータをまとめて処理するため、モデルの精度向上が期待できる。
  • デメリット: 新しいデータが反映されるまでに時間がかかる。

今回の要件ではリアルタイム性が求められるため、オンライン再学習が理想的ですが、データ量がそんなに多くなく、すべてのデータで1から再学習するのが苦にならないため、SageMaker Pipelinesの中で最新のデータをまとめて取得し定期再学習する方法を採用しました。

前提条件

この記事を進めるにあたって、以下のツールとサービスが必要です:

  • GitHubリポジトリ
  • AWSアカウント
  • Amazon SageMaker

また、SageMaker実行ロールを作成し、必要な権限を付与する必要があります。GitHub SecretsにAWS認証情報(AWS_ACCESS_KEY_IDAWS_SECRET_ACCESS_KEY)を設定します。

GitHub Actionsのセットアップ

まず、GitHubリポジトリを作成し、.github/workflowsディレクトリを作成します。次に、以下のような初期のワークフロー設定を行います。

name: Train Model

on:
  workflow_dispatch:
    inputs:
      config_file_name:
        description: 'The name of the configuration file to use'
        required: true
        default: 'default-configuration-filename'
  schedule:
    - cron: '0 21 * * *'  # UTCの午後9時(日本時間の午前6時)

jobs:
  train:
    runs-on: ubuntu-latest
    
    steps:
      - name: Checkout repository
        uses: actions/checkout@v2

AWS認証情報の設定

GitHub SecretsにAWS認証情報を設定します。これにより、GitHub ActionsがAWSにアクセスできるようになります。特に、デフォルトのリージョンを設定するために必要です。次に、aws-actions/configure-aws-credentialsアクションを使用してAWSに認証します。

      - name: Configure AWS credentials
        uses: aws-actions/configure-aws-credentials@v1
        with:
          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          aws-region: ${{ secrets.AWS_DEFAULT_REGION }}

パイプライン実行のための認証設定

SageMaker Pipelinesを実行するために必要な認証設定を行います。以下のようなPythonスクリプトを準備します。

import os
from typing import Optional, Tuple

import boto3
import sagemaker
from dotenv import load_dotenv
from sagemaker.workflow.pipeline_context import PipelineSession

from pipeline.utils.logger import setup_logger


def get_boto_session() -> Tuple[Optional[str], boto3.Session]:
    logger = setup_logger(__name__)
    try:
        load_dotenv()
        role_arn = os.getenv("SAGEMAKER_EXECUTION_ROLE_ARN")
        sts_client = boto3.client("sts")
        assumed_role = sts_client.assume_role(
            RoleArn=role_arn,
            RoleSessionName="my-session",
        )
        credentials = assumed_role["Credentials"]
        boto_session = boto3.Session(
            aws_access_key_id=credentials["AccessKeyId"],
            aws_secret_access_key=credentials["SecretAccessKey"],
            aws_session_token=credentials["SessionToken"],
        )
        logger.info("Boto session created successfully.")
        return role_arn, boto_session
    except Exception as e:
        logger.warn("Boto session creation failed. Trying to create session automatically.")
        return None, boto3.Session()


def get_sagemaker_session(boto_session: boto3.Session, default_bucket) -> sagemaker.Session:
    return sagemaker.Session(boto_session=boto_session, default_bucket=default_bucket)


def get_pipeline_session(boto_session: boto3.Session, default_bucket) -> PipelineSession:
    return PipelineSession(boto_session=boto_session, default_bucket=default_bucket)

コラム:AssumeRoleについて

AssumeRoleは、あるIAMロールの権限を一時的に引き受ける機能です。これにより、より限定された権限を持つIAMユーザーが、高度な権限を持つロールを引き受けて特定の操作を行うことができます。この設計により、セキュリティを保ちながら柔軟な権限管理が可能になります。
具体的には、IAMユーザーに対してRoleをAssumeできるPolicyをアタッチし、SageMakerを操作できるIAMロールを別途定義してそのロールを引き受けるように設計されています。

ワークフローの詳細

GitHub Actionsのワークフローを以下のように定義します。

      - name: Install dependencies
        run: |
          pip install --upgrade --force-reinstall "botocore>1.21.30" "boto3>1.18.30" "awscli>1.20.30" "pydantic>2.6.3" "sagemaker" "python-dotenv"
  
      - name: Run train script
        env:
          CONFIG_FILE_NAME: ${{ inputs.config_file_name || 'default-configuration-filename' }}
          SAGEMAKER_EXECUTION_ROLE_ARN: ${{ secrets.SAGEMAKER_EXECUTION_ROLE_ARN }}
          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          AWS_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
        run: |
          python manage.py runpipeline $CONFIG_FILE_NAME

問題解決とデバッグ

リージョン設定の問題

リージョン設定の問題が発生する場合、以下のポイントを確認してください:

  • GitHub Secretsに設定されたAWS認証情報が正しいことを確認してください。
  • aws-actions/configure-aws-credentialsアクションで設定されたリージョンが、SageMakerを利用可能なリージョンであることを確認してください。

権限エラー

権限エラーが発生する場合、以下のポイントを確認してください:

  • SageMaker実行ロールに必要なポリシーがアタッチされていることを確認してください。
  • AssumeRoleを行うIAMユーザーに対して、RoleをAssumeするための権限が付与されていることを確認してください。

運用とメンテナンス

定期実行と手動実行の併用

この再実行のフローは定期実行されるだけでなく、手動で実行させることもできます。その際には、パイプラインの中身を定義している設定ファイルの名前を指定することで動かせる設計になっています。定期実行と手動実行を併用することで、柔軟な運用が可能です。

on:
  workflow_dispatch:
    inputs:
      config_file_name:
        description: 'The name of the configuration file to use'
        required: true
        default: 'default-configuration-filename'
  schedule:
    - cron: '0 21 * * *'  # UTCの午後9時(日本時間の午前6時)

まとめ

この記事では、GitHub Actionsを使用してAmazon SageMakerパイプラインを定期実行する方法について解説しました。必要なAWS認証情報の設定や、AssumeRoleの設定方法、そして実際のワークフローの詳細について説明しました。この記事の手順を踏むことで、定期的な再学習を自動化し、モデルの精度を維持することができます。

参考資料

SMARTCAMP Engineer Blog

Discussion