AWS CDK で作る AWS Inferentia と Amazon ECS を利用した推論環境（Part 1）

AWS のいくつかのお客様では、OCR （Optical Character Recognition）サービス、AI チャットボット、などの深層学習モデルを利用したアプリケーションを構築されています。これらのお客様の多くが自然言語処理モデルの推論時間の長さに課題を持ち、スループットの向上とレイテンシーの削減に多くの労力を費やしてきました。

AWS では、GPU を搭載した Amazon EC2インスタンス、AWS が設計開発した高機能の機械学習推論チップである AWS Inferentia を搭載した Amazon EC2 Inf1 インスタンスを提供しており、推論環境を構築するためにサーバーレスのための AWS Lambda、あらゆるユースケース向けの機械学習モデルを構築 / トレーニング / デプロイするフルマネージドサービスである Amazon SageMaker、などの選択肢があります。その中でも Amazon EC2 Inf1 インスタンスは、深層学習ワークロードを加速させる高パフォーマンスの機械学習推論環境を低いコストで構築できるため、推論性能及びコストの改善に対する一つのソリューションとなります。

「AWS Inferentia 上の PyTorch 自然言語処理アプリケーションにおいて、12 倍のスループットと最小のレイテンシーを実現」では、HuggingFace Transformers の事前学習された BERT base モデルを修正せずに使用し、PyTorch フレームワークレベルでコードを1行追加するのみで、NLP ベースのソリューションを Amazon EC2 Inf1 インスタンスにデプロイします。このソリューションでは、同じモデルを GPU にデプロイする場合と比較して、AWS Inferentia 上で 70% 低いコストで 12 倍のスループットを実現します。「InfoJobs (Adevinta) が AWS Inferentia と Amazon SageMaker で NLP モデル予測のパフォーマンスをどのように向上させたか」では、AWS Inferentia を使って Amazon SageMaker エンドポイントでホスティングすると、モデル最適化や GPU 搭載インスタンスを使用した垂直スケーリングといった選択肢と比較して、予測レイテンシーを最大 92% 削減し、 75% のコスト削減を実現し、CPU コストで最高の GPU パワーを手に入れるような体験をしています。

現在、Amazon EKS、Amazon ECS、Amazon SageMaker で Amazon EC2 Inf1 インスタンスを利用できます。Amazon SageMaker では、CreateEndpoint API を実行するだけでスケールする推論環境を構築することができ、複雑なデプロイ戦略やスケーリング、A/B テスト環境構築などに頭を悩ます必要がありません。もし柔軟な推論用 API の構築や、Reserved Instances、On-Demand Capacity Reservation によるリソース確保などが必要となる場合には、Amazon ECS によって推論環境を自由に構築することが可能です。Amazon ECS は、フルマネージドコンテナオーケストレーションサービスであり、コンテナ化されたアプリケーションを簡単にデプロイ、管理、およびスケーリングできます。Amazon ECS は推論環境をお客様自身で作り込むことにより、より高い自由度でアプリケーションを構築することができ、Amazon SageMaker では所望の柔軟性を得られない場合に有力な選択肢となり得ます。

本投稿では、推論性能及びコストに課題を持っており AWS Inferentia を利用したいが、柔軟性の観点から私の大好きな Amazon SageMaker を利用できないケースを想定します。Amazon ECS と Amazon EC2 Inf1 インスタンスを利用した推論環境を AWS Cloud Development Kit（AWS CDK）で構築することで、推論速度に課題を持たれる読者の方に実装サンプルを提供し、本ケースにおいて少しでも推論環境の初期構築にかかるコストを下げていただくことを本投稿の目的とします。コンソールから手動で作成する際のオペレーションミスを防ぎ再現性高く実装サンプルを動かしていただくため、AWS CDK による IaC（Infrastructure as Code）を実践します。

図1 に Amazon Inf1 EC2 インスタンスを用いた推論処理の概要図を示します。ブログの Part 1 では、検証環境として Amazon EC2 Inf1 インスタンスを AWS CDK で構築し、構築したインスタンス上でモデルのコンパイル、推論の動作を確認を行います。その後、 Amazon Elastic Container Registry（Amazon ECR）に推論用のコンテナイメージをプッシュします。Part 2 では、AWS Inferentia を利用するために必要な AWS Neuron SDK の機能や制約などの説明を行い、Amazon ECS を AWS CDK を用いて構築し、デプロイされたアプリケーションの動作を確認します。

図1: 推論処理の概要図

事前準備

構築にあたって前提とすることを以下に示します。設定ファイルと認証情報ファイルの設定についてはドキュメントをご参照ください。本投稿のサンプルコードをご利用される際には、セキュリティ観点での確認を必ず行なってください。サンプルコードは、2022年6月21日時点の情報に基づいて作成しており、情報の変更により動かなくなる可能性がありますがご了承ください。

AWS Command Line Interface aws-cli/1.23.2 以降が利用可能な状態であること
AWS CDK 2.23.0 以降が利用可能な状態であること
上記コマンドを実行するために、該当の IAM ロールに AWS 管理ポリシー “AdministratorAccess” がアタッチされること
- 必要に応じて権限を絞ってください

AWS CDK で検証用 Amazon EC2 Inf1 インスタンスを構築する

AWS Inferentia の検証環境を AWS CDK を用いて構築します。カスタムチップである AWS Inferentia を利用するためには、AWS Neuron SDK を用いて学習済みモデルをコンパイルし、推論サーバでコンパイルしたモデルのロードを行う必要があります。AWS Neuron SDK は、他にもプロファイリングツールを提供しています。検証環境では、モデルのコンパイルと推論のどちらも Amazon EC2 Inf1 インスタンス上で行いますが、モデルのコンパイルに関しては、AWS Neuron SDK をインストールさえしていれば、実行環境は Amazon EC2 Inf1 インスタンスに制約されません。AWS Neuron SDK のインストールについては、対応するフレームワークごとにドキュメントに手順が記載されています。この手順に従って導入することはもちろん可能ですが、AWS では AWS Deep Learning AMI（DLAMI）という、深層学習用のAMI があり、PyTorch、TensorFlow といった機械学習フレームワーク、GPU Driver、AWS Neuron SDK、等がプリインストールされています。

構築概要

インスタンスサイズ inf1.xlarge
AWS Deep Learning AMI（DLAMI） Ubuntu 18.04 を利用
デフォルト VPC を利用
ssh ログイン用の keypair を作成
デバイスボリュームを 300 GB に設定
ユーザーデータを利用して python 3.7 インストール等の初期設定を実施
- AWS Neuron SDK はソフトウェアサポートバージョンの依存関係として python 3.7 を要求します
Amazon ECR リポジトリを作成

注意点

Security Group の Inbound Port 22 がオープンされているためセキュリティ観点での必要性に応じてインバウンドルールのソース等の設定をご変更ください。
ap-northeast-1 の利用を前提として AMI ID を設定しているため、別リージョンを利用する場合には適切なAMI IDをご指定ください。DLAMI ID の検索方法はドキュメントをご参照ください。
記事の見やすさの都合上コードが途中で見切れているため、スクロールして最後まで貼り付けてください。

手順

ローカル PC のエディタ、AWS Cloud9 などで事前準備が完了している状態であることをご確認ください。

mkdir ec2 && cd ec2
cdk init sample-app --language typescript
npm install cdk-ec2-key-pair

cdk init を実行すると環境一式が生成されます。 lib/ec2-stack.ts というファイルが生成されているので、エディタで以下のように修正を行います。

import { Stack, StackProps } from "aws-cdk-lib";
import * as ec2 from "aws-cdk-lib/aws-ec2";
import * as ecr from 'aws-cdk-lib/aws-ecr';
import * as cdk from "aws-cdk-lib";
import * as iam from "aws-cdk-lib/aws-iam"
import * as path from "path";
import { Asset } from "aws-cdk-lib/aws-s3-assets";
import { Construct } from "constructs";
import { KeyPair } from "cdk-ec2-key-pair";


export class Ec2Stack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    const region: string = process.env.CDK_DEFAULT_REGION || "ap-northeast-1";
    const imageId: string = "ami-0eeb79c0bc2f4ed75"; // for ap-northeast-1

    const key = new KeyPair(this, "KeyPair", {
      name: "cdk-keypair",
      description: "Key Pair created with CDK Deployment",
      storePublicKey: true,
    });

    const vpc = ec2.Vpc.fromLookup(this, "VPC", {
      isDefault: true
    });

    const securityGroup = new ec2.SecurityGroup(this, "SecurityGroup", {
      vpc,
      description: "Allow SSH (TCP port 22) in",
      allowAllOutbound: true,
    });
    securityGroup.addIngressRule(
      ec2.Peer.anyIpv4(),
      ec2.Port.tcp(22),
      "Allow SSH Access"
    );

    const role = new iam.Role(this, "ec2Role", {
      assumedBy: new iam.ServicePrincipal("ec2.amazonaws.com")
    });

    role.addManagedPolicy(
      iam.ManagedPolicy.fromAwsManagedPolicyName("AmazonSSMManagedInstanceCore")
    );

    // Use Deep Learning AMI(Ubuntu 18.04), Arch: x86_64
    const ami = new ec2.GenericLinuxImage({
      [`${region}`]: imageId,
    });


    const ec2Inf1Instance = new ec2.Instance(this, "Inf1Instance", {
      vpc,
      instanceType: ec2.InstanceType.of(ec2.InstanceClass.INFERENCE1, ec2.InstanceSize.XLARGE),
      machineImage: ami,
      securityGroup: securityGroup,
      keyName: key.keyPairName,
      role: role,
      blockDevices: [
        {
          deviceName: "/dev/sda1",
          volume: ec2.BlockDeviceVolume.ebs(300),
        },
      ],
    });

    const asset = new Asset(this, "Asset", { path: path.join(__dirname, "../config.sh") });
    const localPath = ec2Inf1Instance.userData.addS3DownloadCommand({
      bucket: asset.bucket,
      bucketKey: asset.s3ObjectKey,
    });

    ec2Inf1Instance.userData.addExecuteFileCommand({
      filePath: localPath,
      arguments: "--verbose -y"
    });
    asset.grantRead(ec2Inf1Instance.role);

    new ecr.Repository(this, 'ModelRepo', {
      repositoryName: 'model'
    });

    new ecr.Repository(this, 'WebRepo', {
      repositoryName: 'web'
    });

    new cdk.CfnOutput(this, "EC2 IP Address", { value: ec2Inf1Instance.instancePublicIp });
    new cdk.CfnOutput(this, "Key Name", { value: key.keyPairName })
    new cdk.CfnOutput(this, "Download Key Command", { value: "aws secretsmanager get-secret-value --secret-id ec2-ssh-key/cdk-keypair/private --query SecretString --output text > ~/.ssh/cdk-keypair.pem && chmod 400 ~/.ssh/cdk-keypair.pem" })
    new cdk.CfnOutput(this, "SSH Command", { value: "ssh -i ~/.ssh/cdk-keypair.pem -o IdentitiesOnly=yes ubuntu@" + ec2Inf1Instance.instancePublicIp })
  }    
}

bin/ec2.ts というファイルを以下のように変更します。cdk 実行時の AWS アカウントとリージョンを指定します。

#!/usr/bin/env node
import * as cdk from 'aws-cdk-lib';
import { Ec2Stack } from '../lib/ec2-stack';


const app = new cdk.App();
new Ec2Stack(app, 'Ec2Stack', {
    env: {
        account: process.env.CDK_DEFAULT_ACCOUNT,
        region: process.env.CDK_DEFAULT_REGION
    },
});

新たに config.sh というファイルを以下のように作成します。このシェルスクリプトがインスタンスの初期起動時に実行されます。シェルスクリプトでは、python 3.7 のインストール、スワップ作成、pip での AWS Neuron SDK などのインストールを行います。そして最後に、~/.aws/credentials で認証情報、/tmp/env.hcl に以降の章で利用する環境変数、を設定しておきます。以下のシェルスクリプトの XXX となっている箇所に適切な値を設定してください。（値を設定した上でgitでのconfig.shのファイルの管理をしないようにご注意ください）

#!/bin/sh -x

# Please see the neuron docs
#   https://awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-intro/pytorch-setup/pytorch-install.html
sudo apt-get update -y

export PATH=/opt/aws/neuron/bin:$PATH

sudo apt-get install -y python3.7-venv g++

sudo dd if=/dev/zero of=/var/swpfile8G bs=1M count=8192
sudo mkswap /var/swpfile8G
sudo chmod 600 /var/swpfile8G
sudo swapon /var/swpfile8G
free
sudo echo "/var/swpfile8G swap swap defaults 0 0" >> /etc/fstab

cd /home/ubuntu
sudo -u ubuntu python3.7 -m venv pytorch_venv

. ./pytorch_venv/bin/activate

pip3 install -U pip
pip3 install ipykernel
pip3 install jupyter notebook
pip3 install environment_kernels
pip3 config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
pip3 install torch-neuron==1.10.2.* neuron-cc[tensorflow] "protobuf<4" torchvision
pip3 install transformers==4.19.2 fugashi==1.1.2 ipadic==1.0.0

mkdir -p /home/ubuntu/.aws

cat <<'EOL' > /home/ubuntu/.aws/credentials
[default]
aws_access_key_id = XXX
aws_secret_access_key = XXX
region = XXX
EOL

cat <<'EOL' > /tmp/env.hcl
CDK_DEFAULT_REGION="ap-northeast-1"
CDK_DEFAULT_ACCOUNT="XXX"
REGION="ap-northeast-1"
ACCOUNT_ID="XXX"
EOL

ファイルの修正が全て完了したので、npm run build コマンドでビルドを行います。prettier のエラーが出る場合は、prettier の設定もしくはエラー内容を修正してください。build が成功していれば lib/ec2-stack.d.ts などのファイルが生成されています。そして、以下の環境変数を設定してから cdk deploy コマンドを実行します。

export CDK_DEFAULT_REGION="ap-northeast-1"
export CDK_DEFAULT_ACCOUNT="（ご自身のAWS アカウントID をご指定ください）"

cdk deploy が完了すると以下のような出力が得られます。出力結果の Ec2Stack.DownloadKeyCommand と Ec2Stack.SSHCommand を順次実行し、作成したインスタンスに ssh ログインします。以降は、開発用の Amazon EC2 Inf1 インスタンス上での作業となります。

Outputs:
Ec2Stack.DownloadKeyCommand = aws secretsmanager get-secret-value --secret-id ec2-ssh-key/cdk-keypair/private --query SecretString --output text &gt; ~/.ssh/cdk-keypair.pem &amp;&amp; chmod 400 ~/.ssh/cdk-keypair.pem
Ec2Stack.EC2IPAddress = XXX Ec2Stack.KeyName = cdk-keypair
Ec2Stack.SSHCommand = ssh -i ~/.ssh/cdk-keypair.pem -o IdentitiesOnly=yes ubuntu@XXX
Stack ARN:
XXX

インスタンスにログイン後、/var/log/cloud-init-output.log にユーザーデータの実行ログが残っており、スクリプトの実行状況を監視できます。完了までおよそ10分程度必要です。

アプリケーションを構築する

推論性能とコストの改善をご検討のお客様は既にアプリケーションを構築されている場合があります。例えば、FastAPI フレームワーク、poetry によるパッケージ管理、Gunicorn を用いた複数ワーカーモデル、などを利用されているかもしれません。Part 2 で詳細を説明しますが、AWS Inferentia チップのリソースの共有に関しては制約があります。そして、conda や poetry はサポート対象外です。これらの制約をワークアラウンドとして回避する策は取れますが、既存アプリケーションへの変更が入ってしまいます。そのため、図2 に示すように、既存のアプリケーションサーバ（以降、Web Server）と、新規のAWS Inferentia を用いた推論専用のサーバ（以降、Model Server）に分離する方法を一例として実装します。処理の流れとしては、リクエスタから Web Server に対して invocations API リクエストを実行します。リクエストを受け取った Web Server は、受け取ったリクエストの中身をそのままデータとして、 Model Server に対して inferences API リクエストを実行します。リクエストを受け取った Model Server はリクエストデータの text （例：お爺さんは森に狩りへ出かける）と text に対してマスクしたい箇所を指定する mask_index から お爺さんは森に[MASK]へ出かけるというマスクされた text を作成し、その text に対して推論処理を実行します。推論結果として以下のような json が得られます。出力結果の上位 5 件の token_str をレスポンスとして Web Server に返し、Web Server は Model Server から受け取った結果をそのまま invocations のレスポンスとして返却します。

[
  {
    "score": 0.2534094452857971,
    "token": 8496,
    "token_str": "遊 び",
    "sequence": "お爺さん は 森 に 遊び へ 出かける"
  },
  {
    "score": 0.08862830698490143,
    "token": 6219,
    "token_str": "冒 険",
    "sequence": "お爺さん は 森 に 冒険 へ 出かける"
  },
  {
    "score": 0.0827520489692688,
    "token": 11899,
    "token_str": "狩 り",
    "sequence": "お爺さん は 森 に 狩り へ 出かける"
  },
  {
    "score": 0.0713028535246849,
    "token": 8639,
    "token_str": "釣 り",
    "sequence": "お爺さん は 森 に 釣り へ 出かける"
  },
  {
    "score": 0.03635821118950844,
    "token": 1233,
    "token_str": "調 査",
    "sequence": "お爺さん は 森 に 調査 へ 出かける"
  }
]

DLAMI 同様、AWS では深層学習用の Docker イメージとして、 AWS Deep Learning Containers を提供しています。Model Server では AWS Neuron SDK を利用するため、このDocker イメージを利用します。

図2: 検証環境の全体概要図

Model Server 概要

AWS Deep Learning Containers Neuron Inference Containers をベースイメージとして利用
Hugging Face BERT model cl-tohoku/bert-base-japanese-whole-word-masking を利用
- Licenses: Creative Commons Attribution-ShareAlike 3.0.
FastAPI ウェブフレームワークを利用

Web Server 概要

Gunicorn を用いた複数ワーカーモデル
FastAPI ウェブフレームワークを利用
API リクエストを受けて Model Server に推論要求を出し、その応答をそのまま API レスポンスとして返す

Model Server の構築

Model Server を構築していきましょう。以下にディレクトリ構成を示します。

.
├── app
│   ├── main.py
│   ├── requirements.txt
│   └── run.sh
├── build
│   └── Dockerfile
├── docker-bake.hcl
└── trace
    └── tracer.py

以下のコマンドでディレクトリと空ファイルを作成します。エディタで以降のコードをコピー＆ペーストしてください。

mkdir -p model-server/{app,build,trace} &&\
  touch model-server/{docker-bake.hcl,app/{main.py,requirements.txt,run.sh},build/Dockerfile,trace/tracer.py} &&\
  chmod +x model-server/app/run.sh

model-server/app/

main.py

import os
import sys
from logging import DEBUG, StreamHandler, getLogger
from typing import List

import torch
import torch_neuron
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, constr
from transformers import BertJapaneseTokenizer

app = FastAPI()

logger = getLogger(__name__)
handler = StreamHandler(sys.stdout)
handler.setLevel(DEBUG)
logger.addHandler(handler)
logger.setLevel(DEBUG)

path_prefix = "/app/server/models/"
model_path = os.path.join(path_prefix, "transformers_neuron.pt")
tokenizer_path = os.path.join(path_prefix, "tokenizer")

LENGTH = 512


class ModelInference:
    def __init__(self):
        self.tokenizer = BertJapaneseTokenizer.from_pretrained(tokenizer_path)
        self.model = torch.jit.load(model_path)

    def _conv_to_tokens_from_(self, index):
        return self.tokenizer.convert_ids_to_tokens([index.item()])[0]

    def infer(self, message):

        text = message.text
        mask_index = message.mask_index

        logger.info(f"Input text : {text}")
        tokenized_text = self.tokenizer.tokenize(text)
        logger.info("Tokenized text : " + ",".join(tokenized_text))
        tokenized_text[mask_index] = "[MASK]"
        logger.info("Masked text : " + ",".join(tokenized_text))
        encoding = self.tokenizer.encode_plus(
            text,
            return_tensors="pt",
            max_length=LENGTH,
            padding="max_length",
            truncation=True,
        )
        model_input = (encoding["input_ids"], encoding["attention_mask"])
        with torch.no_grad():
            outputs = self.model(*model_input)
            preds = outputs[0][0, mask_index].topk(5)

        tokens = [self._conv_to_tokens_from_(idx) for idx in preds.indices]

        return tokens


class UserRequestIn(BaseModel):
    text: constr(min_length=1)
    mask_index: int


class MaskedTextOut(BaseModel):
    labels: List[str]


model_class = ModelInference()


@app.head("/", status_code=200)
@app.get("/", status_code=200)
def read_root():
    return {"Status": "Healthy"}


@app.post("/inferences", response_model=MaskedTextOut, status_code=200)
def inferences(message: UserRequestIn):
    try:
        infered = model_class.infer(message)
    except Exception as e:
        msg = f"Internal Server Error, {e}"
        raise HTTPException(status_code=500, detail=msg)

    return {"labels": infered}


if __name__ == "__main__":
    request = {"text": "お爺さんは森に狩りへ出かける", "mask_index": 7}
    message = UserRequestIn(**request)
    print(model_class.predict(message))

requirements.txt

fastapi
uvicorn[standard]

run.sh

#! /usr/bin/env sh
set -e

HOST=${HOST:-0.0.0.0}
PORT=${PORT:-80}
LOG_LEVEL=${LOG_LEVEL:-info}
APP_MODULE=${APP_MODULE:-main:app}

exec uvicorn --reload --host $HOST --port $PORT --log-level $LOG_LEVEL "$APP_MODULE"

model-server/build

Dockerfile

ARG REGION
FROM 763104351884.dkr.ecr.${REGION:-ap-northeast-1}.amazonaws.com/pytorch-inference-neuron:1.10.2-neuron-py37-sdk1.19.0-ubuntu18.04 as base-stage

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV AWS_NEURON_VISIBLE_DEVICES=ALL
ENV PATH=/opt/aws/neuron/bin:$PATH

RUN pip install transformers==4.19.2 fugashi==1.1.2 ipadic==1.0.0

# For Trace ---------------------------------------
FROM base as trace-stage

RUN mkdir -p /app/trace/models

ADD ./trace/tracer.py /app/trace

WORKDIR /app/trace/

# For Model Server --------------------------------
FROM base as model-stage

RUN mkdir -p /app/server

COPY ./app/requirements.txt /app/server/requirements.txt
COPY ./trace/transformers_neuron.pt /app/server/models/transformers_neuron.pt
COPY ./trace/tokenizer /app/server/models/tokenizer
COPY ./app/run.sh /app/server/run.sh
COPY ./app/main.py /app/server/main.py

RUN pip3 install -r /app/server/requirements.txt

WORKDIR /app/server

EXPOSE 8080 80

CMD ["./run.sh"]

model-server/trace

tracer.py

import torch
import torch_neuron
from transformers import BertForMaskedLM, BertJapaneseTokenizer

model_name = "cl-tohoku/bert-base-japanese-whole-word-masking"
ptname = "transformers_neuron.pt"
processor = "inf1"
LENGTH = 512

model = BertForMaskedLM.from_pretrained(model_name, return_dict=False)

model.eval()
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)
tokenizer.save_pretrained("./tokenizer")

text = "お爺さんは森に狩りへ出かける"
inputs = tokenizer.encode_plus(
    text, return_tensors="pt", max_length=LENGTH, padding="max_length", truncation=True
)

example_inputs = (
    torch.cat([inputs["input_ids"]], 0),
    torch.cat([inputs["attention_mask"]], 0),
)

if "inf1" in processor:
    model_traced = torch.neuron.trace(model, example_inputs)
else:
    model_traced = torch.jit.trace(model, example_inputs)

model_traced.save(ptname)

print("Done.")

model-server/docker-bake.hcl

target "base" {
  target = "base-stage"
  dockerfile = "build/Dockerfile"
  args = {
      REGION = "${REGION}"
  }
  tags = ["base"]
}

target "trace" {
  target = "trace-stage"
  inherits = ["base"]
  tags = ["trace"]
}

target "model" {
  target = "model-stage"
  inherits = ["base"]
  tags = ["model"]
}

準備は整ったので、AWS Deep Learning Containers の pull、AWS Neuron SDK を用いたモデルのコンパイル、Model Server の起動、の順に実施していきます。

AWS Deep Learning Containers Base イメージ

model-server/build/Dockerfile で AWS Deep Learning Containers を Base イメージとして指定しています。このイメージを取得するためには Amazon ECR にログインする必要があります。ログイン後に、Dockerfileの base-stage ステージのビルドを行います。

作業ディレクトリ：model-server

FROM 763104351884.dkr.ecr.${REGION:-ap-northeast-1}.amazonaws.com/pytorch-inference-neuron:1.10.2-neuron-py37-sdk1.19.0-ubuntu18.04 as base-stage

export ENVHCL=/tmp/env.hcl &&\
  source $ENVHCL &&\
  aws ecr get-login-password --region $REGION | docker login --username AWS --password-stdin 763104351884.dkr.ecr.$REGION.amazonaws.com &&\
  docker buildx bake base -f docker-bake.hcl -f $ENVHCL

AWS Neuron SDK によるモデルのコンパイル

推論環境を構築することが本投稿の目的であるため、モデル開発やコンパイルオプション等の詳細には触れません。検証用インスタンスは、 jupyter ノートブックを起動できるようにしています。本記事のサンプル実装を元にノートブック上で作業を行いたい場合は、ポートフォワードを利用した一例として、以下の手順でローカル PC のブラウザ上でノートブックを起動できます。モデル開発を行う場合は、CPU / GPU 、フレームワークを柔軟に選択できる Amazon SageMaker Studio がオススメです！

source ~/pytorch_venv/bin/activate

jupyter-notebook --ip='0.0.0.0'

# cdk deploy を実行したローカルPC 等の環境で実行
ssh -i ~/.ssh/cdk-keypair.pem -N -f -L 8888:localhost:8888 ubuntu@XXX

AWS Neuron SDK による AWS Inferentia 向けのモデルのコンパイルのために、tracer.py というファイルを用意しています。tracer.py では torch_neuron.trace API を利用することで AWS Inferentia で実行するための PyTorch モデルを生成します。このモデルは TorchScript としてシリアル化できます。以下のシェルスクリプトを実行することで docker run によりモデルコンパイルを行い、モデルを trace/transformers_neuron.pt に格納します。実行状況は、docker logs -f trace を実行することで確認でき、Done が標準出力されればモデルコンパイル処理は完了です。図2 に示すように Docker ホストと trace コンテナで trace ディレクトリを bind mount しており、コンテナ内で生成されたモデルファイルはホスト側に配置されます。Model Server 起動時に、生成されたこれらのファイルがコンテナ内に取り込まれます。モデルコンパイル処理時に free コマンドによりメモリ使用量を確認すると、8 GB 程度のメモリを利用している場合があり、Out of Memory によるコンパイル処理の中断を防ぐために、スワップを利用しています。モデルによってはコンパイルを実行する上で、より大きめのvCPU、メモリを要求するケースがあります。

作業ディレクトリ：model-server

docker buildx bake trace -f docker-bake.hcl -f /tmp/env.hcl &&\
  cd trace &&\
  docker run -d --rm -m 5G --name trace -v ${PWD}:/app/trace/ --memory-swap -1 --oom-kill-disable trace python tracer.py &&\
  cd -

Model Server の起動

コンパイルされたモデルファイルは Amazon S3 に配置し、サーバ起動時に取得するのが一般的な構成ですが、本サンプルでは簡単のために Docker イメージ内にモデルファイルを配置します。

作業ディレクトリ：model-server

docker network create -d bridge network &&\
　　docker buildx bake model -f docker-bake.hcl -f /tmp/env.hcl &&\
  docker run -d --name model-server -p 80:80 --device=/dev/neuron0 --network=network model &&\
  cd ~

推論結果の確認

以下の post リクエストを実行すると、レスポンスとして{"labels":["狩り","狩","遊び","猟","狩猟"]} が返ってきます。これはお爺さんは森に[MASK]へ出かける の MASK 部分を推論し、推論結果の score の上位5件の token_str を出力しています。

curl -X 'POST' \
  'http://localhost/inferences' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "text": "お爺さんは森に狩りへ出かける",
  "mask_index": 7
}'

Web Server の構築

Web Server を構築していきましょう。以下にディレクトリ構成を示します。

.
├── Dockerfile
├── app
│   ├── main.py
│   ├── requirements.txt
│   └── run.sh
└── docker-bake.hcl

以下のコマンドでディレクトリと空ファイルを作成します。エディタで以降のコードをコピー＆ペーストしてください。

mkdir -p web-server/app &&\
  touch web-server/{Dockerfile,docker-bake.hcl,app/{main.py,run.sh,requirements.txt}} &&\
  chmod +x web-server/app/run.sh

web-server/app

main.py

import json
import os
import sys
from logging import DEBUG, StreamHandler, getLogger
from typing import List

import requests
from fastapi import FastAPI, HTTPException, responses
from pydantic import BaseModel, constr

app = FastAPI()

ENDPOINT_URL = os.environ["ENDPOINT_URL"]

logger = getLogger(__name__)
handler = StreamHandler(sys.stdout)
handler.setLevel(DEBUG)
logger.addHandler(handler)
logger.setLevel(DEBUG)


class Invocation:
    def __init__(self):
        pass

    def invoke(self, message):
        try:
            headers = {
                "accept": "application/json",
                "Content-Type": "application/json",
            }
            payload = {
                "mask_index": message.mask_index,
                "text": message.text,
            }
            data = json.dumps(payload)
            response = requests.post(ENDPOINT_URL, headers=headers, data=data).json()
            logger.info(response)

        except Exception as e:
            raise HTTPException(status_code=500, detail="Model Server invoke exception")
        return response


class UserRequestIn(BaseModel):
    text: constr(min_length=1)
    mask_index: int


class MaskedTextOut(BaseModel):
    labels: List[str]


web_class = Invocation()


@app.get("/")
async def read_root():
    return {"Status": "Healthy"}


@app.post("/invocations", response_model=MaskedTextOut)
async def invocations(message: UserRequestIn):
    return web_class.invoke(message)


if __name__ == "__main__":
    request = {"text": "お爺さんは森に狩りへ出かける", "mask_index": 7}
    message = UserRequestIn(**request)
    print(web_class.invoke(message))

requirements.txt

pydantic==1.9.1
requests==2.27.1

run.sh

#! /usr/bin/env sh
set -e

HOST=${HOST:-0.0.0.0}
PORT=${PORT:-80}
LOG_LEVEL=${LOG_LEVEL:-info}
APP_MODULE=${APP_MODULE:-main:app}
NUM_WORKER=4

exec gunicorn $APP_MODULE --workers $NUM_WORKER --worker-class uvicorn.workers.UvicornWorker --bind ${HOST}:${PORT} --log-level $LOG_LEVEL

web-server/Dockerfile

FROM tiangolo/uvicorn-gunicorn-fastapi:python3.7 as web-stage

WORKDIR /tmp

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE

COPY ./app/run.sh /app/server/run.sh
COPY ./app/main.py /app/server/main.py
COPY ./app/requirements.txt /app/server/requirements.txt

RUN pip install --no-cache-dir --upgrade -r /app/server/requirements.txt

WORKDIR /app/server

EXPOSE 8080

CMD ["./run.sh"]

web-server/docker-bake.hcl

target "web" {
  target = "web-stage"
  dockerfile = "Dockerfile"
  tags = ["web"]
}

準備は整ったので、Web サーバを起動します。

作業ディレクトリ：web-server

docker buildx bake web &&\
  docker run -d --name web-server -p 81:80 --env ENDPOINT_URL=http://model-server/inferences --network=network web

推論結果の確認

以下の post リクエストコマンドを実行すると、レスポンスとして{"labels":["狩り","狩","遊び","猟","狩猟"]} が返ってきます。Web Server はリクエストを Model Server にそのまま渡し、Model Server からのレスポンスをそのまま API のレスポンスとして返却するため、コマンドの結果は Model Server と同じになります。

curl -X 'POST' \
  'http://localhost:81/invocations' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "text": "お爺さんは森に狩りへ出かける",
  "mask_index": 7
}'

Docker Image の Amazon ECR へのプッシュ

作成した web、model Docker イメージを Amazon ECR にプッシュします。

source /tmp/env.hcl &&\
  aws ecr get-login-password --region $REGION \
    | docker login --username AWS --password-stdin ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com &&\
    docker tag web:latest ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/web:latest &&\
    docker push ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/web:latest &&\
    docker tag model:latest ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/model:latest &&\
    docker push ${ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/model:latest

まとめ

Part 1 では、AWS CDK で検証用の Amazon EC2 Inf1 インスタンスを構築し、このインスタンス上で AWS Neuron SDK によるモデルのコンパイル、AWS Inferentia による推論を行う Model Server の構築、既存のアプリケーションサーバを想定した Web Server の構築を行いました。Part 2 では、作成した Docker イメージを利用して AWS CDK で Amazon ECS を構築します。

Part 1 で作業を終了される場合は、cdk destroy で検証用インスタンスを削除してください。

Amazon Web Services ブログ

AWS CDK で作る AWS Inferentia と Amazon ECS を利用した推論環境（Part 1）

目次

事前準備

AWS CDK で 検証用 Amazon EC2 Inf1 インスタンス を構築する

構築概要

注意点

手順

アプリケーションを構築する

Model Server 概要

Web Server 概要

Model Server の構築

model-server/app/

model-server/build

model-server/trace

model-server/docker-bake.hcl

AWS Deep Learning Containers Base イメージ

AWS Neuron SDK によるモデルのコンパイル

Model Server の起動

推論結果の確認

Web Server の構築

web-server/app

web-server/Dockerfile

web-server/docker-bake.hcl

推論結果の確認

Docker Image の Amazon ECR へのプッシュ

まとめ

お役立ちリンク

フォローお願いいたします

AWS CDK で検証用 Amazon EC2 Inf1 インスタンスを構築する