0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@dais-w

Bedrockの課金が暴走したら自動で止まるキルスイッチを作ってみた

0
Posted at

概要

個人で遊んでるアプリでBedrockを呼んでいて、課金が暴走したら怖いなぁと思ったので
予算を超えたら自動でAPIを止める仕組み（キルスイッチ）を作ってみたっていう備忘録
（2026年前半、4月ごろに作ったときの情報です）

結論

Budgets → SNS → Lambda で DynamoDB にフラグを立てて、アプリ側Lambdaが毎リクエストそのフラグを見て503を返す
Budgetsの通知だけでは止まらない。止めるのは自分で作る。
月初はEventBridgeでフラグを戻す
ただBudgetsの反映には遅延があるので、設定額ぴったりで止まるわけじゃない。最後の保険くらいのつもり

前提

  • Bedrockはアプリ側Lambdaから呼んでいる
  • IaCはTerraform、Lambdaは Python
  • DynamoDBのレート制限テーブル（pk/sk、TTL有効、オンデマンド）が既にある。フラグもここに相乗りさせる

そもそも普段の制限

キルスイッチは最後の保険で、普段はその手前で何段か制限してる

  • API Gatewayのステージ全体にスロットリング（秒間リクエスト数とバースト）
  • Lambdaの同時実行数の上限（reserved concurrency）
  • ユーザーごとの回数制限。DynamoDBに日次・月次のカウンタを持って、超えたら429
    • 別で作ってるアプリでは、同じところから短時間に連続で来たら弾く、っていう1分あたりの制限にしてる。仕組みは同じで、期間が違うだけ

連打とかはここで止まるので、キルスイッチが要るのは
これをすり抜けてジワジワ積み上がった総額とか、想定外にユーザーが増えた時とかの話

なぜ作ったか

Bedrockは使った分だけ課金の従量制で、上限が勝手にかかるわけじゃない
リトライ地獄のバグを仕込んだとか、誰かにAPI叩かれまくったとか、なんかの拍子にアクセスが増えたとか
どれでも気づいたら請求が跳ねてそう・・・

Budgetsで予算の通知は設定できるけど、通知はメールが来るだけ
メールだけだと絶対気づかない（寝てる間とか仕事中とか……）
気づいて手で止めるまでは課金が続くので、通知じゃなくて自動で止まる仕組みが欲しかった
なので、アプリ側でフラグを見て止める仕組みを自分で作った。

Budget Actions っていう、IAMのDenyポリシーを当てたりEC2/RDSを止めたりをBudgets側で自動でやってくれるやつもある
アプリのLambdaのロールに bedrock:InvokeModel のDenyを当てれば、それでも止められそう（Converse APIもこの権限で止まる。ストリーミングも使うなら bedrock:InvokeModelWithResponseStream も）

ただ、上に書いたレート制限をもうDynamoDBベースでやってたので、止める仕組みもそこに載せたかった
フラグを同じテーブルに置けば、レート制限の確認と一緒に1回で見られるし、止める・戻すも全部DynamoDBのフラグ1個で済む
なので今回はBudget Actionsは使わず、DynamoDBのフラグで止める方にした

全体はこんな感じ

Budgetsは知らせるだけ、止めるのはLambdaとフラグ、っていう分担。

やったこと

1. Budgets + SNS

月間コストの予算を作って、閾値ごとに通知を出し分けた

  • 50%：予測（FORECASTED）でメールだけ。「このペースだと超えそう」の早めの気づき用
  • 80%：実績（ACTUAL）でメールだけ
  • 100%：実績でメール＋SNS。ここでキルスイッチが動く

予測と実績を使い分けて、SNSに流すのは実績100%だけにしてる（誤発火したくないので）。

で、ここは注意
通知先にSNSトピックを指定しただけだと、BudgetsはそのトピックにPublishできないらしい
トピック側のポリシーで budgets.amazonaws.com に SNS:Publish を許可しないといけない
忘れると通知がSNSに届かなくて、キルスイッチが黙ったまま、ってことになる（一番気づきにくいやつ）

ついでに aws:SourceAccount の条件も付けた
公式のサンプルだと aws:SourceArn の条件も付いてるけど、今回はアカウントだけにした

resource "aws_sns_topic_policy" "budget_alert" {
  arn = aws_sns_topic.budget_alert.arn

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Sid       = "AllowBudgetsPublish"
      Effect    = "Allow"
      Principal = { Service = "budgets.amazonaws.com" }
      Action    = "SNS:Publish"
      Resource  = aws_sns_topic.budget_alert.arn
      Condition = {
        StringEquals = {
          # アカウントIDは直書きせず caller_identity から
          "aws:SourceAccount" = data.aws_caller_identity.current.account_id
        }
      }
    }]
  })
}

aws_sns_topic.budget_alert と data.aws_caller_identity.current は別で定義してる前提
budget側には depends_on = [aws_sns_topic_policy.budget_alert] を付けて、ポリシーが先に効くようにしておいた。

結論：SNSに飛ばすならTopic Policyはセットで書く

2. キルスイッチLambda（フラグを立てる）

SNSから呼ばれて、DynamoDBのフラグ用アイテム（pk=SWITCH, sk=GLOBAL）を locked=true にするだけ
核はここだけ

def lambda_handler(event, context):
    # SNSの中身はログに残すだけ（後で何が起きたか見たいので）
    ...
    now = datetime.now(timezone.utc).isoformat()
    _ddb.update_item(
        TableName=TABLE,
        Key={"pk": {"S": "SWITCH"}, "sk": {"S": "GLOBAL"}},
        UpdateExpression="SET locked = :t, reason = :r, updated_at = :u",
        ExpressionAttributeValues={
            ":t": {"BOOL": True},
            ":r": {"S": "budget exceeded"},
            ":u": {"S": now},
        },
    )
    return {"ok": True}

（TABLE や _ddb、import は省略してるので、このままコピペしても動かないです）

権限は dynamodb:UpdateItem をこのテーブル1つにだけ。SNSから起動できるように aws_lambda_permission とサブスクリプションも定義してる

結論：やることが小さいので権限も小さくできた

3. アプリ側でフラグを見る（503と429）

ここが一番考えたところ
アプリ側Lambdaは Bedrock を呼ぶ前に、1つの transact_write_items でまとめて見てる

  • [0] キルスイッチの ConditionCheck
  • [1] 日次カウンタ
  • [2] 月次カウンタ

どれか1つでも条件違反ならトランザクションごと巻き戻るので、止まってる時とか上限に達してる時はカウンタも進まない。

別々に呼ぶとDynamoDBへの往復が増えるし、間で状態がズレうる
ConditionCheck は書き込まずに条件だけ見るオペレーションなので、レート制限のトランザクションに相乗りさせるのがちょうどよかった。

抜粋するとこんな感じ（日次・月次カウンタの中身は省略）

try:
    _ddb.transact_write_items(
        TransactItems=[
            # [0] キルスイッチが立ってないか見るだけ
            {
                "ConditionCheck": {
                    "TableName": RATE_LIMIT_TABLE,
                    "Key": {"pk": {"S": "SWITCH"}, "sk": {"S": "GLOBAL"}},
                    "ConditionExpression":
                        "attribute_not_exists(locked) OR locked = :false",
                    "ExpressionAttributeValues": {":false": {"BOOL": False}},
                }
            },
            # [1] 日次カウンタ、[2] 月次カウンタ ... 省略
        ]
    )
    return True, ""
except ClientError as e:
    code = e.response.get("Error", {}).get("Code", "")
    if code == "TransactionCanceledException":
        reasons = e.response.get("CancellationReasons", []) or []
        codes = [r.get("Code") for r in reasons]
        # 渡した順に返ってくるので、indexでどこで落ちたか分かる
        if len(codes) >= 1 and codes[0] == "ConditionalCheckFailed":
            return False, "killswitch"
        if len(codes) >= 2 and codes[1] == "ConditionalCheckFailed":
            return False, "daily"
        if len(codes) >= 3 and codes[2] == "ConditionalCheckFailed":
            return False, "monthly"
        return False, "unknown"
    raise

CancellationReasons は渡した TransactItems と同じ順番で返ってくるので、indexを見ればキルスイッチか日次か月次か分かる

ハンドラ側では、この関数（_consume_quota）が返す scope が killswitch なら503（service temporarily disabled）、それ以外は429（rate limit exceeded と上限値）を返してる（トランザクションの競合とかで unknown になった時も429）
クライアント側で「自分が叩きすぎた」のか「サービス全体が止まってる」のか区別できるのでいい感じ。

結論：キルスイッチ判定はレート制限と同じトランザクションに混ぜる

4. 月初の自動リセット

止まったまま放置するとサービスが死んだままになるので、EventBridgeのcronで月初に戻すようにした

# UTCの毎月1日 0:00 = JSTだと毎月1日 9:00
resource "aws_cloudwatch_event_rule" "monthly_reset" {
  name                = "${local.name}-monthly-reset"
  schedule_expression = "cron(0 0 1 * ? *)"
}

ルールだけ作っても何も起きないので、aws_cloudwatch_event_target でreset Lambdaにつないで、aws_lambda_permission（principal events.amazonaws.com）で起動を許可するところまでセットで書く
（ここ抜けると月初にリセットされないまま、ってなる）

呼ばれるreset Lambdaは、さっきの update_item で locked を false、reason を monthly auto reset にするだけ
予算は月間なので、月が替わったらフラグも戻す、っていう素直なやつ。

EventBridgeのルールのcronはUTC基準なので、JSTのつもりで書くとズレる（EventBridge Schedulerならタイムゾーン指定できるらしい）
Budgetsの月の区切りもUTCの0時なので、リセットのタイミングとはちょうど揃う。

結論：戻し忘れ防止は自動化しておく

よかった点

  • 部品が小さい。Lambda数本、DynamoDBオンデマンド、SNS、EventBridgeだけ
  • カウンタは古いものがTTLで勝手に消えるので掃除がいらない

注意点

  • Budgetsの判定はリアルタイムじゃない。公式だとBudgetsの情報の更新は1日最大3回、前回から8〜12時間おき。なので100%を超えてから止まるまで、1日くらい（それ以上も）かかることがある
  • その間のリクエストは止められないので、普段効いてるのは毎リクエスト見てるレート制限のほう
  • 100%の実績通知は1つの予算期間（月）で1回しか来ないらしい。なので月の途中で手でリセットすると、その月はもうキルスイッチが動かない。翌月まではレート制限だけが頼り
  • 前月分の通知が月をまたいで届いたら、リセット直後にまた止まるかも（キルスイッチLambdaは通知がどの月の分か見てないので）。実際に起きるかは未確認
  • 月初リセットは原因を直してる前提。バグや悪用が残ったまま戻ると、また同じことが起きる。くり返すなら EventBridge のルールを一時的に無効にするとか、予算額（budget_limit_usd）を下げてから戻す

感想

仕組み自体はシンプルだけど、遅延のせいで「設定額ぴったりで止まる」わけじゃないのがもどかしい・・・
80%で機能を絞って100%で完全停止、みたいな段階的な落とし方もできたらいいなと思ってる
従量課金が怖い人の参考になれば。

https://docs.aws.amazon.com/cost-management/latest/userguide/budgets-sns-policy.html
https://docs.aws.amazon.com/cost-management/latest/userguide/budgets-managing-costs.html
https://docs.aws.amazon.com/cost-management/latest/userguide/budgets-controls.html

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?