gpu_capacity/prober.py

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)