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が作られています↓
ですので、この記事では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
では権限強すぎる。)
{
"Version": "2012-10-17",
"Statement": [
{
"Sid": "VisualEditor0",
"Effect": "Allow",
"Action": [
"bedrock:InvokeAgent",
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream"
],
"Resource": "*"
}
]
}
以上で、AWSに対する設定は完了です。
認証用のコードを書く
続いて、認証情報を取得するためのコードを書いていきます。
以下のドキュメントを参考にしてGoogle CloudのメータデータサーバーからIDトークンを取得し、その後boto3を使ってAWSの一時認証情報を取得する処理を書きました。
必要なライブラリのインストール
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です。
こちらの記事やコードも参考にさせていただきました
実行結果
ということで、なんとか無事にCloudRunからBedrockを実行することが出来ました。
同様の方法で、 CloudRun → Bedrock に限らず Google Cloud → AWS の認証全般に応用出来るかと思われます。
![NCDCエンジニアブログ](https://storage.googleapis.com/zenn-user-upload/avatar/3a86a409d3.jpeg)
NCDC株式会社( ncdc.co.jp/ )のエンジニアチームです。 募集中のポジションや、採用している技術スタックの紹介などはこちらをご覧ください! github.com/ncdcdev/recruitment
Discussion