LlamaIndexをApp Runner上で動かしてみた

はじめに

CX事業本部アーキテクトチームの佐藤智樹です。

今回はLlamaIndexをApp Runner上で動かすことを試してみます。LlamaIndex自体は以下の記事で記載したコードをローカルPC上で実行すれば動かせますがサービス提供などを考えた場合、クラウド上で扱えた方がより柔軟に扱いやすいので動かしてみます。

アーキテクチャの構成としては以下のようになります。

作成したサンプルコードは以下においてありますのでご参照ください。

今回はインデックスファイルをコンテナ内にビルドして含めます。今後以下の記事のようにS3からインデックスファイルを取得する拡張も検討していきます。

コンテナの作成

まずはOpenAIのAPIを実行するためのコンテナを作成します。今回は以下のソースをほぼそのまま参考にして作成します。

主に変更した部分を記載します。上記のソースのflask周りの部分だけ利用します。index_server.pyにOPENAI_API_KEYを埋め込む部分があるので、実行時にSSM ParameterStoreのSecureStringからAPIキーを取得するよう変更します。

src/index_server.py

import os
import boto3
client = boto3.client('ssm')
response = client.get_parameter(Name="/OpenAiApiKey", WithDecryption=True)
secret_value = response['Parameter']['Value']

os.environ['OPENAI_API_KEY'] = secret_value
...

requirements.txtへ以下のように記載します。元のソースに加えてAWS SDKを使用するためboto3を追加、本番公開用のWSGIサーバとしてGunicornを追加します。

requirements.txt

Flask==2.2.3
langchain==0.0.115
llama-index==0.4.29
boto3==1.26.94
gunicorn==20.1.0

Dockerfileはそのままで、5601を公開ポートに設定しています。

src/Dockerfile

FROM python:3.11.0-slim

WORKDIR /app

COPY . .

RUN pip install -r requirements.txt && pip cache purge

# Flask
CMD ["sh", "launch_app.sh"]
EXPOSE 5601

Dockerfile内のCMDで起動しているlaunch_app.shでApp Runner上でGunicornを使うために末尾の記載を変更します。

src/launch_app.sh

#!/bin/bash

# start backend index server
python ./index_server.py &

# wait for the server to start - if creating a brand new huge index, on startup, increase this further
sleep 60

# start the flask server with gunicorn
gunicorn --workers 2 --bind 0.0.0.0:5601 flask_demo:app

以下のコマンドでSSM ParameterStoreにOpenAI APIキーの追加、Dockerfileをビルドして、Docker上にAWSのリージョンや一時クレデンシャルなど含めてローカルPC上で起動します。

% aws ssm put-parameter --name "/OpenAiApiKey" --value "your key here" --type SecureString
% docker image build -t llamaindex-pack:latest .
% docker run -it \
  -e AWS_DEFAULT_REGION=ap-northeast-1 \
  -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
  -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
  -e AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN \
  -p 8000:5601 llamaindex-pack:latest

上記起動後、以下のようにcurlで動作確認をします。構築がうまくいっていればインデックスに記載されているPaul Grahamの情報が出ます。

% curl "http://127.0.0.1:8000/query?text=Pleasetellmeyourname"
'My name is Paul Graham.'%

App Runner上で起動

前章でコンテナの準備は出来たので、次はApp Runner上にコンテナを乗せていきます。App RunnerはECRにあるコンテナを使うように構築します。実装は全てAWS CDKで実装します。

ECRの作成

以下のように単にリポジトリだけ作成します。ライフサイクルがApp Runnerと別れるためスタックとコンストラクトは分離します。

lib/construct/ecr-construct.ts

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import * as ecr from "aws-cdk-lib/aws-ecr";

export class EcrConstruct extends Construct {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id);

    const repository = new ecr.Repository(this, `${id}-Ecr`, {
      repositoryName: "llama-index-repo",
      imageScanOnPush: true,
    });
  }
}

以下のコマンドでCDKを使ってコンテナリポジトリを作成します。

# コンテナリポジトリをデプロイ
% npx cdk deploy "ContainerStack"

作成が完了したら先ほど作成したコンテナイメージにタグ付けして、リポジトリにpushします。${AWS_ACCOUNT_ID}の部分は自分のAWSアカウントIDが入るように環境変数に設定してください。

% aws ecr get-login-password --region ap-northeast-1 | \
  docker login --username AWS --password-stdin ${AWS_ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com
% docker tag llamaindex-pack:latest ${AWS_ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/llama-index-repo:v1
% docker push ${AWS_ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/llama-index-repo:v1

pushが成功したら最後にApp Runnerを実行します。以下のConstruct部分で実装しています。App Runnerのインスタンスロールには最低限SSM ParameterStoreからOpenAIのAPIキーを取得する権限のみ与えています。もし別のAWSリソースにアクセスしたい場合は、ポリシーを追加してください。

lib/construct/app-runner-construct.ts

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import * as apprunner from "@aws-cdk/aws-apprunner-alpha";
import * as iam from "aws-cdk-lib/aws-iam";
import * as ecr from "aws-cdk-lib/aws-ecr";

export class AppRunnerConstruct extends Construct {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id);

    const instanceRole = new iam.Role(scope, `${id}-InstanceRole`, {
      assumedBy: new iam.ServicePrincipal("tasks.apprunner.amazonaws.com"),
    });
    instanceRole.addToPolicy(
      new iam.PolicyStatement({
        actions: ["ssm:GetParameter"],
        effect: iam.Effect.ALLOW,
        resources: [`arn:aws:ssm:${props?.env?.region}:${props?.env?.account}:parameter/OpenAiApiKey`],
      })
    );

    new apprunner.Service(this, `${id}-Service`, {
      serviceName: `llama-runner`,
      cpu: apprunner.Cpu.ONE_VCPU,
      memory: apprunner.Memory.TWO_GB,
      source: apprunner.Source.fromEcr({
        imageConfiguration: { port: 5601 },
        repository: ecr.Repository.fromRepositoryName(this, `${id}-LlamaIndexRepo`, "llama-index-repo"),
        tagOrDigest: "v1",
      }),
      instanceRole,
    });
  }
}

最後に以下のコマンドでデプロイします。デプロイには大体現時点で5~10分程度かかります。

% npx cdk deploy "LlamaindexRunnerStack"

デプロイが正常に完了すると、App Runnerのデフォルトドメインが作成されるのでデフォルトドメインに向けてcurlでリクエストを飛ばします。正常にデプロイできていると以下のようにレスポンスが返ってきます。

% curl "https://xxxxxxxxxx.ap-northeast-1.awsapprunner.com/query?text=Pleasetellmeyourname"

"My name is Paul Graham. I am the co-founder of the company Viaweb, which was funded by Julian, Idelle's husband. We gave him 10% of the company in return for the seed funding and legal and business advice. I had a negative net worth at the time, so the seed funding was essential for me to live on."%

これでApp Runner上でLlamaIndexを動かせるようになりました。後はインデックス更新のユースケースを整理してFlaskの実装を拡張したり、Cognitoなどで認証認可を組み込むこともできます。そちらについても今後別の記事で紹介します。

おまけ：なぜ App Runner に乗せるのか

AWSでコンピューティングリソースとして主に使われるLambdaやECS/EC2と比較してみます。

LlamaIndexでインデックスを作成・追加する際、データが多くなるとインデックス作成に時間がかかります。これをLambdaで動かすと15分以上かかりタイムアウトしてしまう懸念があります。インデックスを追加書き込みする方式を採用するには、適切な単位でLambdaを構成しエラー処理を組む必要があり少々手間です。ECS/EC2で実行する時と比べた場合、App RunnerはVPCなどの管理はVPC内部のリソースを使わない場合不要なため、NAT GatewayやVPC Endpointなどの常駐リソースの費用を節約できます。また構成もVPC周りのリソースが不要なためかなりシンプルになります。

もちろん利点だけでなく、欠点もあります。App Runner自体のデプロイがコンテナの置き換えに時間が体感5~10分程度かかりAWS上での動作確認には時間がかかります。ただローカルでの開発はDocker上でスピード感をもって作ることは可能なので、欠点はある程度補えます。上記の利点・欠点の比較から今回はApp Runner上で動かしてみています。

参考