Amazon Bedrockを最小権限のIAMポリシーのみで、Lambdaから動かしてみる。

目次

• まえがき
  • Amazon Bedrockはどんなサービス？
• Amazon Bedrock APIでサポートされているアクションについて
  • Bedrock
  • BedrockRuntime
• invoke_modelを最小権限で使えるように設定してみる。
  • 簡単な構成の紹介
  • 1. S3(アウトプット用バケット)を作成する。
  • 2. Lambda関数を作成する。
  • 3. ローカル環境で以下コマンドを実施し、LambdaLayerとしてアップロードする。Lambda関数にも紐付ける。
  • 4. Lambda関数にコードをコピペする。
  • 5. Lambda関数のIAMロールに以下インラインポリシーを作成し付与する。
  • 6. Lambda関数をテストする。
  • 7. コードをコピペしてLambda関数を更新する。
  • 8. Lambda関数のIAMロールに付与している、インラインポリシーを以下内容に修正する。
  • 9. Lambda関数をテストする。
  • 10. 権限が絞られている状態で他のライブラリやアクションが使えない事を確認する
• まとめ

まえがき

ついに日本時間2023年9月29日(金) Amazon Web Services ブログにて、
生成系AIサービスであるAmazon Bedrockが一般利用可能(GA)となった事が発表されました。
私自身Amazon BedrockというサービスがGAとなったことを最近知りましたが、今話題の生成系AIという事で早速触ってみました。

• Amazon Bedrock が一般利用可能に – 基盤モデルを利用した生成系 AI アプリケーションの構築とスケール
  • https://aws.amazon.com/jp/blogs/news/amazon-bedrock-is-now-generally-available-build-and-scale-generative-ai-applications-with-foundation-models/

※IAMポリシーの内容だけ知りたい方は手順8.へ

Amazon Bedrockはどんなサービス？

Amazonおよび主要なAIスタートアップ企業が提供する基盤モデルをAPIを通じて利用可能にする生成系AIサービスになります。

我々開発者にとっては、さまざまなAI基盤モデルを、ユースケースに最適なモデルを選択することが可能になるため、積極的に活用することで新しいビジネスチャンスを生み出すにつながるサービスだと捉えています。

使用方法

後述のように2つの方法があります。
しかし事前手順としてモデルアクセスの許可が必須となります。（AWSコンソール経由）

• API からモデルにアクセスする前に

方法① AWSコンソールを使用する

https://docs.aws.amazon.com/bedrock/latest/userguide/using-console.html

方法② APIを使用する

Amazon Bedrock API には、AWS コマンドラインインターフェイス (AWS CLI)、AWS SDK、または SageMaker Notebook を使用してアクセスが可能でした。
https://docs.aws.amazon.com/bedrock/latest/userguide/using-api.html

API経由でBedrockにアクセスできるという事は、IAMポリシーも最小権限で実装したいという場合も多いかと思います。

Amazon Bedrockに関するAWSドキュメントやAWSコンソールを調査したところ、
現時点(2023/10/03)では、マネージドIAMロールが存在しておりませんでした。
やはりGAされたばかりのサービスのため情報も少ないように見受けられました。
今回は実行したいアクションと必要なIAMポリシーなどについて解説をしたいと思います。

よろしくお願いします。

Amazon Bedrock APIでサポートされているアクションについて

• AWS Python SDK(Boto3)を使用することができ、2種類のクライアントが定義されています。
  
  • Boto3 1.28.57 documentation
    • Bedrock API Reference

2種類のクライアントには、それぞれライブラリが定義されています。

Bedrock

Bedrock
カスタムモデルを設定・管理できたり情報取得が可能です。

BedrockRuntime

BedrockRuntime
Bedrock APIを叩いて生成AIなどのコンテンツを取得することが可能です。

invoke_model
BedrockRuntimeのクライアントの中でも、「invoke_model」が使用頻度の高いライブラリになるかと思います。

invoke_modelを使用してテキストモデルや画像モデルなどを実行し、生成AIコンテンツを取得することが可能となります。

invoke_modelを最小権限で使えるように設定してみる。

簡単な構成の紹介

• Lambda関数
  • Bedrock APIのinvoke_modelライブラリと、AI画像生成するStable Diffusion XL v0.8 – previewモデルを設定する
• Lambda実行IAMロール
  • Bedrock APIのinvoke_modelアクションのみを許可する
  • Bedrock APIのAI画像を生成するStable Diffusion XL v0.8 – previewモデルのみアクセス許可する
• S3(アウトプット用バケット)
  • 生成したAI画像を保存する

1. S3(アウトプット用バケット)を作成する。

一意になるよう命名します。
今回検証バケット名は、「bedrock-serverless-target-us-east-1」

2. Lambda関数を作成する。

3. ローカル環境で以下コマンドを実施し、LambdaLayerとしてアップロードする。Lambda関数にも紐付ける。

mkdir workspace

pip install -t ./workspace boto3

mv ./workspace ./python

zip -r boto3-1.28.57.zip ./python

4. Lambda関数にコードをコピペする。

import os
import boto3
import json
import base64
import datetime

# Lambda環境変数
BUCKET_NAME = os.environ['BUCKET_NAME']

Bucket_Name = BUCKET_NAME

s3 = boto3.client('s3', region_name="us-east-1")
bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-east-1")

def lambda_handler(event, context):
    prompt_data = event.get('prompt')
    modelId = "stability.stable-diffusion-xl-v0" 
    accept = "application/json"
    contentType = "application/json"

    body = json.dumps({
        "text_prompts": [{ "text": prompt_data }],
        "cfg_scale": 10,
        "seed": 20,
        "steps": 50
    })

    response = bedrock.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )
    response_body = json.loads(response.get("body").read())
    print(response_body['result'])

    # 取得した画像データのデコード
    encoded_png_data = response_body.get("artifacts")[0].get("base64")
    decoded_png_data = base64.b64decode(encoded_png_data)

    # ファイル名の付与
    now = datetime.datetime.now()
    formatted_date = now.strftime('bedrock-%Y-%m-%d %H:%M:%S')
    file_name = f"{formatted_date}.png"

    # S3バケットへ出力
    s3.put_object(Bucket=Bucket_Name, Key=file_name, Body=decoded_png_data)

    # Bedrockの画像生成するモデル「stability.stable-diffusion-xl-v0」のARNを調査するためログ出力
    bedrock_c = boto3.client('bedrock', region_name='us-east-1')
    response = bedrock_c.list_foundation_models()

    # Output the response
    print(response)

※特に重要なのが、46行目から51行目「Bedrockの画像生成するモデル「stability.stable-diffusion-xl-v0」のARNを調査するためログ出力」の箇所です。

• IAMポリシー権限を付与するためには、画像生成するモデル「stability.stable-diffusion-xl-v0」のARN(Amazon Resource Name)が必要です。→ ARNについて
  • 今回はカスタムモデルではなく標準のモデルを使用しているため「リソースタイプ 」は「foundation-model」となります。
    • IAMポリシーで権限を絞るためには以下のARNを完成させる必要がありますが、AWSコンソール上では画像生成するモデル「stability.stable-diffusion-xl-v0」のARNが見つけられませんでした。

※画像生成するモデル「stability.stable-diffusion-xl-v0」のARN(Amazon Resource Name)
↓
arn:${Partition}:bedrock:${Region}::foundation-model/${ResourceId}

5. Lambda関数のIAMロールに以下インラインポリシーを作成し付与する。

• 信頼ポリシー

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "Service": "lambda.amazonaws.com"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

• インラインポリシー：BedrockPermissions

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "bedrock:*",
            "Resource": "*",
            "Effect": "Allow"
        }
    ]
}

• インラインポリシー：S3Permissions

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "s3:PutObject",
                "s3:GetObject",
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::<作成したS3バケット名>",
                "arn:aws:s3:::<作成したS3バケット名>/*"
            ],
            "Effect": "Allow"
        }
    ]
}

6. Lambda関数をテストする。

渡すJSON

{
  "prompt": "富士山"
}

テストが完了したことを確認する。ここで失敗する場合は作成したIAMロールに付与したポリシーが意図したものとなっていることを再度確認してください。

S3バケットに生成された画像が存在するか確認

• 「富士山」で生成されていた画像

テスト結果のログから「modelArn」をページ内で検索し、値を参照する。ARN(Amazon Resource Name)になります。

※今回は標準の画像生成モデル「stability.stable-diffusion-xl-v0」のARNを調べたため皆さん共通で使えるARN(Amazon Resource Name)かと思います。
今後カスタムモデルを作成した際にARN(Amazon Resource Name)が特定できなかった場合はログ出力から発見するのがオススメです。

ログ出力コードには、Bedrockクライアントの「list_foundation_models」ライブラリを使用しています。

    # Bedrockの画像生成するモデル「stability.stable-diffusion-xl-v0」のARNを調査するためログ出力
    bedrock_c = boto3.client('bedrock', region_name='us-east-1')
    response = bedrock_c.list_foundation_models()

    # Output the response
    print(response)

7. コードをコピペしてLambda関数を更新する。

import os
import boto3
import json
import base64
import datetime

# Lambda環境変数
BUCKET_NAME = os.environ['BUCKET_NAME']

Bucket_Name = BUCKET_NAME

s3 = boto3.client('s3', region_name="us-east-1")
bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-east-1")

def lambda_handler(event, context):
    prompt_data = event.get('prompt')
    modelId = "stability.stable-diffusion-xl-v0" 
    accept = "application/json"
    contentType = "application/json"

    body = json.dumps({
        "text_prompts": [{ "text": prompt_data }],
        "cfg_scale": 10,
        "seed": 20,
        "steps": 50
    })

    response = bedrock.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )
    response_body = json.loads(response.get("body").read())
    print(response_body['result'])

    # 取得した画像データのデコード
    encoded_png_data = response_body.get("artifacts")[0].get("base64")
    decoded_png_data = base64.b64decode(encoded_png_data)

    # ファイル名の付与
    now = datetime.datetime.now()
    formatted_date = now.strftime('bedrock-%Y-%m-%d %H:%M:%S')
    file_name = f"{formatted_date}.png"

    # S3バケットへ出力
    s3.put_object(Bucket=Bucket_Name, Key=file_name, Body=decoded_png_data)

8. Lambda関数のIAMロールに付与している、インラインポリシーを以下内容に修正する。

• インラインポリシー：BedrockPermissions
  • 変更点
    • Bedrock APIのinvoke_modelアクションのみを許可する
    • Bedrock APIのAI画像を生成するStable Diffusion XL v0.8 – previewモデルのみアクセス許可するようARNを追記

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "bedrock:InvokeModel",
            "Resource": "arn:aws:bedrock:us-east-1::foundation-model/stability.stable-diffusion-xl-v0",
            "Effect": "Allow"
        }
    ]
}

9. Lambda関数をテストする。

渡すJSON

{
  "prompt": "富士山"
}

テストが完了したことを確認する。ここで失敗する場合は作成したIAMロールに付与したポリシーが意図したものとなっていることを再度確認してください。

S3バケットに生成された画像が存在するか確認

• 「富士山」で生成されていた画像

10. 権限が絞られている状態で他のライブラリやアクションが使えない事を確認する

• 以下コード(手順4.で使用したもの)でLambda関数を更新する。
  • 46行目から51行目のログ出力コードには、Bedrockクライアントの「list_foundation_models」ライブラリを使用しているのでエラーとなるはずです。
    • 期待値：AccessDeniedException

import os
import boto3
import json
import base64
import datetime

## Lambda環境変数
BUCKET_NAME = os.environ['BUCKET_NAME']

Bucket_Name = BUCKET_NAME

s3 = boto3.client('s3', region_name="us-east-1")
bedrock = boto3.client(service_name='bedrock-runtime', region_name="us-east-1")

def lambda_handler(event, context):
    prompt_data = event.get('prompt')
    modelId = "stability.stable-diffusion-xl-v0" 
    accept = "application/json"
    contentType = "application/json"

    body = json.dumps({
        "text_prompts": [{ "text": prompt_data }],
        "cfg_scale": 10,
        "seed": 20,
        "steps": 50
    })

    response = bedrock.invoke_model(
        body=body, modelId=modelId, accept=accept, contentType=contentType
    )
    response_body = json.loads(response.get("body").read())
    print(response_body['result'])

    # 取得した画像データのデコード
    encoded_png_data = response_body.get("artifacts")[0].get("base64")
    decoded_png_data = base64.b64decode(encoded_png_data)

    # ファイル名の付与
    now = datetime.datetime.now()
    formatted_date = now.strftime('bedrock-%Y-%m-%d %H:%M:%S')
    file_name = f"{formatted_date}.png"

    # S3バケットへ出力
    s3.put_object(Bucket=Bucket_Name, Key=file_name, Body=decoded_png_data)

    # Bedrockの画像生成するモデル「stability.stable-diffusion-xl-v0」のARNを調査するためログ出力
    bedrock_c = boto3.client('bedrock', region_name='us-east-1')
    response = bedrock_c.list_foundation_models()

    # Output the response
    print(response)

実行すると以下のように「AccessDeniedException」エラーが出力され、期待通りとなりました。

まとめ

Amazon Bedrock使ってみると非常に便利なAWSサービスだなと実感しました。
特にAWSマネージドで生成AI系にAPI経由でアクセスできるのが驚きでした。
今回のようにIAM権限を絞ることでよりセキュアに管理でき、
これまでのサーバレスアーキテクチャに組み込むことでより柔軟なシステムなどが構築できそうでワクワクしています。

最後までご覧頂きありがとうございます。