1# preflight_bedrock.py
2"""Bedrock fail-fast preflight: authorization, region/model availability, quota.
3
4Usage: python preflight_bedrock.py --region <r> --models <id,id,...> [--dataset-size N]
5Prints a JSON verdict to stdout; exit 0 if all models pass, 1 otherwise.
6Always prints JSON — including for missing credentials / bad region — so the
7caller can parse the verdict instead of a traceback. On failure the top level
8carries `reason`/`detail`/`failing_models` lifted from the first failing model.
9Chat models are probed via Converse; embedding models (which don't speak
10Converse) via InvokeModel with their family's request body. Bare OpenAI
11proprietary GPT ids are mantle-served and probed via the OpenAI Responses API;
12GPT-5.6 CRIS profile ids (us./in./global. prefixed) are bedrock-runtime targets
13and take the normal Converse probe.

preflight_bedrock.py on GitHub

The 1-token probe costs a fraction of a cent (noted in output).

"""

import argparse, json, sys

def classify_invoke_error(code: str, message: str) -> dict:

"""Pure: map a botocore error code to a structured preflight verdict."""

if code in ("AccessDeniedException",):

# Bedrock raises AccessDeniedException for two distinct problems:

# (a) model access not enabled in the Bedrock console (message mentions

# model access / "use the model"), fixed in the console, not IAM;

# (b) the IAM principal lacks bedrock:InvokeModel, fixed in IAM.

# Sending the user to the wrong fix wastes their time — disambiguate

# on the message text.

lowered = message.lower()

if ("model access" in lowered or "access to the model" in lowered

or "use the model" in lowered or "model is not" in lowered):

return {"ok": False, "reason": "model_access",

"detail": f"Bedrock model access not enabled for this model — {message}. "

f"Enable it in the Bedrock console (Model access page); "

f"this is separate from IAM."}

return {"ok": False, "reason": "authz",

"detail": f"IAM denies bedrock:InvokeModel — {message}. "

f"Grant bedrock:InvokeModel to this principal."}

if code in ("ValidationException", "ResourceNotFoundException"):

return {"ok": False, "reason": "model_unavailable",

"detail": f"Model not available in this region — {message}. "

f"Try a cross-region inference profile (e.g. us.<model-id>)."}

if code in ("ThrottlingException", "ServiceQuotaExceededException"):

# We got far enough to be throttled => we are authorized.

return {"ok": True, "reason": "throttled_ok",

"detail": "Authorized (probe throttled, which still proves access)."}

if code in ("UnrecognizedClientException", "InvalidSignatureException",

"ExpiredTokenException", "ExpiredToken"):

return {"ok": False, "reason": "credentials",

"detail": f"AWS credentials invalid or expired — {message}. "

f"Run 'aws configure', refresh your SSO session, or set AWS_PROFILE."}

return {"ok": False, "reason": "unknown", "detail": f"{code}: {message}"}

def is_embedding_model(model_id: str) -> bool:

"""Pure: embedding models don't speak the Converse API."""

return "embed" in model_id.lower()

def is_mantle_model(model_id: str) -> bool:

"""Pure: ids that must be probed via the mantle Responses API rather than

Converse/InvokeModel. Matches the BARE proprietary GPT ids (`openai.gpt-5*`),

which are mantle-served. Deliberately does NOT match:

- gpt-oss ids (Converse-capable on bedrock-runtime), and

- GPT-5.6 CRIS profile ids (`us.`/`in.`/`global.` prefixed) — those are

bedrock-runtime targets where Converse IS supported, so falling through to

the standard Converse probe is the correct behavior, not an accident."""

mid = model_id.lower()

return mid.startswith("openai.gpt-5") and "oss" not in mid

def classify_mantle_error(status: int | None, message: str) -> dict:

"""Pure: map an OpenAI-SDK HTTP status from the mantle endpoint to a verdict.

Mirrors classify_invoke_error's contract but for the mantle surface, where the

remedies differ: IAM needs bedrock-mantle:* actions (not bedrock:InvokeModel),

and mantle itself has no cross-region form — though for GPT-5.6 a

bedrock-runtime CRIS id can cover the region instead."""

if status in (401, 403):

lowered = message.lower()

if ("model access" in lowered or "access to the model" in lowered

or "use the model" in lowered or "not authorized to use" in lowered):

return {"ok": False, "reason": "model_access",

"detail": f"Bedrock model access not enabled for this model — {message}. "

f"Enable it in the Bedrock console (Model access page); "

f"this is separate from IAM."}

return {"ok": False, "reason": "authz",

"detail": f"IAM denies mantle inference — {message}. Grant bedrock-mantle "

f"actions (e.g. the AmazonBedrockMantleInferenceAccess managed "

f"policy); bedrock:InvokeModel does NOT authorize these models."}

if status == 404:

return {"ok": False, "reason": "model_unavailable",

"detail": f"Model not available at this mantle endpoint/region — {message}. "

f"Mantle is in-region only; for GPT-5.6 try a bedrock-runtime CRIS "

f"id (us./in./global. prefixed) instead, for GPT-5.5/5.4 switch to "

f"a supported region or a different model."}

if status == 429:

# Reaching a token-per-minute ceiling still proves we are authorized.

return {"ok": True, "reason": "throttled_ok",

"detail": "Authorized (probe throttled on a TPM quota, which still proves access)."}

return {"ok": False, "reason": "unknown", "detail": f"{status}: {message}"}

def probe_mantle_model(model_id: str, region: str) -> dict:

"""Real minimal probe against the mantle Responses API.

Fails CLOSED when the OpenAI SDK is missing. This deliberately does NOT follow

the `embedding_unprobed` precedent: that one passes because an unrecognized

embedding family is a rare edge case, whereas a missing SDK here would be the

normal path for every run if the dependency were absent — returning ok=True

would turn a fail-fast preflight into an unconditional green light and let the

migration proceed with endpoint, model, and IAM access never checked. Both

packages are declared in pyproject.toml, so reaching this branch means the

pinned environment is broken and the operator needs to know."""

try:

from aws_bedrock_token_generator import provide_token

from openai import BedrockOpenAI

except ImportError as e:

return {"ok": False, "reason": "mantle_deps_missing",

"detail": f"Cannot probe mantle-only model {model_id}: {e}. This needs "

f"'openai>=2.45.0' and 'aws-bedrock-token-generator', both declared "

f"in scripts/pyproject.toml — re-sync the pinned environment "

f"(`uv sync --project <scripts dir>`). Access was NOT verified; "

f"preflight fails closed rather than assuming it works."}

try:

client = BedrockOpenAI(

aws_region=region,

bedrock_token_provider=lambda: provide_token(region=region),

max_retries=0,

)

client.responses.create(model=model_id, input="ping", max_output_tokens=16, store=False)

return {"ok": True, "reason": "ok", "detail": "Mantle Responses API authorized."}

except Exception as e: # noqa: BLE001 - SDK raises many types; status is what matters

status = getattr(e, "status_code", None)

if status is None:

resp = getattr(e, "response", None)

status = getattr(resp, "status_code", None)

if status is None:

return {"ok": False, "reason": "credentials",

"detail": f"{type(e).__name__}: {e}. Check AWS credentials, the region "

f"name, and network access to the bedrock-mantle endpoint."}

return classify_mantle_error(int(status), str(e))

def _embed_request_body(model_id: str) -> dict | None:

"""Pure: minimal valid request body per embedding-model family; None if unknown."""

parts = model_id.split(".")

vendor = parts[1] if parts[0] in ("us", "eu", "apac", "global") and len(parts) > 1 else parts[0]

if vendor == "amazon":

return {"inputText": "ping"}

if vendor == "cohere":

return {"texts": ["ping"], "input_type": "search_document"}

return None

def probe_model(client, model_id: str) -> dict:

"""Real minimal probe: Converse for chat models, InvokeModel for embeddings."""

from botocore.exceptions import BotoCoreError, ClientError

try:

if is_embedding_model(model_id):

body = _embed_request_body(model_id)

if body is None:

# Unknown embedding family — probing with a wrong body would

# produce a ValidationException indistinguishable from a real

# availability problem. Pass with an explicit caveat instead.

return {"ok": True, "reason": "embedding_unprobed",

"detail": "Embedding model from an unrecognized family — access not "

"verified by preflight; confirm in the Bedrock console."}

client.invoke_model(modelId=model_id, body=json.dumps(body),

contentType="application/json", accept="application/json")

return {"ok": True, "reason": "ok", "detail": "InvokeModel (embedding) authorized."}

client.converse(

modelId=model_id,

messages=[{"role": "user", "content": [{"text": "ping"}]}],

inferenceConfig={"maxTokens": 1},

)

return {"ok": True, "reason": "ok", "detail": "InvokeModel authorized."}

except ClientError as e:

code = e.response.get("Error", {}).get("Code", "Unknown")

msg = e.response.get("Error", {}).get("Message", str(e))

return classify_invoke_error(code, msg)

except BotoCoreError as e:

# NoCredentialsError, EndpointConnectionError, SSO token errors, etc.

# These are config problems on the caller's machine, not Bedrock verdicts.

return {"ok": False, "reason": "credentials",

"detail": f"{type(e).__name__}: {e}. "

f"Run 'aws configure', refresh your SSO session, or check the region name."}

def aggregate_failure(results: list) -> dict:

"""Pure: lift the first failing model's reason/detail to the top level so the

orchestrator can branch on a single top-level `reason` (its documented contract)."""

failing = [r for r in results if not r["ok"]]

if not failing:

return {}

return {"reason": failing[0]["reason"], "detail": failing[0]["detail"],

"failing_models": [r["model_id"] for r in failing]}

def quota_rpm(quotas: list, model_id: str) -> int | None:

"""Best-effort on-demand requests-per-minute quota for this model from a

pre-fetched quota list; None if no match. Quota names follow the form

'On-demand model inference requests per minute for <Model Display Name>',

so we match per-minute inference quotas whose name shares a token with the

model id (e.g. 'claude', 'nova', 'titan')."""

tokens = [t for t in model_id.lower().replace(":", ".").split(".") if t]

name_tokens = set()

for t in tokens:

name_tokens.update(p for p in t.split("-") if p and not p.isdigit())

lowest = None

for q in quotas:

name = q.get("QuotaName", "").lower()

if "per minute" not in name or "request" not in name:

continue

if not any(tok in name for tok in name_tokens):

continue

v = int(q.get("Value", 0))

lowest = v if lowest is None else min(lowest, v)

return lowest

def fetch_bedrock_quotas(region: str) -> list:

"""Fetch all Bedrock service quotas once; empty list on any failure."""

import boto3

from botocore.exceptions import BotoCoreError, ClientError

try:

sq = boto3.client("service-quotas", region_name=region)

quotas = []

for page in sq.get_paginator("list_service_quotas").paginate(ServiceCode="bedrock"):

quotas.extend(page.get("Quotas", []))

return quotas

except (BotoCoreError, ClientError):

return []

def main(argv=None) -> int:

ap = argparse.ArgumentParser()

ap.add_argument("--region", required=True)

ap.add_argument("--models", required=True, help="comma-separated model ids")

ap.add_argument("--dataset-size", type=int, default=0)

args = ap.parse_args(argv)

model_ids = [m.strip() for m in args.models.split(",") if m.strip()]

if not model_ids:

print(json.dumps({"ok": False, "region": args.region, "models": [],

"reason": "no_models",

"detail": "--models resolved to an empty list; nothing to probe."},

indent=2))

return 1

import boto3

from botocore.exceptions import BotoCoreError

try:

client = boto3.client("bedrock-runtime", region_name=args.region)

except BotoCoreError as e:

print(json.dumps({"ok": False, "region": args.region, "models": [],

"reason": "credentials",

"detail": f"{type(e).__name__}: {e}"}, indent=2))

return 1

quotas = fetch_bedrock_quotas(args.region)

results = []

all_ok = True

for model_id in model_ids:

if is_mantle_model(model_id):

verdict = probe_mantle_model(model_id, args.region)

verdict["model_id"] = model_id

# Mantle enforces per-model input/output TPM quotas and has no RPM

# quota, so an RPM-derived pacing warning would be meaningless here.

verdict["rpm_quota"] = None

verdict["quota_note"] = (

"Mantle quotas are per-model input/output tokens per minute; there is no RPM "

"quota. Pace on token throughput, and note that prompt-cached input tokens are "

"exempt from the input-TPM quota.")

else:

verdict = probe_model(client, model_id)

rpm = quota_rpm(quotas, model_id)

verdict["model_id"] = model_id

verdict["rpm_quota"] = rpm

if rpm is not None and args.dataset_size > rpm:

verdict["quota_warning"] = (

f"Dataset ({args.dataset_size}) exceeds ~{rpm} RPM quota — "

f"Eval will pace with backoff and may be slow.")

all_ok = all_ok and verdict["ok"]

results.append(verdict)

out = {"ok": all_ok, "region": args.region,

"probe_cost_note": "1-token InvokeModel probe per model (~$0.00001 each)",

"models": results}

out.update(aggregate_failure(results))

print(json.dumps(out, indent=2))

return 0 if all_ok else 1

if __name__ == "__main__":

sys.exit(main())