Cloud RunからBedrockを呼び出す(Google Cloud→AWSの認証をIAMで通す)

やりたいこと

VertexAIでもLlama3が使えるようになったのでCloud Run上のアプリから試しに使ってみようと思いました。しかし、デプロイしないといけなくて面倒くさいしデプロイ先のVMの課金が怖いので、簡単に使えそうなAWSのBedrockでLlama3を使ってしまえと思いました。その為にGoogle CloudからAWSの認証突破を試みます。

ですが、実際にやってみたら認証以前にLangChain(langchain-community:0.0.34)からBedrockのLlama3の実行に不具合があるみたいでして、実行は出来ているもののAIの解答の中身がおかしな状態になってしまいました。issueが作られています↓
https://github.com/langchain-ai/langchain/issues/21037
ですので、この記事ではLlama3は後回しにして、Google CloudのCloud RunからAWSのBedrockのClaude 3を実行することを目標にします！！！（当初の目的を見失っていますね。。。）

IAMで、Google Cloud → AWSの認証を通す

方針

AWSとGoogle Cloudは別サービスですので、認証を突破しない限り利用する事はできません。
一番単純なのはAWSにユーザーを追加してアクセスキーとシークレットアクセスキーを発行することですが、セキュリティリスクが高いのでよろしくないですよね。最近もアクセスキーのお漏らしでセキュリティ事故を起こしていた大企業がニュースになっていましたし。
ではどうするかですが、Google CloudのIAMに対してAWSのIAMを紐づければよいわけです。どちらもIAMなのでたぶん出来るでしょう。と思ってやってみました。

やってみた

AWS側でIAMロールを作る

AWSコンソールからIAMロールを新規作成します。信頼されたエンティティタイプとして「ウェブアイデンティティ」を選択してください。

アイデンティティプロバイダーは「Google」を選択してください。
AudienceにはCloud Runに紐づけられているサービスアカウントの「一意のID」を指定します。

許可ポリシーとして、BedrockのInvokeをつけます。
当初はAWSで用意されたpolicyを付けようかと思いったのですが、使えそうなpolicyが用意されてなかったので自分で用意しました。(AmazonBedrockReadOnlyではinvokeできない。AmazonBedrockFullAccessでは権限強すぎる。)

BedrockInvokePolicy

{
	"Version": "2012-10-17",
	"Statement": [
		{
			"Sid": "VisualEditor0",
			"Effect": "Allow",
			"Action": [
				"bedrock:InvokeAgent",
				"bedrock:InvokeModel",
				"bedrock:InvokeModelWithResponseStream"
			],
			"Resource": "*"
		}
	]
}

以上で、AWSに対する設定は完了です。

認証用のコードを書く

続いて、認証情報を取得するためのコードを書いていきます。
以下のドキュメントを参考にしてGoogle CloudのメータデータサーバーからIDトークンを取得し、その後boto3を使ってAWSの一時認証情報を取得する処理を書きました。
https://cloud.google.com/docs/authentication/get-id-token?hl=ja

必要なライブラリのインストール

boto3とgoogle-authをインストールします。

rye add boto3 google-auth
rye sync

AWSの認証情報を取得

Google CloudのメータデータサーバーからIDトークンを取得しそのIDトークンからAWSの一時認証情報を取得する関数を書きます。

from typing import TypedDict

from google.auth import compute_engine
import google.auth.transport.requests as grequests
import boto3


class Credentials(TypedDict):
    '''AWSの一時認証情報の型定義'''
    AccessKeyId: str
    SecretAccessKey: str
    SessionToken: str
    Expiration: str


def get_aws_creds(aws_role_arn: str) -> Credentials | None:
    '''Google CloudのメータデータサーバーからIDトークンを取得し、AWSの一時認証情報を取得する'''

    try:
        # Google CloudのメタデータサーバーからIDトークンを取得
        request = grequests.Request()
        g_credentials = compute_engine.IDTokenCredentials(
            request=request,
            target_audience="https://sts.amazonaws.com/",
            use_metadata_identity_endpoint=True
        )
        g_credentials.refresh(request)
        token: str = g_credentials.token

        # AWSのSTSを使って一時認証情報を取得
        sts = boto3.client('sts')
        assume_role_res = sts.assume_role_with_web_identity(
            RoleArn=aws_role_arn,
            WebIdentityToken=token,
            RoleSessionName="session",
        )
        a_credentials: Credentials = assume_role_res.get('Credentials')  # type: ignore
        return a_credentials
    except Exception as e:
        print(f"Error: {e}")
        return None

実際に使う時は次のように呼び出します。AWS_ROLE_ARNには先ほど作ったAWSのIAMロールのARNを設定してください。

from boto3.session import Session

creds: Credentials | None = get_aws_creds(AWS_ROLE_ARN) if AWS_ROLE_ARN else None
session: Session = Session(
    aws_access_key_id=creds.get("AccessKeyId"),
    aws_secret_access_key=creds.get("SecretAccessKey"),
    aws_session_token=creds.get("SessionToken"),
) if creds else Session()
bedrock_runtime = session.client("bedrock-runtime", region_name="us-west-2")

以降は、通常通りbedrock_runtimeを使えばOKです。

こちらの記事やコードも参考にさせていただきました

https://qiita.com/hyj624117615/items/ca4a8bc8c269098aaeef
https://github.com/cevoaustralia/gcp-sa-to-aws-iam-role

実行結果

ということで、なんとか無事にCloudRunからBedrockを実行することが出来ました。

同様の方法で、 CloudRun → Bedrock に限らず Google Cloud → AWS の認証全般に応用出来るかと思われます。