Coverage for cli/capacity/checker.py: 96.84%
759 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""
2Single-region EC2 capacity checker using real AWS signals.
4This is the core capacity intelligence module (~1265 lines). It queries multiple
5AWS APIs to build a comprehensive picture of GPU/accelerator availability in a
6single region. The MultiRegionCapacityChecker in multi_region.py calls this for
7each region in parallel.
9Data Sources:
10 - EC2 GetSpotPlacementScores: likelihood of getting spot capacity (1-10 score)
11 - EC2 DescribeSpotPriceHistory: current and historical spot prices (7-day window)
12 - EC2 DescribeInstanceTypes: vCPU, memory, GPU count/type/memory, EFA support
13 - EC2 DescribeInstanceTypeOfferings: which instance types are available in the region
14 - EC2 DescribeCapacityBlockOfferings: purchasable Capacity Blocks for ML workloads
15 - EC2 PurchaseCapacityBlock: (optional) purchase a Capacity Block by offering ID
17Key Classes:
18 CapacityChecker: Main class. Instantiated with a region and optional GCOConfig.
19 - check_capacity(instance_type) → CapacityEstimate
20 - get_spot_prices(instance_type) → list[SpotPriceInfo]
21 - get_instance_info(instance_type) → InstanceTypeInfo
22 - check_capacity_blocks(instance_type, count, duration) → list[dict]
24Output Models (defined in models.py):
25 - CapacityEstimate: spot score, price, trend, on-demand price, instance specs
26 - SpotPriceInfo: AZ, price, timestamp
27 - InstanceTypeInfo: vCPU, memory, GPU count/type/memory, EFA, architecture
29The GPU_INSTANCE_SPECS lookup table in models.py provides offline specs for common
30GPU instances so the checker can return useful information even when the EC2 API
31is unavailable or the instance type isn't offered in the region.
32"""
34from __future__ import annotations
36import json
37import logging
38import statistics
39from concurrent.futures import ThreadPoolExecutor, as_completed
40from datetime import UTC, datetime, timedelta
41from typing import Any
43import boto3
44from botocore.config import Config
45from botocore.exceptions import BotoCoreError, ClientError
47from cli.config import GCOConfig, get_config
49from . import blocks
50from .models import (
51 GPU_INSTANCE_SPECS,
52 CapacityCheckError,
53 CapacityEstimate,
54 InstanceTypeInfo,
55 SpotPriceInfo,
56)
58logger = logging.getLogger(__name__)
60# Capacity Block offering API error codes that mean "this instance type / region
61# simply doesn't have Capacity Blocks" rather than a real failure — expected and
62# handled quietly. ``Unsupported`` / ``UnsupportedOperation`` / ``InvalidAction``
63# cover regions where the Capacity Block API isn't available at all (e.g.
64# DescribeCapacityBlockOfferings returns ``InvalidAction`` in eu-west-1).
65_CB_EXPECTED_ERROR_CODES = frozenset(
66 {
67 "Unsupported",
68 "UnsupportedOperation",
69 "InvalidAction",
70 "InvalidParameterValue",
71 "InvalidParameterCombination",
72 }
73)
75# Adaptive client-side retry config for DescribeCapacityBlockOfferings. The
76# find_capacity_blocks sweep fans out region x duration probes in parallel, which
77# can trip the API's per-account request-rate limit (RequestLimitExceeded).
78# Adaptive mode adds client-side rate limiting plus exponential backoff, and a
79# higher attempt budget gives each probe room to succeed under throttling rather
80# than silently returning an empty (misleading) result for a throttled region.
81_CB_RETRY_CONFIG = Config(retries={"max_attempts": 10, "mode": "adaptive"})
83# A conservative cap on the parallel fan-out so a wide region x duration sweep
84# can't open an unbounded number of threads / sockets at once.
85_MAX_SEARCH_WORKERS = 12
88def _instance_desc(instance_type: str, gpu_count: int, gpu_type: str, total_gpu_mem: float) -> str:
89 """Build a human-readable instance description."""
90 if gpu_count > 0 and gpu_type:
91 mem_str = f", {total_gpu_mem:.0f}GB" if total_gpu_mem else ""
92 return f"{instance_type} ({gpu_count}x {gpu_type}{mem_str})"
93 return instance_type
96def _offering_fee(offering: dict[str, Any]) -> float:
97 """Sort key: an offering's upfront fee in USD, with missing fees sorting last."""
98 fee = offering.get("upfront_fee_usd")
99 if fee is None:
100 fee = blocks.parse_upfront_fee(offering.get("upfront_fee"))
101 return fee if fee is not None else float("inf")
104class CapacityChecker:
105 """
106 Checks EC2 capacity availability using real AWS capacity signals.
108 Uses:
109 - Spot Placement Score API for spot capacity estimates
110 - EC2 describe-instance-type-offerings for regional availability
111 - Spot price history for pricing trends
112 - On-demand pricing API
113 """
115 def __init__(self, config: GCOConfig | None = None):
116 self.config = config or get_config()
117 self._session = boto3.Session()
118 self._pricing_cache: dict[str, Any] = {}
119 self._offerings_cache: dict[str, set[str]] = {}
121 def get_instance_info(self, instance_type: str) -> InstanceTypeInfo | None:
122 """Get information about an instance type."""
123 if instance_type in GPU_INSTANCE_SPECS:
124 return GPU_INSTANCE_SPECS[instance_type]
126 # Try to get from EC2 API
127 try:
128 ec2 = self._session.client("ec2", region_name="us-east-1")
129 response = ec2.describe_instance_types(InstanceTypes=[instance_type])
131 if response["InstanceTypes"]:
132 info = response["InstanceTypes"][0]
133 vcpus = info["VCpuInfo"]["DefaultVCpus"]
134 memory = info["MemoryInfo"]["SizeInMiB"] / 1024
136 gpu_count = 0
137 gpu_type = None
138 gpu_memory = 0
140 if "GpuInfo" in info:
141 gpus = info["GpuInfo"].get("Gpus", [])
142 if gpus:
143 gpu_count = gpus[0].get("Count", 0)
144 gpu_type = gpus[0].get("Name")
145 gpu_memory = gpus[0].get("MemoryInfo", {}).get("SizeInMiB", 0) / 1024
147 arch = info["ProcessorInfo"]["SupportedArchitectures"][0]
149 return InstanceTypeInfo(
150 instance_type=instance_type,
151 vcpus=vcpus,
152 memory_gib=memory,
153 gpu_count=gpu_count,
154 gpu_type=gpu_type,
155 gpu_memory_gib=gpu_memory,
156 architecture=arch,
157 )
158 except ClientError as e:
159 logger.debug("Failed to describe instance type %s: %s", instance_type, e)
160 except Exception as e:
161 logger.warning("Unexpected error getting instance info for %s: %s", instance_type, e)
163 return None
165 def validate_instance_type(self, instance_type: str) -> dict[str, Any]:
166 """Classify an instance type as valid / invalid, with friendly normalization.
168 The offering APIs collapse "unknown instance type" and "valid type with
169 zero offerings" into the same empty result. This method separates them so
170 callers can tell a typo from genuine unavailability:
172 * A type in ``GPU_INSTANCE_SPECS`` is valid and ``known`` without any AWS
173 call.
174 * Otherwise EC2 ``DescribeInstanceTypes`` is consulted; an
175 ``InvalidInstanceType`` (or empty result) marks it invalid, while a
176 transient API error leaves it valid-but-unverified with a note.
177 * Friendly aliases are expanded (``p6-b200`` -> ``p6-b200.48xlarge``,
178 ``p6-b300`` -> ``p6-b300.48xlarge``) and UltraServer-only families
179 (the Grace-Blackwell ``gb200``/``gb300`` superchips, sold only as
180 ``P6e-GB`` UltraServers) are flagged invalid-for-``InstanceType`` with
181 guidance toward the UltraServer search flow.
183 Returns a dict: ``requested``, ``instance_type`` (canonical), ``valid``,
184 ``known``, ``note``, ``gpu_count``.
185 """
186 canonical, note = blocks.normalize_instance_type(instance_type)
187 result: dict[str, Any] = {
188 "requested": instance_type,
189 "instance_type": canonical,
190 "valid": True,
191 "known": False,
192 "note": note,
193 "gpu_count": None,
194 }
196 # UltraServer-only families are real accelerators but not standalone EC2
197 # instance types, so InstanceType-based searches can never resolve them.
198 if (instance_type or "").strip().lower() in blocks.NON_STANDALONE_INSTANCE_NOTES:
199 result["valid"] = False
200 return result
202 if canonical in GPU_INSTANCE_SPECS:
203 spec = GPU_INSTANCE_SPECS[canonical]
204 result["known"] = True
205 result["gpu_count"] = spec.gpu_count
206 return result
208 # Not in the offline table — ask EC2, separating "invalid" from "API error".
209 try:
210 ec2 = self._session.client("ec2", region_name="us-east-1")
211 response = ec2.describe_instance_types(InstanceTypes=[canonical])
212 types = response.get("InstanceTypes", [])
213 if not types:
214 result["valid"] = False
215 return result
216 gpu_info = types[0].get("GpuInfo") or {}
217 gpus = gpu_info.get("Gpus", [])
218 result["gpu_count"] = gpus[0].get("Count", 0) if gpus else 0
219 except ClientError as e:
220 code = e.response.get("Error", {}).get("Code", "")
221 if code in ("InvalidInstanceType", "InvalidParameterValue"):
222 result["valid"] = False
223 else:
224 logger.debug("Could not verify instance type %s: %s", canonical, e)
225 result["note"] = note or f"Could not verify instance type via EC2 ({code or e})."
226 except Exception as e:
227 logger.debug("Unexpected error validating instance type %s: %s", canonical, e)
228 result["note"] = note or f"Could not verify instance type: {e}"
229 return result
231 def check_instance_available_in_region(self, instance_type: str, region: str) -> bool:
232 """Check whether an instance type is offered in a region.
234 Returns ``True``/``False`` only from a *successful* offerings lookup, so a
235 ``False`` genuinely means "not offered in this region". If the underlying
236 DescribeInstanceTypeOfferings call fails (throttling, expired/invalid
237 credentials, denied permissions, region not opted in) this raises
238 :class:`CapacityCheckError` instead of silently returning ``False`` — a
239 failed check must not be reported to the user as "not available".
240 """
241 cache_key = f"{region}"
242 if cache_key not in self._offerings_cache:
243 try:
244 ec2 = self._session.client("ec2", region_name=region)
245 paginator = ec2.get_paginator("describe_instance_type_offerings")
246 offerings = set()
247 for page in paginator.paginate(LocationType="region"):
248 for offering in page["InstanceTypeOfferings"]:
249 offerings.add(offering["InstanceType"])
250 self._offerings_cache[cache_key] = offerings
251 except (ClientError, BotoCoreError) as e:
252 logger.warning("Failed to check instance offerings in %s: %s", region, e)
253 raise CapacityCheckError(
254 f"Could not check instance availability in {region}: {e}"
255 ) from e
257 return instance_type in self._offerings_cache[cache_key]
259 def get_availability_zones(self, region: str) -> list[str]:
260 """Get availability zones for a region."""
261 try:
262 ec2 = self._session.client("ec2", region_name=region)
263 response = ec2.describe_availability_zones(
264 Filters=[{"Name": "state", "Values": ["available"]}]
265 )
266 return [az["ZoneName"] for az in response["AvailabilityZones"]]
267 except ClientError as e:
268 logger.warning("Failed to get availability zones for %s: %s", region, e)
269 return []
270 except Exception as e:
271 logger.warning("Unexpected error getting AZs for %s: %s", region, e)
272 return []
274 def get_az_coverage(self, instance_type: str, region: str) -> float | None:
275 """Get the fraction of AZs in a region that offer this instance type.
277 Returns a value between 0.0 and 1.0, or None if we can't determine it.
278 Constrained instances are often available in fewer AZs.
279 """
280 try:
281 ec2 = self._session.client("ec2", region_name=region)
282 total_azs = self.get_availability_zones(region)
283 if not total_azs:
284 return None
286 paginator = ec2.get_paginator("describe_instance_type_offerings")
287 offering_azs = set()
288 for page in paginator.paginate(
289 LocationType="availability-zone",
290 Filters=[{"Name": "instance-type", "Values": [instance_type]}],
291 ):
292 for offering in page["InstanceTypeOfferings"]:
293 offering_azs.add(offering["Location"])
295 return len(offering_azs) / len(total_azs) if total_azs else None
296 except Exception as e:
297 logger.warning("Failed to get AZ coverage for %s in %s: %s", instance_type, region, e)
298 return None
300 def get_spot_placement_score(
301 self, instance_type: str, region: str, target_capacity: int = 1
302 ) -> dict[str, int]:
303 """
304 Get Spot Placement Score for an instance type.
306 The Spot Placement Score (1-10) indicates the likelihood of getting
307 spot capacity. Higher scores mean better availability.
309 Returns:
310 Dict mapping AZ to score (1-10), or empty if not available
311 """
312 try:
313 ec2 = self._session.client("ec2", region_name=region)
315 response = ec2.get_spot_placement_scores(
316 InstanceTypes=[instance_type],
317 TargetCapacity=target_capacity,
318 TargetCapacityUnitType="units",
319 RegionNames=[region],
320 SingleAvailabilityZone=False,
321 )
323 scores = {}
324 for recommendation in response.get("SpotPlacementScores", []):
325 # Regional score
326 if "AvailabilityZoneId" not in recommendation:
327 scores["regional"] = recommendation.get("Score", 0)
328 else:
329 az_id = recommendation["AvailabilityZoneId"]
330 scores[az_id] = recommendation.get("Score", 0)
332 return scores
334 except ClientError as e:
335 error_code = e.response.get("Error", {}).get("Code", "")
336 if error_code in ("InvalidParameterValue", "UnsupportedOperation"):
337 return {}
338 raise
339 except Exception as e:
340 logger.warning(
341 "Failed to get spot placement scores for %s in %s: %s", instance_type, region, e
342 )
343 return {}
345 def get_spot_price_history(
346 self, instance_type: str, region: str, days: int = 7
347 ) -> list[SpotPriceInfo]:
348 """Get spot price history for an instance type."""
349 ec2 = self._session.client("ec2", region_name=region)
351 end_time = datetime.now(UTC)
352 start_time = end_time - timedelta(days=days)
354 try:
355 response = ec2.describe_spot_price_history(
356 InstanceTypes=[instance_type],
357 ProductDescriptions=["Linux/UNIX"],
358 StartTime=start_time,
359 EndTime=end_time,
360 )
362 # Group by availability zone
363 az_prices: dict[str, list[float]] = {}
364 for item in response["SpotPriceHistory"]:
365 az = item["AvailabilityZone"]
366 price = float(item["SpotPrice"])
367 if az not in az_prices:
368 az_prices[az] = []
369 az_prices[az].append(price)
371 results = []
372 for az, prices in az_prices.items():
373 if not prices: 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true
374 continue
376 current = prices[0]
377 avg = statistics.mean(prices)
378 min_price = min(prices)
379 max_price = max(prices)
381 if avg > 0:
382 std_dev = statistics.stdev(prices) if len(prices) > 1 else 0
383 cv = std_dev / avg
384 stability = max(0, 1 - cv)
385 else:
386 stability = 0
388 results.append(
389 SpotPriceInfo(
390 instance_type=instance_type,
391 availability_zone=az,
392 current_price=current,
393 avg_price_7d=avg,
394 min_price_7d=min_price,
395 max_price_7d=max_price,
396 price_stability=stability,
397 )
398 )
400 return results
402 except ClientError as e:
403 if "InvalidParameterValue" in str(e):
404 return []
405 raise
407 def get_on_demand_price(self, instance_type: str, region: str) -> float | None:
408 """Get on-demand price for an instance type."""
409 cache_key = f"{instance_type}:{region}"
410 if cache_key in self._pricing_cache:
411 cached_value = self._pricing_cache[cache_key]
412 return float(cached_value) if cached_value is not None else None
414 try:
415 pricing = self._session.client("pricing", region_name="us-east-1")
417 region_names = {
418 "us-east-1": "US East (N. Virginia)",
419 "us-east-2": "US East (Ohio)",
420 "us-west-1": "US West (N. California)",
421 "us-west-2": "US West (Oregon)",
422 "eu-west-1": "EU (Ireland)",
423 "eu-west-2": "EU (London)",
424 "eu-central-1": "EU (Frankfurt)",
425 "ap-northeast-1": "Asia Pacific (Tokyo)",
426 "ap-southeast-1": "Asia Pacific (Singapore)",
427 "ap-southeast-2": "Asia Pacific (Sydney)",
428 }
430 location = region_names.get(region, region)
432 response = pricing.get_products(
433 ServiceCode="AmazonEC2",
434 Filters=[
435 {"Type": "TERM_MATCH", "Field": "instanceType", "Value": instance_type},
436 {"Type": "TERM_MATCH", "Field": "location", "Value": location},
437 {"Type": "TERM_MATCH", "Field": "operatingSystem", "Value": "Linux"},
438 {"Type": "TERM_MATCH", "Field": "tenancy", "Value": "Shared"},
439 {"Type": "TERM_MATCH", "Field": "preInstalledSw", "Value": "NA"},
440 {"Type": "TERM_MATCH", "Field": "capacitystatus", "Value": "Used"},
441 ],
442 MaxResults=1,
443 )
445 if response["PriceList"]:
446 price_data = json.loads(response["PriceList"][0])
447 terms = price_data.get("terms", {}).get("OnDemand", {})
448 for term in terms.values():
449 for price_dim in term.get("priceDimensions", {}).values():
450 price = float(price_dim["pricePerUnit"]["USD"])
451 self._pricing_cache[cache_key] = price
452 return price
454 except ClientError as e:
455 logger.warning(
456 "Failed to get on-demand price for %s in %s: %s", instance_type, region, e
457 )
458 except Exception as e:
459 logger.warning(
460 "Unexpected error getting pricing for %s in %s: %s", instance_type, region, e
461 )
463 return None
465 def estimate_capacity(
466 self, instance_type: str, region: str, capacity_type: str = "both"
467 ) -> list[CapacityEstimate]:
468 """
469 Estimate capacity availability using real AWS signals.
471 Args:
472 instance_type: EC2 instance type
473 region: AWS region
474 capacity_type: "spot", "on-demand", or "both"
476 Returns:
477 List of CapacityEstimate objects
478 """
479 estimates = []
481 # Check if instance type is available in region
482 if not self.check_instance_available_in_region(instance_type, region):
483 return [
484 CapacityEstimate(
485 instance_type=instance_type,
486 region=region,
487 availability_zone=None,
488 capacity_type="both",
489 availability="unavailable",
490 confidence=1.0,
491 recommendation=f"{instance_type} is not available in {region}",
492 details={"reason": "Instance type not offered in region"},
493 )
494 ]
496 instance_info = self.get_instance_info(instance_type)
498 if capacity_type in ("spot", "both"):
499 spot_estimates = self._estimate_spot_capacity(instance_type, region, instance_info)
500 estimates.extend(spot_estimates)
502 if capacity_type in ("on-demand", "both"):
503 # Pass spot placement scores to on-demand estimator as a scarcity signal
504 spot_scores = (
505 self.get_spot_placement_score(instance_type, region)
506 if capacity_type == "on-demand"
507 else {}
508 )
509 # If we already fetched spot estimates, extract the scores from them
510 if spot_estimates := [e for e in estimates if e.capacity_type == "spot"]:
511 spot_scores = {
512 e.availability_zone or "unknown": e.details.get("spot_placement_score", 0)
513 for e in spot_estimates
514 if e.details.get("spot_placement_score") is not None
515 }
516 # Also gather spot price data for price-ratio signal
517 spot_prices = self.get_spot_price_history(instance_type, region)
518 od_estimate = self._estimate_on_demand_capacity(
519 instance_type, region, instance_info, spot_scores, spot_prices
520 )
521 if od_estimate:
522 estimates.append(od_estimate)
524 return estimates
526 def _estimate_spot_capacity(
527 self, instance_type: str, region: str, instance_info: InstanceTypeInfo | None
528 ) -> list[CapacityEstimate]:
529 """Estimate spot capacity using Spot Placement Score and price history."""
530 estimates = []
532 # Get Spot Placement Score (primary signal)
533 placement_scores = self.get_spot_placement_score(instance_type, region)
535 # Get spot prices for pricing info
536 spot_prices = self.get_spot_price_history(instance_type, region)
537 price_by_az = {sp.availability_zone: sp for sp in spot_prices}
539 on_demand_price = self.get_on_demand_price(instance_type, region)
541 # Get AZs in the region
542 azs = self.get_availability_zones(region)
544 if placement_scores:
545 # Use Spot Placement Score as primary signal
546 regional_score = placement_scores.get("regional", 0)
548 for az in azs:
549 # Try to get AZ-specific score, fall back to regional
550 az_id = az # Note: might need to map zone name to zone ID
551 score = placement_scores.get(az_id, regional_score)
553 # Convert score (1-10) to availability
554 if score >= 8:
555 availability = "high"
556 recommendation = "Excellent spot availability"
557 elif score >= 5:
558 availability = "medium"
559 recommendation = "Good spot availability, some interruption risk"
560 elif score >= 3:
561 availability = "low"
562 recommendation = "Limited spot capacity, consider alternatives"
563 else:
564 availability = "low"
565 recommendation = "Very limited spot capacity"
567 spot_info = price_by_az.get(az)
568 price = spot_info.current_price if spot_info else None
570 details: dict[str, Any] = {
571 "spot_placement_score": score,
572 "score_interpretation": f"{score}/10",
573 }
575 if spot_info:
576 details["current_price"] = spot_info.current_price
577 details["avg_price_7d"] = spot_info.avg_price_7d
578 details["price_stability"] = f"{spot_info.price_stability:.2f}"
580 if on_demand_price and price:
581 savings = (1 - price / on_demand_price) * 100
582 details["savings_vs_on_demand"] = f"{savings:.1f}%"
583 details["on_demand_price"] = on_demand_price
585 estimates.append(
586 CapacityEstimate(
587 instance_type=instance_type,
588 region=region,
589 availability_zone=az,
590 capacity_type="spot",
591 availability=availability,
592 confidence=0.85, # Spot Placement Score is reliable
593 price_per_hour=price,
594 recommendation=recommendation,
595 details=details,
596 )
597 )
599 elif spot_prices:
600 # Fall back to price-based estimation if no placement score
601 for spot_info in spot_prices:
602 # Use price stability as a proxy (less reliable)
603 if spot_info.price_stability > 0.8:
604 availability = "medium"
605 recommendation = "Spot prices stable, likely available"
606 elif spot_info.price_stability > 0.5:
607 availability = "low"
608 recommendation = "Spot prices volatile, capacity uncertain"
609 else:
610 availability = "low"
611 recommendation = "High price volatility, limited capacity likely"
613 details = {
614 "current_price": spot_info.current_price,
615 "avg_price_7d": spot_info.avg_price_7d,
616 "price_stability": f"{spot_info.price_stability:.2f}",
617 "note": "Estimate based on price history (Spot Placement Score unavailable)",
618 }
620 if on_demand_price:
621 savings = (1 - spot_info.current_price / on_demand_price) * 100
622 details["savings_vs_on_demand"] = f"{savings:.1f}%"
624 estimates.append(
625 CapacityEstimate(
626 instance_type=instance_type,
627 region=region,
628 availability_zone=spot_info.availability_zone,
629 capacity_type="spot",
630 availability=availability,
631 confidence=0.5, # Lower confidence without placement score
632 price_per_hour=spot_info.current_price,
633 recommendation=recommendation,
634 details=details,
635 )
636 )
638 if not estimates:
639 estimates.append(
640 CapacityEstimate(
641 instance_type=instance_type,
642 region=region,
643 availability_zone=None,
644 capacity_type="spot",
645 availability="unknown",
646 confidence=0.1,
647 recommendation=f"No spot data available for {instance_type} in {region}",
648 details={"reason": "No spot price history or placement score available"},
649 )
650 )
652 return estimates
654 def _estimate_on_demand_capacity(
655 self,
656 instance_type: str,
657 region: str,
658 instance_info: InstanceTypeInfo | None,
659 spot_placement_scores: dict[str, int] | None = None,
660 spot_prices: list[SpotPriceInfo] | None = None,
661 ) -> CapacityEstimate | None:
662 """Estimate on-demand capacity using live signals for ALL instance types.
664 Uses spot placement scores, instance size (vCPUs, memory, GPUs),
665 pricing, and spot-to-on-demand price ratios as universal scarcity
666 signals — no hardcoded instance families or GPU type lists.
667 """
668 on_demand_price = self.get_on_demand_price(instance_type, region)
670 is_offered = self.check_instance_available_in_region(instance_type, region)
672 if not is_offered:
673 return CapacityEstimate(
674 instance_type=instance_type,
675 region=region,
676 availability_zone=None,
677 capacity_type="on-demand",
678 availability="unavailable",
679 confidence=1.0,
680 recommendation=f"{instance_type} is not offered in {region}",
681 details={"reason": "Instance type not offered in region"},
682 )
684 # Fetch spot placement scores if not provided (on-demand only mode)
685 if spot_placement_scores is None:
686 spot_placement_scores = self.get_spot_placement_score(instance_type, region)
688 if spot_prices is None:
689 spot_prices = self.get_spot_price_history(instance_type, region)
691 az_coverage = self.get_az_coverage(instance_type, region)
693 availability, confidence, recommendation = self._assess_on_demand_availability(
694 instance_type,
695 instance_info,
696 on_demand_price,
697 spot_placement_scores,
698 spot_prices,
699 az_coverage,
700 )
702 if on_demand_price:
703 recommendation += f" Price: ${on_demand_price:.4f}/hr."
704 else:
705 confidence -= 0.1
706 recommendation += " Pricing data unavailable."
708 details: dict[str, Any] = {
709 "price_per_hour": on_demand_price,
710 "is_gpu": instance_info.is_gpu if instance_info else False,
711 }
712 if spot_placement_scores:
713 scores = [s for s in spot_placement_scores.values() if s > 0]
714 if scores:
715 details["avg_spot_placement_score"] = round(sum(scores) / len(scores), 1)
717 return CapacityEstimate(
718 instance_type=instance_type,
719 region=region,
720 availability_zone=None,
721 capacity_type="on-demand",
722 availability=availability,
723 confidence=confidence,
724 price_per_hour=on_demand_price,
725 recommendation=recommendation,
726 details=details,
727 )
729 @staticmethod
730 def _assess_on_demand_availability(
731 instance_type: str,
732 instance_info: InstanceTypeInfo | None,
733 on_demand_price: float | None,
734 spot_placement_scores: dict[str, int] | None = None,
735 spot_prices: list[SpotPriceInfo] | None = None,
736 az_coverage: float | None = None,
737 ) -> tuple[str, float, str]:
738 """Assess on-demand availability using only live market signals.
740 Five live signals, zero hardcoded instance families:
741 1. Spot placement score — AWS's own capacity assessment (1-10)
742 2. Spot-to-on-demand price ratio — when spot approaches on-demand price,
743 the spot market has very little excess capacity
744 3. Spot price volatility — unstable prices reflect capacity fluctuations
745 4. AZ coverage — fraction of AZs that offer this instance type;
746 constrained instances are often available in fewer AZs
747 5. Spot price availability — how many AZs have spot price data;
748 missing price data in some AZs suggests limited capacity there
750 Confidence scales with the number of live signals available.
751 When no signals exist, returns "unknown" rather than guessing.
753 Returns:
754 Tuple of (availability, confidence, recommendation)
755 """
756 price = on_demand_price or 0
757 gpu_count = instance_info.gpu_count if instance_info else 0
758 gpu_type = (instance_info.gpu_type or "") if instance_info else ""
759 total_gpu_mem = instance_info.gpu_memory_gib if instance_info else 0
761 # --- Signal 1: Spot placement score ---
762 avg_spot_score = 0.0
763 has_spot_score = False
764 if spot_placement_scores:
765 scores = [s for s in spot_placement_scores.values() if s > 0]
766 if scores:
767 avg_spot_score = sum(scores) / len(scores)
768 has_spot_score = True
770 # --- Signal 2 & 3: Spot price ratio and volatility ---
771 avg_spot_ratio = 0.0
772 avg_stability = 1.0
773 has_price_signal = False
774 if spot_prices and price > 0:
775 ratios = [sp.current_price / price for sp in spot_prices if sp.current_price > 0]
776 if ratios:
777 avg_spot_ratio = sum(ratios) / len(ratios)
778 has_price_signal = True
779 stabilities = [sp.price_stability for sp in spot_prices]
780 if stabilities: 780 ↛ 784line 780 didn't jump to line 784 because the condition on line 780 was always true
781 avg_stability = sum(stabilities) / len(stabilities)
783 # --- Signal 4: AZ coverage (passed in from caller) ---
784 has_az_signal = az_coverage is not None
786 # --- Combine live signals into scarcity (0.0 - 1.0) ---
787 scarcity = 0.0
788 signal_count = 0
790 if has_spot_score:
791 signal_count += 1
792 if avg_spot_score <= 2:
793 scarcity += 0.5
794 elif avg_spot_score <= 4:
795 scarcity += 0.3
796 elif avg_spot_score <= 6:
797 scarcity += 0.15
799 if has_price_signal:
800 signal_count += 1
801 # Spot price near on-demand = spot market has minimal excess capacity
802 if avg_spot_ratio >= 0.9:
803 scarcity += 0.3
804 elif avg_spot_ratio >= 0.7:
805 scarcity += 0.15
806 elif avg_spot_ratio >= 0.5:
807 scarcity += 0.05
809 # Price instability = capacity fluctuations
810 if avg_stability < 0.6:
811 scarcity += 0.1
812 elif avg_stability < 0.8:
813 scarcity += 0.05
815 if has_az_signal and az_coverage is not None:
816 signal_count += 1
817 # Available in fewer than half the AZs = constrained
818 if az_coverage <= 0.3:
819 scarcity += 0.2
820 elif az_coverage <= 0.5:
821 scarcity += 0.1
823 # --- Confidence scales with signal count ---
824 confidence = min(0.5 + (signal_count * 0.12), 0.9)
826 # --- Map scarcity to availability ---
827 desc = _instance_desc(instance_type, gpu_count, gpu_type, total_gpu_mem)
829 if signal_count == 0:
830 # No live data — be honest about it
831 return (
832 "unknown",
833 0.3,
834 f"No live capacity signals available for {instance_type}."
835 " Unable to assess on-demand availability.",
836 )
838 if scarcity >= 0.6:
839 return (
840 "low",
841 confidence,
842 f"On-demand {desc} is extremely scarce based on live capacity signals."
843 " Capacity reservations or Capacity Blocks are strongly recommended.",
844 )
846 if scarcity >= 0.35:
847 return (
848 "low",
849 confidence,
850 f"On-demand {desc} has limited availability based on live capacity signals."
851 " Consider capacity reservations.",
852 )
854 if scarcity >= 0.15:
855 return (
856 "medium",
857 confidence,
858 f"On-demand {instance_type} may have constrained availability"
859 " based on current market conditions.",
860 )
862 return (
863 "high",
864 confidence,
865 f"On-demand capacity likely available for {instance_type}"
866 " based on live capacity signals.",
867 )
869 def recommend_capacity_type(
870 self, instance_type: str, region: str, fault_tolerance: str = "medium"
871 ) -> tuple[str, str]:
872 """
873 Recommend spot vs on-demand based on actual capacity and requirements.
875 Args:
876 instance_type: EC2 instance type
877 region: AWS region
878 fault_tolerance: "high" (can handle interruptions),
879 "medium" (some tolerance),
880 "low" (needs stability)
882 Returns:
883 Tuple of (recommended_capacity_type, explanation)
884 """
885 estimates = self.estimate_capacity(instance_type, region, "both")
887 spot_estimates = [e for e in estimates if e.capacity_type == "spot"]
888 od_estimates = [e for e in estimates if e.capacity_type == "on-demand"]
890 # Check for unavailable
891 if any(e.availability == "unavailable" for e in estimates):
892 return "unavailable", f"{instance_type} is not available in {region}"
894 # Get best spot option (highest availability)
895 best_spot = None
896 if spot_estimates:
897 available_spots = [e for e in spot_estimates if e.availability != "unknown"]
898 if available_spots:
899 # Sort by availability (high > medium > low) then by price
900 avail_order = {"high": 0, "medium": 1, "low": 2}
901 best_spot = min(
902 available_spots,
903 key=lambda x: (avail_order.get(x.availability, 3), x.price_per_hour or 999),
904 )
906 od_estimate = od_estimates[0] if od_estimates else None
908 # Decision logic based on fault tolerance and actual availability
909 if fault_tolerance == "low":
910 if od_estimate and od_estimate.availability in ("high", "medium"):
911 return "on-demand", "Low fault tolerance requires stable on-demand capacity"
912 return (
913 "on-demand",
914 "On-demand recommended but capacity may be limited; consider capacity reservation",
915 )
917 if best_spot:
918 if best_spot.availability == "high":
919 savings = ""
920 if best_spot.price_per_hour and od_estimate and od_estimate.price_per_hour:
921 pct = (1 - best_spot.price_per_hour / od_estimate.price_per_hour) * 100
922 savings = f" (save ~{pct:.0f}%)"
923 return "spot", f"High spot availability (score-based){savings}"
925 if best_spot.availability == "medium":
926 if fault_tolerance == "high":
927 return "spot", "Medium spot availability acceptable with high fault tolerance"
928 return (
929 "on-demand",
930 "Spot availability is medium; on-demand recommended for reliability",
931 )
933 if best_spot.availability == "low": 933 ↛ 942line 933 didn't jump to line 942 because the condition on line 933 was always true
934 if fault_tolerance == "high":
935 return (
936 "spot",
937 "Low spot availability but acceptable with high fault tolerance",
938 )
939 return "on-demand", "Spot capacity is limited; on-demand recommended"
941 # Default to on-demand
942 return "on-demand", "On-demand recommended (spot availability unknown or limited)"
944 # -------------------------------------------------------------------------
945 # Capacity Reservations (ODCRs) and Capacity Blocks for ML
946 # -------------------------------------------------------------------------
948 def list_capacity_reservations(
949 self,
950 region: str,
951 instance_type: str | None = None,
952 state: str | None = "active",
953 *,
954 include_pricing: bool = False,
955 ) -> list[dict[str, Any]]:
956 """
957 List EC2 On-Demand Capacity Reservations (ODCRs) in a region.
959 Args:
960 region: AWS region to query
961 instance_type: Filter by instance type (optional)
962 state: Filter by state — "active" (default), or None for all
963 include_pricing: Enrich each reservation with On-Demand pricing
964 (per-instance-hour, whole-reservation per-hour, per-GPU-hour).
965 Adds one Pricing API call per distinct instance type (cached), so
966 it is opt-in and off by default to keep the plain list fast.
968 Returns:
969 List of reservation dictionaries with availability details
970 """
971 ec2 = self._session.client("ec2", region_name=region)
973 filters: list[dict[str, Any]] = []
974 if state:
975 filters.append({"Name": "state", "Values": [state]})
976 if instance_type:
977 filters.append({"Name": "instance-type", "Values": [instance_type]})
979 reservations: list[dict[str, Any]] = []
980 try:
981 paginator = ec2.get_paginator("describe_capacity_reservations")
982 page_kwargs: dict[str, Any] = {}
983 if filters:
984 page_kwargs["Filters"] = filters
986 for page in paginator.paginate(**page_kwargs):
987 for cr in page.get("CapacityReservations", []):
988 total = cr.get("TotalInstanceCount", 0)
989 available = cr.get("AvailableInstanceCount", 0)
990 used = total - available
992 entry = {
993 "type": "odcr",
994 "reservation_id": cr.get("CapacityReservationId"),
995 "instance_type": cr.get("InstanceType"),
996 "availability_zone": cr.get("AvailabilityZone"),
997 "region": region,
998 "state": cr.get("State"),
999 "total_instances": total,
1000 "available_instances": available,
1001 "used_instances": used,
1002 "utilization_pct": round(used / total * 100, 1) if total else 0,
1003 "instance_platform": cr.get("InstancePlatform"),
1004 "tenancy": cr.get("Tenancy"),
1005 "instance_match_criteria": cr.get("InstanceMatchCriteria"),
1006 "start_date": (
1007 cr["StartDate"].isoformat() if cr.get("StartDate") else None
1008 ),
1009 "end_date": (cr["EndDate"].isoformat() if cr.get("EndDate") else None),
1010 "end_date_type": cr.get("EndDateType"),
1011 "tags": {t["Key"]: t["Value"] for t in cr.get("Tags", [])},
1012 }
1013 if include_pricing:
1014 self._enrich_reservation_pricing(entry, region)
1015 reservations.append(entry)
1016 except ClientError as e:
1017 logger.debug("Failed to list capacity reservations in %s: %s", region, e)
1019 return reservations
1021 def _enrich_reservation_pricing(self, reservation: dict[str, Any], region: str) -> None:
1022 """Add On-Demand pricing keys to a reservation dict, in place.
1024 ODCRs bill at the On-Demand rate for the reserved instance type, so the
1025 per-instance-hour / per-hour / per-GPU-hour figures mirror the Capacity
1026 Block pricing surface (see :func:`blocks.compute_reservation_pricing`) and
1027 let a caller rank and compare reservations the same way it ranks blocks.
1028 Pricing that can't be resolved is recorded as ``None`` — never fatal.
1029 """
1030 instance_type = reservation.get("instance_type")
1031 if not instance_type: 1031 ↛ 1032line 1031 didn't jump to line 1032 because the condition on line 1031 was never true
1032 return
1033 try:
1034 on_demand = self.get_on_demand_price(instance_type, region)
1035 except Exception as e: # pricing is supplementary; never fail the listing
1036 logger.debug("On-demand price lookup failed for %s in %s: %s", instance_type, region, e)
1037 on_demand = None
1039 spec = GPU_INSTANCE_SPECS.get(instance_type)
1040 gpus_per_instance = spec.gpu_count if spec else None
1042 pricing = blocks.compute_reservation_pricing(
1043 on_demand, reservation.get("total_instances"), gpus_per_instance
1044 )
1045 reservation["on_demand_price_per_hour"] = (
1046 round(float(on_demand), 4) if on_demand is not None else None
1047 )
1048 reservation["gpus_per_instance"] = gpus_per_instance
1049 reservation.update(pricing)
1051 def _build_block_offering(
1052 self,
1053 offering: dict[str, Any],
1054 region: str,
1055 requested_count: int,
1056 gpus_per_instance: int | None,
1057 requested_duration_hours: int,
1058 ) -> dict[str, Any]:
1059 """Shape a raw DescribeCapacityBlockOfferings entry into an enriched dict.
1061 Reports the offering's *actual* duration (the API returns blocks whose
1062 duration is the closest match to the request, not necessarily equal) and
1063 adds per-hour / per-GPU-hour pricing derived from the upfront fee.
1064 ``upfront_fee`` is the raw API value (a string); ``upfront_fee_usd`` is the
1065 parsed float used for ranking and display.
1066 """
1067 start_date = offering.get("StartDate")
1068 end_date = offering.get("EndDate")
1070 actual_duration = offering.get("CapacityBlockDurationHours")
1071 if actual_duration is None:
1072 minutes = offering.get("CapacityBlockDurationMinutes")
1073 actual_duration = round(minutes / 60) if minutes else requested_duration_hours
1075 count = offering.get("InstanceCount")
1076 if count is None: 1076 ↛ 1077line 1076 didn't jump to line 1077 because the condition on line 1076 was never true
1077 count = requested_count
1079 pricing = blocks.compute_offering_pricing(
1080 offering.get("UpfrontFee"), actual_duration, count, gpus_per_instance
1081 )
1083 return {
1084 "type": "capacity_block",
1085 "offering_id": offering.get("CapacityBlockOfferingId"),
1086 "instance_type": offering.get("InstanceType"),
1087 "availability_zone": offering.get("AvailabilityZone"),
1088 "region": region,
1089 "instance_count": count,
1090 "duration_hours": actual_duration,
1091 "duration_days": blocks.hours_to_days(actual_duration),
1092 "start_date": start_date.isoformat() if start_date else None,
1093 "end_date": end_date.isoformat() if end_date else None,
1094 "upfront_fee": offering.get("UpfrontFee"),
1095 "upfront_fee_usd": pricing["upfront_fee_usd"],
1096 "price_per_hour": pricing["price_per_hour"],
1097 "price_per_instance_hour": pricing["price_per_instance_hour"],
1098 "price_per_gpu_hour": pricing["price_per_gpu_hour"],
1099 "gpus_per_instance": gpus_per_instance,
1100 "currency": offering.get("CurrencyCode", "USD"),
1101 "tenancy": offering.get("Tenancy"),
1102 }
1104 def list_capacity_block_offerings(
1105 self,
1106 region: str,
1107 instance_type: str,
1108 instance_count: int = 1,
1109 duration_hours: int = 24,
1110 *,
1111 earliest_start: datetime | None = None,
1112 latest_start: datetime | None = None,
1113 gpus_per_instance: int | None = None,
1114 ) -> list[dict[str, Any]]:
1115 """
1116 List available Capacity Block offerings for ML workloads.
1118 Capacity Blocks provide guaranteed GPU capacity for a fixed duration
1119 at a known price — ideal for training jobs with predictable runtimes.
1120 Queries a single duration in a single region; the date window and the
1121 multi-duration / multi-region sweep are layered on top by
1122 :meth:`find_capacity_blocks`.
1124 Args:
1125 region: AWS region to query
1126 instance_type: GPU instance type (e.g. p5.48xlarge, p4d.24xlarge)
1127 instance_count: Number of instances needed
1128 duration_hours: Desired block duration in hours (must be a supported value)
1129 earliest_start: Only return blocks starting on/after this datetime
1130 (EC2 StartDateRange). Lets callers ask "blocks starting near D1".
1131 latest_start: Only return blocks starting on/before this datetime
1132 (EC2 EndDateRange).
1133 gpus_per_instance: GPUs per instance, used for per-GPU-hour pricing.
1134 Resolved from the instance specs when omitted.
1136 Returns:
1137 List of available capacity block offerings (enriched with pricing).
1138 All matching pages are followed via NextToken.
1139 """
1140 ec2 = self._session.client("ec2", region_name=region, config=_CB_RETRY_CONFIG)
1141 offerings: list[dict[str, Any]] = []
1143 if gpus_per_instance is None:
1144 info = self.get_instance_info(instance_type)
1145 gpus_per_instance = info.gpu_count if info else None
1147 api_kwargs: dict[str, Any] = {
1148 "InstanceType": instance_type,
1149 "InstanceCount": instance_count,
1150 "CapacityDurationHours": duration_hours,
1151 }
1152 if earliest_start is not None:
1153 api_kwargs["StartDateRange"] = earliest_start
1154 if latest_start is not None:
1155 api_kwargs["EndDateRange"] = latest_start
1157 try:
1158 next_token: str | None = None
1159 while True:
1160 if next_token:
1161 api_kwargs["NextToken"] = next_token
1162 response = ec2.describe_capacity_block_offerings(**api_kwargs)
1163 for offering in response.get("CapacityBlockOfferings", []):
1164 offerings.append(
1165 self._build_block_offering(
1166 offering, region, instance_count, gpus_per_instance, duration_hours
1167 )
1168 )
1169 next_token = response.get("NextToken")
1170 if not next_token:
1171 break
1172 except ClientError as e:
1173 error_code = e.response.get("Error", {}).get("Code", "")
1174 if error_code in _CB_EXPECTED_ERROR_CODES:
1175 pass # Type/region doesn't support Capacity Blocks — expected
1176 else:
1177 logger.warning("Failed to list capacity block offerings in %s: %s", region, e)
1178 except BotoCoreError as e:
1179 # Endpoint resolution / connection errors surface here when the
1180 # Capacity Block API isn't available in a region.
1181 logger.warning("Capacity Block API unavailable in %s: %s", region, e)
1183 return offerings
1185 def get_capacity_block_trend(
1186 self,
1187 instance_type: str,
1188 region: str,
1189 ) -> float:
1190 """
1191 Estimate capacity block availability trend via time-series regression.
1193 Queries offerings across the maximum 182-day (26-week) window, buckets
1194 them into weekly bins by start date, and fits a linear regression to
1195 the offering counts per week. The normalized slope indicates whether
1196 capacity is growing or shrinking over time.
1198 Returns:
1199 Trend score from -1.0 to 1.0:
1200 > 0 = capacity growing (offerings increasing week-over-week)
1201 = 0 = stable or no data
1202 < 0 = capacity shrinking (offerings decreasing week-over-week)
1203 """
1204 ec2 = self._session.client("ec2", region_name=region, config=_CB_RETRY_CONFIG)
1206 now = datetime.now(UTC)
1207 far_end = now + timedelta(days=182)
1209 try:
1210 response = ec2.describe_capacity_block_offerings(
1211 InstanceType=instance_type,
1212 InstanceCount=1,
1213 CapacityDurationHours=24, # Minimum duration for broadest results
1214 StartDateRange=now,
1215 EndDateRange=far_end,
1216 )
1217 except ClientError as e:
1218 error_code = e.response.get("Error", {}).get("Code", "")
1219 if error_code not in _CB_EXPECTED_ERROR_CODES:
1220 logger.warning(
1221 "Capacity block trend query failed for %s in %s: %s",
1222 instance_type,
1223 region,
1224 e,
1225 )
1226 return 0.0
1227 except BotoCoreError as e:
1228 logger.warning(
1229 "Capacity block trend query failed for %s in %s: %s", instance_type, region, e
1230 )
1231 return 0.0
1233 offerings = response.get("CapacityBlockOfferings", [])
1234 if not offerings:
1235 return 0.0
1237 # Bucket offerings into weekly bins (week 0 = this week, week 25 = ~6 months out)
1238 num_weeks = 26
1239 bins = [0] * num_weeks
1240 for o in offerings:
1241 start = o.get("StartDate")
1242 if start is None:
1243 continue
1244 delta_days = (start - now).total_seconds() / 86400.0
1245 week_idx = int(delta_days / 7)
1246 if 0 <= week_idx < num_weeks:
1247 bins[week_idx] += 1
1249 # Need at least 2 non-zero bins to detect a meaningful trend
1250 non_zero = sum(1 for b in bins if b > 0)
1251 if non_zero < 2:
1252 return 0.0
1254 # Linear regression: slope of offerings-per-week over time
1255 # Using least-squares: slope = Σ((x-x̄)(y-ȳ)) / Σ((x-x̄)²)
1256 n = len(bins)
1257 x_mean = (n - 1) / 2.0
1258 y_mean = statistics.mean(bins)
1260 numerator = sum((i - x_mean) * (bins[i] - y_mean) for i in range(n))
1261 denominator = sum((i - x_mean) ** 2 for i in range(n))
1263 if denominator == 0: 1263 ↛ 1264line 1263 didn't jump to line 1264 because the condition on line 1263 was never true
1264 return 0.0
1266 slope = numerator / denominator
1268 # Normalize slope to -1..1 range relative to the mean offering count.
1269 # A slope of +y_mean per 26 weeks would be a doubling → maps to ~1.0.
1270 normalized = slope * num_weeks / (y_mean * 2) if y_mean > 0 else 0.0
1272 return round(max(-1.0, min(1.0, normalized)), 4)
1274 def list_all_reservations(
1275 self,
1276 instance_type: str | None = None,
1277 regions: list[str] | None = None,
1278 ) -> dict[str, Any]:
1279 """
1280 List all capacity reservations (ODCRs) across deployed regions.
1282 Args:
1283 instance_type: Filter by instance type (optional)
1284 regions: Regions to query (defaults to deployed GCO regions)
1286 Returns:
1287 Summary dict with reservations grouped by region
1288 """
1289 if not regions:
1290 from cli.aws_client import get_aws_client
1292 aws_client = get_aws_client(self.config)
1293 stacks = aws_client.discover_regional_stacks()
1294 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1296 all_reservations: list[dict[str, Any]] = []
1297 for region in regions:
1298 all_reservations.extend(
1299 self.list_capacity_reservations(region, instance_type=instance_type)
1300 )
1302 total_reserved = sum(r["total_instances"] for r in all_reservations)
1303 total_available = sum(r["available_instances"] for r in all_reservations)
1305 return {
1306 "regions_checked": regions,
1307 "instance_type_filter": instance_type,
1308 "total_reservations": len(all_reservations),
1309 "total_reserved_instances": total_reserved,
1310 "total_available_instances": total_available,
1311 "reservations": all_reservations,
1312 }
1314 def find_capacity_reservations(
1315 self,
1316 instance_type: str | None = None,
1317 regions: list[str] | None = None,
1318 *,
1319 min_count: int = 1,
1320 state: str | None = "active",
1321 include_pricing: bool = True,
1322 max_workers: int | None = None,
1323 ) -> dict[str, Any]:
1324 """Sweep regions for existing ODCRs in one parallel, ranked call.
1326 The ODCR counterpart to :meth:`find_capacity_blocks`. It fans out across
1327 every requested region in parallel, normalizes a friendly instance-type
1328 alias (``p6-b200`` -> ``p6-b200.48xlarge``), enriches each reservation with
1329 On-Demand pricing, and returns a single consolidated report ranked
1330 most-available-first (then cheapest per-GPU-hour). Where
1331 :meth:`list_all_reservations` simply aggregates region by region, this
1332 answers "where do I already have free reserved capacity for this instance
1333 type?" across many regions at once.
1335 Args:
1336 instance_type: Instance type or friendly alias to filter by. Omit to
1337 return every reservation (no type filter).
1338 regions: Regions to search in parallel (any regions, not just
1339 deployed); defaults to the deployed GCO regions when omitted.
1340 min_count: Minimum available instances for the summary to consider the
1341 search satisfied (does not filter the returned list).
1342 state: Reservation state filter ("active" by default; None for all).
1343 include_pricing: Enrich each reservation with On-Demand pricing.
1344 max_workers: Override the parallel fan-out width.
1346 Returns:
1347 A consolidated report dict (see keys assembled below).
1348 """
1349 canonical = instance_type
1350 note: str | None = None
1351 valid = True
1352 known = False
1353 if instance_type:
1354 validation = self.validate_instance_type(instance_type)
1355 canonical = validation["instance_type"]
1356 note = validation["note"]
1357 valid = validation["valid"]
1358 known = validation["known"]
1360 if not regions: 1360 ↛ 1361line 1360 didn't jump to line 1361 because the condition on line 1360 was never true
1361 from cli.aws_client import get_aws_client
1363 aws_client = get_aws_client(self.config)
1364 stacks = aws_client.discover_regional_stacks()
1365 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1367 report: dict[str, Any] = {
1368 "instance_type": canonical,
1369 "requested_instance_type": instance_type,
1370 "valid_instance_type": valid,
1371 "known_instance_type": known,
1372 "note": note,
1373 "min_count": min_count,
1374 "state": state,
1375 "regions_checked": list(regions),
1376 "reservations_found": 0,
1377 "reservations": [],
1378 "ranked": [],
1379 "best": None,
1380 "total_reserved_instances": 0,
1381 "total_available_instances": 0,
1382 "regions_with_reservations": [],
1383 }
1385 if instance_type and not valid:
1386 report["recommendation"] = (
1387 note or f"'{instance_type}' is not a recognized EC2 instance type."
1388 )
1389 return report
1391 def _probe(region: str) -> list[dict[str, Any]]:
1392 try:
1393 return self.list_capacity_reservations(
1394 region,
1395 instance_type=canonical,
1396 state=state,
1397 include_pricing=include_pricing,
1398 )
1399 except Exception as e:
1400 logger.warning("Reservation probe failed for %s in %s: %s", canonical, region, e)
1401 return []
1403 collected: list[dict[str, Any]] = []
1404 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(regions)))
1405 if len(regions) == 1:
1406 collected.extend(_probe(regions[0]))
1407 else:
1408 with ThreadPoolExecutor(max_workers=workers) as executor:
1409 for result in executor.map(_probe, regions):
1410 collected.extend(result)
1412 ranked = blocks.rank_reservations(collected)
1413 best = ranked[0] if ranked else None
1414 total_reserved = sum(r.get("total_instances") or 0 for r in collected)
1415 total_available = sum(r.get("available_instances") or 0 for r in collected)
1417 report.update(
1418 {
1419 "reservations_found": len(collected),
1420 "reservations": blocks.sort_reservations(collected),
1421 "ranked": ranked,
1422 "best": best,
1423 "total_reserved_instances": total_reserved,
1424 "total_available_instances": total_available,
1425 "regions_with_reservations": sorted(
1426 {r["region"] for r in collected if r.get("region")}
1427 ),
1428 "recommendation": self._summarize_reservation_search(
1429 canonical, regions, collected, best, min_count, note
1430 ),
1431 }
1432 )
1433 return report
1435 @staticmethod
1436 def _summarize_reservation_search(
1437 instance_type: str | None,
1438 regions: list[str],
1439 reservations: list[dict[str, Any]],
1440 best: dict[str, Any] | None,
1441 min_count: int,
1442 note: str | None,
1443 ) -> str:
1444 """Build a one-line human recommendation for a consolidated ODCR search."""
1445 label = instance_type or "any instance type"
1446 if not reservations:
1447 msg = (
1448 f"No active On-Demand Capacity Reservations for {label} across "
1449 f"{len(regions)} region(s). Create one with "
1450 "'gco capacity create-reservation', or search purchasable Capacity "
1451 "Blocks with 'gco capacity find-blocks'."
1452 )
1453 return f"{note} {msg}" if note else msg
1455 total_available = sum(r.get("available_instances") or 0 for r in reservations)
1456 regions_with = sorted({r["region"] for r in reservations if r.get("region")})
1457 parts = [
1458 f"Found {len(reservations)} reservation(s) for {label} across "
1459 f"{len(regions_with)} region(s); {total_available} instance(s) available."
1460 ]
1461 if total_available < min_count: 1461 ↛ 1462line 1461 didn't jump to line 1462 because the condition on line 1461 was never true
1462 parts.append(
1463 f"Fewer than the {min_count} requested are free — consider creating "
1464 "another reservation or a Capacity Block."
1465 )
1466 if best and (best.get("available_instances") or 0) > 0: 1466 ↛ 1474line 1466 didn't jump to line 1474 because the condition on line 1466 was always true
1467 gpu_hr = best.get("price_per_gpu_hour")
1468 gpu_hr_str = f", ${gpu_hr}/GPU-hr" if gpu_hr is not None else ""
1469 parts.append(
1470 f"Most available: {best.get('available_instances')}/"
1471 f"{best.get('total_instances')} in {best.get('region')}/"
1472 f"{best.get('availability_zone')} ({best.get('reservation_id')}{gpu_hr_str})."
1473 )
1474 msg = " ".join(parts)
1475 return f"{note} {msg}" if note else msg
1477 @staticmethod
1478 def _summarize_block_search(
1479 instance_type: str,
1480 regions: list[str],
1481 offerings: list[dict[str, Any]],
1482 best: dict[str, Any] | None,
1483 longest: dict[str, Any] | None,
1484 note: str | None,
1485 ) -> str:
1486 """Build a one-line human recommendation for a consolidated block search."""
1487 if not offerings:
1488 msg = (
1489 f"No Capacity Block offerings for {instance_type} across "
1490 f"{len(regions)} region(s) in the requested window. Try a wider "
1491 "date range, a shorter duration, more regions, or check back later."
1492 )
1493 return f"{note} {msg}" if note else msg
1495 regions_with = sorted({o["region"] for o in offerings if o.get("region")})
1496 parts = [
1497 f"Found {len(offerings)} Capacity Block offering(s) for {instance_type} "
1498 f"across {len(regions_with)} region(s)."
1499 ]
1500 if best: 1500 ↛ 1508line 1500 didn't jump to line 1508 because the condition on line 1500 was always true
1501 price = best.get("price_per_gpu_hour")
1502 price_str = f", from ${price}/GPU-hr" if price is not None else ""
1503 parts.append(
1504 f"Cheapest: {best.get('region')}/{best.get('availability_zone')} "
1505 f"{best.get('duration_days')}d starting "
1506 f"{(best.get('start_date') or '')[:16]}{price_str}."
1507 )
1508 if longest and longest is not best:
1509 parts.append(
1510 f"Longest: {longest.get('duration_days')}d in "
1511 f"{longest.get('region')}/{longest.get('availability_zone')}."
1512 )
1513 msg = " ".join(parts)
1514 return f"{note} {msg}" if note else msg
1516 def find_capacity_blocks(
1517 self,
1518 instance_type: str,
1519 regions: list[str] | None = None,
1520 *,
1521 instance_count: int = 1,
1522 duration_hours: int | None = None,
1523 duration_days: int | None = None,
1524 min_duration_hours: int | None = None,
1525 min_duration_days: int | None = None,
1526 max_duration_hours: int | None = None,
1527 max_duration_days: int | None = None,
1528 earliest_start: str | datetime | None = None,
1529 latest_start: str | datetime | None = None,
1530 find_longest: bool = False,
1531 max_workers: int | None = None,
1532 ) -> dict[str, Any]:
1533 """Sweep regions x durations x a date window for Capacity Blocks in one call.
1535 This is the high-level search that answers "find 1x p6-b200.48xlarge
1536 across us-east-1/us-east-2/us-west-2/eu-west-1 for durations 1-63 days
1537 starting 2026-07-01..2026-07-10" without manual multi-call sweeping.
1539 Because EC2 requires an exact ``CapacityDurationHours`` per query, a
1540 duration *range* (or ``find_longest``) expands to every valid Capacity
1541 Block duration in the range, and each (region, duration) pair is probed
1542 in parallel. Offerings are de-duplicated across probes, grouped/sorted by
1543 region + AZ + start date, and ranked cheapest-first by per-GPU-hour price.
1545 Args:
1546 instance_type: GPU instance type or friendly alias (e.g. p6-b200).
1547 regions: Explicit regions to search (any regions, not just deployed).
1548 Defaults to the deployed GCO regions when omitted.
1549 instance_count: Instances per block.
1550 duration_hours / duration_days: A single target duration.
1551 min_duration_hours / min_duration_days: Lower bound of a duration range.
1552 max_duration_hours / max_duration_days: Upper bound of a duration range.
1553 earliest_start / latest_start: Date (YYYY-MM-DD) or ISO datetime window
1554 for the block start, threaded to StartDateRange / EndDateRange.
1555 find_longest: Sweep the full duration ladder (within any range given)
1556 and surface the longest available block.
1557 max_workers: Override the parallel fan-out width.
1559 Returns:
1560 A consolidated report dict (see keys assembled below).
1561 """
1562 validation = self.validate_instance_type(instance_type)
1563 canonical = validation["instance_type"]
1564 gpus_per_instance = validation["gpu_count"]
1566 parsed_earliest = blocks.parse_date_input(earliest_start)
1567 parsed_latest = blocks.parse_date_input(latest_start)
1569 durations = blocks.resolve_search_durations(
1570 duration_hours=blocks.coerce_hours(duration_hours, duration_days),
1571 min_duration_hours=blocks.coerce_hours(min_duration_hours, min_duration_days),
1572 max_duration_hours=blocks.coerce_hours(max_duration_hours, max_duration_days),
1573 find_longest=find_longest,
1574 )
1576 if not regions:
1577 from cli.aws_client import get_aws_client
1579 aws_client = get_aws_client(self.config)
1580 stacks = aws_client.discover_regional_stacks()
1581 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1583 report: dict[str, Any] = {
1584 "instance_type": canonical,
1585 "requested_instance_type": instance_type,
1586 "valid_instance_type": validation["valid"],
1587 "known_instance_type": validation["known"],
1588 "note": validation["note"],
1589 "instance_count": instance_count,
1590 "regions_checked": list(regions),
1591 "durations_probed_hours": durations,
1592 "durations_probed_days": [blocks.hours_to_days(h) for h in durations],
1593 "date_window": {
1594 "earliest_start": parsed_earliest.isoformat() if parsed_earliest else None,
1595 "latest_start": parsed_latest.isoformat() if parsed_latest else None,
1596 },
1597 "offerings_found": 0,
1598 "offerings": [],
1599 "ranked": [],
1600 "best": None,
1601 "longest": None,
1602 "regions_with_offerings": [],
1603 }
1605 if not validation["valid"]:
1606 report["recommendation"] = (
1607 validation["note"]
1608 or f"'{instance_type}' is not a valid standalone EC2 instance type "
1609 "for Capacity Blocks."
1610 )
1611 return report
1613 probes = [(region, dur) for region in regions for dur in durations]
1614 collected: list[dict[str, Any]] = []
1615 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(probes)))
1617 def _probe(region: str, dur: int) -> list[dict[str, Any]]:
1618 try:
1619 return self.list_capacity_block_offerings(
1620 region,
1621 canonical,
1622 instance_count=instance_count,
1623 duration_hours=dur,
1624 earliest_start=parsed_earliest,
1625 latest_start=parsed_latest,
1626 gpus_per_instance=gpus_per_instance,
1627 )
1628 except Exception as e:
1629 logger.warning(
1630 "Capacity block probe failed for %s in %s (%sh): %s",
1631 canonical,
1632 region,
1633 dur,
1634 e,
1635 )
1636 return []
1638 with ThreadPoolExecutor(max_workers=workers) as executor:
1639 futures = [executor.submit(_probe, region, dur) for region, dur in probes]
1640 for future in as_completed(futures):
1641 collected.extend(future.result())
1643 unique = blocks.dedupe_offerings(collected)
1644 ranked = blocks.rank_offerings(unique)
1645 best = ranked[0] if ranked else None
1646 longest = blocks.longest_offering(unique)
1648 report.update(
1649 {
1650 "offerings_found": len(unique),
1651 "offerings": blocks.sort_offerings(unique),
1652 "ranked": ranked,
1653 "best": best,
1654 "longest": longest,
1655 "regions_with_offerings": sorted({o["region"] for o in unique if o.get("region")}),
1656 "recommendation": self._summarize_block_search(
1657 canonical, regions, unique, best, longest, validation["note"]
1658 ),
1659 }
1660 )
1661 return report
1663 def check_reservation_availability(
1664 self,
1665 instance_type: str,
1666 min_count: int = 1,
1667 include_capacity_blocks: bool = True,
1668 block_duration_hours: int = 24,
1669 *,
1670 regions: list[str] | None = None,
1671 block_duration_days: int | None = None,
1672 earliest_start: str | datetime | None = None,
1673 latest_start: str | datetime | None = None,
1674 max_workers: int | None = None,
1675 ) -> dict[str, Any]:
1676 """
1677 Check if capacity reservations or blocks have available instances.
1679 Checks both ODCRs (existing reservations) and Capacity Block offerings
1680 (purchasable guaranteed capacity) for a given instance type, across one
1681 or many regions queried in parallel.
1683 Args:
1684 instance_type: EC2 instance type to check
1685 min_count: Minimum number of available instances needed
1686 include_capacity_blocks: Also check Capacity Block offerings
1687 block_duration_hours: Duration for capacity block search (hours)
1688 regions: Regions to check in parallel (any regions, not just
1689 deployed); falls back to the deployed regions when omitted.
1690 block_duration_days: Block duration in days (overrides hours when set).
1691 earliest_start / latest_start: Date window for block start
1692 (StartDateRange / EndDateRange).
1693 max_workers: Override the parallel fan-out width.
1695 Returns:
1696 Dictionary with ODCR availability and capacity block offerings
1697 """
1698 if regions:
1699 target_regions = list(regions)
1700 else:
1701 from cli.aws_client import get_aws_client
1703 aws_client = get_aws_client(self.config)
1704 stacks = aws_client.discover_regional_stacks()
1705 target_regions = list(stacks.keys()) if stacks else [self.config.default_region]
1707 effective_duration = blocks.snap_duration_hours(
1708 blocks.coerce_hours(block_duration_hours, block_duration_days) or 24
1709 )
1710 parsed_earliest = blocks.parse_date_input(earliest_start)
1711 parsed_latest = blocks.parse_date_input(latest_start)
1713 def _check_region(r: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
1714 reservations = self.list_capacity_reservations(r, instance_type=instance_type)
1715 region_blocks: list[dict[str, Any]] = []
1716 if include_capacity_blocks:
1717 region_blocks = self.list_capacity_block_offerings(
1718 r,
1719 instance_type=instance_type,
1720 instance_count=min_count,
1721 duration_hours=effective_duration,
1722 earliest_start=parsed_earliest,
1723 latest_start=parsed_latest,
1724 )
1725 return reservations, region_blocks
1727 odcr_results: list[dict[str, Any]] = []
1728 block_offerings: list[dict[str, Any]] = []
1729 total_available = 0
1730 total_reserved = 0
1732 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(target_regions)))
1733 if len(target_regions) == 1:
1734 region_results = [_check_region(target_regions[0])]
1735 else:
1736 with ThreadPoolExecutor(max_workers=workers) as executor:
1737 region_results = list(executor.map(_check_region, target_regions))
1739 for reservations, region_blocks in region_results:
1740 for res in reservations:
1741 avail = res["available_instances"]
1742 total_available += avail
1743 total_reserved += res["total_instances"]
1744 if avail > 0:
1745 odcr_results.append(res)
1746 block_offerings.extend(region_blocks)
1748 block_offerings = blocks.sort_offerings(block_offerings)
1749 has_odcr = total_available >= min_count
1750 has_blocks = len(block_offerings) > 0
1752 # Build recommendation
1753 if has_odcr:
1754 recommendation = (
1755 f"ODCR capacity available: {total_available} instances "
1756 f"across {len(odcr_results)} reservation(s)"
1757 )
1758 elif has_blocks:
1759 cheapest = min(block_offerings, key=_offering_fee)
1760 fee_display = cheapest.get("upfront_fee_usd")
1761 if fee_display is None:
1762 fee_display = cheapest.get("upfront_fee", "?")
1763 recommendation = (
1764 f"No ODCR capacity, but {len(block_offerings)} Capacity Block offering(s) "
1765 f"available (from ${fee_display} for {effective_duration}h)"
1766 )
1767 else:
1768 recommendation = (
1769 "No reserved capacity or block offerings found. "
1770 "Consider on-demand or spot, or request a Capacity Block "
1771 "for a different duration/region."
1772 )
1774 return {
1775 "instance_type": instance_type,
1776 "min_count_requested": min_count,
1777 "regions_checked": target_regions,
1778 "odcr": {
1779 "total_reserved_instances": total_reserved,
1780 "total_available_instances": total_available,
1781 "has_availability": has_odcr,
1782 "reservations": odcr_results,
1783 },
1784 "capacity_blocks": {
1785 "offerings_found": len(block_offerings),
1786 "has_offerings": has_blocks,
1787 "duration_hours": effective_duration,
1788 "date_window": {
1789 "earliest_start": parsed_earliest.isoformat() if parsed_earliest else None,
1790 "latest_start": parsed_latest.isoformat() if parsed_latest else None,
1791 },
1792 "offerings": block_offerings,
1793 },
1794 "recommendation": recommendation,
1795 }
1797 def purchase_capacity_block(
1798 self,
1799 offering_id: str,
1800 region: str,
1801 dry_run: bool = False,
1802 ) -> dict[str, Any]:
1803 """
1804 Purchase a Capacity Block offering by its ID.
1806 Args:
1807 offering_id: Capacity Block offering ID (cb-xxx) from list_capacity_block_offerings
1808 region: AWS region where the offering exists
1809 dry_run: If True, validate the offering without purchasing
1811 Returns:
1812 Dictionary with the created capacity reservation details
1813 """
1814 ec2 = self._session.client("ec2", region_name=region)
1816 if dry_run:
1817 # Validate the offering exists by describing capacity block offerings
1818 # and matching the ID
1819 try:
1820 # Use EC2 DryRun to validate permissions without purchasing
1821 ec2.purchase_capacity_block(
1822 CapacityBlockOfferingId=offering_id,
1823 InstancePlatform="Linux/UNIX",
1824 DryRun=True,
1825 )
1826 except ClientError as e:
1827 error_code = e.response.get("Error", {}).get("Code", "")
1828 if error_code == "DryRunOperation":
1829 # DryRunOperation means the request would have succeeded
1830 return {
1831 "success": True,
1832 "dry_run": True,
1833 "offering_id": offering_id,
1834 "region": region,
1835 "message": "Dry run succeeded — offering is valid and purchasable",
1836 }
1837 error_msg = e.response.get("Error", {}).get("Message", str(e))
1838 return {
1839 "success": False,
1840 "dry_run": True,
1841 "offering_id": offering_id,
1842 "region": region,
1843 "error_code": error_code,
1844 "error": error_msg,
1845 }
1847 try:
1848 response = ec2.purchase_capacity_block(
1849 CapacityBlockOfferingId=offering_id,
1850 InstancePlatform="Linux/UNIX",
1851 )
1853 reservation = response.get("CapacityReservation", {})
1854 reservation_id = reservation.get("CapacityReservationId", "")
1855 instance_type = reservation.get("InstanceType", "")
1856 az = reservation.get("AvailabilityZone", "")
1857 total = reservation.get("TotalInstanceCount", 0)
1858 start = reservation.get("StartDate")
1859 end = reservation.get("EndDate")
1861 return {
1862 "success": True,
1863 "dry_run": False,
1864 "reservation_id": reservation_id,
1865 "offering_id": offering_id,
1866 "instance_type": instance_type,
1867 "availability_zone": az,
1868 "region": region,
1869 "total_instances": total,
1870 "start_date": start.isoformat() if start else None,
1871 "end_date": end.isoformat() if end else None,
1872 "state": reservation.get("State", ""),
1873 }
1874 except ClientError as e:
1875 error_code = e.response.get("Error", {}).get("Code", "")
1876 error_msg = e.response.get("Error", {}).get("Message", str(e))
1877 return {
1878 "success": False,
1879 "dry_run": False,
1880 "offering_id": offering_id,
1881 "region": region,
1882 "error_code": error_code,
1883 "error": error_msg,
1884 }
1886 def create_capacity_reservation(
1887 self,
1888 instance_type: str,
1889 region: str,
1890 availability_zone: str,
1891 instance_count: int = 1,
1892 *,
1893 instance_platform: str = "Linux/UNIX",
1894 tenancy: str = "default",
1895 instance_match_criteria: str = "open",
1896 end_date: str | datetime | None = None,
1897 ebs_optimized: bool = False,
1898 dry_run: bool = False,
1899 ) -> dict[str, Any]:
1900 """Create a new On-Demand Capacity Reservation (ODCR).
1902 The ODCR counterpart to :meth:`purchase_capacity_block`: it reserves
1903 On-Demand capacity for an instance type in a specific Availability Zone.
1904 Unlike a Capacity Block (a fixed-term block bought by offering id), an ODCR
1905 is open-ended by default and billed at the On-Demand rate until cancelled.
1906 Friendly instance-type aliases are normalized (``p6-b200`` ->
1907 ``p6-b200.48xlarge``).
1909 Args:
1910 instance_type: EC2 instance type or friendly alias.
1911 region: AWS region.
1912 availability_zone: Target AZ (e.g. us-east-1a).
1913 instance_count: Number of instances to reserve.
1914 instance_platform: Platform/OS (default "Linux/UNIX").
1915 tenancy: "default" or "dedicated".
1916 instance_match_criteria: "open" (any matching instance) or "targeted".
1917 end_date: Optional end date (ISO string or datetime). When set the
1918 reservation is "limited" and auto-releases then; omit for
1919 "unlimited".
1920 ebs_optimized: Reserve EBS-optimized capacity.
1921 dry_run: Validate permissions/parameters without creating (no cost).
1923 Returns:
1924 Dict with the created reservation's details, or an error payload.
1925 """
1926 validation = self.validate_instance_type(instance_type)
1927 canonical = validation["instance_type"]
1928 if not validation["valid"]:
1929 return {
1930 "success": False,
1931 "dry_run": dry_run,
1932 "instance_type": canonical,
1933 "region": region,
1934 "error_code": "InvalidInstanceType",
1935 "error": (
1936 validation["note"] or f"'{instance_type}' is not a valid EC2 instance type."
1937 ),
1938 }
1940 ec2 = self._session.client("ec2", region_name=region)
1941 parsed_end = blocks.parse_date_input(end_date)
1942 api_kwargs: dict[str, Any] = {
1943 "InstanceType": canonical,
1944 "InstancePlatform": instance_platform,
1945 "AvailabilityZone": availability_zone,
1946 "InstanceCount": instance_count,
1947 "Tenancy": tenancy,
1948 "InstanceMatchCriteria": instance_match_criteria,
1949 "EbsOptimized": ebs_optimized,
1950 }
1951 if parsed_end is not None:
1952 api_kwargs["EndDate"] = parsed_end
1953 api_kwargs["EndDateType"] = "limited"
1954 else:
1955 api_kwargs["EndDateType"] = "unlimited"
1957 if dry_run:
1958 try:
1959 ec2.create_capacity_reservation(DryRun=True, **api_kwargs)
1960 except ClientError as e:
1961 error_code = e.response.get("Error", {}).get("Code", "")
1962 if error_code == "DryRunOperation": 1962 ↛ 1972line 1962 didn't jump to line 1972 because the condition on line 1962 was always true
1963 return {
1964 "success": True,
1965 "dry_run": True,
1966 "instance_type": canonical,
1967 "availability_zone": availability_zone,
1968 "region": region,
1969 "instance_count": instance_count,
1970 "message": "Dry run succeeded — reservation parameters are valid.",
1971 }
1972 error_msg = e.response.get("Error", {}).get("Message", str(e))
1973 return {
1974 "success": False,
1975 "dry_run": True,
1976 "instance_type": canonical,
1977 "region": region,
1978 "error_code": error_code,
1979 "error": error_msg,
1980 }
1981 return {
1982 "success": True,
1983 "dry_run": True,
1984 "instance_type": canonical,
1985 "availability_zone": availability_zone,
1986 "region": region,
1987 "instance_count": instance_count,
1988 "message": "Dry run completed.",
1989 }
1991 try:
1992 response = ec2.create_capacity_reservation(**api_kwargs)
1993 cr = response.get("CapacityReservation", {})
1994 return {
1995 "success": True,
1996 "dry_run": False,
1997 "reservation_id": cr.get("CapacityReservationId", ""),
1998 "instance_type": cr.get("InstanceType", canonical),
1999 "availability_zone": cr.get("AvailabilityZone", availability_zone),
2000 "region": region,
2001 "total_instances": cr.get("TotalInstanceCount", instance_count),
2002 "available_instances": cr.get("AvailableInstanceCount"),
2003 "state": cr.get("State", ""),
2004 "tenancy": cr.get("Tenancy", tenancy),
2005 "instance_match_criteria": cr.get("InstanceMatchCriteria", instance_match_criteria),
2006 "start_date": cr["StartDate"].isoformat() if cr.get("StartDate") else None,
2007 "end_date": cr["EndDate"].isoformat() if cr.get("EndDate") else None,
2008 "end_date_type": cr.get("EndDateType"),
2009 }
2010 except ClientError as e:
2011 error_code = e.response.get("Error", {}).get("Code", "")
2012 error_msg = e.response.get("Error", {}).get("Message", str(e))
2013 return {
2014 "success": False,
2015 "dry_run": False,
2016 "instance_type": canonical,
2017 "region": region,
2018 "error_code": error_code,
2019 "error": error_msg,
2020 }
2022 def cancel_capacity_reservation(
2023 self,
2024 reservation_id: str,
2025 region: str,
2026 dry_run: bool = False,
2027 ) -> dict[str, Any]:
2028 """Cancel an On-Demand Capacity Reservation by id.
2030 Releases reserved capacity so it stops incurring On-Demand charges. Only
2031 ODCRs can be cancelled this way; a Capacity Block runs for its fixed term.
2032 Cancelling does not terminate instances already running against the
2033 reservation — they simply revert to normal On-Demand billing.
2035 Args:
2036 reservation_id: Capacity Reservation id (cr-...).
2037 region: AWS region where the reservation exists.
2038 dry_run: Validate permissions without cancelling.
2040 Returns:
2041 Dict describing the outcome.
2042 """
2043 ec2 = self._session.client("ec2", region_name=region)
2045 if dry_run:
2046 try:
2047 ec2.cancel_capacity_reservation(CapacityReservationId=reservation_id, DryRun=True)
2048 except ClientError as e:
2049 error_code = e.response.get("Error", {}).get("Code", "")
2050 if error_code == "DryRunOperation": 2050 ↛ 2058line 2050 didn't jump to line 2058 because the condition on line 2050 was always true
2051 return {
2052 "success": True,
2053 "dry_run": True,
2054 "reservation_id": reservation_id,
2055 "region": region,
2056 "message": "Dry run succeeded — reservation can be cancelled.",
2057 }
2058 error_msg = e.response.get("Error", {}).get("Message", str(e))
2059 return {
2060 "success": False,
2061 "dry_run": True,
2062 "reservation_id": reservation_id,
2063 "region": region,
2064 "error_code": error_code,
2065 "error": error_msg,
2066 }
2067 return {
2068 "success": True,
2069 "dry_run": True,
2070 "reservation_id": reservation_id,
2071 "region": region,
2072 "message": "Dry run completed.",
2073 }
2075 try:
2076 response = ec2.cancel_capacity_reservation(CapacityReservationId=reservation_id)
2077 return {
2078 "success": bool(response.get("Return", True)),
2079 "dry_run": False,
2080 "reservation_id": reservation_id,
2081 "region": region,
2082 "message": f"Capacity reservation {reservation_id} cancelled.",
2083 }
2084 except ClientError as e:
2085 error_code = e.response.get("Error", {}).get("Code", "")
2086 error_msg = e.response.get("Error", {}).get("Message", str(e))
2087 return {
2088 "success": False,
2089 "dry_run": False,
2090 "reservation_id": reservation_id,
2091 "region": region,
2092 "error_code": error_code,
2093 "error": error_msg,
2094 }
2096 def recommend_region_for_job(
2097 self,
2098 gpu_required: bool = False,
2099 min_gpus: int = 0,
2100 instance_type: str | None = None,
2101 gpu_count: int = 0,
2102 ) -> dict[str, Any]:
2103 """
2104 Recommend the optimal region for job placement.
2106 Delegates to MultiRegionCapacityChecker for cross-region analysis.
2108 Args:
2109 gpu_required: Whether the job requires GPUs
2110 min_gpus: Minimum number of GPUs required
2111 instance_type: Specific instance type for workload-aware scoring
2112 gpu_count: Number of GPUs required
2114 Returns:
2115 Dictionary with recommended region and justification
2116 """
2117 from .multi_region import MultiRegionCapacityChecker
2119 checker = MultiRegionCapacityChecker(self.config)
2120 return checker.recommend_region_for_job(
2121 gpu_required, min_gpus, instance_type=instance_type, gpu_count=gpu_count
2122 )
2125def get_capacity_checker(config: GCOConfig | None = None) -> CapacityChecker:
2126 """Get a configured capacity checker instance."""
2127 return CapacityChecker(config)