Setting the file. One moment.
Deploy Model · Amazon DynamoDB · aws/agent-toolkit-for-aws · Skills Docs
ContentsBack to the top of the page Reference Architecture
93
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Creating API Gateway Stage
def _build_create_kwargs
— line 393
This file
Number 65.10
Position 10 of 14
Type Python
Size 47 KB
Lines 1,106 scripts/ deploy_model.py
Python · 1,106 lines · 47 KB
16 These defaults match SKILL.md Data modeling #3's own logic for ephemeral /
17 cache-like tables. Production designs still get Best Practices defaults —
18 this is a tooling choice for the benchmark, not a rule change.
19
20 Safety contract:
21 - Refuses unless --yes-deploy is passed.
22 - Aborts on caller-identity ARN / alias containing prod, production, prd, live.
23 - On an interactive terminal, prints the target account and waits a few
24 seconds for a Ctrl-C abort before creating anything; non-interactive
25 (agent/CI) runs proceed immediately on the --yes-deploy consent.
26 - resource_prefix is optional: if absent it is auto-generated as
27 ddb-skill-bench-<date>-<uuid8>; if supplied it MUST start with
28 ddb-skill-bench- (so teardown can scope deletions) or the deploy refuses.
29 - Refuses if boto3 is missing, credentials are missing/expired,
30 access_patterns is empty, any pattern has missing/zero peak_rps, or any
31 table is missing key_schema.
32
33 Usage:
34 python3 deploy_model.py \\
35 --model dynamodb_data_model.json \\
36 --config benchmark_config.json \\
37 --manifest-out created_resources.json \\
38 --yes-deploy
39 """
40 from __future__ import annotations
41
42 import argparse
43 import datetime as _dt
44 import io
45 import json
46 import sys
47 import time
48 import uuid
49 import zipfile
50 from concurrent.futures import ThreadPoolExecutor, as_completed
51 from pathlib import Path
52
53 PROD_MARKERS = ( "prod" , "production" , "prd" , "live" )
54 REQUIRED_PREFIX = "ddb-skill-bench-"
55 TYPE_MAP = { "S" : "S" , "N" : "N" , "B" : "B" }
56
57
58 # Lambda role trust policy — the only principal allowed to assume is the Lambda
59 # service itself, and only when the source account matches the deploying account
60 # (aws:SourceAccount guards against confused-deputy assumption from other
61 # accounts). The account is known from the preflight identity check, so we build
62 # the document at role-creation time rather than hardcoding it.
63 def _build_lambda_trust_policy (account: str ) -> dict :
64 return {
65 "Version" : "2012-10-17" ,
66 "Statement" : [
67 {
68 "Effect" : "Allow" ,
69 "Principal" : { "Service" : "lambda.amazonaws.com" },
70 "Action" : "sts:AssumeRole" ,
71 "Condition" : { "StringEquals" : { "aws:SourceAccount" : str (account)}},
72 }
73 ],
74 }
75
76
77 # 1769 MB is the threshold where Lambda allocates a full vCPU. The benchmark
78 # driver runs up to concurrency_per_pattern (default 32) I/O-bound threads; a
79 # full vCPU gives headroom for the GIL-bound boto3/JSON work between round trips
80 # so the open-loop scheduler can sustain the higher single-Lambda rps ceiling.
81 # Cost impact is negligible (a benchmark runs minutes). Override via config.
82 DEFAULT_LAMBDA_MEMORY_MB = 1769
83 DEFAULT_LAMBDA_TIMEOUT_S = 900
84 LAMBDA_RUNTIME = "python3.12"
85 LAMBDA_HANDLER = "benchmark.handler"
86
87 # Brief abort window shown ONLY on an interactive terminal, after the
88 # caller-identity banner and before any resource is created. Non-interactive
89 # (agent/CI) runs skip it — consent was already given via --yes-deploy.
90 _ABORT_WINDOW_S = 5
91
92
93 def _die (msg: str , code: int = 2 ) -> None :
94 print ( f "ERROR: { msg } " , file = sys.stderr)
95 sys.exit(code)
96
97
98 def _require_boto3 ():
99 try :
100 import boto3 # noqa: F401
101 from botocore.exceptions import ClientError # noqa: F401
102
103 return boto3
104 except ImportError :
105 _die( "boto3 not installed. Run: pip install boto3>=1.34" )
106
107
108 def _load_json (path: Path) -> dict :
109 if not path.exists():
110 _die( f "file not found: { path } " )
111 with path.open() as f:
112 return json.load(f)
113
114
115 def _validate_design (model: dict ) -> None :
116 if not model.get( "tables" ):
117 _die( "design JSON has no tables" )
118 aps = model.get( "access_patterns" ) or []
119 if not aps:
120 _die(
121 "design JSON has no access_patterns — refusing (Mechanics #2: "
122 "unknown RPS is a design gap, not a benchmark input)."
123 )
124 for ap in aps:
125 if not ap.get( "peak_rps" ):
126 _die(
127 f "access pattern { ap.get( 'pattern_id' , '?' ) } has missing or "
128 "zero peak_rps — refusing per Mechanics #2."
129 )
130 for t in model[ "tables" ]:
131 ks = t.get( "key_schema" ) or {}
132 if not ks.get( "partition_key" ):
133 _die(
134 f "table { t.get( 'table_name' , '?' ) } has no "
135 "key_schema.partition_key. Live deploy requires an explicit "
136 '{"key_schema": {"partition_key": "<attr>", "sort_key": '
137 '"<attr>?" }} block per table. Add it to the JSON and re-run. '
138 "(The cost calculator does not require this; live deploy does.)"
139 )
140 _validate_pattern_refs(model)
141
142
143 def _validate_pattern_refs (model: dict ) -> None :
144 """Every access pattern must point at a table that exists, and any Query/Scan
145 `index` must name a GSI defined on that table. Catching this up front turns a
146 silent benchmark failure (the operation errors on 100% of calls, observed CU
147 reads as 0, and a naive report calls it cheap) into a clear, pre-deploy
148 refusal. This guards the structural-reference class of mistakes the real-AWS
149 run surfaced (Query on a missing GSI, pattern on a missing table)."""
150 tables_by_name = {t.get( "table_name" ): t for t in model.get( "tables" , [])}
151 for ap in model.get( "access_patterns" ) or []:
152 pid = ap.get( "pattern_id" , "?" )
153 tn = ap.get( "table" )
154 if not tn:
155 _die(
156 f 'access pattern { pid } has no "table" — every pattern must '
157 "name the table it runs against."
158 )
159 td = tables_by_name.get(tn)
160 if td is None :
161 _die(
162 f "access pattern { pid } references table { tn !r} , which is not "
163 f "defined in tables[]. Defined tables: "
164 f ' { sorted (tables_by_name) } . Fix the "table" field or add the '
165 "table."
166 )
167 idx = ap.get( "index" )
168 op = ap.get( "operation" )
169 vec_names = {v.get( "index_name" ) for v in ((td or {}).get( "vector_indexes" ) or [])}
170 gsi_names = {g.get( "index_name" ) for g in ((td or {}).get( "gsis" ) or [])}
171
172 # SearchVectors reads a VECTOR index, never a GSI — and the reverse is also
173 # true, so each operation is checked against the right index family. Getting
174 # this wrong fails every call at runtime while the observed CU reads as 0,
175 # which a naive report then calls cheap.
176 if op == "SearchVectors" :
177 if not idx:
178 _die(
179 f 'access pattern { pid } is a SearchVectors but has no "index" — '
180 "it must name the vector index to search."
181 )
182 if idx not in vec_names:
183 _die(
184 f "access pattern { pid } searches vector index { idx !r} on table "
185 f " { tn !r} , but that table defines no such vector index. Defined "
186 f "vector indexes on { tn } : { sorted (n for n in vec_names if n) } . "
187 'Add it to vector_indexes[] or correct the "index" field.'
188 )
189 elif idx:
190 if idx in vec_names:
191 _die(
192 f "access pattern { pid } targets { idx !r} on table { tn !r} with "
193 f "operation { op !r} , but { idx !r} is a VECTOR index. Query and Scan "
194 "are rejected against a vector index with "
195 "'ValidationException: Query operation not supported on this index "
196 'type. \' Use operation "SearchVectors" to read it.'
197 )
198 if idx not in gsi_names:
199 _die(
200 f "access pattern { pid } uses index { idx !r} on table { tn !r} , "
201 f "but that table defines no such GSI. Defined GSIs on { tn } : "
202 f " { sorted (n for n in gsi_names if n) } . A Query/Scan against a "
203 "non-existent index fails every call at runtime — add the GSI "
204 'to the table or correct the "index" field.'
205 )
206
207
208 def _ensure_resource_prefix (cfg: dict , today: str ) -> bool :
209 """Fill in resource_prefix when the caller did not supply one.
210
211 The agent should NOT have to hand-build the prefix (date + uuid8) and get
212 its shape exactly right only to be refused — the script can generate a
213 valid, run-unique prefix itself. `today` is the YYYYMMDD already derived
214 from this run's clock so the prefix and the manifest's created_at agree.
215
216 Returns True if a prefix was auto-generated (so the caller can surface it),
217 False if the caller supplied one explicitly. An explicitly supplied prefix
218 is still validated for the required namespace by _validate_config.
219 """
220 if cfg.get( "resource_prefix" ):
221 return False
222 cfg[ "resource_prefix" ] = f " { REQUIRED_PREFIX }{ today } - { uuid.uuid4().hex[: 8 ] } "
223 return True
224
225
226 def _validate_config (cfg: dict ) -> None :
227 for k in ( "aws_profile" , "region" ):
228 if not cfg.get(k):
229 _die( f "benchmark_config.json missing required field: { k } " )
230 # resource_prefix is auto-generated by _ensure_resource_prefix when absent,
231 # so by the time we get here it is always set; we only police its namespace
232 # (which matters most for caller-supplied values).
233 if not cfg.get( "resource_prefix" ):
234 _die( "benchmark_config.json missing required field: resource_prefix" )
235 if not cfg[ "resource_prefix" ].startswith( REQUIRED_PREFIX ):
236 _die(
237 f "resource_prefix must start with { REQUIRED_PREFIX !r} so teardown "
238 f "can scope deletions safely. Got: { cfg[ 'resource_prefix' ] !r} . "
239 "Tip: omit resource_prefix entirely and the deploy generates a "
240 f "valid one ( { REQUIRED_PREFIX } <date>-<uuid8>) for you."
241 )
242
243
244 def _preflight (boto3_mod, cfg: dict , dry_run: bool = False ) -> dict :
245 """Run sts get-caller-identity, refuse on prod markers, return identity dict.
246
247 On an interactive terminal (and only for a real deploy, not a dry run) this
248 prints the target account and waits a brief window for a Ctrl-C abort before
249 returning. Non-interactive runs and dry runs return immediately.
250 """
251 from botocore.exceptions import (
252 ClientError,
253 EndpointConnectionError,
254 NoCredentialsError,
255 )
256
257 try :
258 session = boto3_mod.Session( profile_name = cfg[ "aws_profile" ], region_name = cfg[ "region" ])
259 sts = session.client( "sts" )
260 ident = sts.get_caller_identity()
261 except NoCredentialsError:
262 _die(
263 "AWS credentials not found for profile "
264 f " { cfg[ 'aws_profile' ] !r} . Try: aws sso login --profile "
265 f " { cfg[ 'aws_profile' ] } "
266 )
267 except ClientError as e:
268 code = e.response.get( "Error" , {}).get( "Code" , "Unknown" )
269 _die(
270 f "sts get-caller-identity failed ( { code } ): { e } . "
271 f "Credentials may be expired — try: aws sso login --profile "
272 f " { cfg[ 'aws_profile' ] } "
273 )
274 except EndpointConnectionError as e:
275 _die( f "cannot reach AWS endpoint: { e } " )
276 except Exception as e:
277 _die( f "AWS credential check failed: { type (e). __name__ } : { e } " )
278
279 arn = ident.get( "Arn" , "" )
280 account = ident.get( "Account" , "" )
281 account_alias = ""
282 try :
283 iam = session.client( "iam" )
284 aliases = iam.list_account_aliases().get( "AccountAliases" , [])
285 account_alias = aliases[ 0 ] if aliases else ""
286 except Exception :
287 # list_account_aliases can fail with no permission; prod-check still
288 # works off the ARN alone.
289 pass
290
291 lower = f " { arn } { account_alias } " .lower()
292 matched = [m for m in PROD_MARKERS if m in lower]
293 if matched:
294 _die(
295 f "caller identity appears to be production ( { matched } ). "
296 f "ARN: { arn } . Alias: { account_alias or '<none>' } . "
297 "REFUSING to deploy benchmark resources against a prod account."
298 )
299
300 bar = "#" * 72
301 print (bar)
302 print ( "# WARNING — this will create REAL AWS resources in the account shown" )
303 print ( "# below. A benchmark run typically costs single-digit cents but" )
304 print ( "# involves:" )
305 print ( "# - Multiple DynamoDB tables and GSIs with live capacity." )
306 print ( "# - A Lambda function and IAM role." )
307 print ( "# - Seed + warmup + measurement traffic (hundreds to thousands" )
308 print ( "# of reads/writes)." )
309 print ( "#" )
310 print ( "# USE AN AWS ACCOUNT DEDICATED TO TESTING." )
311 print ( "# Do NOT run this against production or any account holding real" )
312 print ( "# user data. The teardown script is generated separately and must" )
313 print ( "# be run by you." )
314 print (bar)
315 print ()
316 print ( "=" * 72 )
317 print ( "Caller identity confirmation:" )
318 print ( f " Account: { account } " )
319 print ( f " Alias: { account_alias or '<none>' } " )
320 print ( f " ARN: { arn } " )
321 print ( f " Region: { cfg[ 'region' ] } " )
322 print ( "=" * 72 )
323 # Consent was already given via --yes-deploy (validated in main before we
324 # got here), so the deploy proceeds immediately. Only when this is an
325 # interactive terminal AND a real deploy do we offer a real, brief abort
326 # window — in an agent/CI run (no TTY) there is no human to press Ctrl-C, so
327 # we must not print a "cancel now" prompt that nobody can act on, and a dry
328 # run creates nothing so there is nothing to abort.
329 if dry_run:
330 return { "account" : account, "alias" : account_alias, "arn" : arn}
331 interactive = False
332 try :
333 interactive = sys.stdin.isatty() and sys.stdout.isatty()
334 except Exception :
335 interactive = False
336 if interactive:
337 print (
338 f " \n This account is shown above. Deploying in { _ABORT_WINDOW_S } s — "
339 "press Ctrl-C now to abort if it is NOT a testing account."
340 )
341 try :
342 time.sleep( _ABORT_WINDOW_S )
343 except KeyboardInterrupt :
344 _die( "aborted by user before any resource was created." , code = 130 )
345 else :
346 print (
347 " \n Proceeding (consent given via --yes-deploy; non-interactive run, "
348 "no abort prompt). Verify the account above is a testing account."
349 )
350 return { "account" : account, "alias" : account_alias, "arn" : arn}
351
352
353 def _collect_attr_types (tables: list[ dict ]) -> dict[ str , dict[ str , str ]]:
354 """Per-table map of attribute-name → DynamoDB type letter ("S"/"N"/"B").
355
356 Types are read from TWO sources, in increasing precedence:
357 1. entities[].attributes[] as {"name","type"} — the canonical schema form
358 documented in references/cost-model-schema.md.
359 2. a table-level "attribute_definitions" block — the raw-CreateTable-API
360 spelling an author (or LLM) naturally reaches for. Accepts BOTH
361 {"attribute_name","attribute_type"} (API style) and the {"name","type"}
362 shorthand. An explicit attribute_definitions entry WINS over an
363 entities-derived type for the same attribute.
364
365 A key attribute whose type is never declared falls through to "S" in
366 _build_create_kwargs. Declaring a numeric/binary key as anything other than
367 its true type (or leaving it to default to "S") makes CreateTable build the
368 wrong AttributeType, so writes of the real value fail with
369 ValidationException in production. scripts/benchmark_lambda.py builds the
370 identical map with the same precedence — keep the two in sync.
371 """
372
373 def _norm_t (v) -> str :
374 return TYPE_MAP .get((v or "S" ).upper(), "S" )
375
376 by_table: dict[ str , dict[ str , str ]] = {}
377 for t in tables:
378 tm = by_table.setdefault(t[ "table_name" ], {})
379 # 1. entities[].attributes[] (lower precedence)
380 for e in t.get( "entities" ) or []:
381 for a in e.get( "attributes" ) or []:
382 name = a.get( "name" )
383 if name:
384 tm[name] = _norm_t(a.get( "type" ))
385 # 2. table-level attribute_definitions (higher precedence)
386 for a in t.get( "attribute_definitions" ) or []:
387 name = a.get( "attribute_name" ) or a.get( "name" )
388 if name:
389 tm[name] = _norm_t(a.get( "attribute_type" ) or a.get( "type" ))
390 return by_table
391
392
393 def _build_create_kwargs (
394 table_def: dict ,
395 prefix: str ,
396 tags: dict ,
397 attr_types: dict[ str , str ],
398 _global_provisioned: dict | None = None ,
399 ) -> dict :
400 ks = table_def[ "key_schema" ]
401 pk = ks[ "partition_key" ]
402 sk = ks.get( "sort_key" )
403
404 referenced_attrs: dict[ str , str ] = {pk: attr_types.get(pk, "S" )}
405 if sk:
406 referenced_attrs[sk] = attr_types.get(sk, "S" )
407
408 key_schema = [{ "AttributeName" : pk, "KeyType" : "HASH" }]
409 if sk:
410 key_schema.append({ "AttributeName" : sk, "KeyType" : "RANGE" })
411
412 gsis = []
413 for g in table_def.get( "gsis" ) or []:
414 g_pk = g.get( "partition_key" )
415 g_sk = g.get( "sort_key" )
416 if not g_pk:
417 _die(
418 f "GSI { g.get( 'index_name' , '?' ) } on table "
419 f " { table_def[ 'table_name' ] } has no partition_key"
420 )
421 referenced_attrs[g_pk] = attr_types.get(g_pk, "S" )
422 g_key_schema = [{ "AttributeName" : g_pk, "KeyType" : "HASH" }]
423 if g_sk:
424 referenced_attrs[g_sk] = attr_types.get(g_sk, "S" )
425 g_key_schema.append({ "AttributeName" : g_sk, "KeyType" : "RANGE" })
426
427 proj = g.get( "projection" ) or { "type" : "ALL" }
428 p_type = (proj.get( "type" ) or "ALL" ).upper()
429 proj_kwargs: dict = { "ProjectionType" : p_type}
430 if p_type == "INCLUDE" :
431 # Accept any of the three spellings the rest of the toolchain reads
432 # (`attributes` is canonical per cost-model-schema.md; calculate_costs
433 # and iterate_design's fingerprint also accept `non_key_attributes` /
434 # `NonKeyAttributes`). Reading only `attributes` here would silently
435 # create the GSI with an EMPTY include list — projecting nothing —
436 # when the design used another spelling, while the calculator and
437 # fingerprint happily saw the attributes. Stay consistent.
438 include_attrs = (
439 proj.get( "attributes" )
440 or proj.get( "non_key_attributes" )
441 or proj.get( "NonKeyAttributes" )
442 or []
443 )
444 if not include_attrs:
445 _die(
446 f "GSI { g[ 'index_name' ] } on table { table_def[ 'table_name' ] } "
447 "has projection type INCLUDE but no projected attributes. "
448 'Add a non-empty "attributes" list to the projection, or '
449 'use "ALL"/"KEYS_ONLY".'
450 )
451 proj_kwargs[ "NonKeyAttributes" ] = list (include_attrs)
452
453 gsis.append(
454 {
455 "IndexName" : g[ "index_name" ],
456 "KeySchema" : g_key_schema,
457 "Projection" : proj_kwargs,
458 }
459 )
460
461 # Vector indexes. Two things bite here and both are hard failures at CreateTable
462 # rather than at write time, so they are handled before AttributeDefinitions is built:
463 # * every SearchSchema attribute must ALSO be declared in AttributeDefinitions,
464 # exactly as a GSI key attribute must ("One element in SearchSchema is not
465 # defined in attribute definitions"), and
466 # * vector indexes require on-demand capacity, so a provisioned bench run cannot
467 # carry one.
468 vector_indexes = []
469 for v in table_def.get( "vector_indexes" ) or []:
470 name = v.get( "index_name" )
471 attr = v.get( "vector_attribute" )
472 dims = v.get( "dimensions" )
473 fn = (v.get( "distance_function" ) or "" ).upper()
474 if not (name and attr and dims and fn):
475 _die(
476 f "vector index { name or '?' } on table { table_def[ 'table_name' ] } needs "
477 "index_name, vector_attribute, dimensions and distance_function"
478 )
479 if fn not in { "COSINE" , "EUCLIDEAN" , "DOT_PRODUCT" }:
480 _die(
481 f "vector index { name } : distance_function must be COSINE, EUCLIDEAN or "
482 f "DOT_PRODUCT, got { v.get( 'distance_function' ) !r} "
483 )
484 if not 1 <= int (dims) <= 4096 :
485 _die( f "vector index { name } : dimensions must be 1-4096, got { dims !r} " )
486
487 vproj = v.get( "projection" ) or { "type" : "ALL" }
488 # Accept the bare-string shorthand: "KEYS_ONLY" == {"type": "KEYS_ONLY"}. Writing the
489 # string is the natural mistake (it is how the API-level value reads), and without this
490 # the next line raised a bare AttributeError instead of anything actionable. Mirrors the
491 # attribute_definitions shorthand accepted elsewhere in these scripts.
492 if isinstance (vproj, str ):
493 vproj = { "type" : vproj}
494 vp_type = (vproj.get( "type" ) or "ALL" ).upper()
495 vproj_kwargs: dict = { "ProjectionType" : vp_type}
496 if vp_type == "INCLUDE" :
497 inc = (
498 vproj.get( "attributes" )
499 or vproj.get( "non_key_attributes" )
500 or vproj.get( "NonKeyAttributes" )
501 or []
502 )
503 if not inc:
504 _die(
505 f "vector index { name } has projection INCLUDE but no projected "
506 'attributes. Add a non-empty "attributes" list, or use '
507 '"ALL"/"KEYS_ONLY".'
508 )
509 vproj_kwargs[ "NonKeyAttributes" ] = list (inc)
510
511 spec: dict = {
512 "IndexName" : name,
513 "VectorAttribute" : { "AttributeName" : attr},
514 "Projection" : vproj_kwargs,
515 "Dimensions" : int (dims),
516 "DistanceFunction" : fn,
517 }
518
519 schema = v.get( "search_schema" ) or {}
520 elements = []
521 v_pk = schema.get( "partition_key" )
522 if v_pk:
523 referenced_attrs[v_pk] = attr_types.get(v_pk, "S" )
524 elements.append({ "AttributeName" : v_pk, "SearchSchemaElementType" : "HASH" })
525 for f in schema.get( "inline_filters" ) or []:
526 referenced_attrs[f] = attr_types.get(f, "S" )
527 elements.append({ "AttributeName" : f, "SearchSchemaElementType" : "INLINE_FILTER" })
528 if elements:
529 spec[ "SearchSchema" ] = elements
530 vector_indexes.append(spec)
531
532 attr_defs = [
533 { "AttributeName" : name, "AttributeType" : t} for name, t in sorted (referenced_attrs.items())
534 ]
535
536 kwargs: dict = {
537 "TableName" : f " { prefix }{ table_def[ 'table_name' ] } " ,
538 "AttributeDefinitions" : attr_defs,
539 "KeySchema" : key_schema,
540 "BillingMode" : "PAY_PER_REQUEST" ,
541 # Bench-only: deletion protection OFF (see module docstring).
542 "DeletionProtectionEnabled" : False ,
543 # Encryption at rest: DynamoDB ALWAYS encrypts at rest. The default is an
544 # AWS-owned key (no cost, no key management) — correct for ephemeral bench
545 # tables — and is the implicit state when no SSESpecification is sent. We
546 # intentionally do NOT pass SSESpecification here: DynamoDB's SSEType only
547 # accepts "KMS" (a customer-/AWS-managed CMK); there is no "AES256" SSEType
548 # on DynamoDB (that is an S3 spelling) and sending one is rejected with a
549 # ValidationException. The at-rest state is surfaced in the manifest from
550 # DescribeTable instead. Production designs that need a customer-managed CMK
551 # for audit/compliance set SSESpecification={"Enabled": true, "SSEType":
552 # "KMS", "KMSMasterKeyId": "<arn>"}; see SKILL.md "Security considerations".
553 "Tags" : [{ "Key" : k, "Value" : str (v)} for k, v in tags.items()],
554 }
555
556 # Optional PROVISIONED capacity for a deliberate capacity-ceiling test. On
557 # on-demand, adaptive capacity auto-scales the table and absorbs a hot key,
558 # so a single load generator rarely produces throttles. A low provisioned
559 # total imposes a HARD table-wide ceiling adaptive capacity cannot exceed —
560 # the reliable way to observe hot-partition throttling (Mechanics #3) and to
561 # validate a planned provisioned capacity (Mechanics #19). Set via the
562 # design's table.provisioned_capacity = {"read": N, "write": M}, or globally
563 # via benchmark_config.provisioned_capacity. Bench-only; production designs
564 # still default to on-demand.
565 prov = table_def.get( "provisioned_capacity" ) or _global_provisioned
566 if prov and vector_indexes:
567 _die(
568 f "table { table_def[ 'table_name' ] } declares a vector index AND provisioned "
569 "capacity. Vector indexes are supported only on on-demand "
570 "(PAY_PER_REQUEST) tables, so this table cannot be created as declared. "
571 "Remove provisioned_capacity for this table (or drop the vector index) and "
572 "re-run. Note a provisioned-ceiling throttle experiment is therefore not "
573 "available on a vector-indexed table."
574 )
575 if prov:
576 kwargs[ "BillingMode" ] = "PROVISIONED"
577 kwargs[ "ProvisionedThroughput" ] = {
578 "ReadCapacityUnits" : int (prov.get( "read" , 5 )),
579 "WriteCapacityUnits" : int (prov.get( "write" , 5 )),
580 }
581
582 if gsis:
583 if prov:
584 for g in gsis:
585 g[ "ProvisionedThroughput" ] = {
586 "ReadCapacityUnits" : int (prov.get( "read" , 5 )),
587 "WriteCapacityUnits" : int (prov.get( "write" , 5 )),
588 }
589 kwargs[ "GlobalSecondaryIndexes" ] = gsis
590
591 if vector_indexes:
592 # A single CreateTable may define several vector indexes (up to the per-table
593 # limit of 5), and creating them inline avoids the UpdateTable backfill path
594 # entirely — measured at 20s inline versus 8m33s for an add-to-existing-table on
595 # a six-item table.
596 if len (vector_indexes) > 5 :
597 _die(
598 f "table { table_def[ 'table_name' ] } declares { len (vector_indexes) } vector "
599 "indexes; the per-table limit is 5"
600 )
601 kwargs[ "VectorIndexes" ] = vector_indexes
602
603 streams = table_def.get( "streams" ) or {}
604 if streams.get( "enabled" ):
605 view = streams.get( "view_type" , "NEW_AND_OLD_IMAGES" )
606 kwargs[ "StreamSpecification" ] = {
607 "StreamEnabled" : True ,
608 "StreamViewType" : view,
609 }
610 return kwargs
611
612
613 def _create_table_one (client, kwargs: dict ) -> None :
614 """Issue CreateTable. Swallow ResourceInUseException as a rerun-safety."""
615 from botocore.exceptions import ClientError
616
617 try :
618 client.create_table( ** kwargs)
619 except ClientError as e:
620 code = e.response.get( "Error" , {}).get( "Code" , "" )
621 if code == "ResourceInUseException" :
622 print ( f " note: { kwargs[ 'TableName' ] } already exists — reusing." )
623 else :
624 raise
625
626
627 def _wait_table_active (client, table_name: str ) -> dict :
628 waiter = client.get_waiter( "table_exists" )
629 waiter.wait( TableName = table_name)
630 desc = client.describe_table( TableName = table_name)[ "Table" ]
631 if desc.get( "VectorIndexes" ):
632 desc = _wait_vector_indexes_ready(client, table_name)
633 return desc
634
635
636 def _wait_vector_indexes_ready (client, table_name: str , timeout_s: int = 1800 ) -> dict :
637 """Wait until every vector index is ACTIVE and not backfilling.
638
639 The `table_exists` waiter returns when the TABLE is ACTIVE and says nothing about a
640 vector index, so returning there hands the benchmark an index that rejects
641 SearchVectors. There is no `BACKFILLING` index status: the correct predicate is
642 `IndexStatus == "ACTIVE" and not Backfilling`, and it is the only one correct on both
643 creation paths — measured, the two report `Backfilling` differently:
644
645 inline CreateTable : CREATING with the Backfilling key ABSENT, then ACTIVE (~20s)
646 UpdateTable add : CREATING/Backfilling=false, then CREATING/Backfilling=true,
647 then ACTIVE (8m33s on a six-item table)
648
649 So "wait for the Backfilling key to disappear" returns immediately-but-wrongly on the
650 inline path, and "wait for Backfilling == false" returns minutes early on the other.
651 """
652 deadline = time.time() + timeout_s
653 announced = False
654 while True :
655 desc = client.describe_table( TableName = table_name)[ "Table" ]
656 vis = desc.get( "VectorIndexes" ) or []
657 pending = [v for v in vis if v.get( "IndexStatus" ) != "ACTIVE" or v.get( "Backfilling" )]
658 if not pending:
659 return desc
660 if not announced:
661 names = ", " .join(v.get( "IndexName" , "?" ) for v in pending)
662 print ( f " waiting on vector index/indexes: { names } " f "(ACTIVE and not backfilling)" )
663 announced = True
664 if time.time() > deadline:
665 states = {v.get( "IndexName" ): (v.get( "IndexStatus" ), v.get( "Backfilling" )) for v in vis}
666 _die(
667 f "vector index on { table_name } not ready after { timeout_s } s: { states } . "
668 "Backfill duration is driven by index construction, not item count, so a "
669 "small table can still take many minutes."
670 )
671 time.sleep( 5 )
672
673
674 def _deploy_tables_concurrent (client, table_specs: list[tuple[ dict , str ]]):
675 """Fire CreateTable for each spec in parallel, then wait on all concurrently.
676
677 Each spec is (create_kwargs, original_table_name); the name element is unused
678 here (the kwargs already carry the prefixed TableName) and is kept only so the
679 caller can correlate results back to the source table.
680 """
681 # Phase 1: fire all CreateTable calls.
682 for kwargs, _ in table_specs:
683 print ( f " CreateTable: { kwargs[ 'TableName' ] } " )
684 _create_table_one(client, kwargs)
685
686 # Phase 2: wait on all waiters concurrently.
687 descriptions: dict[ str , dict ] = {}
688 with ThreadPoolExecutor( max_workers = min ( 16 , len (table_specs) or 1 )) as pool:
689 futures = {
690 pool.submit(_wait_table_active, client, kwargs[ "TableName" ]): kwargs[ "TableName" ]
691 for kwargs, _ in table_specs
692 }
693 for fut in as_completed(futures):
694 tn = futures[fut]
695 descriptions[tn] = fut.result()
696 print ( f " ACTIVE: { tn } " )
697 return descriptions
698
699
700 # DynamoDB vector operations (SearchVectors and friends) ship in botocore 1.43.64. The
701 # managed Lambda python3.12 runtime bundles a much older boto3, so a vector benchmark run
702 # MUST carry its own SDK or every SearchVectors call dies with AttributeError while the
703 # observed capacity reads as zero — i.e. it looks like a free pattern rather than a broken
704 # one.
705 #
706 # Shipping only the newer dynamodb service model via AWS_DATA_PATH does NOT work: the new
707 # endpoint rule set uses a `stringArray` parameter type the old botocore cannot parse
708 # (`EndpointResolutionError: Unknown parameter type: stringArray`), and SearchVectors is
709 # served by a dedicated endpoint that must be resolved from those rules. Verified.
710 #
711 # So the whole SDK is vendored from the LOCAL install — which the skill already requires —
712 # rather than pip-installed at deploy time. No network, no new dependency. The cost (~18 MB
713 # zipped, a few seconds of upload) is paid ONLY by designs that declare a vector index;
714 # every other design gets the same few-KB package as before.
715 VECTOR_SDK_MIN = ( 1 , 43 , 64 )
716
717
718 def _local_sdk_paths () -> tuple[Path, Path]:
719 import boto3 as _b3
720 import botocore as _bc
721
722 ver = tuple ( int (x) for x in _bc. __version__ .split( "." )[: 3 ])
723 if ver < VECTOR_SDK_MIN :
724 _die(
725 f "this design declares a vector index, which the benchmark Lambda can only "
726 f "drive with botocore >= { '.' .join( map ( str , VECTOR_SDK_MIN )) } . The locally "
727 f "installed botocore is { _bc. __version__ } , and it is what gets vendored into "
728 f "the Lambda package. Upgrade locally first: \n "
729 f " pip install 'boto3>= { '.' .join( map ( str , VECTOR_SDK_MIN )) } '"
730 )
731 return Path(_b3. __file__ ).resolve().parent, Path(_bc. __file__ ).resolve().parent
732
733
734 def _build_lambda_zip (bundle_vector_sdk: bool = False ) -> bytes :
735 """Zip the handler source into an in-memory Lambda deployment package.
736
737 With bundle_vector_sdk, also vendor the local boto3/botocore so the handler can call
738 SearchVectors. Lambda puts /var/task ahead of the runtime's site-packages on sys.path,
739 so the vendored copy wins.
740 """
741 here = Path( __file__ ).resolve().parent
742 handler_src = here / "benchmark_lambda.py"
743 if not handler_src.exists():
744 _die(
745 f "lambda handler not found at { handler_src } . The skill is missing "
746 "the scripts/benchmark_lambda.py file — reinstall the skill."
747 )
748 buf = io.BytesIO()
749 with zipfile.ZipFile(buf, "w" , zipfile. ZIP_DEFLATED ) as z:
750 # The deployed module is named benchmark.py inside the zip so the Lambda
751 # Handler config (LAMBDA_HANDLER = "benchmark.handler") resolves; the
752 # local source file is scripts/benchmark_lambda.py.
753 z.writestr( "benchmark.py" , handler_src.read_text())
754 if bundle_vector_sdk:
755 b3_dir, bc_dir = _local_sdk_paths()
756 n = 0
757 for pkg_dir in (b3_dir, bc_dir):
758 root = pkg_dir.parent
759 for f in pkg_dir.rglob( "*" ):
760 if not f.is_file():
761 continue
762 # Skip caches and test trees; they are dead weight in a Lambda.
763 if "__pycache__" in f.parts or f.suffix in ( ".pyc" , ".pyo" ):
764 continue
765 z.write(f, str (f.relative_to(root)))
766 n += 1
767 print (
768 f " vendoring local boto3/botocore into the Lambda package "
769 f "( { n } files) so SearchVectors can be driven"
770 )
771 return buf.getvalue()
772
773
774 def _build_role_policy (prefix: str , region: str , account: str ) -> dict :
775 table_arn = f "arn:aws:dynamodb: { region } : { account } :table/ { prefix } *"
776 return {
777 "Version" : "2012-10-17" ,
778 "Statement" : [
779 {
780 "Effect" : "Allow" ,
781 "Action" : [
782 "dynamodb:DescribeTable" ,
783 "dynamodb:GetItem" ,
784 "dynamodb:Query" ,
785 "dynamodb:Scan" ,
786 "dynamodb:BatchGetItem" ,
787 "dynamodb:PutItem" ,
788 "dynamodb:UpdateItem" ,
789 "dynamodb:DeleteItem" ,
790 "dynamodb:BatchWriteItem" ,
791 "dynamodb:TransactWriteItems" ,
792 "dynamodb:TransactGetItems" ,
793 ],
794 "Resource" : [
795 table_arn,
796 f " { table_arn } /index/*" ,
797 ],
798 },
799 {
800 # dynamodb:SearchVectors gets its OWN statement, scoped to the index
801 # resource, and deliberately carries NO condition block.
802 #
803 # This is not stylistic. DynamoDB's fine-grained access control condition
804 # keys (dynamodb:LeadingKeys, dynamodb:Attributes, dynamodb:Select) are
805 # absent from the SearchVectors request context, so a statement whose
806 # Condition references one does not match a SearchVectors call and the
807 # result is a DENIAL rather than a grant. Folding this action into a
808 # conditioned statement silently breaks vector search while that
809 # statement's other actions keep working — which presents as a broken
810 # index rather than a policy fault. See SKILL.md "Security
811 # considerations" #6 and references/vector-search.md.
812 #
813 # The `/index/*` wildcard is deliberate HERE and only here: this is a
814 # throwaway benchmark role whose table ARN is already bounded by the bench
815 # prefix, and whose indexes this same run creates and destroys. Do NOT copy
816 # the wildcard into a production policy — name the specific index(es) the
817 # workload searches, or a vector index added to the table later inherits
818 # the grant silently. references/vector-search.md shows the scoped form.
819 "Effect" : "Allow" ,
820 "Action" : [ "dynamodb:SearchVectors" ],
821 "Resource" : [ f " { table_arn } /index/*" ],
822 },
823 {
824 "Effect" : "Allow" ,
825 "Action" : [
826 "logs:CreateLogGroup" ,
827 "logs:CreateLogStream" ,
828 "logs:PutLogEvents" ,
829 ],
830 # Least privilege: scope to this bench run's own Lambda log
831 # groups (prefix-namespaced) — never account-wide "*".
832 "Resource" : (
833 f "arn:aws:logs: { region } : { account } :" f "log-group:/aws/lambda/ { prefix } *:*"
834 ),
835 },
836 ],
837 }
838
839
840 def _deploy_lambda (
841 session, cfg: dict , ident: dict , prefix: str , tags: dict , bundle_vector_sdk: bool = False
842 ) -> dict :
843 """Create the IAM role and the Lambda function. Returns their ARNs."""
844 from botocore.exceptions import ClientError
845
846 iam = session.client( "iam" )
847 lam = session.client( "lambda" )
848
849 role_name = f " { prefix } bench-role" [: 64 ] # IAM role name max 64 chars
850 fn_name = f " { prefix } bench" [: 64 ]
851
852 # --- IAM role ---
853 try :
854 role = iam.create_role(
855 RoleName = role_name,
856 AssumeRolePolicyDocument = json.dumps(_build_lambda_trust_policy(ident[ "account" ])),
857 Description = f "ddb-skill-bench role for run { tags[ 'run_id' ] } " ,
858 Tags = [{ "Key" : k, "Value" : str (v)} for k, v in tags.items()],
859 )[ "Role" ]
860 role_arn = role[ "Arn" ]
861 print ( f " CreateRole: { role_name } " )
862 except ClientError as e:
863 code = e.response.get( "Error" , {}).get( "Code" , "" )
864 if code == "EntityAlreadyExists" :
865 role = iam.get_role( RoleName = role_name)[ "Role" ]
866 role_arn = role[ "Arn" ]
867 print ( f " note: role { role_name } already exists — reusing." )
868 else :
869 raise
870
871 iam.put_role_policy(
872 RoleName = role_name,
873 PolicyName = f " { prefix } bench-policy" ,
874 PolicyDocument = json.dumps(_build_role_policy(prefix, cfg[ "region" ], ident[ "account" ])),
875 )
876 print ( f " PutRolePolicy: { role_name } / { prefix } bench-policy" )
877
878 # --- Lambda function ---
879 zip_bytes = _build_lambda_zip(bundle_vector_sdk)
880 memory = int (cfg.get( "lambda_memory_mb" ) or DEFAULT_LAMBDA_MEMORY_MB )
881 timeout = int (cfg.get( "lambda_timeout_seconds" ) or DEFAULT_LAMBDA_TIMEOUT_S )
882
883 # IAM role is eventually consistent — CreateFunction can fail on "cannot
884 # be assumed by Lambda" even after CreateRole returns. Retry with backoff.
885 create_kwargs = dict (
886 FunctionName = fn_name,
887 Runtime = LAMBDA_RUNTIME ,
888 Role = role_arn,
889 Handler = LAMBDA_HANDLER ,
890 Code = { "ZipFile" : zip_bytes},
891 Timeout = timeout,
892 MemorySize = memory,
893 Architectures = [ "arm64" ],
894 Tags = {k: str (v) for k, v in tags.items()},
895 Description = f "ddb-skill-bench Lambda for run { tags[ 'run_id' ] } " ,
896 )
897
898 deadline = time.monotonic() + 45.0 # up to ~45s of retries
899 delay = 2.0
900 last_err = None
901 fn_arn = None
902 while time.monotonic() < deadline:
903 try :
904 resp = lam.create_function( ** create_kwargs)
905 fn_arn = resp[ "FunctionArn" ]
906 print ( f " CreateFunction: { fn_name } " )
907 break
908 except ClientError as e:
909 code = e.response.get( "Error" , {}).get( "Code" , "" )
910 msg = str (e)
911 last_err = e
912 if code == "ResourceConflictException" :
913 print ( f " note: function { fn_name } already exists — updating code + config." )
914 lam.update_function_code( FunctionName = fn_name, ZipFile = zip_bytes)
915 waiter = lam.get_waiter( "function_updated" )
916 waiter.wait( FunctionName = fn_name)
917 # Keep memory/timeout/role in sync with the current config on
918 # a same-prefix rerun.
919 lam.update_function_configuration(
920 FunctionName = fn_name,
921 Role = role_arn,
922 MemorySize = memory,
923 Timeout = timeout,
924 )
925 waiter.wait( FunctionName = fn_name)
926 fn_arn = lam.get_function( FunctionName = fn_name)[ "Configuration" ][ "FunctionArn" ]
927 break
928 if code == "InvalidParameterValueException" and (
929 "cannot be assumed" in msg or "role defined" in msg
930 ):
931 print ( f " waiting for IAM role to be assumable… ( { delay :.0f} s)" )
932 time.sleep(delay)
933 delay = min (delay * 1.5 , 8.0 )
934 continue
935 raise
936 if fn_arn is None :
937 _die(
938 "Lambda CreateFunction kept failing after IAM-consistency backoff. "
939 f "Last error: { last_err } "
940 )
941
942 # Wait for Active state before returning.
943 waiter = lam.get_waiter( "function_active_v2" )
944 waiter.wait( FunctionName = fn_name)
945
946 return {
947 "function_name" : fn_name,
948 "function_arn" : fn_arn,
949 "role_name" : role_name,
950 "role_arn" : role_arn,
951 "policy_name" : f " { prefix } bench-policy" ,
952 }
953
954
955 def main ():
956 p = argparse.ArgumentParser( description = __doc__ .splitlines()[ 0 ])
957 p.add_argument( "--model" , required = True , help = "path to dynamodb_data_model.json" )
958 p.add_argument( "--config" , required = True , help = "path to benchmark_config.json" )
959 p.add_argument(
960 "--manifest-out" , required = True , help = "path to write the created-resources manifest"
961 )
962 p.add_argument(
963 "--yes-deploy" ,
964 action = "store_true" ,
965 help = "required: explicit consent to create real AWS resources" ,
966 )
967 args = p.parse_args()
968
969 if not args.yes_deploy:
970 _die( "refusing to deploy without --yes-deploy" )
971
972 boto3_mod = _require_boto3()
973 model = _load_json(Path(args.model))
974 cfg = _load_json(Path(args.config))
975
976 # One clock for the whole run: the manifest's created_at and any
977 # auto-generated resource_prefix share the same UTC timestamp.
978 run_now = _dt.datetime.now(_dt.timezone.utc).replace( microsecond = 0 )
979 created_at = run_now.isoformat().replace( "+00:00" , "Z" )
980 run_id = uuid.uuid4().hex
981
982 _validate_design(model)
983 prefix_was_generated = _ensure_resource_prefix(cfg, run_now.strftime( "%Y%m %d " ))
984 _validate_config(cfg)
985 if prefix_was_generated:
986 print (
987 f "No resource_prefix supplied — generated { cfg[ 'resource_prefix' ] !r} "
988 "for this run (override by setting resource_prefix in the config)."
989 )
990
991 dry_run = bool (cfg.get( "dry_run" ))
992 ident = _preflight(boto3_mod, cfg, dry_run = dry_run)
993 session = boto3_mod.Session( profile_name = cfg[ "aws_profile" ], region_name = cfg[ "region" ])
994 ddb = session.client( "dynamodb" )
995
996 prefix = cfg[ "resource_prefix" ].rstrip( "-" ) + "-"
997
998 built_in_tags = {
999 "purpose" : "ddb-skill-bench" ,
1000 "run_id" : run_id,
1001 "created_at" : created_at,
1002 "prefix" : prefix,
1003 }
1004 tags = { ** (cfg.get( "tags" ) or {}), ** built_in_tags}
1005
1006 attr_types_by_table = _collect_attr_types(model[ "tables" ])
1007
1008 # Optional global provisioned-capacity override from the benchmark config
1009 # (a per-table table.provisioned_capacity still wins). Used for deliberate
1010 # capacity-ceiling tests where on-demand adaptive capacity would mask
1011 # hot-partition throttling.
1012 global_provisioned = cfg.get( "provisioned_capacity" )
1013 if global_provisioned:
1014 print (
1015 f "Provisioned-capacity mode: { global_provisioned } "
1016 "(hard ceiling; bench-only — production designs default to "
1017 "on-demand)."
1018 )
1019
1020 # Build specs. Each spec is (create_kwargs, original_table_name).
1021 specs: list[tuple[ dict , str ]] = []
1022 for t in model[ "tables" ]:
1023 kwargs = _build_create_kwargs(
1024 t,
1025 prefix,
1026 tags,
1027 attr_types_by_table.get(t[ "table_name" ], {}),
1028 global_provisioned,
1029 )
1030 specs.append((kwargs, t[ "table_name" ]))
1031
1032 if dry_run:
1033 print ( " \n DRY RUN — intended resources:" )
1034 for k, orig in specs:
1035 print ( f " table: { k[ 'TableName' ] } (from { orig } )" )
1036 print ( f " lambda: { prefix } bench (not built in dry run)" )
1037 print ( f " role: { prefix } bench-role" )
1038 return
1039
1040 # --- Deploy tables concurrently ---
1041 print ( f " \n Deploying { len (specs) } table(s) concurrently …" )
1042 descs = _deploy_tables_concurrent(ddb, specs)
1043
1044 deployed_tables = []
1045 for kwargs, orig in specs:
1046 tn = kwargs[ "TableName" ]
1047 desc = descs.get(tn) or {}
1048 sse_desc = desc.get( "SSEDescription" ) or {}
1049 deployed_tables.append(
1050 {
1051 "name" : tn,
1052 "original_name" : orig,
1053 "arn" : desc.get( "TableArn" ),
1054 "gsis" : [g[ "IndexName" ] for g in kwargs.get( "GlobalSecondaryIndexes" ) or []],
1055 "stream_arn" : desc.get( "LatestStreamArn" ),
1056 # Surface encryption-at-rest from DescribeTable so the user can verify
1057 # it. DynamoDB always encrypts at rest; with the AWS-owned-key default
1058 # DescribeTable returns NO SSEDescription block, so an empty/absent
1059 # SSEDescription means "encrypted with the AWS-owned key" (the implicit
1060 # default), and a present block means a customer-/AWS-managed CMK is in
1061 # use (Status/SSEType/KMSMasterKeyArn populated).
1062 "encryption" : {
1063 "sse_type" : sse_desc.get( "SSEType" , "AWS_OWNED" ),
1064 "status" : sse_desc.get( "Status" , "ENABLED (AWS-owned key, default)" ),
1065 "kms_key_arn" : sse_desc.get( "KMSMasterKeyArn" ),
1066 },
1067 "key_schema" : {
1068 "partition_key" : kwargs[ "KeySchema" ][ 0 ][ "AttributeName" ],
1069 "sort_key" : (
1070 kwargs[ "KeySchema" ][ 1 ][ "AttributeName" ]
1071 if len (kwargs[ "KeySchema" ]) > 1
1072 else None
1073 ),
1074 },
1075 }
1076 )
1077
1078 # --- Deploy Lambda + IAM role ---
1079 print ( f " \n Deploying benchmark Lambda and IAM role …" )
1080 has_vectors = any (t.get( "vector_indexes" ) for t in (model.get( "tables" ) or []))
1081 lambda_info = _deploy_lambda(session, cfg, ident, prefix, tags, bundle_vector_sdk = has_vectors)
1082
1083 manifest = {
1084 "account" : ident[ "account" ],
1085 "alias" : ident[ "alias" ],
1086 "aws_profile" : cfg[ "aws_profile" ],
1087 "region" : cfg[ "region" ],
1088 "run_id" : run_id,
1089 "created_at" : created_at,
1090 "prefix" : prefix,
1091 "tables" : deployed_tables,
1092 "lambda" : lambda_info,
1093 }
1094 out_path = Path(args.manifest_out)
1095 with out_path.open( "w" ) as f:
1096 json.dump(manifest, f, indent = 2 , default = str )
1097 print ( f " \n Manifest written to { out_path } " )
1098 print (
1099 f "Created { len (deployed_tables) } table(s), 1 Lambda "
1100 f "( { lambda_info[ 'function_name' ] } ), 1 IAM role "
1101 f "( { lambda_info[ 'role_name' ] } ). Run benchmark_model.py next."
1102 )
1103
1104
1105 if __name__ == "__main__" :
1106 main()