Amazon EKS Pod Identity を利用した Amazon Bedrock アプリケーションの実行

Amazon EKS Pod Identity を利用すると Amazon EKS の Pod 単位で簡単に IAM ロール権限を与えることができます。このブログでは Amazon EKS Pod Identity を利用して Amazon Bedrock アクセスするアプリケーションを実行する例をご紹介します。

Amazon EKS Pod Identity の開始方法や詳細につきましては以下のブログをご確認ください。
ブログ：Amazon EKS Pod Identity は、Amazon EKS クラスター上のアプリケーションの IAM 許可を簡素化します

前提条件

このブログでは、既存の Amazon EKS クラスターを使用します。
Amazon Bedrock は東京リージョン ( ap-northeast-1 ) で anthropic.claude-instant-v1 のモデルが有効化している状態とします。
Amazon EKS Pod Identity の Add-on が既存の Amazon EKS クラスターにインストールしている状態とします。
本手順で公開するアプリケーションは 0.0.0.0:80 でパブリックに公開されますので、必要に応じてセキュリティグループを利用してアクセス制限を行ってください。
特に本番環境で利用する上では Amazon EKS のセキュリティベストプラクティスをご確認の上デプロイを行ってください。

Amazon Bedrock アクセスするアプリケーションの実行

まずは IAM ロール権限がない状態で以下のソースコードで実行されるアプリケーションを build して Amazon EKS 上で実行し、Amazon Bedrock アクセスできないことを確認します。

コンテナイメージの作成

Amazon Bedrock アクセスするアプリケーションのコードの例 ( app.py )

import streamlit as st
import boto3
import json

st.title("Bedrock Chat")

# Bedrock Runtimeサービス用のクライアント
bedrock = boto3.client(
    service_name='bedrock-runtime',
    region_name='ap-northeast-1')
    
def format_chat_history(messages):
    formatted_history = ""
    for message in messages:
        # ユーザーかアシスタントかに応じてロールを設定
        role = "Human:" if message["role"] == "user" else "Assistant:"
        # メッセージを整形して追加
        formatted_history += f"{role} {message['content']}\n"
    return formatted_history

# チャット履歴の初期化
if "messages" not in st.session_state:
    st.session_state.messages = []

# アプリ再実行時にチャットメッセージを表示
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# ユーザー入力の受け取り
if prompt := st.chat_input("What is up?"):
    # ユーザーメッセージをチャットメッセージコンテナに表示
    with st.chat_message("user"):
        st.markdown(prompt)
    # ユーザーメッセージをチャット履歴に追加
    st.session_state.messages.append({"role": "user", "content": prompt})

    # 対話履歴を整形してモデルの入力用に準備
    chat_history = format_chat_history(st.session_state.messages)

    # アシスタントの応答をチャットメッセージコンテナに表示
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""

        # 整形された対話履歴と新しいユーザー入力をモデルに送信
        body = json.dumps({
            'prompt': chat_history + 'Human: ' + prompt + '\n\nAssistant:',
            'max_tokens_to_sample': 1000,
            "stop_sequences": ["\n\nHuman:"],
            "anthropic_version": "bedrock-2023-05-31"
        })
        accept = 'application/json'
        contentType = 'application/json'
        model_id = 'anthropic.claude-instant-v1'

        response = bedrock.invoke_model_with_response_stream(
            body=body, modelId=model_id, accept=accept, contentType=contentType)
        stream = response.get('body')

        for event in stream:
            chunk = event.get('chunk')
            if chunk:
                # 取得したチャンクを追加し、Streamlitに表示
                full_response += json.loads(chunk.get('bytes').decode()
                                            )["completion"]
            message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
        # アシスタントの応答をチャット履歴に追加
        st.session_state.messages.append(
            {"role": "assistant", "content": full_response})

Dockerfileの例

FROM python:3.11-slim

WORKDIR /app

RUN apt-get update && apt-get install -y \
    curl \
    && rm -rf /var/lib/apt/lists/*

COPY ./requirements.txt /app/requirements.txt

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

COPY ./app.py /app

RUN adduser --disabled-password --gecos '' streamlit

USER streamlit

EXPOSE 8501

HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health

ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

requirements.txtの例

boto3==1.29.4
botocore==1.32.4
streamlit==1.28.2

Amazon EKS Pod Identity 使用する場合、 AWS SDK は一定以上のバージョンが必要になります。

以下のようにディレクトリにソースコードを配置してコンテナ image を build し Amazon EKS から利用できるコンテナレジストリーに push を行います。

$ tree
.
├── app.py
├── Dockerfile
├── requirements.txt

このブログでは Amazon ECR を利用した手順をご紹介します。
以下の手順では事前に Amazon ECR 上に llm-app というプライベートレジストリが作成が必要です。

# ご自身のAWSアカウントIDを設定
export AWS_ACCOUNT_ID=XXXXXXXXXXXX

# 認証トークンを取得し、レジストリに対して Docker クライアントを認証
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 イメージを構築
docker build -t llm-app:sample .

# リポジトリにイメージをプッシュできるように、イメージにタグ付けを行う
docker tag llm-app:sample ${AWS_ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/llm-app:sample

# イメージをプッシュ
docker push ${AWS_ACCOUNT_ID}.dkr.ecr.ap-northeast-1.amazonaws.com/llm-app:sample

Amazon EKS 上での実行

以下の manifest ファイルのようにしてアプリケーションを実行します。
この手順では sample-app という Namespace を作成し、Service は type: LoadBalancer を利用してアプリケーションを公開します。
※Amazon EKS でロードバランサーを取り扱う場合は AWS Load Balancer Controller の利用を推奨します

Namespace の manifest ファイルの例

apiVersion: v1
kind: Namespace
metadata:
  name: sample-app

Deployment の manifest ファイルの例
※ Deployment のコンテナイメージはご自身のコンテナレジストリのイメージに変更してください。

apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: llm-app
  name: llm-app
  namespace: sample-app
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llm-app
  template:
    metadata:
      labels:
        app: llm-app
    spec:
      containers:
      # ご自身のコンテナレジストリのイメージに変更してください
      - image: XXXXXXXXXXXX.dkr.ecr.ap-northeast-1.amazonaws.com/llm-app:sample
        name: sampleapp

Service の manifest ファイルの例

apiVersion: v1
kind: Service
metadata:
  name: llm-app-lb
  namespace: sample-app
spec:
  type: LoadBalancer
  selector:
    app: llm-app
  ports:
  - protocol: TCP
    port: 80
    targetPort: 8501

アプリケーションを実行すると以下のようなアプリケーションが実行されます。
しかしながら、 Amazon Bedrock への IAM ロール権限が無いため画面のように AccessDeniedException となるはずです。

アプリケーションへ Amazon Bedrock へのアクセス権限を追加する

Amazon Bedrock へのアクセス権限を持つ IAM ロールを作成し、 Amazon EKS 上で実行されている Pod に IAM ロールの追加を行います。

Amazon Bedrock へのアクセス権限を持つ IAM ロールの作成

IAM ロールの作成からユースケースで EKS – Pod Identity を選択し、 AmazonBedrockFullAccess の許可ポリシーを追加したロールを作成します。

Amazon EKS Pod Identity を利用して IAM ロールと ServiceAccount のアソシエーションを行う

以下の manifest ファイルのようにして ServiceAccount を作成し、Amazon EKS Pod Identity を利用して IAM ロールと ServiceAccount のアソシエーションを行います。アソシエーションは AWS Management Console 、もしくは CLI を利用して行うことができます。

ServiceAccount の manifest ファイルの例

apiVersion: v1
kind: ServiceAccount
metadata:
  name: bedrock-sa
  namespace: sample-app

CLI を利用したアソシエーションの例
※ —role-arn ではご自身で作成した IAM ロールの arn を指定してください。

# IAM ロールと ServiceAccount の association
aws eks create-pod-identity-association \
    --cluster-name <CLUSTER_NAME> \
    --role-arn arn:aws:iam::XXXXXXXXXXXX:role/eks-pod-identity-bedrock \
    --namespace sample-app \
    --service-account bedrock-sa
    
# association の確認
eksctl get podidentityassociation --cluster <CLUSTER_NAME>

ASSOCIATION ARN                                                                                 NAMESPACE       SERVICE ACCOUNT NAME    IAM ROLE ARN
arn:aws:eks:ap-northeast-1:XXXXXXXXXXXX:podidentityassociation/develop/a-zxckaesu4um17oi7e      sample-app      bedrock-sa              arn:aws:iam::XXXXXXXXXXXX:role/eks-pod-identity-bedrock

アプリケーションへ ServiceAccount を追加する
IAM ロールと ServiceAccount のアソシエーションが完了したら、あとはアプリケーションが定義されている Pod へ ServiceAccount の追加を行うだけです。 Deployment に serviceAccountName を追加し、先程の手順でIAM ロールとアソシエーションした ServiceAccount を指定します。

Deployment の manifest ファイルの例

apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: llm-app
  name: llm-app
  namespace: sample-app
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llm-app
  template:
    metadata:
      labels:
        app: llm-app
    spec:
      # serviceAccountName を追加
      serviceAccountName: bedrock-sa
      containers:
      # ご自身のコンテナレジストリのイメージに変更してください
      - image: XXXXXXXXXXXX.dkr.ecr.ap-northeast-1.amazonaws.com/llm-app:sample
        name: sampleapp

アプリケーションを更新すると以下の様に、Amazon Bedrock へのアクセス権限が追加されていることがわかります。
Amazon EKS Pod Identity の良さを感じて頂いたのではないでしょうか。

まとめ

このブログでは、 Amazon EKS 上で動作するアプリケーションに対して、Amazon EKS Pod Identity を利用して Amazon Bedrock への IAM ロール権限を追加する方法をご紹介しました
Amazon EKS Pod Identity を利用して ServiceAccount との IAM ロールのアソシエーションを行うと、 ServiceAccount を Pod に指定するだけで、 IAM ロールベースのアクセス権限を簡単に付与することができます。
サービスアカウントの IAM ロール (IRSA) と比較すると、 IAM ロールに Amazon EKS クラスター毎に信頼ポリシーを更新する必要がなく、複数のクラスターから IAM ロールを簡単に再利用できます。また ServiceAccount にアノテーションを付与する必要がないためよりシンプルな構成にすることができます。

このブログが皆様の Amazon EKS を利用したビジネスにお役立ちできれば幸いです。

作者情報

ソリューションアーキテクト瀧田直斗 (Takita Naoto)
中堅中小企業様を中心に技術的な側面からお客様のビジネスを支援させて頂いております。

Amazon Web Services ブログ