aws-bench の診断タスクを AWS 利用料ほぼゼロ、Kiro クレジットで試してみた

はじめに

先日、AWS が公開した AI エージェント向けベンチマーク aws-bench について紹介しました。

https://dev.classmethod.jp/articles/aws-bench-overview/

公式の Getting started の手順でベンチマークを実行するには、新しい AWS アカウントを払い出せる AWS Organizations 環境が必要です。また、シナリオによっては高額な AWS リソースを使用するものもあるため、気軽に試せるものではありませんでした。

そこで、公式環境の再現は目指さず、4つの診断タスクを選び、既存の AWS アカウント1つで動作するミニテスト環境を用意しました。

ミニテストの方針

診断対象の AWS 環境に意図的な異常を作り、隔離コンテナ内で read-only ロールを使用するエージェントに調査させて、回答ファイルを作成します。

回答ファイルは、判定担当の別の Kiro が評価します。

Organizations によるアカウント払い出し、公式のオーケストレーション、Bedrock の judge、ドリフト検証、変更系タスクのリソース状態検証は省略しました。

4つの診断タスク環境を CDK で作る

公式シナリオでは、CDK スタックと必要な setup で、意図的な異常を含む AWS リソースを作成し、30スタックを7リージョンに展開します。今回は、4つのタスク分のリソースと共通の read-only ロールだけに絞り、5スタックを3リージョンに直接展開しました。

4つのタスク環境の構築には、公式データセットの資材を使いました。

https://github.com/aws-bench/aws-bench-datasets

今回の5スタックを定義するコードは environment.ts です。

// 全タスク共通（read-only ロール）
new QARolesStack(app, `${envId}-QARoles-us-east-1`, { env: { account, region: 'us-east-1' } });

// Task: check-vpc-flow-log-destinations
new Cloudformation_vesgrw3ay(app, `${envId}-cloudformation-vesgrw3ay-us-east-1`, { env: { account, region: 'us-east-1' } });

// Task: diagnose-step-functions-failing-executions
new Lambda_mw9wjm2q7(app, `${envId}-lambda-mw9wjm2q7-ap-southeast-2`, { env: { account, region: 'ap-southeast-2' } });

// Task: fix-cloudformation-update-failed-stack
new Cloudformation_t9dx4pgqw(app, `${envId}-cloudformation-t9dx4pgqw-us-east-1`, { env: { account, region: 'us-east-1' } });

// Task: lambda-appconfig-stale-value-troubleshoot
new Lambda_1sscp39dx(app, `${envId}-lambda-1sscp39dx-us-west-2`, { env: { account, region: 'us-west-2' } });

4つのタスクの主なリソース、仕込まれた異常、診断内容を表にまとめます。

fix-cloudformation-update-failed-stack という名称には fix が含まれますが、このミニテストでは復旧操作を実行させません。失敗原因と必要な復旧操作を、診断レポートに記載させます。

CDK のスタック定義と異常状態を作るロジックは、公式資材を流用したものです。単一アカウントで動作するように、Dockerfile、デプロイ手順、setup における認証プロファイルの解決方法を手直ししました。

fix-cloudformation-update-failed-stack で作る異常状態

4つのタスクのうち、デプロイ後に異常状態を作成するタスクである fix-cloudformation-update-failed-stack を例にします。CDK では、2つの Lambda 関数と現在のバージョンを指す alias を作ります。

const simpleEmailServiceAlias = new lambda.Alias(this, 'SimpleEmailServiceLambdaAlias', {
    aliasName: 'LATEST',
    version: simpleEmailServiceLambda.currentVersion,
});

const getDetectorOutcomeAlias = new lambda.Alias(this, 'GetDetectorOutcomeLambdaAlias', {
    aliasName: 'LATEST',
    version: getDetectorOutcomeLambda.currentVersion,
});

デプロイ後、setup スクリプトが CloudFormation の元テンプレートを書き換えます。AWS::Lambda::Alias の FunctionVersion を存在しないバージョン番号に差し替え、スタックを更新します。DisableRollback=True を指定するため、更新失敗後も UPDATE_FAILED 状態のスタックを調査できます。

for resource in template.get("Resources", {}).values():
    if resource.get("Type") != "AWS::Lambda::Alias":
        continue

    ref = resource["Properties"].get("FunctionName", {}).get("Ref", "")
    resource["Properties"]["FunctionVersion"] = (
        "593" if "SimpleEmailService" in ref else "268"
    )

update_and_wait(cfn, template, disable_rollback=True)

final_status = cfn.describe_stacks(StackName=STACK_NAME)["Stacks"][0]["StackStatus"]
if final_status != "UPDATE_FAILED":
    raise RuntimeError(f"Expected stack to reach UPDATE_FAILED but got {final_status}")

このタスクでは、エージェントは UPDATE_FAILED の原因と復旧に必要な操作を診断レポートに記載します。

コンテナに隔離した Kiro に診断させる

当初は、公式ベンチマークの資材を置いたプロジェクトディレクトリで、Kiro をヘッドレスかつ全ツールを信頼するモード（--trust-all-tools）で実行していました。しかし実行ログから、エージェントが正答ファイルなど採点に使用する資材を参照した痕跡を確認しました。この状態では、AWS 環境を調査する能力ではなく、ローカルに置かれた評価用資材へのアクセスが回答に影響し得ます。

そこで、公式 aws-bench と同様に、エージェントを Docker コンテナへ隔離して実行する構成を今回のミニテストでも踏襲しました。エージェントコンテナには、Kiro の認証情報、タスク指示、読み取り専用の AWS 認証情報、回答の出力先だけを渡します。以降の例では、実際の ARN と参照回答の原文を伏せます。正答やルーブリックなどの採点資材にはアクセスさせません。

共通の QARoles スタックは、エージェントの AWS 調査用に read-only ロールを作ります。Kiro は各タスクの指示を受け、このロールを使用して AWS の構成、ログ、実行履歴を参照し、診断レポートを書きます。

エージェント用コンテナには、AWS CLI v2 と Kiro CLI を含む Python slim イメージを使用し、その中で Kiro をヘッドレス実行します。ホストにあるベンチマーク資材はマウントしません。コンテナ側には実行環境と必要な入力だけを置き、判定はホスト側で行います。

例えば Step Functions のタスクでは、エージェントは読み取り専用の AWS 操作で実行履歴と Lambda のエラーログをたどり、Lambda へ渡される入力に targetBucket がないことを原因として特定します。

Step Functions タスクを隔離コンテナの Kiro に渡す

ランナーは、タスク指示中のプレースホルダを単一アカウント環境内の識別子に置き換えてから、コンテナ内で Kiro をヘッドレスモードで実行します。エージェントには、次の指示を渡します。

My Step Functions state machine <resolved StateMachine ARN>
in ap-southeast-2 has failing executions. Why?

コンテナ起動時には、認証ボリュームとタスクごとの一時出力先だけをマウントします。AWS の一時認証情報は環境変数で渡します。

docker run --rm \
  -v <kiro-auth-volume>:/root/.local/share/kiro-cli \
  -v <per-task-output-directory>:/logs/agent \
  -e TASK_INSTRUCTION="<resolved instruction>" \
  <isolation-image> \
  'kiro-cli chat \
    --model <model-id> \
    --trust-all-tools \
    --no-interactive \
    "$TASK_INSTRUCTION

Benchmark safety rule: investigate with read-only AWS operations only.
Do NOT create, update, delete, deploy, invoke, or otherwise modify any AWS resource.

Write only your final answer to:
/logs/agent/agent-output.txt"'

タスク指示としてエージェントに渡すのは、この問題文と制約のみです。ホストのベンチマーク資材、参照回答 ground_truth.json、公式ルーブリック judge_prompt.md はコンテナに渡しません。参照回答とルーブリックは、エージェントの回答を判定する後段の Kiro だけが読みます。

公式の採点資材をローカル Kiro で使う

公式データセットには、タスクの指示、参照回答、ルーブリック（judge_prompt.md）が含まれます。タスク指示と参照回答に含まれるプレースホルダを、単一アカウント環境内の識別子に置き換えて使います。

Step Functions タスクを採点する

この Step Functions タスクでは、採点側の Kiro がエージェントの診断結果 agent-output.txt を判定します。判定時には、解決済みの参照回答と公式ルーブリックを参照します。参照回答には、タスク指示と期待される診断内容が含まれます。

公式ルーブリックの原文は、GitHub 上の次のファイルで確認できます。

https://github.com/aws-bench/aws-bench-datasets/blob/2daf77d2d41c21bae00bf8227fc463be51f721d0/tasks/troubleshooting-multiservice/diagnose-step-functions-failing-executions/tests/judge_prompt.md

この例では、前節で示した Step Functions のタスク指示に対する回答を判定します。参照回答では、InvokeProcessingLambda に渡される入力に targetBucket がなく、Lambda が失敗することを期待します。

採点側の Kiro には agent-output.txt、ground_truth-resolved.json、judge_prompt.md を読ませます。それぞれ、エージェントの診断結果、解決済みの参照回答、公式ルーブリックです。

参照回答が示す原因・障害の状態・対象リソースと一致するかを確認します。表現の違いと追加された正しい情報は許容します。

一方、原因が異なる場合や主要な原因が抜けている場合は不合格です。復旧策を求めるタスクでは、実行不可能な復旧策を示す場合も不合格とします。

各判定は、次の JSON 形式で保存します。

{
  "criterion": "answers_equivalent",
  "result": "<pass|fail>",
  "reason": "<brief explanation>",
  "task": "diagnose-step-functions-failing-executions"
}

コスト要素

テスト環境では VPC Flow Logs、CloudWatch Logs、Amazon S3 などを使用しますが、AWS 利用料はほぼゼロです。AI エージェントの実行には Kiro のクレジットを使用し、付帯クレジットの範囲内であれば、1クレジットあたり $0.02 で利用できます。

まとめ

AWS Organizations や高額なリソースを用意しにくい場合でも、単一アカウントの小規模な環境であれば、AWS 利用料をほぼゼロに抑えられます。Kiro クレジットでエージェントに診断させ、その回答をローカルの Kiro で判定する流れを試せるようになりました。実行結果は別の記事で紹介する予定です。

参考リンク

• aws-bench mini tasks local runner