438 lines
17 KiB
Python
438 lines
17 KiB
Python
"""EC2 runtime adapter for managing Spot instances."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import http.client
|
|
import json
|
|
import logging
|
|
import random
|
|
import socket
|
|
import threading
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
import boto3
|
|
from botocore.exceptions import ClientError
|
|
|
|
from ..config import AwsConfig
|
|
from .base import RuntimeAdapter
|
|
from .base import RuntimeError as RuntimeAdapterError
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# EC2 ClientError code → normalized error category
|
|
_ERROR_CATEGORIES: dict[str, str] = {
|
|
"InsufficientInstanceCapacity": "capacity_unavailable",
|
|
"SpotMaxPriceTooLow": "price_too_low",
|
|
"RequestLimitExceeded": "throttled",
|
|
}
|
|
|
|
_RETRYABLE_CODES: frozenset[str] = frozenset({"RequestLimitExceeded"})
|
|
|
|
|
|
class _UnixSocketHTTPConnection(http.client.HTTPConnection):
|
|
"""HTTP connection over a Unix domain socket."""
|
|
|
|
def __init__(self, socket_path: str, timeout: float = 1.0) -> None:
|
|
super().__init__("local-tailscaled.sock", timeout=timeout)
|
|
self._socket_path = socket_path
|
|
|
|
def connect(self) -> None:
|
|
self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
self.sock.connect(self._socket_path)
|
|
|
|
|
|
class EC2Runtime(RuntimeAdapter):
|
|
"""EC2 Spot instance runtime adapter.
|
|
|
|
Args:
|
|
config: AWS configuration dataclass.
|
|
environment: Environment tag value (e.g. ``"dev"``, ``"prod"``).
|
|
_client: Optional pre-configured boto3 EC2 client (for testing).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: AwsConfig,
|
|
environment: str = "dev",
|
|
*,
|
|
_client: Any = None,
|
|
_tailscale_socket_path: str = "/run/tailscale/tailscaled.sock",
|
|
) -> None:
|
|
self._base_client: Any = _client or boto3.client("ec2", region_name=config.region)
|
|
self._region = config.region
|
|
self._assume_role_arn = config.assume_role_arn.strip()
|
|
self._assume_role_session_name = "nix-builder-autoscaler"
|
|
self._assumed_client: Any | None = None
|
|
self._assumed_client_expiration: datetime | None = None
|
|
self._assume_role_lock = threading.Lock()
|
|
self._sts_client: Any | None = None
|
|
if self._assume_role_arn != "":
|
|
self._sts_client = boto3.client("sts", region_name=config.region)
|
|
self._instance_type = config.instance_type
|
|
self._launch_template_id = config.launch_template_id
|
|
self._on_demand_launch_template_id = config.on_demand_launch_template_id
|
|
self._subnet_ids = list(config.subnet_ids)
|
|
self._security_group_ids = list(config.security_group_ids)
|
|
self._instance_profile_arn = config.instance_profile_arn
|
|
self._environment = environment
|
|
self._subnet_index = 0
|
|
self._tailscale_socket_path = _tailscale_socket_path
|
|
|
|
def preflight_validate(self) -> None:
|
|
"""Check that the configured instance type is available in the configured subnets' AZs.
|
|
|
|
Logs a clear error if the instance type is absent from the region or missing from
|
|
any subnet AZ so misconfigurations are surfaced at startup rather than discovered
|
|
silently on every failed launch attempt. Never raises; API failures are logged as
|
|
warnings so a transient permissions issue does not prevent startup.
|
|
"""
|
|
try:
|
|
ec2_client = self._get_ec2_client()
|
|
target_azs: set[str] = set()
|
|
if self._subnet_ids:
|
|
resp = self._call_with_backoff(
|
|
ec2_client.describe_subnets, SubnetIds=self._subnet_ids
|
|
)
|
|
for subnet in resp.get("Subnets", []):
|
|
az = subnet.get("AvailabilityZone")
|
|
if az:
|
|
target_azs.add(az)
|
|
|
|
filters: list[dict[str, Any]] = [
|
|
{"Name": "instance-type", "Values": [self._instance_type]},
|
|
]
|
|
if target_azs:
|
|
filters.append({"Name": "location", "Values": list(target_azs)})
|
|
|
|
resp = self._call_with_backoff(
|
|
ec2_client.describe_instance_type_offerings,
|
|
LocationType="availability-zone",
|
|
Filters=filters,
|
|
)
|
|
available_azs = {o["Location"] for o in resp.get("InstanceTypeOfferings", [])}
|
|
|
|
if not available_azs:
|
|
region = ec2_client.meta.region_name
|
|
log.error(
|
|
"preflight_misconfiguration",
|
|
extra={
|
|
"error": (
|
|
f"instance type {self._instance_type!r} is not available in"
|
|
f" region {region!r} - all launches will fail with Unsupported"
|
|
),
|
|
"category": "misconfiguration",
|
|
},
|
|
)
|
|
return
|
|
|
|
missing_azs = target_azs - available_azs
|
|
if missing_azs:
|
|
log.warning(
|
|
"preflight_misconfiguration",
|
|
extra={
|
|
"error": (
|
|
f"instance type {self._instance_type!r} is not available in"
|
|
f" AZs {sorted(missing_azs)} - launches into those subnets will"
|
|
f" fail with Unsupported"
|
|
),
|
|
"category": "misconfiguration",
|
|
},
|
|
)
|
|
else:
|
|
log.info(
|
|
"preflight_ok",
|
|
extra={
|
|
"error": None,
|
|
"category": None,
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
log.warning(
|
|
"preflight_validate_failed",
|
|
extra={"error": str(exc), "category": "unknown"},
|
|
)
|
|
|
|
def launch_instance(
|
|
self, slot_id: str, user_data: str, *, nested_virtualization: bool = False
|
|
) -> str:
|
|
"""Launch an instance for *slot_id*. Return instance ID.
|
|
|
|
When nested_virtualization is True, an on-demand instance is launched using the
|
|
on-demand launch template (cpu_options nested virt enabled, no spot market options).
|
|
When False (default), a spot instance is launched using the spot launch template.
|
|
"""
|
|
if nested_virtualization:
|
|
if not self._on_demand_launch_template_id:
|
|
raise RuntimeAdapterError(
|
|
"nested_virtualization=True but on_demand_launch_template_id is not configured",
|
|
category="misconfiguration",
|
|
)
|
|
lt_id = self._on_demand_launch_template_id
|
|
else:
|
|
lt_id = self._launch_template_id
|
|
|
|
params: dict[str, Any] = {
|
|
"MinCount": 1,
|
|
"MaxCount": 1,
|
|
"LaunchTemplate": {
|
|
"LaunchTemplateId": lt_id,
|
|
"Version": "$Latest",
|
|
},
|
|
"UserData": user_data,
|
|
"TagSpecifications": [
|
|
{
|
|
"ResourceType": "instance",
|
|
"Tags": [
|
|
{"Key": "Name", "Value": f"nix-builder-{slot_id}"},
|
|
{"Key": "AutoscalerSlot", "Value": slot_id},
|
|
{"Key": "ManagedBy", "Value": "nix-builder-autoscaler"},
|
|
{"Key": "Service", "Value": "nix-builder"},
|
|
{"Key": "Environment", "Value": self._environment},
|
|
],
|
|
}
|
|
],
|
|
}
|
|
|
|
if not nested_virtualization:
|
|
params["InstanceMarketOptions"] = {
|
|
"MarketType": "spot",
|
|
"SpotOptions": {
|
|
"SpotInstanceType": "one-time",
|
|
"InstanceInterruptionBehavior": "terminate",
|
|
},
|
|
}
|
|
|
|
if self._subnet_ids:
|
|
subnet = self._subnet_ids[self._subnet_index % len(self._subnet_ids)]
|
|
self._subnet_index += 1
|
|
params["SubnetId"] = subnet
|
|
|
|
ec2_client = self._get_ec2_client()
|
|
resp = self._call_with_backoff(ec2_client.run_instances, **params)
|
|
return resp["Instances"][0]["InstanceId"]
|
|
|
|
def describe_instance(self, instance_id: str) -> dict:
|
|
"""Return normalized instance info dict."""
|
|
try:
|
|
ec2_client = self._get_ec2_client()
|
|
resp = self._call_with_backoff(ec2_client.describe_instances, InstanceIds=[instance_id])
|
|
except RuntimeAdapterError:
|
|
return {"state": "terminated", "tailscale_ip": None, "launch_time": None}
|
|
|
|
reservations = resp.get("Reservations", [])
|
|
if not reservations or not reservations[0].get("Instances"):
|
|
return {"state": "terminated", "tailscale_ip": None, "launch_time": None}
|
|
|
|
inst = reservations[0]["Instances"][0]
|
|
tags = inst.get("Tags", [])
|
|
slot_id = self._get_tag(tags, "AutoscalerSlot")
|
|
state = inst["State"]["Name"]
|
|
tailscale_ip: str | None = None
|
|
if state == "running" and slot_id:
|
|
tailscale_ip = self._discover_tailscale_ip(slot_id, instance_id)
|
|
|
|
launch_time = inst.get("LaunchTime")
|
|
return {
|
|
"state": state,
|
|
"tailscale_ip": tailscale_ip,
|
|
"launch_time": launch_time.isoformat() if launch_time else None,
|
|
}
|
|
|
|
def terminate_instance(self, instance_id: str) -> None:
|
|
"""Terminate the instance."""
|
|
ec2_client = self._get_ec2_client()
|
|
self._call_with_backoff(ec2_client.terminate_instances, InstanceIds=[instance_id])
|
|
|
|
def list_managed_instances(self) -> list[dict]:
|
|
"""Return list of managed instances."""
|
|
ec2_client = self._get_ec2_client()
|
|
resp = self._call_with_backoff(
|
|
ec2_client.describe_instances,
|
|
Filters=[
|
|
{"Name": "tag:ManagedBy", "Values": ["nix-builder-autoscaler"]},
|
|
{
|
|
"Name": "instance-state-name",
|
|
"Values": ["pending", "running", "shutting-down", "stopping"],
|
|
},
|
|
],
|
|
)
|
|
|
|
result: list[dict] = []
|
|
for reservation in resp.get("Reservations", []):
|
|
for inst in reservation.get("Instances", []):
|
|
tags = inst.get("Tags", [])
|
|
result.append(
|
|
{
|
|
"instance_id": inst["InstanceId"],
|
|
"state": inst["State"]["Name"],
|
|
"slot_id": self._get_tag(tags, "AutoscalerSlot"),
|
|
}
|
|
)
|
|
return result
|
|
|
|
def _get_ec2_client(self) -> Any:
|
|
if self._assume_role_arn == "":
|
|
return self._base_client
|
|
return self._get_assumed_ec2_client()
|
|
|
|
def _get_assumed_ec2_client(self) -> Any:
|
|
now = datetime.now(UTC)
|
|
with self._assume_role_lock:
|
|
if (
|
|
self._assumed_client is not None
|
|
and self._assumed_client_expiration is not None
|
|
and (self._assumed_client_expiration - now).total_seconds() > 300
|
|
):
|
|
return self._assumed_client
|
|
|
|
if self._sts_client is None:
|
|
raise RuntimeAdapterError(
|
|
"assume_role_arn is set but STS client is not initialized",
|
|
category="misconfiguration",
|
|
)
|
|
|
|
try:
|
|
response = self._sts_client.assume_role(
|
|
RoleArn=self._assume_role_arn,
|
|
RoleSessionName=self._assume_role_session_name,
|
|
)
|
|
except ClientError as exc:
|
|
raise RuntimeAdapterError(str(exc), category="misconfiguration") from exc
|
|
credentials = response["Credentials"]
|
|
self._assumed_client_expiration = credentials["Expiration"]
|
|
self._assumed_client = boto3.client(
|
|
"ec2",
|
|
region_name=self._region,
|
|
aws_access_key_id=credentials["AccessKeyId"],
|
|
aws_secret_access_key=credentials["SecretAccessKey"],
|
|
aws_session_token=credentials["SessionToken"],
|
|
)
|
|
return self._assumed_client
|
|
|
|
def _call_with_backoff(self, fn: Any, *args: Any, max_retries: int = 3, **kwargs: Any) -> Any:
|
|
"""Call *fn* with exponential backoff and full jitter on retryable errors."""
|
|
delay = 0.5
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
except ClientError as e:
|
|
code = e.response["Error"]["Code"]
|
|
if code in _RETRYABLE_CODES and attempt < max_retries:
|
|
jitter = random.uniform(0, min(delay, 10.0))
|
|
time.sleep(jitter)
|
|
delay *= 2
|
|
log.warning(
|
|
"Retryable EC2 error (attempt %d/%d): %s",
|
|
attempt + 1,
|
|
max_retries,
|
|
code,
|
|
)
|
|
continue
|
|
category = _ERROR_CATEGORIES.get(code, "unknown")
|
|
raise RuntimeAdapterError(str(e), category=category) from e
|
|
|
|
# Unreachable — loop always returns or raises on every path
|
|
msg = "Retries exhausted"
|
|
raise RuntimeAdapterError(msg, category="unknown")
|
|
|
|
def _discover_tailscale_ip(self, slot_id: str, instance_id: str) -> str | None:
|
|
"""Resolve Tailscale IP for instance identity via local tailscaled LocalAPI."""
|
|
status = self._read_tailscale_status()
|
|
if status is None:
|
|
return None
|
|
|
|
peers_obj = status.get("Peer")
|
|
if not isinstance(peers_obj, dict):
|
|
return None
|
|
|
|
online_candidates: list[tuple[str, str]] = []
|
|
for peer in peers_obj.values():
|
|
if not isinstance(peer, dict):
|
|
continue
|
|
if not self._peer_is_online(peer):
|
|
continue
|
|
hostname = self._peer_hostname(peer)
|
|
if hostname is None:
|
|
continue
|
|
ip = self._peer_tailscale_ip(peer)
|
|
if ip is None:
|
|
continue
|
|
online_candidates.append((hostname, ip))
|
|
|
|
identity = f"nix-builder-{slot_id}-{instance_id}".lower()
|
|
identity_matches = [ip for host, ip in online_candidates if identity in host]
|
|
if len(identity_matches) == 1:
|
|
return identity_matches[0]
|
|
if len(identity_matches) > 1:
|
|
log.warning(
|
|
"tailscale_identity_ambiguous",
|
|
extra={"slot_id": slot_id, "instance_id": instance_id},
|
|
)
|
|
return None
|
|
|
|
slot_identity = f"nix-builder-{slot_id}".lower()
|
|
slot_matches = [ip for host, ip in online_candidates if slot_identity in host]
|
|
if len(slot_matches) == 1:
|
|
return slot_matches[0]
|
|
if len(slot_matches) > 1:
|
|
log.warning("tailscale_slot_ambiguous", extra={"slot_id": slot_id})
|
|
return None
|
|
return None
|
|
|
|
def _read_tailscale_status(self) -> dict[str, Any] | None:
|
|
"""Query local tailscaled LocalAPI status endpoint over Unix socket."""
|
|
conn = _UnixSocketHTTPConnection(self._tailscale_socket_path, timeout=1.0)
|
|
try:
|
|
conn.request(
|
|
"GET",
|
|
"/localapi/v0/status",
|
|
headers={"Host": "local-tailscaled.sock", "Accept": "application/json"},
|
|
)
|
|
response = conn.getresponse()
|
|
if response.status != 200:
|
|
return None
|
|
payload = response.read()
|
|
parsed = json.loads(payload.decode())
|
|
if isinstance(parsed, dict):
|
|
return parsed
|
|
return None
|
|
except (OSError, PermissionError, TimeoutError, json.JSONDecodeError, UnicodeDecodeError):
|
|
return None
|
|
except http.client.HTTPException:
|
|
return None
|
|
finally:
|
|
conn.close()
|
|
|
|
@staticmethod
|
|
def _peer_is_online(peer: dict[str, Any]) -> bool:
|
|
return bool(peer.get("Online"))
|
|
|
|
@staticmethod
|
|
def _peer_hostname(peer: dict[str, Any]) -> str | None:
|
|
host = peer.get("HostName") or peer.get("DNSName")
|
|
if not isinstance(host, str) or not host:
|
|
return None
|
|
return host.strip(".").lower()
|
|
|
|
@staticmethod
|
|
def _peer_tailscale_ip(peer: dict[str, Any]) -> str | None:
|
|
ips = peer.get("TailscaleIPs")
|
|
if not isinstance(ips, list):
|
|
return None
|
|
ipv4 = [ip for ip in ips if isinstance(ip, str) and "." in ip]
|
|
if ipv4:
|
|
return ipv4[0]
|
|
for ip in ips:
|
|
if isinstance(ip, str) and ip:
|
|
return ip
|
|
return None
|
|
|
|
@staticmethod
|
|
def _get_tag(tags: list[dict[str, str]], key: str) -> str | None:
|
|
"""Extract a tag value from an EC2 tag list."""
|
|
for tag in tags:
|
|
if tag.get("Key") == key:
|
|
return tag.get("Value")
|
|
return None
|