Terraformで構築する機械学習ワークロード（Lambda編）

こんちには。

データアナリティクス事業本部 インテグレーション部 機械学習チームの中村です。

今回は「Terraformで構築する機械学習ワークロード」ということで、Lambdaを使って物体検出モデルの１つであるYOLOXの推論環境を構築していこうと思います。

構成イメージ

構成としては以下のようなものを作成していきます。

物体検出はLambda上でコンテナイメージを動かすことで実現します。

このコンテナイメージ内にMMDetectionというフレームワークをインストールしておき、その中で物体検出モデルの一つであるYOLOXを動かしていきます。

MMDetectionの説明については少しコードが古い部分もありますが、以下が参考となります。

動作環境

Docker、Terraformはインストール済みとします。

Terraformを実行する際のAWSリソースへの権限は、aws-vaultで環境構築をしておきます。

aws-vaultについては以下も参考にされてください。

ホストPCとしてのバージョン情報配下です。

• OS: Windows 10 バージョン22H2
• docker: Docker version 24.0.2-rd, build e63f5fa
• terraform: Terraform v1.4.6 (on windows_amd64)

またコンテナ内のバージョン情報配下となります。

• Python: 3.10.12 (main, Aug 15 2023, 15:43:05) [GCC 7.3.1 20180712 (Red Hat 7.3.1-15)]
• PyTorch: 2.0.1+cu117
• OpenCV: 4.8.0
• MMEngine: 0.8.4
• MMDetection: 3.1.0+

成果物

成果物はGitHub上にあげておきましたので、詳細は必要に応じてこちらを参照ください。

フォルダ構成は以下のようになっています。

├─asset
│      yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth
├─docker
│  └─lambda
│          docker-compose.yml
│          push_ecr.ps1
│
├─python
│      Dockerfile
│      lambda_handler.py
│      mmdetect_handler.py
│      s3_handler.py
│
└─terraform
    ├─environments
    │  └─dev
    │          main.tf
    │
    └─modules
        ├─ecr
        │      main.tf
        │
        ├─iam
        │      main.tf
        │      outputs.tf
        │
        ├─lambda
        │      main.tf
        │      outputs.tf
        │
        └─s3
                main.tf

assetフォルダはモデル置き場として使用しています。

dockerフォルダはコンテナイメージのビルド時やECRへのpush時に使用しますが、Dockerfile自体はpythonフォルダに置いています。

pythonフォルダにはDockerfileに加えて使用するPythonスクリプトを置いています。

terraformフォルダにはmoduleに分割したtfファイルが置かれており、作業フォルダはterraform/environments/devとなります。

コードの説明

ここからはコードを説明します。構築手順のみ知りたい方はスキップして手順に移動してください。

コンテナイメージのビルド

まず以下のような内容でpython/Dockerfileを準備しておきます。

FROM public.ecr.aws/lambda/python:3.10

WORKDIR /opt

RUN pip3 install torch

RUN pip install openmim && \
    mim install "mmengine>=0.7.1" "mmcv>=2.0.0rc4"

RUN yum install -y git
RUN git clone https://github.com/open-mmlab/mmdetection.git
WORKDIR /opt/mmdetection
RUN pip3 install --no-cache-dir -e .

RUN yum install -y tar
RUN yum install -y mesa-libGL.x86_64

RUN mkdir -p /opt/mmdetection/checkpoints

COPY *.py /opt/mmdetection

CMD ["sh"]

ベースイメージはLambdaのものを使用しており、ここにPyTorchやMMDetection関連のライブラリをインストールします。

こちらはMMDetectionのレポジトリのDockerfileを参考にしています。

バージョンを固定されたい場合は、pipによるインストールやgit cloneコマンドで固定化を実施しておく必要があります。

その他COPY *.py /opt/mmdetectionにより、pythonフォルダに置かれているスクリプトをコンテナ内に複製します。

デバッグ等をする際はVSCodeのDev Containerなどを使用すると捗るかと思います。

Pythonファイルについて

Pythonファイルは以下の３種類あります。

│      lambda_handler.py
│      mmdetect_handler.py
│      s3_handler.py

lambda_handler.py

lambda_handler.pyは以下のような内容となります。

import os
from mmdetect_handler import collect_env, find_checkpoint, inference
from s3_handler import download_directory, download_file, upload_file

def handler(event: dict, context):

    print(f"{event=}")
    print(f"{context=}")

    bucket_name = os.getenv("BUCKET_NAME")
    input_prefix = os.getenv("OBJECT_INPUT_PREFIX")
    output_prefix = os.getenv("OBJECT_OUTPUT_PREFIX")
    print(f"{bucket_name=}")
    print(f"{input_prefix=}")
    print(f"{output_prefix=}")

    # 処理対象のオブジェクトキーを取得
    target_object_key = event['Records'][0]['s3']['object']['key']

    # 処理対象のオブジェクトを取得
    download_file("/tmp/input.jpg", bucket_name, target_object_key)

    # モデル等をダウンロード
    download_directory(destination_path="/tmp/asset",
        bucket_name=bucket_name, prefix="asset/")

    # 環境ログを出力
    for name, val in collect_env().items():
        print(f"{name}: {val}")

    # モデルファイルを探索
    checkpoint_file = find_checkpoint(model_name="yolox_l_8x8_300e_coco",
        checkpoints_dir="/tmp/asset")

    # 推論処理
    inference(checkpoint_file=str(checkpoint_file),
        model_name="yolox_l_8x8_300e_coco",
        device="cpu",
        input_image_file="/tmp/input.jpg",
        output_image_file="/tmp/output.jpg")

    # 結果をS3にupload
    output_object_key = output_prefix + target_object_key[len(input_prefix):]
    print(f"{output_object_key=}")
    upload_file("/tmp/output.jpg", bucket_name, output_object_key)

    return {
        "statusCode": 200,
        "body": "OK"
    }

以下の流れで処理をします。

• S3イベントでLambdaを起動するのでeventからオブジェクトキーを取得
• S3から対象のオブジェクトを/tmpにダウンロード
• モデルファイルをS3にあらかじめ置いているので/tmpにダウンロード
• /tmpからモデルファイルを探索（ここはべた書きでも良いと思います）
• 推論処理
• 推論で出力された画像をS3にアップロード

S3関連の処理やMMDetectionに関する処理は、それぞれのハンドラを呼び出して処理をします。

s3_handler.py

S3関連の処理はs3_handler.pyに集めています。

import boto3
import pathlib

def download_directory(destination_path: str, bucket_name: str , prefix: str=""):

    client = boto3.client('s3')
    paginator = client.get_paginator('list_objects')

    for result in paginator.paginate(Bucket=bucket_name, Prefix=prefix):
        for file in result.get('Contents', []):
            target = (file.get('Key')[len(prefix):])
            size = file.get('Size')
            if size == 0:
                continue
            print(target)
            dest_file = pathlib.Path(destination_path).joinpath(target)
            dest_file.parent.mkdir(parents=True, exist_ok=True)
            client.download_file(bucket_name, file.get('Key'), str(dest_file))

def download_file(destination_path: str, bucket_name: str, object_key: str):
    client = boto3.client('s3')
    client.download_file(bucket_name, object_key, str(destination_path))

def upload_file(source_path: str, bucket_name: str, object_key: str):
    client = boto3.client('s3')
    client.upload_file(str(source_path), bucket_name, object_key)

こちらは特筆するような内容はありません。

mmdetect_handler.py

MMDetectionに関する処理はmmdetect_handler.pyに集めています。

from rich.pretty import pprint

def collect_env():
    import mmdet
    from mmengine.utils import get_git_hash
    from mmengine.utils.dl_utils import collect_env as collect_base_env

    """Collect the information of the running environments."""
    env_info = collect_base_env()
    env_info['MMDetection'] = f'{mmdet.__version__}+{get_git_hash()[:7]}'
    return env_info

def find_checkpoint(model_name: str,
                         checkpoints_dir: str="./checkpoints"):

    import pathlib
    checkpoints_path = pathlib.Path(checkpoints_dir)
    pattern = f"{model_name}_*.pth"
    checkpoints = sorted([ f for f in checkpoints_path.glob(pattern)])
    print(f"{checkpoints=}")
    if len(checkpoints) == 0:
        assert False, f"Cannot find checkpoint file: {checkpoints=}, {checkpoints_path=}, {pattern=}"

    checkpoint_file = checkpoints[0]

    print(f"{checkpoint_file=}")
    return checkpoint_file

def download_checkpoint(model_name: str,
                         checkpoints_dir: str="./checkpoints"):

    print(f"{model_name=}")
    print(f"{checkpoints_dir=}")

    import subprocess

    subprocess.run(["mkdir", "-p", checkpoints_dir])
    subprocess.run(["mim", "download", "mmdet", "--config", model_name, "--dest", checkpoints_dir])

    return find_checkpoint(model_name, checkpoints_dir)

def inference(model_name: str,
        checkpoint_file: str,
        input_image_file: str,
        output_image_file: str,
        device: str="cpu"):

    print(f"{model_name=}")
    print(f"{checkpoint_file=}")
    print(f"{input_image_file=}")
    print(f"{output_image_file=}")
    print(f"{device=}")

    # Set the device to be used for evaluation
    # device = 'cuda:0'
    device = 'cpu'

    # Initialize the DetInferencer
    from mmdet.apis import DetInferencer
    inferencer = DetInferencer(model_name, str(checkpoint_file), device)

    # Use the detector to do inference
    result = inferencer(input_image_file, no_save_vis=True, return_vis=True)
    pprint(result, max_length=4)

    from PIL import Image
    out_img = Image.fromarray(result["visualization"][0])
    out_img.save(output_image_file)

    return

メインとなるのは、inference関数です。

こちらの処理はMMDetectionのレポジトリのノートブックを参考にしています。

また今回はlambda_handler.pyからdownload_checkpoint関数を呼び出していませんが、以下のようにしてダウンロードしたモデルをレポジトリのassetフォルダに格納して使用しています。

download_checkpoint(model_name="yolox_l_8x8_300e_coco", checkpoints_dir="./checkpoints")

こちらが手間と感じられる場合は、以下のリンクから取得されてください。

tfファイル

tfファイルは以下のようなフォルダ構成となっています。

    ├─environments
    │  └─dev
    │          main.tf
    │
    └─modules
        ├─ecr
        │      main.tf
        │
        ├─iam
        │      main.tf
        │      outputs.tf
        │
        ├─lambda
        │      main.tf
        │      outputs.tf
        │
        └─s3
                main.tf

environments/dev/main.tf

起点となるのはこちらのtfファイルです。以下のようになっています。

variable "project_prefix" {}

data "aws_caller_identity" "current" {}
data "aws_region" "current" {}

locals {
  account_id           = data.aws_caller_identity.current.account_id
  region               = data.aws_region.current.name
  bucket_name          = "${var.project_prefix}-${local.account_id}"
  object_input_prefix  = "input/"
  object_output_prefix = "output/"
  function_name        = "${var.project_prefix}"
  iam_role_name        = "${var.project_prefix}-iam-role"
  iam_policy_name      = "${var.project_prefix}-iam-policy"
  repository_name      = "${var.project_prefix}"
  image_uri            = "${local.account_id}.dkr.ecr.${local.region}.amazonaws.com/${local.repository_name}"
}

module ecr {
  source="../../modules/ecr"
  repository_name=local.repository_name
}

module iam {
  source="../../modules/iam"
  iam_role_name=local.iam_role_name
  iam_policy_name=local.iam_policy_name
}

module lambda {
  source="../../modules/lambda"
  function_name=local.function_name
  image_uri=local.image_uri
  iam_role_arn=module.iam.iam_role_arn
  bucket_name=local.bucket_name
  object_input_prefix=local.object_input_prefix
  object_output_prefix=local.object_output_prefix
}

module s3 {
  source="../../modules/s3"
  bucket_name=local.bucket_name
  object_input_prefix=local.object_input_prefix
  object_output_prefix=local.object_output_prefix
  lamba_function_arn=module.lambda.lamba_function_arn
}

定数の定義とモジュールの呼び出しを行っています。

また実行時にproject_prefixを変数として与えられるようにしています。

modules/ecr/

ecrは以下のようなシンプルな構成です。

variable repository_name {}

resource "aws_ecr_repository" "main" {
  name = var.repository_name
  force_delete = true
}

modules/iam/

iamは以下のようにLambda用のIAMロールを作成します。（正直なところあまり権限を絞れてはいませんので参考までに）

variable iam_role_name {}
variable iam_policy_name {}

data "aws_iam_policy_document" "main" {
  statement {
    effect = "Allow"
    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }
    actions = ["sts:AssumeRole"]
  }
}

resource "aws_iam_role" "main" {
  name = var.iam_role_name
  assume_role_policy = data.aws_iam_policy_document.main.json
  inline_policy {
    name = var.iam_policy_name
    policy = jsonencode({
      "Version" : "2012-10-17"
      "Statement" : [
        {
          "Effect": "Allow",
          "Action": [
            "logs:CreateLogGroup",
            "logs:CreateLogStream",
            "logs:PutLogEvents"
          ],
          "Resource": "*"
        },
        {
          "Effect": "Allow",
          "Action": [
            "s3:*",
            "s3-object-lambda:*"
          ],
          "Resource": "*"
        }
      ]
    })
  }
}

outputs.tfも以下のように記載して、Lambdaに情報を渡します。

output "iam_role_arn" {
  value = aws_iam_role.main.arn
}

modules/lambda/

lambdaは以下のように構築しています。

variable function_name {}
variable image_uri {}
variable iam_role_arn {}
variable bucket_name {}
variable object_input_prefix {}
variable object_output_prefix {}

resource "aws_lambda_function" "function" {
  function_name    = var.function_name
  role             = var.iam_role_arn
  image_uri        = "${var.image_uri}:latest"
  package_type     = "Image"
  timeout          = 600
  memory_size      = 8192

  ephemeral_storage {
    size = 8192
  }

  image_config {
    command = ["lambda_handler.handler"]
  }

  environment {
    variables = {
      BUCKET_NAME = var.bucket_name
      OBJECT_INPUT_PREFIX = var.object_input_prefix
      OBJECT_OUTPUT_PREFIX = var.object_output_prefix
    }
  }

  depends_on = [
    aws_cloudwatch_log_group.log
  ]
}

resource "aws_cloudwatch_log_group" "log" {
  name = "/aws/lambda/${var.function_name}"
}

ECRに登録してあるイメージやIAMロールの情報、その他Lambda内で使用する環境変数などを入力としています。

スペックは最適ではないかもしれませんが、少し余裕めな構成にしています。

image_configのcommandで、呼び出す際のエントリポイントを指定しています。

LambdaのARNがS3イベントに必要なので、outputs.tfに以下を記載しています。

output lamba_function_arn {
    value = aws_lambda_function.function.arn
}

modules/s3/

s3のリソース構成は以下のようにしています。

variable bucket_name {}
variable object_input_prefix {}
variable object_output_prefix {}
variable lamba_function_arn {}

resource "aws_s3_bucket" "main" {
  bucket = var.bucket_name
  force_destroy = true
}

resource "aws_lambda_permission" "permission" {
  statement_id  = "AllowExecutionFromS3Bucket"
  action        = "lambda:InvokeFunction"
  function_name = var.lamba_function_arn
  principal     = "s3.amazonaws.com"
  source_arn    = aws_s3_bucket.main.arn
}

resource "aws_s3_bucket_notification" "bucket_notification" {
  bucket = aws_s3_bucket.main.id

  lambda_function {
    lambda_function_arn = var.lamba_function_arn
    events              = ["s3:ObjectCreated:*"]
    filter_prefix       = "input/"
    filter_suffix       = ".jpg"
  }

  depends_on = [aws_lambda_permission.permission]
}

resource "aws_s3_object" "object_input" {
  bucket = var.bucket_name
  key    = var.object_input_prefix
}

resource "aws_s3_object" "object_output" {
  bucket = var.bucket_name
  key    = var.object_output_prefix
}

resource "aws_s3_object" "object_model_file" {
  bucket = var.bucket_name
  key    = "asset/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth"
  source = "../../../asset/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth"
}

aws_lambda_permissionの記載場所は今回S3側にしています。aws_s3_bucket_notificationと関連すると考えたためですが、今後見直す可能性もあります。

aws_s3_bucket_notificationは、prefixに"input/"を指定してイベント実行をするようにしています。

空のフォルダを作成するためにinputとoutputのaws_s3_objectを作成しておきます。

またローカルにあるモデルファイルをsourceとしてaws_s3_objectをアップロードしています。

手順

TerraformでECRのみ構築

まずはECRのみを構築します。理由はECRレポジトリが先に無いとコンテナイメージをpushできないためです。

またコンテナイメージが無いとTerraformでLambdaを作成することができないため、先にECRレポジトリを作成します。

# 作業フォルダ: terraform/environments/dev/
terraform init # 初回のみでOK
aws-vault exec {プロファイル名} -- terraform apply -target="module.ecr" -var 'project_prefix={任意のプレフィックス}' # aws-vault経由で実行

工夫すればこの辺りの依存関係を自動化できる可能性もありますが、今回は手動でやっています。

イメージのビルド

まずdocker/lambda/.envというファイルを作成して、環境変数を入力しておきます。

PROJECT_PREFIX="{任意のプレフィックス}"

{任意のプレフィックス}は、terraform applyの際に指定したものと合致するようにしておいてください。

その後は以下でビルドができます。

# 作業フォルダ: docker/lambda/
docker compose build

ビルドには時間がかかりました。PyTorchをインストールしているためと考えられます。

またイメージの容量もかなり大きいためご注意ください。

# 確認
docker images

# REPOSITORY             TAG       IMAGE ID       CREATED         SIZE
# sample-yolox-lambda    latest    b0b91c72e91d   2 hours ago     9.81GB

ECRへコンテナイメージをpush

push_ecr.ps1というスクリプトを準備していますのでそちらを実行してください。

.\push_ecr.ps1 -ProfileName {プロファイル名} -ProjectPrefix {任意のプレフィックス}

{任意のプレフィックス}は、terraform applyの際に指定したものと合致するようにしておいてください。

push_ecr.ps1の内容は以下です。

param(
    [Parameter(Mandatory)]
    [string]$ProfileName,
    [Parameter(Mandatory)]
    [string]$ProjectPrefix
)

$REGION = $(aws configure get region --profile $ProfileName)
$ACCOUNT_ID = $(aws sts get-caller-identity --query 'Account' --output text --profile $ProfileName)

$REPOSITORY_NAME = $ProjectPrefix
$ECR_BASE_URL = "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"
$ECR_IMAGE_URI = "${ECR_BASE_URL}/${REPOSITORY_NAME}"

Write-Output "ECR_BASE_URL: ${ECR_BASE_URL}"
Write-Output "ECR_IMAGE_URI: ${ECR_IMAGE_URI}"

# ECRへのログイン
aws ecr get-login-password --region ${REGION} --profile $ProfileName | docker login --username AWS --password-stdin ${ECR_BASE_URL}

# tagの付け替え
docker tag "${REPOSITORY_NAME}:latest" "${ECR_IMAGE_URI}:latest"

# ECRへのpush
docker push "${ECR_IMAGE_URI}:latest"

# Lambdaを更新
aws lambda update-function-code --function-name $ProjectPrefix --image-uri "${ECR_IMAGE_URI}:latest" --profile $ProfileName | Out-Null

なお、この時点ではpush_ecr.ps1スクリプトの最後のaws lambda update-function-codeのみ失敗します。リソースがまだ作られてないためです。

こちらは、更新を兼ねているためこのようなスクリプトとなっています。

また同様のことが可能なシェルもpush_ecr.shとして置いてありますのでご活用ください。

push_ecr.sh {プロファイル名} {任意のプレフィックス}

push_ecr.shの内容は以下となっています。

ProfileName=$1
ProjectPrefix=$2

REGION=$(aws configure get region --profile $ProfileName)
ACCOUNT_ID=$(aws sts get-caller-identity --query 'Account' --output text --profile $ProfileName)

REPOSITORY_NAME=$ProjectPrefix
ECR_BASE_URL="${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"
ECR_IMAGE_URI="${ECR_BASE_URL}/${REPOSITORY_NAME}"

echo "ECR_BASE_URL: ${ECR_BASE_URL}"
echo "ECR_IMAGE_URI: ${ECR_IMAGE_URI}"

# ECRへのログイン
aws ecr get-login-password --region ${REGION} --profile $ProfileName | docker login --username AWS --password-stdin ${ECR_BASE_URL}

# tagの付け替え
docker tag "${REPOSITORY_NAME}:latest" "${ECR_IMAGE_URI}:latest"

# ECRへのpush
docker push "${ECR_IMAGE_URI}:latest"

# Lambdaを更新
aws lambda update-function-code --function-name $ProjectPrefix --image-uri "${ECR_IMAGE_URI}:latest" --profile $ProfileName > /dev/null

Terraformでリソースを全て構築

最後にECR以外のリソースを作成します。

# 作業フォルダ: terraform/environments/dev/
aws-vault exec {プロファイル名} -- terraform apply -var 'project_prefix={任意のプレフィックス}' # aws-vault経由で実行

動作確認

AWS CLIでjpgファイルをアップロードしてみます。

asset/demo.jpgにサンプル画像を配置してありますので良ければお試しください。

aws s3 cp asset/demo.jpg s3://{バケット名}/input/demo.jpg --profile {プロファイル名}

処理が終わると、output/に結果は配置されます。

aws s3 ls s3://{バケット名}/output/ --profile {プロファイル名}
# 2023-09-08 20:28:20          0
# 2023-09-12 12:06:11      78343 demo.jpg

以下のような画像が出力されます。

まとめ

いかがでしたでしょうか。

今後はこの環境をベースに様々なワークロードを流すための環境を構築していきたいと思いますのでご期待ください。

本記事が皆様のお役に立てば幸いです。