scripts/aws_image_tools.py
Rama master · réplica de lectura del repositorio original.
"""
Tools for managing dstack AWS AMIs across regions.
dstack publishes public AMIs (see scripts/packer/aws-image.json) to all regions
listed in scripts/packer/aws-vars-prod.json. Over time these accumulate and hit
the per-region AMI service quota (the AWS error looks like "the maximum number of
AMIs has been reached"). This script helps to:
1. request-quota Request a service quota increase across regions (e.g. the EC2
"AMIs" / "Public AMIs" quota).
2. list-quotas Discover quota codes and/or quota names
(e.g. search for "AMI") to use with request-quota.
3. delete-amis Deregister AMIs older than a date and delete their snapshots.
Dry-run by default — pass --yes to actually delete.
"""
import logging
import sys
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import List, Optional
import boto3
# Regions dstack copies AMIs to, kept in sync with scripts/packer/aws-vars-prod.json.
PROD_REGIONS = [
"us-east-2",
"us-east-1",
"us-west-1",
"us-west-2",
"ca-central-1",
"eu-central-1",
"eu-west-1",
"eu-west-2",
"eu-west-3",
"eu-north-1",
"ap-southeast-1",
]
# Default name prefix of dstack AMIs (e.g. dstack-0.18, dstack-cuda-0.18).
DEFAULT_NAME_PREFIX = "dstack-"
EC2_SERVICE_CODE = "ec2"
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s",
datefmt="%Y-%m-%dT%H:%M:%S",
)
logger = logging.getLogger("aws_image_tools")
class ScriptError(Exception):
pass
@dataclass
class RequestQuotaCommandArgs:
regions: List[str]
service_code: str
quota_code: str
value: float
yes: bool
@classmethod
def setup_parser(cls, parser: ArgumentParser) -> None:
parser.add_argument(
"--regions",
metavar="REGION",
nargs="*",
help="Regions to request the increase in (default: dstack prod regions)",
)
parser.add_argument("--service-code", default=EC2_SERVICE_CODE)
parser.add_argument(
"--quota-code",
required=True,
help="Quota code, e.g. L-0E3CBAB9 (same across regions; find it with list-quotas)",
)
parser.add_argument("--value", type=float, required=True, help="Desired new quota value")
parser.add_argument(
"--yes",
action="store_true",
help="Actually submit the requests (default: preview only)",
)
parser.set_defaults(to_struct=cls.from_namespace, run_command=request_quota_command)
@staticmethod
def from_namespace(args: Namespace) -> "RequestQuotaCommandArgs":
return RequestQuotaCommandArgs(
regions=args.regions or list(PROD_REGIONS),
service_code=args.service_code,
quota_code=args.quota_code,
value=args.value,
yes=args.yes,
)
@dataclass
class ListQuotasCommandArgs:
region: str
service_code: str
search: Optional[str]
@classmethod
def setup_parser(cls, parser: ArgumentParser) -> None:
parser.add_argument("--region", default=PROD_REGIONS[0], help="Region to list quotas in")
parser.add_argument("--service-code", default=EC2_SERVICE_CODE)
parser.add_argument("--search", help="Case-insensitive substring to filter quota names by")
parser.set_defaults(to_struct=cls.from_namespace, run_command=list_quotas_command)
@staticmethod
def from_namespace(args: Namespace) -> "ListQuotasCommandArgs":
return ListQuotasCommandArgs(
region=args.region,
service_code=args.service_code,
search=args.search,
)
@dataclass
class DeleteAmisCommandArgs:
regions: List[str]
before: datetime
name_prefix: str
name_contains: Optional[str]
keep_latest: int
yes: bool
@classmethod
def setup_parser(cls, parser: ArgumentParser) -> None:
parser.add_argument(
"--regions",
metavar="REGION",
nargs="*",
help="Regions to delete AMIs in (default: dstack prod regions)",
)
parser.add_argument(
"--before",
required=True,
metavar="YYYY-MM-DD",
help="Delete AMIs created strictly before this date (UTC)",
)
parser.add_argument(
"--name-prefix",
default=DEFAULT_NAME_PREFIX,
help=f"Only consider AMIs whose name starts with this (default: {DEFAULT_NAME_PREFIX!r})",
)
parser.add_argument(
"--name-contains",
help="Further restrict to AMIs whose name contains this substring "
"(case-insensitive), e.g. a version like 0.18",
)
parser.add_argument(
"--keep-latest",
type=int,
default=0,
help="Always keep this many newest matching AMIs per region, "
"regardless of --before (default: 0)",
)
parser.add_argument(
"--yes",
action="store_true",
help="Actually deregister AMIs and delete snapshots (default: preview only)",
)
parser.set_defaults(to_struct=cls.from_namespace, run_command=delete_amis_command)
@staticmethod
def from_namespace(args: Namespace) -> "DeleteAmisCommandArgs":
try:
before = datetime.strptime(args.before, "%Y-%m-%d").replace(tzinfo=timezone.utc)
except ValueError:
raise ScriptError(f"Invalid --before date {args.before!r}, expected YYYY-MM-DD")
if args.keep_latest < 0:
raise ScriptError("--keep-latest must be >= 0")
return DeleteAmisCommandArgs(
regions=args.regions or list(PROD_REGIONS),
before=before,
name_prefix=args.name_prefix,
name_contains=args.name_contains,
keep_latest=args.keep_latest,
yes=args.yes,
)
def main() -> None:
parser = ArgumentParser(description="Tools for managing dstack AWS AMIs")
subparsers = parser.add_subparsers()
request_quota_parser = subparsers.add_parser(
name="request-quota",
description="Request a service quota increase across regions.",
)
RequestQuotaCommandArgs.setup_parser(request_quota_parser)
list_quotas_parser = subparsers.add_parser(
name="list-quotas",
description="List service quotas (and their codes) in a region.",
)
ListQuotasCommandArgs.setup_parser(list_quotas_parser)
delete_amis_parser = subparsers.add_parser(
name="delete-amis",
description=(
"Deregister AMIs older than a date and delete their snapshots. "
"Dry-run by default; pass --yes to actually delete."
),
)
DeleteAmisCommandArgs.setup_parser(delete_amis_parser)
args = parser.parse_args()
if not hasattr(args, "run_command"):
parser.print_help()
sys.exit(1)
try:
args.run_command(args.to_struct(args))
except ScriptError as e:
logger.error("%s", e)
sys.exit(1)
def request_quota_command(args: RequestQuotaCommandArgs) -> None:
failed = False
quota_code = args.quota_code
for region in args.regions:
client = boto3.client("service-quotas", region_name=region)
current = _get_quota_value(client, args.service_code, quota_code)
if current is not None and current >= args.value:
logger.info(
"[%s] %s already at %g >= %g, skipping",
region,
quota_code,
current,
args.value,
)
continue
pending = _get_pending_request_value(client, args.service_code, quota_code)
if pending is not None and pending >= args.value:
logger.info(
"[%s] %s already has a pending request for %g, skipping",
region,
quota_code,
pending,
)
continue
if not args.yes:
logger.info(
"[%s] would request %s increase to %g (current: %s)",
region,
quota_code,
args.value,
"unknown" if current is None else f"{current:g}",
)
continue
try:
client.request_service_quota_increase(
ServiceCode=args.service_code,
QuotaCode=quota_code,
DesiredValue=args.value,
)
logger.info("[%s] requested %s increase to %g", region, quota_code, args.value)
except Exception as e:
logger.error("[%s] failed to request %s increase: %s", region, quota_code, e)
failed = True
if not args.yes:
logger.info("Preview only. Re-run with --yes to submit the requests.")
if failed:
raise ScriptError("Some quota requests failed or were skipped")
def _get_quota_value(client, service_code: str, quota_code: str) -> Optional[float]:
try:
resp = client.get_service_quota(ServiceCode=service_code, QuotaCode=quota_code)
return resp["Quota"]["Value"]
except Exception:
return None
def _get_pending_request_value(client, service_code: str, quota_code: str) -> Optional[float]:
try:
paginator = client.get_paginator("list_requested_service_quota_change_history_by_quota")
latest_value = None
for page in paginator.paginate(ServiceCode=service_code, QuotaCode=quota_code):
for req in page["RequestedQuotas"]:
if req["Status"] in ("PENDING", "CASE_OPENED"):
value = req["DesiredValue"]
if latest_value is None or value > latest_value:
latest_value = value
return latest_value
except Exception:
return None
def list_quotas_command(args: ListQuotasCommandArgs) -> None:
client = boto3.client("service-quotas", region_name=args.region)
needle = args.search.lower() if args.search else None
paginator = client.get_paginator("list_service_quotas")
rows = []
for page in paginator.paginate(ServiceCode=args.service_code):
for quota in page["Quotas"]:
if needle and needle not in quota["QuotaName"].lower():
continue
rows.append((quota["QuotaCode"], quota["Value"], quota["QuotaName"]))
rows.sort(key=lambda r: r[2])
if not rows:
logger.info("No quotas found matching the filter.")
return
print(f"{'QUOTA CODE':<16} {'VALUE':>10} NAME")
for code, value, name in rows:
print(f"{code:<16} {value:>10g} {name}")
def delete_amis_command(args: DeleteAmisCommandArgs) -> None:
total_deleted = 0
for region in args.regions:
ec2 = boto3.client("ec2", region_name=region)
images = _find_self_owned_images(ec2, args.name_prefix, args.name_contains)
# Sort newest first so --keep-latest preserves the most recent images.
images.sort(key=lambda img: img["_created"], reverse=True)
to_delete = []
for index, image in enumerate(images):
if index < args.keep_latest:
continue
if image["_created"] < args.before:
to_delete.append(image)
keep = [img for img in images if img not in to_delete]
logger.info(
"[%s] %d matching AMIs: %d to delete, %d to keep",
region,
len(images),
len(to_delete),
len(keep),
)
for image in keep:
logger.info(
"[%s] KEEP %s %s (%s)",
region,
image["ImageId"],
image["Name"],
image["CreationDate"],
)
for image in to_delete:
snapshot_ids = _image_snapshot_ids(image)
logger.info(
"[%s] DELETE %s %s (%s) snapshots=%s",
region,
image["ImageId"],
image["Name"],
image["CreationDate"],
",".join(snapshot_ids) or "none",
)
if args.yes:
_deregister_image(ec2, region, image, snapshot_ids)
total_deleted += 1
if not args.yes:
logger.info("Preview only. Re-run with --yes to deregister AMIs and delete snapshots.")
else:
logger.info("Deleted %d AMIs.", total_deleted)
def _find_self_owned_images(ec2, name_prefix: str, name_contains: Optional[str]) -> List[dict]:
resp = ec2.describe_images(
Owners=["self"],
Filters=[{"Name": "name", "Values": [f"{name_prefix}*"]}],
)
images = resp["Images"]
if name_contains:
needle = name_contains.lower()
images = [img for img in images if needle in img["Name"].lower()]
for image in images:
image["_created"] = datetime.strptime(
image["CreationDate"], "%Y-%m-%dT%H:%M:%S.%fZ"
).replace(tzinfo=timezone.utc)
return images
def _image_snapshot_ids(image: dict) -> List[str]:
snapshot_ids = []
for mapping in image.get("BlockDeviceMappings", []):
ebs = mapping.get("Ebs")
if ebs and ebs.get("SnapshotId"):
snapshot_ids.append(ebs["SnapshotId"])
return snapshot_ids
def _deregister_image(ec2, region: str, image: dict, snapshot_ids: List[str]) -> None:
try:
ec2.deregister_image(ImageId=image["ImageId"])
logger.info("[%s] deregistered %s", region, image["ImageId"])
except Exception as e:
logger.error("[%s] failed to deregister %s: %s", region, image["ImageId"], e)
return
for snapshot_id in snapshot_ids:
try:
ec2.delete_snapshot(SnapshotId=snapshot_id)
logger.info("[%s] deleted snapshot %s", region, snapshot_id)
except Exception as e:
logger.error("[%s] failed to delete snapshot %s: %s", region, snapshot_id, e)
if __name__ == "__main__":
main()