Bedrock の利用額に上限を設けて超過ユーザーを SCP で制限してみた

"sub": "a1b2c3d4-5678-90ab-cdef-1234567890ab"

aws sts assume-role-with-web-identity \
  --role-arn arn:aws:iam::111122223333:role/ClaudeCodeBedrockUser \
  --role-session-name a1b2c3d4-5678-90ab-cdef-1234567890ab \
  --web-identity-token "$ID_TOKEN"

"Condition": {
  "StringEquals": {
    "sts:RoleSessionName": "${<issuer>:sub}"
  }
}

SELECT
  regexp_extract(line_item_iam_principal, '[^/]+$') AS sub,
  sum(line_item_unblended_cost)                     AS cost_usd
FROM cur2_table
WHERE billing_period = date_format(current_date, '%Y-%m')
  AND line_item_line_item_type = 'Usage'
  AND line_item_usage_type LIKE '%TokenCount%'
  AND regexp_like(line_item_iam_principal,
        ':assumed-role/BedrockUserRole/[0-9a-f-]{36}$')
GROUP BY 1
HAVING sum(line_item_unblended_cost) > 100   -- 月の上限（USD）

{
  "Effect": "Deny",
  "Action": [
    "bedrock:InvokeModel",
    "bedrock:InvokeModelWithResponseStream",
    "bedrock:Converse",
    "bedrock:ConverseStream"
  ],
  "Resource": "*",
  "Condition": {
    "StringLike": {
      "aws:userid": ["*:a1b2c3d4-5678-90ab-cdef-1234567890ab"]
    }
  }
}

import json
import os
import time
import boto3

athena = boto3.client("athena")
org = boto3.client("organizations")

POLICY_ID = os.environ["POLICY_ID"]
DB = os.environ.get("ATHENA_DB", "cur2")
TABLE = os.environ.get("ATHENA_TABLE", "cur2.test_export_bedrock")
ROLE = os.environ.get("ROLE_NAME", "ClaudeCodeBedrockUser")
OUTPUT = os.environ["ATHENA_OUTPUT"]

PLACEHOLDER = "__placeholder__"
SCP_MAX = 10240

ACTIONS = [
    "bedrock:InvokeModel",
    "bedrock:InvokeModelWithResponseStream",
    "bedrock:Converse",
    "bedrock:ConverseStream"
]

QUERY = """
SELECT
    element_at(split(line_item_iam_principal, '/'), -1) AS user_key,
    round(sum(line_item_unblended_cost), 6) AS cost_usd
FROM {table}
WHERE billing_period = date_format(current_date, '%Y-%m')
  AND line_item_line_item_type = 'Usage'
  AND line_item_usage_type LIKE '%TokenCount%'
  AND line_item_iam_principal LIKE 'arn:aws:sts::%:assumed-role/{role}/%'
  AND regexp_like(
      element_at(split(line_item_iam_principal, '/'), -1),
      '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
  )
GROUP BY 1
"""

def run_query():
    """Athenaで集計し、全ページの結果を取得する。"""
    qid = athena.start_query_execution(
        QueryString=QUERY.format(
            table=TABLE,
            role=ROLE,
        ),
        QueryExecutionContext={
            "Database": DB,
        },
        ResultConfiguration={
            "OutputLocation": OUTPUT,
        },
    )["QueryExecutionId"]

    for _ in range(60):
        status = athena.get_query_execution(
            QueryExecutionId=qid,
        )["QueryExecution"]["Status"]

        if status["State"] in ("SUCCEEDED", "FAILED", "CANCELLED"):
            break

        time.sleep(2)

    if status["State"] != "SUCCEEDED":
        raise RuntimeError(
            f"Athena {status['State']}: "
            f"{status.get('StateChangeReason')}"
        )

    costs = {}
    next_token = None
    first_page = True

    while True:
        params = {
            "QueryExecutionId": qid,
            "MaxResults": 1000,
        }

        if next_token:
            params["NextToken"] = next_token

        response = athena.get_query_results(**params)
        rows = response["ResultSet"]["Rows"]

        # 列名行があるのは最初のページだけ。
        if first_page:
            rows = rows[1:]
            first_page = False

        for row in rows:
            cells = [
                cell.get("VarCharValue")
                for cell in row["Data"]
            ]

            if cells[0]:
                costs[cells[0]] = float(cells[1] or 0)

        next_token = response.get("NextToken")

        if not next_token:
            break

    return costs, qid

def scp_content(subs):
    """拒否対象のsub一覧からSCP全体のJSONを生成する。"""
    return json.dumps(
        {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Sid": "DenyBedrockForOverQuotaUsers",
                    "Effect": "Deny",
                    "Action": ACTIONS,
                    "Resource": "*",
                    "Condition": {
                        "StringLike": {
                            "aws:userid": [
                                f"*:{sub}"
                                for sub in subs
                            ],
                        },
                    },
                },
            ],
        },
        separators=(",", ":"),
    )

def current_deny_list():
    """現在のSCPから、"*:<sub>"形式の拒否対象を取り出す。"""
    response = org.describe_policy(
        PolicyId=POLICY_ID,
    )
    document = json.loads(
        response["Policy"]["Content"]
    )

    for statement in document.get("Statement", []):
        value = (
            statement.get("Condition", {})
            .get("StringLike", {})
            .get("aws:userid")
        )

        if value:
            values = value if isinstance(value, list) else [value]

            return [
                item.split(":", 1)[1]
                for item in values
                if ":" in item
            ]

    return []

def lambda_handler(event, context):
    thresholds = json.loads(
        os.environ.get(
            "THRESHOLDS",
            '{"default_usd": 0.002}',
        )
    )

    default_threshold = float(
        thresholds.get("default_usd", 0.002)
    )

    costs, qid = run_query()

    # しきい値を超えたユーザーを拒否対象にする。
    over_quota = sorted(
        user_key
        for user_key, cost in costs.items()
        if cost > default_threshold
    )

    before = sorted(
        user_key
        for user_key in current_deny_list()
        if user_key != PLACEHOLDER
    )

    added = [
        user_key
        for user_key in over_quota
        if user_key not in before
    ]
    removed = [
        user_key
        for user_key in before
        if user_key not in over_quota
    ]

    # 拒否対象がいない場合は、空リストの代わりにダミー値を設定する。
    desired = over_quota or [PLACEHOLDER]
    content = scp_content(desired)

    result = {
        "query_execution_id": qid,
        "users_seen": len(costs),
        "costs": costs,
        "threshold_default_usd": default_threshold,
        "over_quota": over_quota,
        "deny_list_before": before,
        "added": added,
        "removed": removed,
        "scp_bytes": len(content),
        "updated": False,
    }

    if len(content) > SCP_MAX:
        result["error"] = (
            f"SCPが上限を超える"
            f"（{len(content)} > {SCP_MAX}）。更新しない"
        )
        print(json.dumps(result, ensure_ascii=False))
        return result

    if not added and not removed:
        result["note"] = "変更なし"
    else:
        # 今回は、指定したSCPの本文を丸ごと更新する。
        org.update_policy(
            PolicyId=POLICY_ID,
            Content=content,
        )
        result["updated"] = True

    print(json.dumps(result, ensure_ascii=False))
    return result

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "DenyBedrockForOverQuotaUsers",
      "Effect": "Deny",
      "Action": [
        "bedrock:InvokeModel",
        "bedrock:InvokeModelWithResponseStream",
        "bedrock:Converse",
        "bedrock:ConverseStream"
      ],
      "Resource": "*",
      "Condition": {
        "StringLike": {
          "aws:userid": [
            "*:a1b2c3d4-5678-90ab-cdef-1234567890ab"
          ]
        }
      }
    }
  ]
}

An error occurred (AccessDeniedException) when calling the Converse operation:
User: arn:aws:sts::111122223333:assumed-role/BedrockUserRole/a1b2c3d4-5678-90ab-cdef-1234567890ab
is not authorized to perform: bedrock:InvokeModel ...
with an explicit deny in a service control policy