Coverage for cli / capacity / models.py: 100.00%
128 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Data classes for capacity checking.
3Instance-type characteristics are resolved live from
4``ec2:DescribeInstanceTypes`` — see :func:`instance_type_info_from_ec2`. This
5module intentionally holds no checked-in instance specification table: one used
6to live here (``GPU_INSTANCE_SPECS``, 25 GPU types) and short-circuited the API,
7which meant a new accelerator family was invisible until someone remembered to
8hand-edit it, and a wrong number propagated silently into NodePool sizing and
9capacity scores. One API call per lookup is cheaper than that failure mode.
10"""
12from __future__ import annotations
14from dataclasses import dataclass, field
15from typing import Any
18class CapacityCheckError(Exception):
19 """A primary capacity availability check failed at the AWS API level.
21 Raised when an availability lookup (e.g. DescribeInstanceTypeOfferings)
22 fails due to throttling, expired/invalid credentials, denied permissions,
23 or a region that isn't opted in. This is distinct from a *successful* lookup
24 that reports an instance type as genuinely not offered — masking such
25 failures as "unavailable" hides real, actionable errors from the caller.
26 """
29@dataclass
30class InstanceTypeInfo:
31 """Compute characteristics of an EC2 instance type, as EC2 reports them.
33 Every field is resolved live from ``ec2:DescribeInstanceTypes``. There is
34 deliberately no checked-in specification table behind this: a hand-maintained
35 catalog silently goes stale, and a stale accelerator count feeds directly
36 into NodePool sizing and capacity scoring. Asking EC2 costs one API call and
37 is always right.
39 The first seven fields keep their historical names and order — the capacity
40 heuristics and the recommender read them positionally — and everything added
41 since carries a default, so a partial record from an older API shape still
42 constructs.
44 Memory conventions, which are easy to get wrong:
46 * ``memory_gib`` is host RAM.
47 * ``gpu_memory_gib`` is the **total** across all accelerators, taken from
48 ``GpuInfo.TotalGpuMemoryInMiB``. It is not per-device: a heterogeneous
49 ``Gpus[]`` list would make ``count x per_device`` wrong, and the
50 on-demand availability heuristics read this field as a total.
51 * ``gpu_devices`` carries the per-model breakdown when you need it.
52 """
54 instance_type: str
55 vcpus: int
56 memory_gib: float
57 gpu_count: int = 0
58 gpu_type: str | None = None
59 gpu_memory_gib: float = 0
60 architecture: str = "x86_64"
62 # --- Region the description came from -------------------------------
63 # DescribeInstanceTypes is region-scoped: a type absent here may exist
64 # elsewhere, so the answer is only meaningful alongside its region.
65 region: str | None = None
67 # --- Processor ------------------------------------------------------
68 cores: int | None = None
69 threads_per_core: int | None = None
70 architectures: list[str] = field(default_factory=list)
71 sustained_clock_speed_ghz: float | None = None
72 processor_manufacturer: str | None = None
73 current_generation: bool | None = None
74 bare_metal: bool | None = None
75 hypervisor: str | None = None
76 burstable: bool | None = None
77 free_tier_eligible: bool | None = None
79 # --- Accelerators ---------------------------------------------------
80 #: Per-model GPU breakdown: name, manufacturer, count, memory_gib.
81 gpu_devices: list[dict[str, Any]] = field(default_factory=list)
82 gpu_manufacturer: str | None = None
83 #: AWS Neuron (Trainium / Inferentia2) devices, same shape as gpu_devices.
84 neuron_devices: list[dict[str, Any]] = field(default_factory=list)
85 neuron_count: int = 0
86 neuron_memory_gib: float = 0
87 #: First-generation Inferentia, reported separately by EC2.
88 inference_accelerators: list[dict[str, Any]] = field(default_factory=list)
89 inference_accelerator_count: int = 0
90 media_accelerators: list[dict[str, Any]] = field(default_factory=list)
91 fpgas: list[dict[str, Any]] = field(default_factory=list)
93 # --- Network --------------------------------------------------------
94 efa_supported: bool | None = None
95 efa_max_interfaces: int | None = None
96 network_performance: str | None = None
97 maximum_network_interfaces: int | None = None
98 maximum_network_cards: int | None = None
99 ipv6_supported: bool | None = None
100 ena_support: str | None = None
101 encryption_in_transit_supported: bool | None = None
103 # --- Storage --------------------------------------------------------
104 instance_storage_supported: bool | None = None
105 instance_storage_total_gb: int | None = None
106 instance_storage_disks: list[dict[str, Any]] = field(default_factory=list)
107 instance_storage_nvme: str | None = None
108 ebs_optimized_support: str | None = None
109 ebs_encryption_support: str | None = None
110 ebs_nvme_support: str | None = None
111 ebs_baseline_iops: int | None = None
112 ebs_maximum_iops: int | None = None
113 ebs_baseline_throughput_mbps: float | None = None
114 ebs_maximum_throughput_mbps: float | None = None
116 # --- Purchasing and placement ---------------------------------------
117 supported_usage_classes: list[str] = field(default_factory=list)
118 supported_placement_strategies: list[str] = field(default_factory=list)
119 dedicated_hosts_supported: bool | None = None
121 # --- Platform capabilities ------------------------------------------
122 supported_root_device_types: list[str] = field(default_factory=list)
123 supported_virtualization_types: list[str] = field(default_factory=list)
124 supported_boot_modes: list[str] = field(default_factory=list)
125 hibernation_supported: bool | None = None
126 auto_recovery_supported: bool | None = None
127 nitro_enclaves_support: str | None = None
128 nitro_tpm_support: str | None = None
130 @property
131 def is_gpu(self) -> bool:
132 return self.gpu_count > 0
134 @property
135 def is_accelerated(self) -> bool:
136 """True when the type carries any accelerator, not just an NVIDIA GPU.
138 ``is_gpu`` deliberately stays GPU-only because the capacity heuristics
139 use it to reason about GPU scarcity specifically; a Trainium node is a
140 different supply pool.
141 """
142 return bool(
143 self.gpu_count
144 or self.neuron_count
145 or self.inference_accelerator_count
146 or self.media_accelerators
147 or self.fpgas
148 )
150 @property
151 def spot_supported(self) -> bool:
152 return "spot" in self.supported_usage_classes
154 @property
155 def capacity_block_supported(self) -> bool:
156 return "capacity-block" in self.supported_usage_classes
159@dataclass
160class SpotPriceInfo:
161 """Spot price information for an instance type."""
163 instance_type: str
164 availability_zone: str
165 current_price: float
166 avg_price_7d: float
167 min_price_7d: float
168 max_price_7d: float
169 price_stability: float # 0-1, higher is more stable
172@dataclass
173class CapacityEstimate:
174 """Capacity availability estimate."""
176 instance_type: str
177 region: str
178 availability_zone: str | None
179 capacity_type: str # "spot" or "on-demand"
180 availability: str # "high", "medium", "low", "unavailable", "unknown"
181 confidence: float # 0-1
182 estimated_wait_time: str | None = None
183 price_per_hour: float | None = None
184 recommendation: str = ""
185 details: dict[str, Any] = field(default_factory=dict)
186 error: str | None = None # Explicit error string for partial/degraded results
189def _mib_to_gib(value: Any) -> float | None:
190 """Convert an EC2 MiB quantity to GiB, or None when absent."""
191 if value is None:
192 return None
193 try:
194 return round(float(value) / 1024, 2)
195 except TypeError, ValueError:
196 return None
199def instance_type_info_from_ec2(
200 record: dict[str, Any], region: str | None = None
201) -> InstanceTypeInfo:
202 """Build an :class:`InstanceTypeInfo` from one ``DescribeInstanceTypes`` entry.
204 A pure function so the mapping can be tested against recorded API shapes
205 without AWS credentials, and so the CLI and any future caller derive the
206 same fields from the same payload.
208 Every read is defensive. EC2 omits whole field groups depending on the
209 instance family (``GpuInfo`` only on GPU types, ``NeuronInfo`` only on
210 Trainium/Inferentia2, ``InstanceStorageInfo`` only when local disks exist),
211 and it adds new groups over time, so a missing key is normal and must not
212 raise. Lists are read in full rather than indexed at ``[0]``: a family with
213 two accelerator models or two supported architectures would otherwise be
214 silently misreported.
215 """
216 vcpu_info = record.get("VCpuInfo") or {}
217 memory_info = record.get("MemoryInfo") or {}
218 processor_info = record.get("ProcessorInfo") or {}
219 gpu_info = record.get("GpuInfo") or {}
220 neuron_info = record.get("NeuronInfo") or {}
221 inference_info = record.get("InferenceAcceleratorInfo") or {}
222 media_info = record.get("MediaAcceleratorInfo") or {}
223 fpga_info = record.get("FpgaInfo") or {}
224 network_info = record.get("NetworkInfo") or {}
225 efa_info = network_info.get("EfaInfo") or {}
226 storage_info = record.get("InstanceStorageInfo") or {}
227 ebs_info = record.get("EbsInfo") or {}
228 ebs_throughput = ebs_info.get("EbsOptimizedInfo") or {}
229 placement_info = record.get("PlacementGroupInfo") or {}
231 architectures = list(processor_info.get("SupportedArchitectures") or [])
233 gpu_devices = [
234 {
235 "name": device.get("Name"),
236 "manufacturer": device.get("Manufacturer"),
237 "count": device.get("Count"),
238 "memory_gib": _mib_to_gib((device.get("MemoryInfo") or {}).get("SizeInMiB")),
239 }
240 for device in (gpu_info.get("Gpus") or [])
241 ]
242 # Sum across models rather than trusting Gpus[0].Count — a heterogeneous
243 # list would otherwise undercount the node's real accelerator budget.
244 gpu_count = sum(int(device.get("count") or 0) for device in gpu_devices)
246 # Prefer EC2's own total. Fall back to count x per-device across models when
247 # the field is absent, because reporting 0 GiB for a node that plainly has
248 # GPUs is worse than a derived figure — the on-demand availability
249 # heuristics read this as a total and would score the node as CPU-only.
250 total_gpu_memory_gib = _mib_to_gib(gpu_info.get("TotalGpuMemoryInMiB"))
251 if total_gpu_memory_gib is None and gpu_devices:
252 derived = sum(
253 float(device.get("memory_gib") or 0) * int(device.get("count") or 0)
254 for device in gpu_devices
255 )
256 total_gpu_memory_gib = round(derived, 2) if derived else None
258 neuron_devices = [
259 {
260 "name": device.get("Name"),
261 "count": device.get("Count"),
262 "core_count": (device.get("CoreInfo") or {}).get("Count"),
263 "core_version": (device.get("CoreInfo") or {}).get("Version"),
264 "memory_gib": _mib_to_gib((device.get("MemoryInfo") or {}).get("SizeInMiB")),
265 }
266 for device in (neuron_info.get("NeuronDevices") or [])
267 ]
269 inference_accelerators = [
270 {
271 "name": device.get("Name"),
272 "manufacturer": device.get("Manufacturer"),
273 "count": device.get("Count"),
274 "memory_gib": _mib_to_gib((device.get("MemoryInfo") or {}).get("SizeInMiB")),
275 }
276 for device in (inference_info.get("Accelerators") or [])
277 ]
279 return InstanceTypeInfo(
280 instance_type=str(record.get("InstanceType") or ""),
281 vcpus=int(vcpu_info.get("DefaultVCpus") or 0),
282 memory_gib=_mib_to_gib(memory_info.get("SizeInMiB")) or 0.0,
283 gpu_count=gpu_count,
284 # Join multiple models so the scalar field stays informative on a
285 # heterogeneous type instead of naming only the first.
286 gpu_type=(
287 "+".join(str(device["name"]) for device in gpu_devices if device.get("name")) or None
288 ),
289 # Total, not per-device — see the InstanceTypeInfo docstring.
290 gpu_memory_gib=total_gpu_memory_gib or 0.0,
291 architecture=architectures[0] if architectures else "x86_64",
292 region=region,
293 cores=vcpu_info.get("DefaultCores"),
294 threads_per_core=vcpu_info.get("DefaultThreadsPerCore"),
295 architectures=architectures,
296 sustained_clock_speed_ghz=processor_info.get("SustainedClockSpeedInGhz"),
297 processor_manufacturer=processor_info.get("Manufacturer"),
298 current_generation=record.get("CurrentGeneration"),
299 bare_metal=record.get("BareMetal"),
300 hypervisor=record.get("Hypervisor"),
301 burstable=record.get("BurstablePerformanceSupported"),
302 free_tier_eligible=record.get("FreeTierEligible"),
303 gpu_devices=gpu_devices,
304 gpu_manufacturer=next(
305 (str(device["manufacturer"]) for device in gpu_devices if device.get("manufacturer")),
306 None,
307 ),
308 neuron_devices=neuron_devices,
309 neuron_count=sum(int(device.get("count") or 0) for device in neuron_devices),
310 neuron_memory_gib=_mib_to_gib(neuron_info.get("TotalNeuronDeviceMemoryInMiB")) or 0.0,
311 inference_accelerators=inference_accelerators,
312 inference_accelerator_count=sum(
313 int(device.get("count") or 0) for device in inference_accelerators
314 ),
315 media_accelerators=[
316 {
317 "name": device.get("Name"),
318 "manufacturer": device.get("Manufacturer"),
319 "count": device.get("Count"),
320 "memory_gib": _mib_to_gib((device.get("MemoryInfo") or {}).get("SizeInMiB")),
321 }
322 for device in (media_info.get("Accelerators") or [])
323 ],
324 fpgas=[
325 {
326 "name": device.get("Name"),
327 "manufacturer": device.get("Manufacturer"),
328 "count": device.get("Count"),
329 "memory_gib": _mib_to_gib((device.get("MemoryInfo") or {}).get("SizeInMiB")),
330 }
331 for device in (fpga_info.get("Fpgas") or [])
332 ],
333 efa_supported=network_info.get("EfaSupported"),
334 efa_max_interfaces=efa_info.get("MaximumEfaInterfaces"),
335 network_performance=network_info.get("NetworkPerformance"),
336 maximum_network_interfaces=network_info.get("MaximumNetworkInterfaces"),
337 maximum_network_cards=network_info.get("MaximumNetworkCards"),
338 ipv6_supported=network_info.get("Ipv6Supported"),
339 ena_support=network_info.get("EnaSupport"),
340 encryption_in_transit_supported=network_info.get("EncryptionInTransitSupported"),
341 instance_storage_supported=record.get("InstanceStorageSupported"),
342 instance_storage_total_gb=storage_info.get("TotalSizeInGB"),
343 instance_storage_disks=[
344 {
345 "size_gb": disk.get("SizeInGB"),
346 "count": disk.get("Count"),
347 "type": disk.get("Type"),
348 }
349 for disk in (storage_info.get("Disks") or [])
350 ],
351 instance_storage_nvme=storage_info.get("NvmeSupport"),
352 ebs_optimized_support=ebs_info.get("EbsOptimizedSupport"),
353 ebs_encryption_support=ebs_info.get("EncryptionSupport"),
354 ebs_nvme_support=ebs_info.get("NvmeSupport"),
355 ebs_baseline_iops=ebs_throughput.get("BaselineIops"),
356 ebs_maximum_iops=ebs_throughput.get("MaximumIops"),
357 ebs_baseline_throughput_mbps=ebs_throughput.get("BaselineThroughputInMBps"),
358 ebs_maximum_throughput_mbps=ebs_throughput.get("MaximumThroughputInMBps"),
359 supported_usage_classes=list(record.get("SupportedUsageClasses") or []),
360 supported_placement_strategies=list(placement_info.get("SupportedStrategies") or []),
361 dedicated_hosts_supported=record.get("DedicatedHostsSupported"),
362 supported_root_device_types=list(record.get("SupportedRootDeviceTypes") or []),
363 supported_virtualization_types=list(record.get("SupportedVirtualizationTypes") or []),
364 supported_boot_modes=list(record.get("SupportedBootModes") or []),
365 hibernation_supported=record.get("HibernationSupported"),
366 auto_recovery_supported=record.get("AutoRecoverySupported"),
367 nitro_enclaves_support=record.get("NitroEnclavesSupport"),
368 nitro_tpm_support=record.get("NitroTpmSupport"),
369 )