Coverage for cli / capacity / checker.py: 100.00%
813 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""
2EC2 capacity checker using real AWS signals.
4This is the core capacity intelligence module. It queries multiple AWS APIs to
5build a picture of GPU/accelerator availability in a region, and manages
6reserved capacity end to end: On-Demand Capacity Reservations (ODCRs) and
7Capacity Blocks. Every query takes the region it should ask; the multi-region
8sweeps (``find_capacity_reservations``, ``find_capacity_blocks``,
9``check_reservation_availability``) fan out across regions in parallel
10themselves. ``MultiRegionCapacityChecker`` in multi_region.py reuses one
11instance, one region after another, for the instance-specific signals in its
12weighted region scoring.
14Data Sources:
15 - EC2 GetSpotPlacementScores: likelihood of getting spot capacity (1-10 score)
16 - EC2 DescribeSpotPriceHistory: current and historical spot prices (7-day window)
17 - AWS Price List GetProducts: Linux on-demand prices (through us-east-1)
18 - EC2 DescribeInstanceTypes: full compute characteristics — vCPU/cores/threads,
19 memory, every accelerator class, EFA and network limits, local NVMe and EBS,
20 placement-group support, purchase options, platform capabilities
21 - EC2 DescribeInstanceTypeOfferings and DescribeAvailabilityZones: which
22 instance types a region offers, and in how many of its AZs
23 - EC2 DescribeCapacityBlockOfferings: purchasable Capacity Blocks for ML workloads
24 - EC2 DescribeCapacityReservations: existing ODCRs and their free capacity
25 - EC2 PurchaseCapacityBlock, CreateCapacityReservation, CancelCapacityReservation:
26 the explicit, opt-in actions that spend money or release capacity
28Key Classes:
29 CapacityChecker: Main class. Instantiated with an optional GCOConfig; each
30 method takes the region (or regions) to query.
31 - estimate_capacity(instance_type, region) → list[CapacityEstimate]
32 - get_spot_price_history(instance_type, region) → list[SpotPriceInfo]
33 - get_instance_info(instance_type, region) → InstanceTypeInfo | None
34 - list_capacity_block_offerings(region, instance_type, ...) → list[dict]
35 - find_capacity_reservations(...) / find_capacity_blocks(...) → one
36 ranked report across the given regions (default: deployed regions)
38Output Models (defined in models.py):
39 - CapacityEstimate: availability, confidence, wait estimate, hourly price,
40 and a recommendation for one capacity type (spot or on-demand)
41 - SpotPriceInfo: one AZ's current spot price with its 7-day average, min,
42 max, and stability
43 - InstanceTypeInfo: every compute characteristic EC2 reports for a type
45Instance-type characteristics are always resolved live. A checked-in
46GPU_INSTANCE_SPECS table used to short-circuit DescribeInstanceTypes here; it was
47removed because a hand-maintained catalog hid new accelerator families until
48someone edited it, and a stale number fed straight into NodePool sizing and
49capacity scoring. The trade is one EC2 call per lookup, and that lookups now
50require credentials and a region that offers the type.
51"""
53from __future__ import annotations
55import functools
56import importlib.util
57import json
58import logging
59import statistics
60import sys
61from concurrent.futures import ThreadPoolExecutor, as_completed
62from datetime import UTC, datetime, timedelta
63from pathlib import Path
64from typing import Any
66import boto3
67from botocore.config import Config
68from botocore.exceptions import BotoCoreError, ClientError
70from cli.config import GCOConfig, get_config
72from . import blocks
73from .models import (
74 CapacityCheckError,
75 CapacityEstimate,
76 InstanceTypeInfo,
77 SpotPriceInfo,
78 instance_type_info_from_ec2,
79)
81logger = logging.getLogger(__name__)
84class SpotPlacementConfigLimitError(Exception):
85 """GetSpotPlacementScores refused the request with ``MaxConfigLimitExceeded``.
87 The account has asked about too many distinct Spot placement
88 configurations in the rolling 24-hour window. This is a refusal, not a
89 capacity signal: it must never be collapsed into an empty score mapping,
90 which callers could mistake for "this instance type has no capacity
91 data". Callers degrade to non-SPS signals and say the score was
92 unavailable.
93 """
96# Checkout-relative location of the pool catalog module, used when a plain
97# ``import scripts`` fails because the repository root is not on sys.path
98# (the installed ``gco`` console entrypoint). Patchable in tests.
99_POOL_CATALOG_PATH = Path(__file__).resolve().parents[2] / "scripts" / "accelerator_catalog.py"
102@functools.cache
103def _pool_catalog_lookup() -> Any | None:
104 """Return ``scripts.accelerator_catalog.pool_for_instance_type`` or ``None``.
106 The pool catalog lives beside the repository's maintenance tooling, not in
107 the installed ``cli``/``gco`` packages. Two situations need handling:
109 - The ``gco`` console entrypoint does not put the repository root on
110 ``sys.path``, so a plain ``import scripts`` fails even when running from
111 a checkout. The catalog module is self-contained (stdlib + PyYAML), so
112 load it directly from the file that sits two levels above this package.
113 - An installed wheel outside a checkout has no ``scripts/`` directory at
114 all; return ``None`` and let callers treat every type as unpooled
115 rather than issue a known-invalid single-type request.
117 Cached because the catalog is static for the life of the process and the
118 lookup runs on every capacity check.
119 """
120 try:
121 from scripts.accelerator_catalog import pool_for_instance_type
123 return pool_for_instance_type
124 except ImportError:
125 pass
126 catalog_path = _POOL_CATALOG_PATH
127 if not catalog_path.is_file():
128 return None
129 module_name = "_gco_pool_catalog"
130 try:
131 spec = importlib.util.spec_from_file_location(module_name, catalog_path)
132 if spec is None or spec.loader is None:
133 return None
134 module = importlib.util.module_from_spec(spec)
135 # dataclass creation resolves cls.__module__ through sys.modules, so
136 # the module must be registered before its body executes.
137 sys.modules[module_name] = module
138 try:
139 spec.loader.exec_module(module)
140 except BaseException:
141 sys.modules.pop(module_name, None)
142 raise
143 return module.pool_for_instance_type
144 except Exception as exc:
145 logger.debug("failed to load the instance pool catalog from %s: %s", catalog_path, exc)
146 return None
149# Capacity Block offering API error codes that mean "this instance type / region
150# simply doesn't have Capacity Blocks" rather than a real failure — expected and
151# handled quietly. ``Unsupported`` / ``UnsupportedOperation`` / ``InvalidAction``
152# cover regions where the Capacity Block API isn't available at all (e.g.
153# DescribeCapacityBlockOfferings returns ``InvalidAction`` in eu-west-1).
154_CB_EXPECTED_ERROR_CODES = frozenset(
155 {
156 "Unsupported",
157 "UnsupportedOperation",
158 "InvalidAction",
159 "InvalidParameterValue",
160 "InvalidParameterCombination",
161 }
162)
164# Adaptive client-side retry config for DescribeCapacityBlockOfferings. The
165# find_capacity_blocks sweep fans out region x duration probes in parallel, which
166# can trip the API's per-account request-rate limit (RequestLimitExceeded).
167# Adaptive mode adds client-side rate limiting plus exponential backoff, and a
168# higher attempt budget gives each probe room to succeed under throttling rather
169# than silently returning an empty (misleading) result for a throttled region.
170_CB_RETRY_CONFIG = Config(retries={"max_attempts": 10, "mode": "adaptive"})
172# A conservative cap on the parallel fan-out so a wide region x duration sweep
173# can't open an unbounded number of threads / sockets at once.
174_MAX_SEARCH_WORKERS = 12
177def _instance_desc(instance_type: str, gpu_count: int, gpu_type: str, total_gpu_mem: float) -> str:
178 """Build a human-readable instance description."""
179 if gpu_count > 0 and gpu_type:
180 mem_str = f", {total_gpu_mem:.0f}GB" if total_gpu_mem else ""
181 return f"{instance_type} ({gpu_count}x {gpu_type}{mem_str})"
182 return instance_type
185def _offering_fee(offering: dict[str, Any]) -> float:
186 """Sort key: an offering's upfront fee in USD, with missing fees sorting last."""
187 fee = offering.get("upfront_fee_usd")
188 if fee is None:
189 fee = blocks.parse_upfront_fee(offering.get("upfront_fee"))
190 return fee if fee is not None else float("inf")
193class CapacityChecker:
194 """
195 Checks EC2 capacity availability using real AWS capacity signals.
197 Uses:
198 - Spot Placement Score API for spot capacity estimates
199 - EC2 describe-instance-type-offerings for regional availability
200 - Spot price history for pricing trends
201 - On-demand pricing API
202 """
204 def __init__(self, config: GCOConfig | None = None):
205 self.config = config or get_config()
206 self._session = boto3.Session()
207 self._pricing_cache: dict[str, Any] = {}
208 self._offerings_cache: dict[str, set[str]] = {}
209 # Instance descriptions are now always resolved from EC2, and callers
210 # like _enrich_reservation_pricing ask per row. Memoize per
211 # (type, region) so listing 50 reservations of one type stays one call.
212 self._instance_info_cache: dict[tuple[str, str], InstanceTypeInfo | None] = {}
214 def get_instance_info(
215 self, instance_type: str, region: str | None = None
216 ) -> InstanceTypeInfo | None:
217 """Describe an instance type's compute characteristics from EC2.
219 Always resolved live. There is no checked-in specification table behind
220 this any more: the previous 25-entry GPU catalog short-circuited the API,
221 so a new accelerator family stayed invisible until someone hand-edited
222 it, and a wrong number propagated silently into NodePool sizing and
223 capacity scores. Every field now comes from
224 ``ec2:DescribeInstanceTypes`` via :func:`instance_type_info_from_ec2`.
226 ``DescribeInstanceTypes`` is region-scoped and returns a type only where
227 it is offered, so ``region`` matters. It defaults to the configured
228 default region rather than a hardcoded one; a type that region does not
229 offer returns ``None`` even though it may exist elsewhere. Use
230 ``check_instance_available_in_region`` or ``recommend-region`` to find
231 where a type is offered.
233 Friendly aliases are normalized first (``p6-b200`` ->
234 ``p6-b200.48xlarge``) so this agrees with the rest of the capacity
235 surface.
237 Returns ``None`` when the type is unknown, is not offered in the target
238 region, or cannot be described. The distinction is logged; callers
239 needing to tell those cases apart should use
240 :meth:`validate_instance_type`.
241 """
242 canonical, _ = blocks.normalize_instance_type(instance_type)
243 target_region = region or self.config.default_region
245 try:
246 ec2 = self._session.client("ec2", region_name=target_region)
247 response = ec2.describe_instance_types(InstanceTypes=[canonical])
248 except ClientError as e:
249 logger.debug(
250 "Failed to describe instance type %s in %s: %s", canonical, target_region, e
251 )
252 return None
253 except Exception as e:
254 logger.warning(
255 "Unexpected error getting instance info for %s in %s: %s",
256 canonical,
257 target_region,
258 e,
259 )
260 return None
262 records = response.get("InstanceTypes") or []
263 if not records:
264 logger.debug(
265 "Instance type %s is not offered in %s (it may exist in another region)",
266 canonical,
267 target_region,
268 )
269 return None
270 return instance_type_info_from_ec2(records[0], region=target_region)
272 def validate_instance_type(self, instance_type: str) -> dict[str, Any]:
273 """Classify an instance type as valid / invalid, with friendly normalization.
275 The offering APIs collapse "unknown instance type" and "valid type with
276 zero offerings" into the same empty result. This method separates them so
277 callers can tell a typo from genuine unavailability:
279 * EC2 ``DescribeInstanceTypes`` is consulted for every type; an
280 ``InvalidInstanceType`` marks it invalid, an empty result means the
281 type is not offered in the region consulted, and a transient API error
282 leaves it valid-but-unverified with a note.
283 * Friendly aliases are expanded (``p6-b200`` -> ``p6-b200.48xlarge``,
284 ``p6-b300`` -> ``p6-b300.48xlarge``) and UltraServer-only families
285 (the Grace-Blackwell ``gb200``/``gb300`` superchips, sold only as
286 ``P6e-GB`` UltraServers) are flagged invalid-for-``InstanceType`` with
287 guidance toward the UltraServer search flow.
289 Returns a dict: ``requested``, ``instance_type`` (canonical), ``valid``,
290 ``known``, ``note``, ``gpu_count``.
292 ``known`` means "EC2 returned a description for this type in the region
293 consulted". It used to mean "present in our hardcoded GPU catalog"; that
294 catalog is gone, so the flag now reflects live confirmation. It is
295 reported for information (as ``known_instance_type``) and never gates
296 anything.
298 Note the three-way outcome, which callers do depend on: a bogus type is
299 ``valid=False``; a real type the consulted region does not offer stays
300 ``valid=True`` with an explanatory ``note`` and ``known=False``; and a
301 confirmed type is ``valid=True, known=True``. Collapsing the middle case
302 into ``valid=False`` would reject a p5 in a region that merely lacks it,
303 which is what the removed catalog used to paper over.
304 """
305 canonical, note = blocks.normalize_instance_type(instance_type)
306 result: dict[str, Any] = {
307 "requested": instance_type,
308 "instance_type": canonical,
309 "valid": True,
310 "known": False,
311 "note": note,
312 "gpu_count": None,
313 }
315 # UltraServer-only families are real accelerators but not standalone EC2
316 # instance types, so InstanceType-based searches can never resolve them.
317 if (instance_type or "").strip().lower() in blocks.NON_STANDALONE_INSTANCE_NOTES:
318 result["valid"] = False
319 return result
321 # Ask EC2, separating "invalid type" from "not offered here" from
322 # "API error". An empty result is the middle case: the type is real but
323 # this region does not offer it, so it must not be reported invalid.
324 region = self.config.default_region
325 try:
326 ec2 = self._session.client("ec2", region_name=region)
327 response = ec2.describe_instance_types(InstanceTypes=[canonical])
328 types = response.get("InstanceTypes", [])
329 if not types:
330 result["note"] = note or (
331 f"Instance type is not offered in {region}; it may exist in another region."
332 )
333 return result
334 result["known"] = True
335 gpu_info = types[0].get("GpuInfo") or {}
336 # Sum across models rather than reading Gpus[0] so a heterogeneous
337 # accelerator list is not undercounted.
338 result["gpu_count"] = sum(
339 int(gpu.get("Count") or 0) for gpu in (gpu_info.get("Gpus") or [])
340 )
341 except ClientError as e:
342 code = e.response.get("Error", {}).get("Code", "")
343 if code in ("InvalidInstanceType", "InvalidParameterValue"):
344 result["valid"] = False
345 else:
346 logger.debug("Could not verify instance type %s: %s", canonical, e)
347 result["note"] = note or f"Could not verify instance type via EC2 ({code or e})."
348 except Exception as e:
349 logger.debug("Unexpected error validating instance type %s: %s", canonical, e)
350 result["note"] = note or f"Could not verify instance type: {e}"
351 return result
353 def check_instance_available_in_region(self, instance_type: str, region: str) -> bool:
354 """Check whether an instance type is offered in a region.
356 Returns ``True``/``False`` only from a *successful* offerings lookup, so a
357 ``False`` genuinely means "not offered in this region". If the underlying
358 DescribeInstanceTypeOfferings call fails (throttling, expired/invalid
359 credentials, denied permissions, region not opted in) this raises
360 :class:`CapacityCheckError` instead of silently returning ``False`` — a
361 failed check must not be reported to the user as "not available".
362 """
363 cache_key = f"{region}"
364 if cache_key not in self._offerings_cache:
365 try:
366 ec2 = self._session.client("ec2", region_name=region)
367 paginator = ec2.get_paginator("describe_instance_type_offerings")
368 offerings = set()
369 for page in paginator.paginate(LocationType="region"):
370 for offering in page["InstanceTypeOfferings"]:
371 offerings.add(offering["InstanceType"])
372 self._offerings_cache[cache_key] = offerings
373 except (ClientError, BotoCoreError) as e:
374 logger.warning("Failed to check instance offerings in %s: %s", region, e)
375 raise CapacityCheckError(
376 f"Could not check instance availability in {region}: {e}"
377 ) from e
379 return instance_type in self._offerings_cache[cache_key]
381 def get_availability_zones(self, region: str) -> list[str]:
382 """Get availability zones for a region."""
383 try:
384 ec2 = self._session.client("ec2", region_name=region)
385 response = ec2.describe_availability_zones(
386 Filters=[{"Name": "state", "Values": ["available"]}]
387 )
388 return [az["ZoneName"] for az in response["AvailabilityZones"]]
389 except ClientError as e:
390 logger.warning("Failed to get availability zones for %s: %s", region, e)
391 return []
392 except Exception as e:
393 logger.warning("Unexpected error getting AZs for %s: %s", region, e)
394 return []
396 def get_az_coverage(self, instance_type: str, region: str) -> float | None:
397 """Get the fraction of AZs in a region that offer this instance type.
399 Returns a value between 0.0 and 1.0, or None if we can't determine it.
400 Constrained instances are often available in fewer AZs.
401 """
402 try:
403 ec2 = self._session.client("ec2", region_name=region)
404 total_azs = self.get_availability_zones(region)
405 if not total_azs:
406 return None
408 paginator = ec2.get_paginator("describe_instance_type_offerings")
409 offering_azs = set()
410 for page in paginator.paginate(
411 LocationType="availability-zone",
412 Filters=[{"Name": "instance-type", "Values": [instance_type]}],
413 ):
414 for offering in page["InstanceTypeOfferings"]:
415 offering_azs.add(offering["Location"])
417 return len(offering_azs) / len(total_azs) if total_azs else None
418 except Exception as e:
419 logger.warning("Failed to get AZ coverage for %s in %s: %s", instance_type, region, e)
420 return None
422 @staticmethod
423 def instance_pool_for(instance_type: str) -> tuple[str, tuple[str, ...]] | None:
424 """Return (pool name, member types) for the first catalog pool containing the type.
426 AWS documents that ``GetSpotPlacementScores`` needs at least three
427 instance types to return meaningful scores, so every SPS request this
428 checker makes is scoped to a reviewed instance pool from
429 ``scripts/accelerator_catalog.py`` (the same catalog the capacity
430 poller uses). Returns ``None`` when no pool contains the type — and,
431 as a deliberate degradation, when the pool catalog is unavailable (an
432 installed wheel running outside a repository checkout ships without
433 ``scripts/``): an honest "no score obtained" beats a known-invalid
434 single-type request in both cases.
435 """
436 lookup = _pool_catalog_lookup()
437 if lookup is None:
438 logger.debug(
439 "instance pool catalog (scripts/accelerator_catalog.py) is unavailable; "
440 "treating %s as unpooled and skipping Spot Placement Scores",
441 instance_type,
442 )
443 return None
444 pool = lookup(instance_type)
445 if pool is None:
446 return None
447 return (pool.name, pool.members)
449 def get_spot_placement_score(
450 self, instance_type: str, region: str, target_capacity: int = 1
451 ) -> dict[str, int]:
452 """
453 Get the Spot Placement Score for the instance pool containing a type.
455 The Spot Placement Score (1-10) indicates the likelihood of getting
456 spot capacity. Higher scores mean better availability. The request
457 carries the full member list of the first catalog pool containing
458 ``instance_type`` — never the single type, which AWS documents as
459 returning misleadingly low scores — so the score describes the pool
460 at ``target_capacity``, not one instance type.
462 Returns:
463 Dict mapping AZ to score (1-10). Empty when no pool contains the
464 type (no request is made) or the API reports the type/region as
465 unsupported.
467 Raises:
468 SpotPlacementConfigLimitError: the account has queried too many
469 distinct SPS configurations in the rolling window
470 (``MaxConfigLimitExceeded``). Surfaced as a distinct condition
471 because "refused" must never be mistaken for "no capacity
472 data exists".
473 """
474 pool = self.instance_pool_for(instance_type)
475 if pool is None:
476 logger.info(
477 "no instance pool contains %s; skipping the Spot Placement Score request "
478 "(a single-type request returns misleadingly low scores)",
479 instance_type,
480 )
481 return {}
482 pool_name, members = pool
483 try:
484 ec2 = self._session.client("ec2", region_name=region)
486 response = ec2.get_spot_placement_scores(
487 InstanceTypes=list(members),
488 TargetCapacity=target_capacity,
489 TargetCapacityUnitType="units",
490 RegionNames=[region],
491 SingleAvailabilityZone=False,
492 )
494 scores = {}
495 for recommendation in response.get("SpotPlacementScores", []):
496 # Regional score
497 if "AvailabilityZoneId" not in recommendation:
498 scores["regional"] = recommendation.get("Score", 0)
499 else:
500 az_id = recommendation["AvailabilityZoneId"]
501 scores[az_id] = recommendation.get("Score", 0)
503 return scores
505 except ClientError as e:
506 error_code = e.response.get("Error", {}).get("Code", "")
507 if error_code == "MaxConfigLimitExceeded":
508 logger.warning(
509 "Spot Placement Score request refused (MaxConfigLimitExceeded) for "
510 "pool %s (%s) in %s: the account has queried too many distinct SPS "
511 "configurations in the rolling 24h window",
512 pool_name,
513 instance_type,
514 region,
515 )
516 raise SpotPlacementConfigLimitError(
517 f"Spot Placement Score request for pool {pool_name} in {region} was "
518 "refused: the account reached its Spot placement configuration limit "
519 "(MaxConfigLimitExceeded)"
520 ) from e
521 if error_code in ("InvalidParameterValue", "UnsupportedOperation"):
522 return {}
523 raise
524 except Exception as e:
525 logger.warning(
526 "Failed to get spot placement scores for %s in %s: %s", instance_type, region, e
527 )
528 return {}
530 def get_spot_price_history(
531 self, instance_type: str, region: str, days: int = 7
532 ) -> list[SpotPriceInfo]:
533 """Get spot price history for an instance type."""
534 ec2 = self._session.client("ec2", region_name=region)
536 end_time = datetime.now(UTC)
537 start_time = end_time - timedelta(days=days)
539 try:
540 response = ec2.describe_spot_price_history(
541 InstanceTypes=[instance_type],
542 ProductDescriptions=["Linux/UNIX"],
543 StartTime=start_time,
544 EndTime=end_time,
545 )
547 # Group by availability zone
548 az_prices: dict[str, list[float]] = {}
549 for item in response["SpotPriceHistory"]:
550 az = item["AvailabilityZone"]
551 price = float(item["SpotPrice"])
552 if az not in az_prices:
553 az_prices[az] = []
554 az_prices[az].append(price)
556 results = []
557 for az, prices in az_prices.items():
558 current = prices[0]
559 avg = statistics.mean(prices)
560 min_price = min(prices)
561 max_price = max(prices)
563 if avg > 0:
564 std_dev = statistics.stdev(prices) if len(prices) > 1 else 0
565 cv = std_dev / avg
566 stability = max(0, 1 - cv)
567 else:
568 stability = 0
570 results.append(
571 SpotPriceInfo(
572 instance_type=instance_type,
573 availability_zone=az,
574 current_price=current,
575 avg_price_7d=avg,
576 min_price_7d=min_price,
577 max_price_7d=max_price,
578 price_stability=stability,
579 )
580 )
582 return results
584 except ClientError as e:
585 if "InvalidParameterValue" in str(e):
586 return []
587 raise
589 def get_on_demand_price(self, instance_type: str, region: str) -> float | None:
590 """Get on-demand price for an instance type."""
591 cache_key = f"{instance_type}:{region}"
592 if cache_key in self._pricing_cache:
593 cached_value = self._pricing_cache[cache_key]
594 return float(cached_value) if cached_value is not None else None
596 try:
597 pricing = self._session.client("pricing", region_name="us-east-1")
599 # Filter on the regionCode product attribute — the AWS region
600 # identifier itself — never on the human-readable location name.
601 # The former ten-entry region-name map silently returned None
602 # for every region outside it, and Price List location strings
603 # are not derivable from any published source (eu-north-1 is
604 # "EU (Stockholm)", not "Europe (Stockholm)"), so a name map is
605 # inherently fragile. regionCode removes the class of bug:
606 # verified to return identical prices for all ten formerly
607 # mapped regions and correct prices for previously failing ones.
608 response = pricing.get_products(
609 ServiceCode="AmazonEC2",
610 Filters=[
611 {"Type": "TERM_MATCH", "Field": "instanceType", "Value": instance_type},
612 {"Type": "TERM_MATCH", "Field": "regionCode", "Value": region},
613 {"Type": "TERM_MATCH", "Field": "operatingSystem", "Value": "Linux"},
614 {"Type": "TERM_MATCH", "Field": "tenancy", "Value": "Shared"},
615 {"Type": "TERM_MATCH", "Field": "preInstalledSw", "Value": "NA"},
616 {"Type": "TERM_MATCH", "Field": "capacitystatus", "Value": "Used"},
617 ],
618 MaxResults=1,
619 )
621 if response["PriceList"]:
622 price_data = json.loads(response["PriceList"][0])
623 terms = price_data.get("terms", {}).get("OnDemand", {})
624 for term in terms.values():
625 for price_dim in term.get("priceDimensions", {}).values():
626 price = float(price_dim["pricePerUnit"]["USD"])
627 self._pricing_cache[cache_key] = price
628 return price
630 except ClientError as e:
631 logger.warning(
632 "Failed to get on-demand price for %s in %s: %s", instance_type, region, e
633 )
634 except Exception as e:
635 logger.warning(
636 "Unexpected error getting pricing for %s in %s: %s", instance_type, region, e
637 )
639 return None
641 def estimate_capacity(
642 self, instance_type: str, region: str, capacity_type: str = "both"
643 ) -> list[CapacityEstimate]:
644 """
645 Estimate capacity availability using real AWS signals.
647 Args:
648 instance_type: EC2 instance type
649 region: AWS region
650 capacity_type: "spot", "on-demand", or "both"
652 Returns:
653 List of CapacityEstimate objects
654 """
655 estimates = []
657 # Check if instance type is available in region
658 if not self.check_instance_available_in_region(instance_type, region):
659 return [
660 CapacityEstimate(
661 instance_type=instance_type,
662 region=region,
663 availability_zone=None,
664 capacity_type="both",
665 availability="unavailable",
666 confidence=1.0,
667 recommendation=f"{instance_type} is not available in {region}",
668 details={"reason": "Instance type not offered in region"},
669 )
670 ]
672 instance_info = self.get_instance_info(instance_type)
674 if capacity_type in ("spot", "both"):
675 spot_estimates = self._estimate_spot_capacity(instance_type, region, instance_info)
676 estimates.extend(spot_estimates)
678 if capacity_type in ("on-demand", "both"):
679 # Pass spot placement scores to on-demand estimator as a scarcity signal
680 spot_scores = {}
681 if capacity_type == "on-demand":
682 try:
683 spot_scores = self.get_spot_placement_score(instance_type, region)
684 except SpotPlacementConfigLimitError:
685 # The refusal is already logged distinctly; the on-demand
686 # estimate degrades to its other scarcity signals.
687 spot_scores = {}
688 # If we already fetched spot estimates, extract the scores from them
689 if spot_estimates := [e for e in estimates if e.capacity_type == "spot"]:
690 spot_scores = {
691 e.availability_zone or "unknown": e.details.get("spot_placement_score", 0)
692 for e in spot_estimates
693 if e.details.get("spot_placement_score") is not None
694 }
695 # Also gather spot price data for price-ratio signal
696 spot_prices = self.get_spot_price_history(instance_type, region)
697 od_estimate = self._estimate_on_demand_capacity(
698 instance_type, region, instance_info, spot_scores, spot_prices
699 )
700 if od_estimate:
701 estimates.append(od_estimate)
703 return estimates
705 def _estimate_spot_capacity(
706 self, instance_type: str, region: str, instance_info: InstanceTypeInfo | None
707 ) -> list[CapacityEstimate]:
708 """Estimate spot capacity using pooled Spot Placement Scores and price history."""
709 estimates = []
711 # Spot Placement Score (primary signal), requested for the instance
712 # pool containing the type. Three distinct non-score outcomes exist
713 # and must stay distinguishable: the type is in no pool (no request
714 # made), the account's configuration limit refused the request, or
715 # the API had nothing to say. None of them may borrow the confidence
716 # reserved for a real score.
717 pool = self.instance_pool_for(instance_type)
718 sps_target_capacity = 1
719 placement_scores: dict[str, int] = {}
720 if pool is None:
721 sps_note = (
722 f"no placement score was obtained: {instance_type} belongs to no "
723 "instance pool, and a single-type request returns misleadingly "
724 "low scores"
725 )
726 else:
727 sps_note = "no placement score was obtained"
728 try:
729 placement_scores = self.get_spot_placement_score(
730 instance_type, region, sps_target_capacity
731 )
732 except SpotPlacementConfigLimitError:
733 sps_note = (
734 "no placement score was obtained: the account reached its Spot "
735 "placement configuration limit (MaxConfigLimitExceeded)"
736 )
738 # Get spot prices for pricing info
739 spot_prices = self.get_spot_price_history(instance_type, region)
740 price_by_az = {sp.availability_zone: sp for sp in spot_prices}
742 on_demand_price = self.get_on_demand_price(instance_type, region)
744 # Get AZs in the region
745 azs = self.get_availability_zones(region)
747 if placement_scores:
748 # Use Spot Placement Score as primary signal
749 pool_name = pool[0] if pool is not None else "unknown"
750 regional_score = placement_scores.get("regional", 0)
752 for az in azs:
753 # Try to get AZ-specific score, fall back to regional
754 az_id = az # Note: might need to map zone name to zone ID
755 score = placement_scores.get(az_id, regional_score)
757 # Convert score (1-10) to availability
758 if score >= 8:
759 availability = "high"
760 recommendation = "Excellent spot availability"
761 elif score >= 5:
762 availability = "medium"
763 recommendation = "Good spot availability, some interruption risk"
764 elif score >= 3:
765 availability = "low"
766 recommendation = "Limited spot capacity, consider alternatives"
767 else:
768 availability = "low"
769 recommendation = "Very limited spot capacity"
771 spot_info = price_by_az.get(az)
772 price = spot_info.current_price if spot_info else None
774 details: dict[str, Any] = {
775 "spot_placement_score": score,
776 # The score describes the pool's fleet at the stated
777 # target capacity, not this one instance type.
778 "score_interpretation": (
779 f"{score}/10 for instance pool {pool_name} "
780 f"at target capacity {sps_target_capacity}"
781 ),
782 "spot_pool": pool_name,
783 "sps_target_capacity": sps_target_capacity,
784 }
786 if spot_info:
787 details["current_price"] = spot_info.current_price
788 details["avg_price_7d"] = spot_info.avg_price_7d
789 details["price_stability"] = f"{spot_info.price_stability:.2f}"
791 if on_demand_price and price:
792 savings = (1 - price / on_demand_price) * 100
793 details["savings_vs_on_demand"] = f"{savings:.1f}%"
794 details["on_demand_price"] = on_demand_price
796 estimates.append(
797 CapacityEstimate(
798 instance_type=instance_type,
799 region=region,
800 availability_zone=az,
801 capacity_type="spot",
802 availability=availability,
803 confidence=0.85, # Spot Placement Score is reliable
804 price_per_hour=price,
805 recommendation=recommendation,
806 details=details,
807 )
808 )
810 elif spot_prices:
811 # Fall back to price-based estimation if no placement score
812 for spot_info in spot_prices:
813 # Use price stability as a proxy (less reliable)
814 if spot_info.price_stability > 0.8:
815 availability = "medium"
816 recommendation = "Spot prices stable, likely available"
817 elif spot_info.price_stability > 0.5:
818 availability = "low"
819 recommendation = "Spot prices volatile, capacity uncertain"
820 else:
821 availability = "low"
822 recommendation = "High price volatility, limited capacity likely"
824 details = {
825 "current_price": spot_info.current_price,
826 "avg_price_7d": spot_info.avg_price_7d,
827 "price_stability": f"{spot_info.price_stability:.2f}",
828 "note": f"Estimate based on price history; {sps_note}",
829 }
831 if on_demand_price:
832 savings = (1 - spot_info.current_price / on_demand_price) * 100
833 details["savings_vs_on_demand"] = f"{savings:.1f}%"
835 estimates.append(
836 CapacityEstimate(
837 instance_type=instance_type,
838 region=region,
839 availability_zone=spot_info.availability_zone,
840 capacity_type="spot",
841 availability=availability,
842 confidence=0.5, # Lower confidence without placement score
843 price_per_hour=spot_info.current_price,
844 recommendation=f"{recommendation} ({sps_note})",
845 details=details,
846 )
847 )
849 if not estimates:
850 estimates.append(
851 CapacityEstimate(
852 instance_type=instance_type,
853 region=region,
854 availability_zone=None,
855 capacity_type="spot",
856 availability="unknown",
857 confidence=0.1,
858 recommendation=(
859 f"No spot data available for {instance_type} in {region}; {sps_note}"
860 ),
861 details={"reason": f"No spot price history, and {sps_note}"},
862 )
863 )
865 return estimates
867 def _estimate_on_demand_capacity(
868 self,
869 instance_type: str,
870 region: str,
871 instance_info: InstanceTypeInfo | None,
872 spot_placement_scores: dict[str, int] | None = None,
873 spot_prices: list[SpotPriceInfo] | None = None,
874 ) -> CapacityEstimate | None:
875 """Estimate on-demand capacity using live signals for ALL instance types.
877 Uses spot placement scores, instance size (vCPUs, memory, GPUs),
878 pricing, and spot-to-on-demand price ratios as universal scarcity
879 signals — no hardcoded instance families or GPU type lists.
880 """
881 on_demand_price = self.get_on_demand_price(instance_type, region)
883 is_offered = self.check_instance_available_in_region(instance_type, region)
885 if not is_offered:
886 return CapacityEstimate(
887 instance_type=instance_type,
888 region=region,
889 availability_zone=None,
890 capacity_type="on-demand",
891 availability="unavailable",
892 confidence=1.0,
893 recommendation=f"{instance_type} is not offered in {region}",
894 details={"reason": "Instance type not offered in region"},
895 )
897 # Fetch spot placement scores if not provided (on-demand only mode)
898 if spot_placement_scores is None:
899 try:
900 spot_placement_scores = self.get_spot_placement_score(instance_type, region)
901 except SpotPlacementConfigLimitError:
902 # Logged distinctly at the source; degrade to the other
903 # universal scarcity signals rather than failing the estimate.
904 spot_placement_scores = {}
906 if spot_prices is None:
907 spot_prices = self.get_spot_price_history(instance_type, region)
909 az_coverage = self.get_az_coverage(instance_type, region)
911 availability, confidence, recommendation = self._assess_on_demand_availability(
912 instance_type,
913 instance_info,
914 on_demand_price,
915 spot_placement_scores,
916 spot_prices,
917 az_coverage,
918 )
920 if on_demand_price:
921 recommendation += f" Price: ${on_demand_price:.4f}/hr."
922 else:
923 confidence -= 0.1
924 recommendation += " Pricing data unavailable."
926 details: dict[str, Any] = {
927 "price_per_hour": on_demand_price,
928 "is_gpu": instance_info.is_gpu if instance_info else False,
929 }
930 if spot_placement_scores:
931 scores = [s for s in spot_placement_scores.values() if s > 0]
932 if scores:
933 details["avg_spot_placement_score"] = round(sum(scores) / len(scores), 1)
935 return CapacityEstimate(
936 instance_type=instance_type,
937 region=region,
938 availability_zone=None,
939 capacity_type="on-demand",
940 availability=availability,
941 confidence=confidence,
942 price_per_hour=on_demand_price,
943 recommendation=recommendation,
944 details=details,
945 )
947 @staticmethod
948 def _assess_on_demand_availability(
949 instance_type: str,
950 instance_info: InstanceTypeInfo | None,
951 on_demand_price: float | None,
952 spot_placement_scores: dict[str, int] | None = None,
953 spot_prices: list[SpotPriceInfo] | None = None,
954 az_coverage: float | None = None,
955 ) -> tuple[str, float, str]:
956 """Assess on-demand availability using only live market signals.
958 Five live signals, zero hardcoded instance families:
959 1. Spot placement score — AWS's own capacity assessment (1-10)
960 2. Spot-to-on-demand price ratio — when spot approaches on-demand price,
961 the spot market has very little excess capacity
962 3. Spot price volatility — unstable prices reflect capacity fluctuations
963 4. AZ coverage — fraction of AZs that offer this instance type;
964 constrained instances are often available in fewer AZs
965 5. Spot price availability — how many AZs have spot price data;
966 missing price data in some AZs suggests limited capacity there
968 Confidence scales with the number of live signals available.
969 When no signals exist, returns "unknown" rather than guessing.
971 Returns:
972 Tuple of (availability, confidence, recommendation)
973 """
974 price = on_demand_price or 0
975 gpu_count = instance_info.gpu_count if instance_info else 0
976 gpu_type = (instance_info.gpu_type or "") if instance_info else ""
977 total_gpu_mem = instance_info.gpu_memory_gib if instance_info else 0
979 # --- Signal 1: Spot placement score ---
980 avg_spot_score = 0.0
981 has_spot_score = False
982 if spot_placement_scores:
983 scores = [s for s in spot_placement_scores.values() if s > 0]
984 if scores:
985 avg_spot_score = sum(scores) / len(scores)
986 has_spot_score = True
988 # --- Signal 2 & 3: Spot price ratio and volatility ---
989 avg_spot_ratio = 0.0
990 avg_stability = 1.0
991 has_price_signal = False
992 if spot_prices and price > 0:
993 ratios = [sp.current_price / price for sp in spot_prices if sp.current_price > 0]
994 if ratios:
995 avg_spot_ratio = sum(ratios) / len(ratios)
996 has_price_signal = True
997 stabilities = [sp.price_stability for sp in spot_prices]
998 avg_stability = sum(stabilities) / len(stabilities)
1000 # --- Signal 4: AZ coverage (passed in from caller) ---
1001 has_az_signal = az_coverage is not None
1003 # --- Combine live signals into scarcity (0.0 - 1.0) ---
1004 scarcity = 0.0
1005 signal_count = 0
1007 if has_spot_score:
1008 signal_count += 1
1009 if avg_spot_score <= 2:
1010 scarcity += 0.5
1011 elif avg_spot_score <= 4:
1012 scarcity += 0.3
1013 elif avg_spot_score <= 6:
1014 scarcity += 0.15
1016 if has_price_signal:
1017 signal_count += 1
1018 # Spot price near on-demand = spot market has minimal excess capacity
1019 if avg_spot_ratio >= 0.9:
1020 scarcity += 0.3
1021 elif avg_spot_ratio >= 0.7:
1022 scarcity += 0.15
1023 elif avg_spot_ratio >= 0.5:
1024 scarcity += 0.05
1026 # Price instability = capacity fluctuations
1027 if avg_stability < 0.6:
1028 scarcity += 0.1
1029 elif avg_stability < 0.8:
1030 scarcity += 0.05
1032 if has_az_signal and az_coverage is not None:
1033 signal_count += 1
1034 # Available in fewer than half the AZs = constrained
1035 if az_coverage <= 0.3:
1036 scarcity += 0.2
1037 elif az_coverage <= 0.5:
1038 scarcity += 0.1
1040 # --- Confidence scales with signal count ---
1041 confidence = min(0.5 + (signal_count * 0.12), 0.9)
1043 # --- Map scarcity to availability ---
1044 desc = _instance_desc(instance_type, gpu_count, gpu_type, total_gpu_mem)
1046 if signal_count == 0:
1047 # No live data — be honest about it
1048 return (
1049 "unknown",
1050 0.3,
1051 f"No live capacity signals available for {instance_type}."
1052 " Unable to assess on-demand availability.",
1053 )
1055 if scarcity >= 0.6:
1056 return (
1057 "low",
1058 confidence,
1059 f"On-demand {desc} is extremely scarce based on live capacity signals."
1060 " Capacity reservations or Capacity Blocks are strongly recommended.",
1061 )
1063 if scarcity >= 0.35:
1064 return (
1065 "low",
1066 confidence,
1067 f"On-demand {desc} has limited availability based on live capacity signals."
1068 " Consider capacity reservations.",
1069 )
1071 if scarcity >= 0.15:
1072 return (
1073 "medium",
1074 confidence,
1075 f"On-demand {instance_type} may have constrained availability"
1076 " based on current market conditions.",
1077 )
1079 return (
1080 "high",
1081 confidence,
1082 f"On-demand capacity likely available for {instance_type}"
1083 " based on live capacity signals.",
1084 )
1086 def recommend_capacity_type(
1087 self, instance_type: str, region: str, fault_tolerance: str = "medium"
1088 ) -> tuple[str, str]:
1089 """
1090 Recommend spot vs on-demand based on actual capacity and requirements.
1092 Args:
1093 instance_type: EC2 instance type
1094 region: AWS region
1095 fault_tolerance: "high" (can handle interruptions),
1096 "medium" (some tolerance),
1097 "low" (needs stability)
1099 Returns:
1100 Tuple of (recommended_capacity_type, explanation)
1101 """
1102 estimates = self.estimate_capacity(instance_type, region, "both")
1104 spot_estimates = [e for e in estimates if e.capacity_type == "spot"]
1105 od_estimates = [e for e in estimates if e.capacity_type == "on-demand"]
1107 # Check for unavailable
1108 if any(e.availability == "unavailable" for e in estimates):
1109 return "unavailable", f"{instance_type} is not available in {region}"
1111 # Get best spot option (highest availability)
1112 best_spot = None
1113 if spot_estimates:
1114 available_spots = [e for e in spot_estimates if e.availability != "unknown"]
1115 if available_spots:
1116 # Sort by availability (high > medium > low) then by price
1117 avail_order = {"high": 0, "medium": 1, "low": 2}
1118 best_spot = min(
1119 available_spots,
1120 key=lambda x: (avail_order.get(x.availability, 3), x.price_per_hour or 999),
1121 )
1123 od_estimate = od_estimates[0] if od_estimates else None
1125 # Decision logic based on fault tolerance and actual availability
1126 if fault_tolerance == "low":
1127 if od_estimate and od_estimate.availability in ("high", "medium"):
1128 return "on-demand", "Low fault tolerance requires stable on-demand capacity"
1129 return (
1130 "on-demand",
1131 "On-demand recommended but capacity may be limited; consider capacity reservation",
1132 )
1134 if best_spot:
1135 if best_spot.availability == "high":
1136 savings = ""
1137 if best_spot.price_per_hour and od_estimate and od_estimate.price_per_hour:
1138 pct = (1 - best_spot.price_per_hour / od_estimate.price_per_hour) * 100
1139 savings = f" (save ~{pct:.0f}%)"
1140 return "spot", f"High spot availability (score-based){savings}"
1142 if best_spot.availability == "medium":
1143 if fault_tolerance == "high":
1144 return "spot", "Medium spot availability acceptable with high fault tolerance"
1145 return (
1146 "on-demand",
1147 "Spot availability is medium; on-demand recommended for reliability",
1148 )
1150 # ``unknown`` values were filtered above; high and medium already
1151 # returned, so low is the only remaining availability in the model.
1152 if fault_tolerance == "high":
1153 return (
1154 "spot",
1155 "Low spot availability but acceptable with high fault tolerance",
1156 )
1157 return "on-demand", "Spot capacity is limited; on-demand recommended"
1159 # Default to on-demand
1160 return "on-demand", "On-demand recommended (spot availability unknown or limited)"
1162 # -------------------------------------------------------------------------
1163 # Capacity Reservations (ODCRs) and Capacity Blocks for ML
1164 # -------------------------------------------------------------------------
1166 def list_capacity_reservations(
1167 self,
1168 region: str,
1169 instance_type: str | None = None,
1170 state: str | None = "active",
1171 *,
1172 include_pricing: bool = False,
1173 ) -> list[dict[str, Any]]:
1174 """
1175 List EC2 On-Demand Capacity Reservations (ODCRs) in a region.
1177 Args:
1178 region: AWS region to query
1179 instance_type: Filter by instance type (optional)
1180 state: Filter by state — "active" (default), or None for all
1181 include_pricing: Enrich each reservation with On-Demand pricing
1182 (per-instance-hour, whole-reservation per-hour, per-GPU-hour).
1183 Adds one Pricing API call per distinct instance type (cached), so
1184 it is opt-in and off by default to keep the plain list fast.
1186 Returns:
1187 List of reservation dictionaries with availability details
1188 """
1189 ec2 = self._session.client("ec2", region_name=region)
1191 filters: list[dict[str, Any]] = []
1192 if state:
1193 filters.append({"Name": "state", "Values": [state]})
1194 if instance_type:
1195 filters.append({"Name": "instance-type", "Values": [instance_type]})
1197 reservations: list[dict[str, Any]] = []
1198 try:
1199 paginator = ec2.get_paginator("describe_capacity_reservations")
1200 page_kwargs: dict[str, Any] = {}
1201 if filters:
1202 page_kwargs["Filters"] = filters
1204 for page in paginator.paginate(**page_kwargs):
1205 for cr in page.get("CapacityReservations", []):
1206 total = cr.get("TotalInstanceCount", 0)
1207 available = cr.get("AvailableInstanceCount", 0)
1208 used = total - available
1210 entry = {
1211 "type": "odcr",
1212 "reservation_id": cr.get("CapacityReservationId"),
1213 "instance_type": cr.get("InstanceType"),
1214 "availability_zone": cr.get("AvailabilityZone"),
1215 "region": region,
1216 "state": cr.get("State"),
1217 "total_instances": total,
1218 "available_instances": available,
1219 "used_instances": used,
1220 "utilization_pct": round(used / total * 100, 1) if total else 0,
1221 "instance_platform": cr.get("InstancePlatform"),
1222 "tenancy": cr.get("Tenancy"),
1223 "instance_match_criteria": cr.get("InstanceMatchCriteria"),
1224 "start_date": (
1225 cr["StartDate"].isoformat() if cr.get("StartDate") else None
1226 ),
1227 "end_date": (cr["EndDate"].isoformat() if cr.get("EndDate") else None),
1228 "end_date_type": cr.get("EndDateType"),
1229 "tags": {t["Key"]: t["Value"] for t in cr.get("Tags", [])},
1230 }
1231 if include_pricing:
1232 self._enrich_reservation_pricing(entry, region)
1233 reservations.append(entry)
1234 except ClientError as e:
1235 logger.debug("Failed to list capacity reservations in %s: %s", region, e)
1237 return reservations
1239 def _gpus_per_instance(self, instance_type: str, region: str) -> int | None:
1240 """GPUs per instance for per-GPU-hour pricing, memoized per type+region.
1242 Returns ``None`` when the type cannot be described, which callers treat
1243 as "omit the per-GPU figure" rather than as zero GPUs — reporting a
1244 per-GPU-hour price of infinity, or silently dividing by a wrong count,
1245 would be worse than omitting it.
1246 """
1247 key = (instance_type, region)
1248 if key not in self._instance_info_cache:
1249 self._instance_info_cache[key] = self.get_instance_info(instance_type, region=region)
1250 info = self._instance_info_cache[key]
1251 return info.gpu_count if info else None
1253 def _enrich_reservation_pricing(self, reservation: dict[str, Any], region: str) -> None:
1254 """Add On-Demand pricing keys to a reservation dict, in place.
1256 ODCRs bill at the On-Demand rate for the reserved instance type, so the
1257 per-instance-hour / per-hour / per-GPU-hour figures mirror the Capacity
1258 Block pricing surface (see :func:`blocks.compute_reservation_pricing`) and
1259 let a caller rank and compare reservations the same way it ranks blocks.
1260 Pricing that can't be resolved is recorded as ``None`` — never fatal.
1261 """
1262 instance_type = reservation.get("instance_type")
1263 if not instance_type:
1264 return
1265 try:
1266 on_demand = self.get_on_demand_price(instance_type, region)
1267 except Exception as e: # pricing is supplementary; never fail the listing
1268 logger.debug("On-demand price lookup failed for %s in %s: %s", instance_type, region, e)
1269 on_demand = None
1271 gpus_per_instance = self._gpus_per_instance(str(instance_type), region)
1273 pricing = blocks.compute_reservation_pricing(
1274 on_demand, reservation.get("total_instances"), gpus_per_instance
1275 )
1276 reservation["on_demand_price_per_hour"] = (
1277 round(float(on_demand), 4) if on_demand is not None else None
1278 )
1279 reservation["gpus_per_instance"] = gpus_per_instance
1280 reservation.update(pricing)
1282 def _build_block_offering(
1283 self,
1284 offering: dict[str, Any],
1285 region: str,
1286 requested_count: int,
1287 gpus_per_instance: int | None,
1288 requested_duration_hours: int,
1289 ) -> dict[str, Any]:
1290 """Shape a raw DescribeCapacityBlockOfferings entry into an enriched dict.
1292 Reports the offering's *actual* duration (the API returns blocks whose
1293 duration is the closest match to the request, not necessarily equal) and
1294 adds per-hour / per-GPU-hour pricing derived from the upfront fee.
1295 ``upfront_fee`` is the raw API value (a string); ``upfront_fee_usd`` is the
1296 parsed float used for ranking and display.
1297 """
1298 start_date = offering.get("StartDate")
1299 end_date = offering.get("EndDate")
1301 actual_duration = offering.get("CapacityBlockDurationHours")
1302 if actual_duration is None:
1303 minutes = offering.get("CapacityBlockDurationMinutes")
1304 actual_duration = round(minutes / 60) if minutes else requested_duration_hours
1306 count = offering.get("InstanceCount")
1307 if count is None:
1308 count = requested_count
1310 pricing = blocks.compute_offering_pricing(
1311 offering.get("UpfrontFee"), actual_duration, count, gpus_per_instance
1312 )
1314 return {
1315 "type": "capacity_block",
1316 "offering_id": offering.get("CapacityBlockOfferingId"),
1317 "instance_type": offering.get("InstanceType"),
1318 "availability_zone": offering.get("AvailabilityZone"),
1319 "region": region,
1320 "instance_count": count,
1321 "duration_hours": actual_duration,
1322 "duration_days": blocks.hours_to_days(actual_duration),
1323 "start_date": start_date.isoformat() if start_date else None,
1324 "end_date": end_date.isoformat() if end_date else None,
1325 "upfront_fee": offering.get("UpfrontFee"),
1326 "upfront_fee_usd": pricing["upfront_fee_usd"],
1327 "price_per_hour": pricing["price_per_hour"],
1328 "price_per_instance_hour": pricing["price_per_instance_hour"],
1329 "price_per_gpu_hour": pricing["price_per_gpu_hour"],
1330 "gpus_per_instance": gpus_per_instance,
1331 "currency": offering.get("CurrencyCode", "USD"),
1332 "tenancy": offering.get("Tenancy"),
1333 }
1335 def list_capacity_block_offerings(
1336 self,
1337 region: str,
1338 instance_type: str,
1339 instance_count: int = 1,
1340 duration_hours: int = 24,
1341 *,
1342 earliest_start: datetime | None = None,
1343 latest_start: datetime | None = None,
1344 gpus_per_instance: int | None = None,
1345 ) -> list[dict[str, Any]]:
1346 """
1347 List available Capacity Block offerings for ML workloads.
1349 Capacity Blocks provide guaranteed GPU capacity for a fixed duration
1350 at a known price — ideal for training jobs with predictable runtimes.
1351 Queries a single duration in a single region; the date window and the
1352 multi-duration / multi-region sweep are layered on top by
1353 :meth:`find_capacity_blocks`.
1355 Args:
1356 region: AWS region to query
1357 instance_type: GPU instance type (e.g. p5.48xlarge, p4d.24xlarge)
1358 instance_count: Number of instances needed
1359 duration_hours: Desired block duration in hours (must be a supported value)
1360 earliest_start: Only return blocks starting on/after this datetime
1361 (EC2 StartDateRange). Lets callers ask "blocks starting near D1".
1362 latest_start: Only return blocks starting on/before this datetime
1363 (EC2 EndDateRange).
1364 gpus_per_instance: GPUs per instance, used for per-GPU-hour pricing.
1365 Resolved from the instance specs when omitted.
1367 Returns:
1368 List of available capacity block offerings (enriched with pricing).
1369 All matching pages are followed via NextToken.
1370 """
1371 ec2 = self._session.client("ec2", region_name=region, config=_CB_RETRY_CONFIG)
1372 offerings: list[dict[str, Any]] = []
1374 if gpus_per_instance is None:
1375 info = self.get_instance_info(instance_type)
1376 gpus_per_instance = info.gpu_count if info else None
1378 api_kwargs: dict[str, Any] = {
1379 "InstanceType": instance_type,
1380 "InstanceCount": instance_count,
1381 "CapacityDurationHours": duration_hours,
1382 }
1383 if earliest_start is not None:
1384 api_kwargs["StartDateRange"] = earliest_start
1385 if latest_start is not None:
1386 api_kwargs["EndDateRange"] = latest_start
1388 try:
1389 next_token: str | None = None
1390 while True:
1391 if next_token:
1392 api_kwargs["NextToken"] = next_token
1393 response = ec2.describe_capacity_block_offerings(**api_kwargs)
1394 for offering in response.get("CapacityBlockOfferings", []):
1395 offerings.append(
1396 self._build_block_offering(
1397 offering, region, instance_count, gpus_per_instance, duration_hours
1398 )
1399 )
1400 next_token = response.get("NextToken")
1401 if not next_token:
1402 break
1403 except ClientError as e:
1404 error_code = e.response.get("Error", {}).get("Code", "")
1405 if error_code in _CB_EXPECTED_ERROR_CODES:
1406 pass # Type/region doesn't support Capacity Blocks — expected
1407 else:
1408 logger.warning("Failed to list capacity block offerings in %s: %s", region, e)
1409 except BotoCoreError as e:
1410 # Endpoint resolution / connection errors surface here when the
1411 # Capacity Block API isn't available in a region.
1412 logger.warning("Capacity Block API unavailable in %s: %s", region, e)
1414 return offerings
1416 def get_capacity_block_trend(
1417 self,
1418 instance_type: str,
1419 region: str,
1420 ) -> float:
1421 """
1422 Estimate capacity block availability trend via time-series regression.
1424 Queries offerings across the maximum 182-day (26-week) window, buckets
1425 them into weekly bins by start date, and fits a linear regression to
1426 the offering counts per week. The normalized slope indicates whether
1427 capacity is growing or shrinking over time.
1429 Returns:
1430 Trend score from -1.0 to 1.0:
1431 > 0 = capacity growing (offerings increasing week-over-week)
1432 = 0 = stable or no data
1433 < 0 = capacity shrinking (offerings decreasing week-over-week)
1434 """
1435 ec2 = self._session.client("ec2", region_name=region, config=_CB_RETRY_CONFIG)
1437 now = datetime.now(UTC)
1438 far_end = now + timedelta(days=182)
1440 try:
1441 response = ec2.describe_capacity_block_offerings(
1442 InstanceType=instance_type,
1443 InstanceCount=1,
1444 CapacityDurationHours=24, # Minimum duration for broadest results
1445 StartDateRange=now,
1446 EndDateRange=far_end,
1447 )
1448 except ClientError as e:
1449 error_code = e.response.get("Error", {}).get("Code", "")
1450 if error_code not in _CB_EXPECTED_ERROR_CODES:
1451 logger.warning(
1452 "Capacity block trend query failed for %s in %s: %s",
1453 instance_type,
1454 region,
1455 e,
1456 )
1457 return 0.0
1458 except BotoCoreError as e:
1459 logger.warning(
1460 "Capacity block trend query failed for %s in %s: %s", instance_type, region, e
1461 )
1462 return 0.0
1464 offerings = response.get("CapacityBlockOfferings", [])
1465 if not offerings:
1466 return 0.0
1468 # Bucket offerings into weekly bins (week 0 = this week, week 25 = ~6 months out)
1469 num_weeks = 26
1470 bins = [0] * num_weeks
1471 for o in offerings:
1472 start = o.get("StartDate")
1473 if start is None:
1474 continue
1475 delta_days = (start - now).total_seconds() / 86400.0
1476 week_idx = int(delta_days / 7)
1477 if 0 <= week_idx < num_weeks:
1478 bins[week_idx] += 1
1480 # Need at least 2 non-zero bins to detect a meaningful trend
1481 non_zero = sum(1 for b in bins if b > 0)
1482 if non_zero < 2:
1483 return 0.0
1485 # Linear regression: slope of offerings-per-week over time
1486 # Using least-squares: slope = Σ((x-x̄)(y-ȳ)) / Σ((x-x̄)²)
1487 n = len(bins)
1488 x_mean = (n - 1) / 2.0
1489 y_mean = statistics.mean(bins)
1491 numerator = sum((i - x_mean) * (bins[i] - y_mean) for i in range(n))
1492 denominator = sum((i - x_mean) ** 2 for i in range(n))
1493 slope = numerator / denominator
1495 # Normalize slope to -1..1 range relative to the mean offering count.
1496 # A slope of +y_mean per 26 weeks would be a doubling → maps to ~1.0.
1497 normalized = slope * num_weeks / (y_mean * 2) if y_mean > 0 else 0.0
1499 return round(max(-1.0, min(1.0, normalized)), 4)
1501 def list_all_reservations(
1502 self,
1503 instance_type: str | None = None,
1504 regions: list[str] | None = None,
1505 ) -> dict[str, Any]:
1506 """
1507 List all capacity reservations (ODCRs) across deployed regions.
1509 Args:
1510 instance_type: Filter by instance type (optional)
1511 regions: Regions to query (defaults to deployed GCO regions)
1513 Returns:
1514 Summary dict with reservations grouped by region
1515 """
1516 if not regions:
1517 from cli.aws_client import get_aws_client
1519 aws_client = get_aws_client(self.config)
1520 stacks = aws_client.discover_regional_stacks()
1521 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1523 all_reservations: list[dict[str, Any]] = []
1524 for region in regions:
1525 all_reservations.extend(
1526 self.list_capacity_reservations(region, instance_type=instance_type)
1527 )
1529 total_reserved = sum(r["total_instances"] for r in all_reservations)
1530 total_available = sum(r["available_instances"] for r in all_reservations)
1532 return {
1533 "regions_checked": regions,
1534 "instance_type_filter": instance_type,
1535 "total_reservations": len(all_reservations),
1536 "total_reserved_instances": total_reserved,
1537 "total_available_instances": total_available,
1538 "reservations": all_reservations,
1539 }
1541 def find_capacity_reservations(
1542 self,
1543 instance_type: str | None = None,
1544 regions: list[str] | None = None,
1545 *,
1546 min_count: int = 1,
1547 state: str | None = "active",
1548 include_pricing: bool = True,
1549 max_workers: int | None = None,
1550 ) -> dict[str, Any]:
1551 """Sweep regions for existing ODCRs in one parallel, ranked call.
1553 The ODCR counterpart to :meth:`find_capacity_blocks`. It fans out across
1554 every requested region in parallel, normalizes a friendly instance-type
1555 alias (``p6-b200`` -> ``p6-b200.48xlarge``), enriches each reservation with
1556 On-Demand pricing, and returns a single consolidated report ranked
1557 most-available-first (then cheapest per-GPU-hour). Where
1558 :meth:`list_all_reservations` simply aggregates region by region, this
1559 answers "where do I already have free reserved capacity for this instance
1560 type?" across many regions at once.
1562 Args:
1563 instance_type: Instance type or friendly alias to filter by. Omit to
1564 return every reservation (no type filter).
1565 regions: Regions to search in parallel (any regions, not just
1566 deployed); defaults to the deployed GCO regions when omitted.
1567 min_count: Minimum available instances for the summary to consider the
1568 search satisfied (does not filter the returned list).
1569 state: Reservation state filter ("active" by default; None for all).
1570 include_pricing: Enrich each reservation with On-Demand pricing.
1571 max_workers: Override the parallel fan-out width.
1573 Returns:
1574 A consolidated report dict (see keys assembled below).
1575 """
1576 canonical = instance_type
1577 note: str | None = None
1578 valid = True
1579 known = False
1580 if instance_type:
1581 validation = self.validate_instance_type(instance_type)
1582 canonical = validation["instance_type"]
1583 note = validation["note"]
1584 valid = validation["valid"]
1585 known = validation["known"]
1587 if not regions:
1588 from cli.aws_client import get_aws_client
1590 aws_client = get_aws_client(self.config)
1591 stacks = aws_client.discover_regional_stacks()
1592 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1594 report: dict[str, Any] = {
1595 "instance_type": canonical,
1596 "requested_instance_type": instance_type,
1597 "valid_instance_type": valid,
1598 "known_instance_type": known,
1599 "note": note,
1600 "min_count": min_count,
1601 "state": state,
1602 "regions_checked": list(regions),
1603 "reservations_found": 0,
1604 "reservations": [],
1605 "ranked": [],
1606 "best": None,
1607 "total_reserved_instances": 0,
1608 "total_available_instances": 0,
1609 "regions_with_reservations": [],
1610 }
1612 if instance_type and not valid:
1613 report["recommendation"] = (
1614 note or f"'{instance_type}' is not a recognized EC2 instance type."
1615 )
1616 return report
1618 def _probe(region: str) -> list[dict[str, Any]]:
1619 try:
1620 return self.list_capacity_reservations(
1621 region,
1622 instance_type=canonical,
1623 state=state,
1624 include_pricing=include_pricing,
1625 )
1626 except Exception as e:
1627 logger.warning("Reservation probe failed for %s in %s: %s", canonical, region, e)
1628 return []
1630 collected: list[dict[str, Any]] = []
1631 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(regions)))
1632 if len(regions) == 1:
1633 collected.extend(_probe(regions[0]))
1634 else:
1635 with ThreadPoolExecutor(max_workers=workers) as executor:
1636 for result in executor.map(_probe, regions):
1637 collected.extend(result)
1639 ranked = blocks.rank_reservations(collected)
1640 best = ranked[0] if ranked else None
1641 total_reserved = sum(r.get("total_instances") or 0 for r in collected)
1642 total_available = sum(r.get("available_instances") or 0 for r in collected)
1644 report.update(
1645 {
1646 "reservations_found": len(collected),
1647 "reservations": blocks.sort_reservations(collected),
1648 "ranked": ranked,
1649 "best": best,
1650 "total_reserved_instances": total_reserved,
1651 "total_available_instances": total_available,
1652 "regions_with_reservations": sorted(
1653 {r["region"] for r in collected if r.get("region")}
1654 ),
1655 "recommendation": self._summarize_reservation_search(
1656 canonical, regions, collected, best, min_count, note
1657 ),
1658 }
1659 )
1660 return report
1662 @staticmethod
1663 def _summarize_reservation_search(
1664 instance_type: str | None,
1665 regions: list[str],
1666 reservations: list[dict[str, Any]],
1667 best: dict[str, Any] | None,
1668 min_count: int,
1669 note: str | None,
1670 ) -> str:
1671 """Build a one-line human recommendation for a consolidated ODCR search."""
1672 label = instance_type or "any instance type"
1673 if not reservations:
1674 msg = (
1675 f"No active On-Demand Capacity Reservations for {label} across "
1676 f"{len(regions)} region(s). Create one with "
1677 "'gco capacity create-reservation', or search purchasable Capacity "
1678 "Blocks with 'gco capacity find-blocks'."
1679 )
1680 return f"{note} {msg}" if note else msg
1682 total_available = sum(r.get("available_instances") or 0 for r in reservations)
1683 regions_with = sorted({r["region"] for r in reservations if r.get("region")})
1684 parts = [
1685 f"Found {len(reservations)} reservation(s) for {label} across "
1686 f"{len(regions_with)} region(s); {total_available} instance(s) available."
1687 ]
1688 if total_available < min_count:
1689 parts.append(
1690 f"Fewer than the {min_count} requested are free — consider creating "
1691 "another reservation or a Capacity Block."
1692 )
1693 if best and (best.get("available_instances") or 0) > 0:
1694 gpu_hr = best.get("price_per_gpu_hour")
1695 gpu_hr_str = f", ${gpu_hr}/GPU-hr" if gpu_hr is not None else ""
1696 parts.append(
1697 f"Most available: {best.get('available_instances')}/"
1698 f"{best.get('total_instances')} in {best.get('region')}/"
1699 f"{best.get('availability_zone')} ({best.get('reservation_id')}{gpu_hr_str})."
1700 )
1701 msg = " ".join(parts)
1702 return f"{note} {msg}" if note else msg
1704 @staticmethod
1705 def _summarize_block_search(
1706 instance_type: str,
1707 regions: list[str],
1708 offerings: list[dict[str, Any]],
1709 best: dict[str, Any] | None,
1710 longest: dict[str, Any] | None,
1711 note: str | None,
1712 ) -> str:
1713 """Build a one-line human recommendation for a consolidated block search."""
1714 if not offerings:
1715 msg = (
1716 f"No Capacity Block offerings for {instance_type} across "
1717 f"{len(regions)} region(s) in the requested window. Try a wider "
1718 "date range, a shorter duration, more regions, or check back later."
1719 )
1720 return f"{note} {msg}" if note else msg
1722 regions_with = sorted({o["region"] for o in offerings if o.get("region")})
1723 parts = [
1724 f"Found {len(offerings)} Capacity Block offering(s) for {instance_type} "
1725 f"across {len(regions_with)} region(s)."
1726 ]
1727 if best:
1728 price = best.get("price_per_gpu_hour")
1729 price_str = f", from ${price}/GPU-hr" if price is not None else ""
1730 parts.append(
1731 f"Cheapest: {best.get('region')}/{best.get('availability_zone')} "
1732 f"{best.get('duration_days')}d starting "
1733 f"{(best.get('start_date') or '')[:16]}{price_str}."
1734 )
1735 if longest and longest is not best:
1736 parts.append(
1737 f"Longest: {longest.get('duration_days')}d in "
1738 f"{longest.get('region')}/{longest.get('availability_zone')}."
1739 )
1740 msg = " ".join(parts)
1741 return f"{note} {msg}" if note else msg
1743 def find_capacity_blocks(
1744 self,
1745 instance_type: str,
1746 regions: list[str] | None = None,
1747 *,
1748 instance_count: int = 1,
1749 duration_hours: int | None = None,
1750 duration_days: int | None = None,
1751 min_duration_hours: int | None = None,
1752 min_duration_days: int | None = None,
1753 max_duration_hours: int | None = None,
1754 max_duration_days: int | None = None,
1755 earliest_start: str | datetime | None = None,
1756 latest_start: str | datetime | None = None,
1757 find_longest: bool = False,
1758 max_workers: int | None = None,
1759 ) -> dict[str, Any]:
1760 """Sweep regions x durations x a date window for Capacity Blocks in one call.
1762 This is the high-level search that answers "find 1x p6-b200.48xlarge
1763 across us-east-1/us-east-2/us-west-2/eu-west-1 for durations 1-63 days
1764 starting 2026-07-01..2026-07-10" without manual multi-call sweeping.
1766 Because EC2 requires an exact ``CapacityDurationHours`` per query, a
1767 duration *range* (or ``find_longest``) expands to every valid Capacity
1768 Block duration in the range, and each (region, duration) pair is probed
1769 in parallel. Offerings are de-duplicated across probes, grouped/sorted by
1770 region + AZ + start date, and ranked cheapest-first by per-GPU-hour price.
1772 Args:
1773 instance_type: GPU instance type or friendly alias (e.g. p6-b200).
1774 regions: Explicit regions to search (any regions, not just deployed).
1775 Defaults to the deployed GCO regions when omitted.
1776 instance_count: Instances per block.
1777 duration_hours / duration_days: A single target duration.
1778 min_duration_hours / min_duration_days: Lower bound of a duration range.
1779 max_duration_hours / max_duration_days: Upper bound of a duration range.
1780 earliest_start / latest_start: Date (YYYY-MM-DD) or ISO datetime window
1781 for the block start, threaded to StartDateRange / EndDateRange.
1782 find_longest: Sweep the full duration ladder (within any range given)
1783 and surface the longest available block.
1784 max_workers: Override the parallel fan-out width.
1786 Returns:
1787 A consolidated report dict (see keys assembled below).
1788 """
1789 validation = self.validate_instance_type(instance_type)
1790 canonical = validation["instance_type"]
1791 gpus_per_instance = validation["gpu_count"]
1793 parsed_earliest = blocks.parse_date_input(earliest_start)
1794 parsed_latest = blocks.parse_date_input(latest_start)
1796 durations = blocks.resolve_search_durations(
1797 duration_hours=blocks.coerce_hours(duration_hours, duration_days),
1798 min_duration_hours=blocks.coerce_hours(min_duration_hours, min_duration_days),
1799 max_duration_hours=blocks.coerce_hours(max_duration_hours, max_duration_days),
1800 find_longest=find_longest,
1801 )
1803 if not regions:
1804 from cli.aws_client import get_aws_client
1806 aws_client = get_aws_client(self.config)
1807 stacks = aws_client.discover_regional_stacks()
1808 regions = list(stacks.keys()) if stacks else [self.config.default_region]
1810 report: dict[str, Any] = {
1811 "instance_type": canonical,
1812 "requested_instance_type": instance_type,
1813 "valid_instance_type": validation["valid"],
1814 "known_instance_type": validation["known"],
1815 "note": validation["note"],
1816 "instance_count": instance_count,
1817 "regions_checked": list(regions),
1818 "durations_probed_hours": durations,
1819 "durations_probed_days": [blocks.hours_to_days(h) for h in durations],
1820 "date_window": {
1821 "earliest_start": parsed_earliest.isoformat() if parsed_earliest else None,
1822 "latest_start": parsed_latest.isoformat() if parsed_latest else None,
1823 },
1824 "offerings_found": 0,
1825 "offerings": [],
1826 "ranked": [],
1827 "best": None,
1828 "longest": None,
1829 "regions_with_offerings": [],
1830 }
1832 if not validation["valid"]:
1833 report["recommendation"] = (
1834 validation["note"]
1835 or f"'{instance_type}' is not a valid standalone EC2 instance type "
1836 "for Capacity Blocks."
1837 )
1838 return report
1840 probes = [(region, dur) for region in regions for dur in durations]
1841 collected: list[dict[str, Any]] = []
1842 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(probes)))
1844 def _probe(region: str, dur: int) -> list[dict[str, Any]]:
1845 try:
1846 return self.list_capacity_block_offerings(
1847 region,
1848 canonical,
1849 instance_count=instance_count,
1850 duration_hours=dur,
1851 earliest_start=parsed_earliest,
1852 latest_start=parsed_latest,
1853 gpus_per_instance=gpus_per_instance,
1854 )
1855 except Exception as e:
1856 logger.warning(
1857 "Capacity block probe failed for %s in %s (%sh): %s",
1858 canonical,
1859 region,
1860 dur,
1861 e,
1862 )
1863 return []
1865 with ThreadPoolExecutor(max_workers=workers) as executor:
1866 futures = [executor.submit(_probe, region, dur) for region, dur in probes]
1867 for future in as_completed(futures):
1868 collected.extend(future.result())
1870 unique = blocks.dedupe_offerings(collected)
1871 ranked = blocks.rank_offerings(unique)
1872 best = ranked[0] if ranked else None
1873 longest = blocks.longest_offering(unique)
1875 report.update(
1876 {
1877 "offerings_found": len(unique),
1878 "offerings": blocks.sort_offerings(unique),
1879 "ranked": ranked,
1880 "best": best,
1881 "longest": longest,
1882 "regions_with_offerings": sorted({o["region"] for o in unique if o.get("region")}),
1883 "recommendation": self._summarize_block_search(
1884 canonical, regions, unique, best, longest, validation["note"]
1885 ),
1886 }
1887 )
1888 return report
1890 def check_reservation_availability(
1891 self,
1892 instance_type: str,
1893 min_count: int = 1,
1894 include_capacity_blocks: bool = True,
1895 block_duration_hours: int = 24,
1896 *,
1897 regions: list[str] | None = None,
1898 block_duration_days: int | None = None,
1899 earliest_start: str | datetime | None = None,
1900 latest_start: str | datetime | None = None,
1901 max_workers: int | None = None,
1902 ) -> dict[str, Any]:
1903 """
1904 Check if capacity reservations or blocks have available instances.
1906 Checks both ODCRs (existing reservations) and Capacity Block offerings
1907 (purchasable guaranteed capacity) for a given instance type, across one
1908 or many regions queried in parallel.
1910 Args:
1911 instance_type: EC2 instance type to check
1912 min_count: Minimum number of available instances needed
1913 include_capacity_blocks: Also check Capacity Block offerings
1914 block_duration_hours: Duration for capacity block search (hours)
1915 regions: Regions to check in parallel (any regions, not just
1916 deployed); falls back to the deployed regions when omitted.
1917 block_duration_days: Block duration in days (overrides hours when set).
1918 earliest_start / latest_start: Date window for block start
1919 (StartDateRange / EndDateRange).
1920 max_workers: Override the parallel fan-out width.
1922 Returns:
1923 Dictionary with ODCR availability and capacity block offerings
1924 """
1925 if regions:
1926 target_regions = list(regions)
1927 else:
1928 from cli.aws_client import get_aws_client
1930 aws_client = get_aws_client(self.config)
1931 stacks = aws_client.discover_regional_stacks()
1932 target_regions = list(stacks.keys()) if stacks else [self.config.default_region]
1934 effective_duration = blocks.snap_duration_hours(
1935 blocks.coerce_hours(block_duration_hours, block_duration_days) or 24
1936 )
1937 parsed_earliest = blocks.parse_date_input(earliest_start)
1938 parsed_latest = blocks.parse_date_input(latest_start)
1940 def _check_region(r: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
1941 reservations = self.list_capacity_reservations(r, instance_type=instance_type)
1942 region_blocks: list[dict[str, Any]] = []
1943 if include_capacity_blocks:
1944 region_blocks = self.list_capacity_block_offerings(
1945 r,
1946 instance_type=instance_type,
1947 instance_count=min_count,
1948 duration_hours=effective_duration,
1949 earliest_start=parsed_earliest,
1950 latest_start=parsed_latest,
1951 )
1952 return reservations, region_blocks
1954 odcr_results: list[dict[str, Any]] = []
1955 block_offerings: list[dict[str, Any]] = []
1956 total_available = 0
1957 total_reserved = 0
1959 workers = max(1, min(max_workers or _MAX_SEARCH_WORKERS, len(target_regions)))
1960 if len(target_regions) == 1:
1961 region_results = [_check_region(target_regions[0])]
1962 else:
1963 with ThreadPoolExecutor(max_workers=workers) as executor:
1964 region_results = list(executor.map(_check_region, target_regions))
1966 for reservations, region_blocks in region_results:
1967 for res in reservations:
1968 avail = res["available_instances"]
1969 total_available += avail
1970 total_reserved += res["total_instances"]
1971 if avail > 0:
1972 odcr_results.append(res)
1973 block_offerings.extend(region_blocks)
1975 block_offerings = blocks.sort_offerings(block_offerings)
1976 has_odcr = total_available >= min_count
1977 has_blocks = len(block_offerings) > 0
1979 # Build recommendation
1980 if has_odcr:
1981 recommendation = (
1982 f"ODCR capacity available: {total_available} instances "
1983 f"across {len(odcr_results)} reservation(s)"
1984 )
1985 elif has_blocks:
1986 cheapest = min(block_offerings, key=_offering_fee)
1987 fee_display = cheapest.get("upfront_fee_usd")
1988 if fee_display is None:
1989 fee_display = cheapest.get("upfront_fee", "?")
1990 recommendation = (
1991 f"No ODCR capacity, but {len(block_offerings)} Capacity Block offering(s) "
1992 f"available (from ${fee_display} for {effective_duration}h)"
1993 )
1994 else:
1995 recommendation = (
1996 "No reserved capacity or block offerings found. "
1997 "Consider on-demand or spot, or request a Capacity Block "
1998 "for a different duration/region."
1999 )
2001 return {
2002 "instance_type": instance_type,
2003 "min_count_requested": min_count,
2004 "regions_checked": target_regions,
2005 "odcr": {
2006 "total_reserved_instances": total_reserved,
2007 "total_available_instances": total_available,
2008 "has_availability": has_odcr,
2009 "reservations": odcr_results,
2010 },
2011 "capacity_blocks": {
2012 "offerings_found": len(block_offerings),
2013 "has_offerings": has_blocks,
2014 "duration_hours": effective_duration,
2015 "date_window": {
2016 "earliest_start": parsed_earliest.isoformat() if parsed_earliest else None,
2017 "latest_start": parsed_latest.isoformat() if parsed_latest else None,
2018 },
2019 "offerings": block_offerings,
2020 },
2021 "recommendation": recommendation,
2022 }
2024 def purchase_capacity_block(
2025 self,
2026 offering_id: str,
2027 region: str,
2028 dry_run: bool = False,
2029 ) -> dict[str, Any]:
2030 """
2031 Purchase a Capacity Block offering by its ID.
2033 Args:
2034 offering_id: Capacity Block offering ID (cb-xxx) from list_capacity_block_offerings
2035 region: AWS region where the offering exists
2036 dry_run: If True, validate the offering without purchasing
2038 Returns:
2039 Dictionary with the created capacity reservation details
2040 """
2041 ec2 = self._session.client("ec2", region_name=region)
2043 if dry_run:
2044 try:
2045 # AWS signals an authorized DryRun with DryRunOperation. A
2046 # normal return violates that protocol and must never fall
2047 # through to the real purchase call.
2048 ec2.purchase_capacity_block(
2049 CapacityBlockOfferingId=offering_id,
2050 InstancePlatform="Linux/UNIX",
2051 DryRun=True,
2052 )
2053 except ClientError as e:
2054 error_code = e.response.get("Error", {}).get("Code", "")
2055 if error_code == "DryRunOperation":
2056 return {
2057 "success": True,
2058 "dry_run": True,
2059 "offering_id": offering_id,
2060 "region": region,
2061 "message": "Dry run succeeded — offering is valid and purchasable",
2062 }
2063 error_msg = e.response.get("Error", {}).get("Message", str(e))
2064 return {
2065 "success": False,
2066 "dry_run": True,
2067 "offering_id": offering_id,
2068 "region": region,
2069 "error_code": error_code,
2070 "error": error_msg,
2071 }
2072 return {
2073 "success": False,
2074 "dry_run": True,
2075 "offering_id": offering_id,
2076 "region": region,
2077 "error_code": "DryRunProtocolError",
2078 "error": "AWS returned normally for a DryRun request; no purchase was attempted.",
2079 }
2081 try:
2082 response = ec2.purchase_capacity_block(
2083 CapacityBlockOfferingId=offering_id,
2084 InstancePlatform="Linux/UNIX",
2085 )
2087 reservation = response.get("CapacityReservation", {})
2088 reservation_id = reservation.get("CapacityReservationId", "")
2089 instance_type = reservation.get("InstanceType", "")
2090 az = reservation.get("AvailabilityZone", "")
2091 total = reservation.get("TotalInstanceCount", 0)
2092 start = reservation.get("StartDate")
2093 end = reservation.get("EndDate")
2095 return {
2096 "success": True,
2097 "dry_run": False,
2098 "reservation_id": reservation_id,
2099 "offering_id": offering_id,
2100 "instance_type": instance_type,
2101 "availability_zone": az,
2102 "region": region,
2103 "total_instances": total,
2104 "start_date": start.isoformat() if start else None,
2105 "end_date": end.isoformat() if end else None,
2106 "state": reservation.get("State", ""),
2107 }
2108 except ClientError as e:
2109 error_code = e.response.get("Error", {}).get("Code", "")
2110 error_msg = e.response.get("Error", {}).get("Message", str(e))
2111 return {
2112 "success": False,
2113 "dry_run": False,
2114 "offering_id": offering_id,
2115 "region": region,
2116 "error_code": error_code,
2117 "error": error_msg,
2118 }
2120 def create_capacity_reservation(
2121 self,
2122 instance_type: str,
2123 region: str,
2124 availability_zone: str,
2125 instance_count: int = 1,
2126 *,
2127 instance_platform: str = "Linux/UNIX",
2128 tenancy: str = "default",
2129 instance_match_criteria: str = "open",
2130 end_date: str | datetime | None = None,
2131 ebs_optimized: bool = False,
2132 dry_run: bool = False,
2133 ) -> dict[str, Any]:
2134 """Create a new On-Demand Capacity Reservation (ODCR).
2136 The ODCR counterpart to :meth:`purchase_capacity_block`: it reserves
2137 On-Demand capacity for an instance type in a specific Availability Zone.
2138 Unlike a Capacity Block (a fixed-term block bought by offering id), an ODCR
2139 is open-ended by default and billed at the On-Demand rate until cancelled.
2140 Friendly instance-type aliases are normalized (``p6-b200`` ->
2141 ``p6-b200.48xlarge``).
2143 Args:
2144 instance_type: EC2 instance type or friendly alias.
2145 region: AWS region.
2146 availability_zone: Target AZ (e.g. us-east-1a).
2147 instance_count: Number of instances to reserve.
2148 instance_platform: Platform/OS (default "Linux/UNIX").
2149 tenancy: "default" or "dedicated".
2150 instance_match_criteria: "open" (any matching instance) or "targeted".
2151 end_date: Optional end date (ISO string or datetime). When set the
2152 reservation is "limited" and auto-releases then; omit for
2153 "unlimited".
2154 ebs_optimized: Reserve EBS-optimized capacity.
2155 dry_run: Validate permissions/parameters without creating (no cost).
2157 Returns:
2158 Dict with the created reservation's details, or an error payload.
2159 """
2160 validation = self.validate_instance_type(instance_type)
2161 canonical = validation["instance_type"]
2162 if not validation["valid"]:
2163 return {
2164 "success": False,
2165 "dry_run": dry_run,
2166 "instance_type": canonical,
2167 "region": region,
2168 "error_code": "InvalidInstanceType",
2169 "error": (
2170 validation["note"] or f"'{instance_type}' is not a valid EC2 instance type."
2171 ),
2172 }
2174 ec2 = self._session.client("ec2", region_name=region)
2175 parsed_end = blocks.parse_date_input(end_date)
2176 api_kwargs: dict[str, Any] = {
2177 "InstanceType": canonical,
2178 "InstancePlatform": instance_platform,
2179 "AvailabilityZone": availability_zone,
2180 "InstanceCount": instance_count,
2181 "Tenancy": tenancy,
2182 "InstanceMatchCriteria": instance_match_criteria,
2183 "EbsOptimized": ebs_optimized,
2184 }
2185 if parsed_end is not None:
2186 api_kwargs["EndDate"] = parsed_end
2187 api_kwargs["EndDateType"] = "limited"
2188 else:
2189 api_kwargs["EndDateType"] = "unlimited"
2191 if dry_run:
2192 try:
2193 ec2.create_capacity_reservation(DryRun=True, **api_kwargs)
2194 except ClientError as e:
2195 error_code = e.response.get("Error", {}).get("Code", "")
2196 if error_code == "DryRunOperation":
2197 return {
2198 "success": True,
2199 "dry_run": True,
2200 "instance_type": canonical,
2201 "availability_zone": availability_zone,
2202 "region": region,
2203 "instance_count": instance_count,
2204 "message": "Dry run succeeded — reservation parameters are valid.",
2205 }
2206 error_msg = e.response.get("Error", {}).get("Message", str(e))
2207 return {
2208 "success": False,
2209 "dry_run": True,
2210 "instance_type": canonical,
2211 "region": region,
2212 "error_code": error_code,
2213 "error": error_msg,
2214 }
2215 return {
2216 "success": False,
2217 "dry_run": True,
2218 "instance_type": canonical,
2219 "availability_zone": availability_zone,
2220 "region": region,
2221 "instance_count": instance_count,
2222 "error_code": "DryRunProtocolError",
2223 "error": "AWS returned normally for a DryRun request; no reservation was created.",
2224 }
2226 try:
2227 response = ec2.create_capacity_reservation(**api_kwargs)
2228 cr = response.get("CapacityReservation", {})
2229 return {
2230 "success": True,
2231 "dry_run": False,
2232 "reservation_id": cr.get("CapacityReservationId", ""),
2233 "instance_type": cr.get("InstanceType", canonical),
2234 "availability_zone": cr.get("AvailabilityZone", availability_zone),
2235 "region": region,
2236 "total_instances": cr.get("TotalInstanceCount", instance_count),
2237 "available_instances": cr.get("AvailableInstanceCount"),
2238 "state": cr.get("State", ""),
2239 "tenancy": cr.get("Tenancy", tenancy),
2240 "instance_match_criteria": cr.get("InstanceMatchCriteria", instance_match_criteria),
2241 "start_date": cr["StartDate"].isoformat() if cr.get("StartDate") else None,
2242 "end_date": cr["EndDate"].isoformat() if cr.get("EndDate") else None,
2243 "end_date_type": cr.get("EndDateType"),
2244 }
2245 except ClientError as e:
2246 error_code = e.response.get("Error", {}).get("Code", "")
2247 error_msg = e.response.get("Error", {}).get("Message", str(e))
2248 return {
2249 "success": False,
2250 "dry_run": False,
2251 "instance_type": canonical,
2252 "region": region,
2253 "error_code": error_code,
2254 "error": error_msg,
2255 }
2257 def cancel_capacity_reservation(
2258 self,
2259 reservation_id: str,
2260 region: str,
2261 dry_run: bool = False,
2262 ) -> dict[str, Any]:
2263 """Cancel an On-Demand Capacity Reservation by id.
2265 Releases reserved capacity so it stops incurring On-Demand charges. Only
2266 ODCRs can be cancelled this way; a Capacity Block runs for its fixed term.
2267 Cancelling does not terminate instances already running against the
2268 reservation — they simply revert to normal On-Demand billing.
2270 Args:
2271 reservation_id: Capacity Reservation id (cr-...).
2272 region: AWS region where the reservation exists.
2273 dry_run: Validate permissions without cancelling.
2275 Returns:
2276 Dict describing the outcome.
2277 """
2278 ec2 = self._session.client("ec2", region_name=region)
2280 if dry_run:
2281 try:
2282 ec2.cancel_capacity_reservation(CapacityReservationId=reservation_id, DryRun=True)
2283 except ClientError as e:
2284 error_code = e.response.get("Error", {}).get("Code", "")
2285 if error_code == "DryRunOperation":
2286 return {
2287 "success": True,
2288 "dry_run": True,
2289 "reservation_id": reservation_id,
2290 "region": region,
2291 "message": "Dry run succeeded — reservation can be cancelled.",
2292 }
2293 error_msg = e.response.get("Error", {}).get("Message", str(e))
2294 return {
2295 "success": False,
2296 "dry_run": True,
2297 "reservation_id": reservation_id,
2298 "region": region,
2299 "error_code": error_code,
2300 "error": error_msg,
2301 }
2302 return {
2303 "success": False,
2304 "dry_run": True,
2305 "reservation_id": reservation_id,
2306 "region": region,
2307 "error_code": "DryRunProtocolError",
2308 "error": "AWS returned normally for a DryRun request; no reservation was cancelled.",
2309 }
2311 try:
2312 response = ec2.cancel_capacity_reservation(CapacityReservationId=reservation_id)
2313 return {
2314 "success": bool(response.get("Return", True)),
2315 "dry_run": False,
2316 "reservation_id": reservation_id,
2317 "region": region,
2318 "message": f"Capacity reservation {reservation_id} cancelled.",
2319 }
2320 except ClientError as e:
2321 error_code = e.response.get("Error", {}).get("Code", "")
2322 error_msg = e.response.get("Error", {}).get("Message", str(e))
2323 return {
2324 "success": False,
2325 "dry_run": False,
2326 "reservation_id": reservation_id,
2327 "region": region,
2328 "error_code": error_code,
2329 "error": error_msg,
2330 }
2332 def recommend_region_for_job(
2333 self,
2334 gpu_required: bool = False,
2335 min_gpus: int = 0,
2336 instance_type: str | None = None,
2337 gpu_count: int = 0,
2338 ) -> dict[str, Any]:
2339 """
2340 Recommend the optimal region for job placement.
2342 Delegates to MultiRegionCapacityChecker for cross-region analysis.
2344 Args:
2345 gpu_required: Whether the job requires GPUs
2346 min_gpus: Minimum number of GPUs required
2347 instance_type: Specific instance type for workload-aware scoring
2348 gpu_count: Number of GPUs required
2350 Returns:
2351 Dictionary with recommended region and justification
2352 """
2353 from .multi_region import MultiRegionCapacityChecker
2355 checker = MultiRegionCapacityChecker(self.config)
2356 return checker.recommend_region_for_job(
2357 gpu_required, min_gpus, instance_type=instance_type, gpu_count=gpu_count
2358 )
2361def get_capacity_checker(config: GCOConfig | None = None) -> CapacityChecker:
2362 """Get a configured capacity checker instance."""
2363 return CapacityChecker(config)