import logging import threading import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone import boto3 from botocore.config import Config from botocore.exceptions import ClientError from . import config log = logging.getLogger(__name__) BOTO_CONFIG = Config(retries={"mode": "adaptive", "max_attempts": 5}) # For probe creates only: botocore treats InsufficientInstanceCapacity as # retryable, but here that error IS the answer — retrying turns a 1s probe # into 20s of backoff. Throttling is retried manually in probe_az. PROBE_BOTO_CONFIG = Config(retries={"mode": "standard", "max_attempts": 1}) # CreateCapacityReservation error code -> probe status STATUS_BY_CODE = { "InsufficientInstanceCapacity": "unavailable", "InstanceLimitExceeded": "quota", "ReservationCapacityExceeded": "quota", "MaxConfigLimitExceededException": "quota", "Unsupported": "unsupported", "UnsupportedOperation": "unsupported", "UnauthorizedOperation": "forbidden", } _clients: dict[str, object] = {} _clients_lock = threading.Lock() # (region, instance_type) -> (azs, fetched_at); describes are free, cached 24h _offerings: dict[tuple[str, str], tuple[list[str], float]] = {} _offerings_lock = threading.Lock() OFFERINGS_TTL = 24 * 3600 def ec2(region: str): with _clients_lock: if region not in _clients: _clients[region] = boto3.client("ec2", region_name=region, config=BOTO_CONFIG) return _clients[region] def probe_ec2(region: str): with _clients_lock: key = f"probe:{region}" if key not in _clients: _clients[key] = boto3.client("ec2", region_name=region, config=PROBE_BOTO_CONFIG) return _clients[key] def offered_azs(region: str, instance_type: str) -> list[str]: """AZs of `region` where `instance_type` is offered (cached 24h). Works for any instance type, not just configured ones — an unknown or unoffered type simply matches nothing and yields []. """ key = (region, instance_type) with _offerings_lock: cached = _offerings.get(key) if cached and time.time() - cached[1] < OFFERINGS_TTL: return cached[0] resp = ec2(region).describe_instance_type_offerings( LocationType="availability-zone", Filters=[{"Name": "instance-type", "Values": [instance_type]}], ) azs = sorted(off["Location"] for off in resp["InstanceTypeOfferings"]) with _offerings_lock: _offerings[key] = (azs, time.time()) return azs def _create_with_throttle_retry(region: str, **kwargs): for attempt in range(3): try: return probe_ec2(region).create_capacity_reservation(**kwargs) except ClientError as e: if e.response["Error"]["Code"] == "RequestLimitExceeded" and attempt < 2: time.sleep(1 + attempt) continue raise # instance_type -> {"name", "count", "memory_gb"} or None; resolved once per process _gpu_info: dict[str, dict | None] = {} _gpu_lock = threading.Lock() def gpu_info(instance_types: list[str]) -> dict[str, dict | None]: """GPU spec per type via DescribeInstanceTypes, from the first region that knows the type (newer families only exist in a few regions).""" with _gpu_lock: missing = [t for t in instance_types if t not in _gpu_info] for region in config.REGIONS: if not missing: break try: # Filters are applied per page, so matches must be collected across # all pages (a bare call would return first-page matches only). found = [] paginator = ec2(region).get_paginator("describe_instance_types") for page in paginator.paginate( Filters=[{"Name": "instance-type", "Values": missing}] ): found.extend(page["InstanceTypes"]) except ClientError as e: log.warning("describe_instance_types in %s failed: %s", region, e) continue with _gpu_lock: for it in found: gpus = (it.get("GpuInfo") or {}).get("Gpus") or [] _gpu_info[it["InstanceType"]] = { "name": f"{gpus[0]['Manufacturer']} {gpus[0]['Name']}", "count": gpus[0]["Count"], "memory_gb": round(gpus[0]["MemoryInfo"]["SizeInMiB"] / 1024), } if gpus else None missing = [t for t in instance_types if t not in _gpu_info] with _gpu_lock: for t in missing: # unknown everywhere: cache None to avoid re-querying _gpu_info[t] = None return {t: _gpu_info[t] for t in instance_types} def probe_az(region: str, instance_type: str, az: str) -> dict: """Create a 1-instance targeted capacity reservation in `az`, then cancel it. Success means real on-demand capacity exists right now. The reservation carries an EndDate dead-man switch so a crash between create and cancel can't leak an open-ended reservation. """ client = ec2(region) result = {"az": az, "status": "error", "error": None} cr_id = None try: resp = _create_with_throttle_retry( region, InstanceType=instance_type, InstancePlatform="Linux/UNIX", AvailabilityZone=az, InstanceCount=1, InstanceMatchCriteria="targeted", EndDateType="limited", EndDate=datetime.now(timezone.utc) + timedelta(minutes=config.PROBE_END_DATE_MINUTES), TagSpecifications=[ { "ResourceType": "capacity-reservation", "Tags": [ {"Key": config.PROBE_TAG_KEY, "Value": config.PROBE_TAG_VALUE} ], } ], ) cr_id = resp["CapacityReservation"]["CapacityReservationId"] result["status"] = "available" except ClientError as e: code = e.response["Error"]["Code"] result["status"] = STATUS_BY_CODE.get(code, "error") result["error"] = code if result["status"] == "error": log.warning("probe %s %s/%s: %s", instance_type, region, az, e) finally: if cr_id: _cancel(client, cr_id, region) return result def _cancel(client, cr_id: str, region: str): for attempt in range(4): try: client.cancel_capacity_reservation(CapacityReservationId=cr_id) return except ClientError as e: log.warning("cancel %s in %s failed (attempt %d): %s", cr_id, region, attempt + 1, e) time.sleep(2**attempt) log.error( "LEAKED reservation %s in %s — will auto-expire in %d min", cr_id, region, config.PROBE_END_DATE_MINUTES, ) def probe_region(region: str, instance_type: str) -> dict: """Probe every offered AZ of a region. Region is 'available' if any AZ is.""" started = time.time() try: azs = offered_azs(region, instance_type) except ClientError as e: code = e.response["Error"]["Code"] return {"status": "error", "error": code, "azs": {}, "checked_at": time.time()} if not azs: return {"status": "not_offered", "azs": {}, "checked_at": time.time()} with ThreadPoolExecutor(max_workers=min(len(azs), 6)) as pool: az_results = list(pool.map(lambda az: probe_az(region, instance_type, az), azs)) statuses = [r["status"] for r in az_results] if "available" in statuses: status = "available" elif "unavailable" in statuses: status = "unavailable" else: status = statuses[0] # quota / unsupported / forbidden / error return { "status": status, "azs": {r["az"]: {"status": r["status"], "error": r["error"]} for r in az_results}, "checked_at": time.time(), "duration_ms": int((time.time() - started) * 1000), } def cleanup_leaked(regions: list[str] | None = None): """Cancel any active probe-tagged reservation left over from a crash.""" for region in regions or config.REGIONS: try: resp = ec2(region).describe_capacity_reservations( Filters=[ {"Name": f"tag:{config.PROBE_TAG_KEY}", "Values": [config.PROBE_TAG_VALUE]}, {"Name": "state", "Values": ["active"]}, ] ) for cr in resp["CapacityReservations"]: log.warning("cleaning leaked reservation %s in %s", cr["CapacityReservationId"], region) _cancel(ec2(region), cr["CapacityReservationId"], region) except Exception as e: # ClientError, but also expired-SSO credential errors log.warning("cleanup in %s failed: %s", region, e)