Coverage for cli/capacity/blocks.py: 96.92%

1"""

2Capacity Block search primitives: duration math, instance-type normalization,

3offering pricing, and de-dup / sort helpers.

5EC2 Capacity Blocks for ML (the ``DescribeCapacityBlockOfferings`` API) accept

6reservation durations in **1-day increments up to 14 days, then 7-day increments

7up to 182 days** (26 weeks). ``CapacityDurationHours`` is a *required* parameter,

8so searching a duration *range* — or asking "what's the longest block available?"

9— means probing several discrete duration values and merging the results. These

10helpers centralize that math, plus the offering enrichment (per-hour and

11per-GPU-hour pricing) and the de-dup / sort the search layer applies across the

12region x duration probe matrix.

14Everything here is pure (no boto3, no I/O) so it is cheap to unit-test in

15isolation. The single-region API call, the parallel fan-out, and the instance

16validation that needs AWS live in ``checker.py``.

17"""

19from __future__ import annotations

21from datetime import UTC, datetime

22from typing import Any

24# ---------------------------------------------------------------------------

25# Duration math

26# ---------------------------------------------------------------------------

28#: Default Capacity Block search duration (24h = 1 day), the smallest valid block.

29DEFAULT_DURATION_HOURS = 24

31#: AWS caps a single Capacity Block at 182 days (26 weeks).

32MAX_DURATION_DAYS = 182

34#: Durations are expressed in 1-day increments up to this many days...

35_DAILY_INCREMENT_MAX_DAYS = 14

37#: ...then in 7-day increments beyond that, up to ``MAX_DURATION_DAYS``.

38_WEEKLY_INCREMENT_DAYS = 7

40HOURS_PER_DAY = 24

43def capacity_block_duration_days() -> tuple[int, ...]:

44 """Return every Capacity Block duration AWS allows, in days, ascending.

46 1..14 in 1-day steps, then 21, 28, ... 182 in 7-day steps.

47 """

48 daily = list(range(1, _DAILY_INCREMENT_MAX_DAYS + 1))

49 weekly = list(

50 range(

51 _DAILY_INCREMENT_MAX_DAYS + _WEEKLY_INCREMENT_DAYS,

52 MAX_DURATION_DAYS + 1,

53 _WEEKLY_INCREMENT_DAYS,

54 )

55 )

56 return tuple(daily + weekly)

59def capacity_block_duration_hours() -> tuple[int, ...]:

60 """Return every allowed Capacity Block duration, in hours, ascending."""

61 return tuple(d * HOURS_PER_DAY for d in capacity_block_duration_days())

64def is_valid_duration_hours(hours: int) -> bool:

65 """True if ``hours`` is exactly one of the AWS-allowed Capacity Block durations."""

66 return hours in capacity_block_duration_hours()

69def snap_duration_hours(hours: int) -> int:

70 """Snap an arbitrary hour count to the nearest valid Capacity Block duration.

72 Ties (equidistant between two valid durations) round up to the longer one so

73 callers never under-request. Values below/above the allowed range clamp to

74 the smallest/largest valid duration.

75 """

76 options = capacity_block_duration_hours()

77 if hours <= options[0]:

78 return options[0]

79 if hours >= options[-1]:

80 return options[-1]

81 # Nearest, ties → longer.

82 return min(options, key=lambda opt: (abs(opt - hours), -opt))

85def coerce_hours(hours: int | None, days: int | None) -> int | None:

86 """Resolve a duration to hours from an hours-or-days pair.

88 ``days`` wins when both are supplied so a caller that passes ``duration_days``

89 always gets day-granularity. Returns ``None`` when neither is set.

90 """

91 if days is not None:

92 return days * HOURS_PER_DAY

93 return hours

96def resolve_search_durations(

97 *,

98 duration_hours: int | None = None,

99 min_duration_hours: int | None = None,

100 max_duration_hours: int | None = None,

101 find_longest: bool = False,

102) -> list[int]:

103 """Resolve the set of Capacity Block durations (in hours) to probe.

104

105 Precedence:

106

107 * A min/max range (either bound) — or ``find_longest`` — expands to *every*

108 valid duration within the (possibly open-ended) range. ``find_longest``

109 without a range sweeps the full 1-day..182-day ladder.

110 * A single ``duration_hours`` snaps to the nearest valid duration.

111 * Otherwise the default 24h block.

112

113 The result is always non-empty and ascending so the caller can probe each

114 value and let the search layer pick the longest / cheapest.

115 """

116 options = capacity_block_duration_hours()

117 has_range = min_duration_hours is not None or max_duration_hours is not None

118

119 if has_range or find_longest:

120 low = min_duration_hours if min_duration_hours is not None else options[0]

121 high = max_duration_hours if max_duration_hours is not None else options[-1]

122 if low > high:

123 low, high = high, low

124 chosen = [h for h in options if low <= h <= high]

125 # A tight range that straddles no valid duration (e.g. 30h..40h) still

126 # gets a single best-effort probe rather than an empty search.

127 return chosen or [snap_duration_hours((low + high) // 2)]

128

129 if duration_hours is not None:

130 return [snap_duration_hours(duration_hours)]

131

132 return [DEFAULT_DURATION_HOURS]

133

134

135def hours_to_days(hours: int | None) -> float | None:

136 """Convert a duration in hours to days (float), or ``None`` passthrough."""

137 if hours is None:

138 return None

139 return round(hours / HOURS_PER_DAY, 2)

140

141

142# ---------------------------------------------------------------------------

143# Date parsing

144# ---------------------------------------------------------------------------

145

146

147def parse_date_input(value: str | datetime | None) -> datetime | None:

148 """Parse a date/datetime input into a timezone-aware UTC ``datetime``.

149

150 Accepts ``None`` (passthrough), a ``datetime`` (naive values are assumed

151 UTC), an ISO-8601 string (``2026-07-01`` or ``2026-07-01T11:30:00Z``), and

152 tolerates a trailing ``Z``. Raises ``ValueError`` with a friendly message on

153 anything else so the CLI can surface a clear error.

154 """

155 if value is None:

156 return None

157 if isinstance(value, datetime):

158 return value if value.tzinfo else value.replace(tzinfo=UTC)

159 text = str(value).strip()

160 if not text:

161 return None

162 try:

163 parsed = datetime.fromisoformat(text.replace("Z", "+00:00"))

164 except ValueError as exc:

165 raise ValueError(

166 f"Invalid date {value!r}; use YYYY-MM-DD or an ISO-8601 timestamp "

167 "(e.g. 2026-07-01 or 2026-07-01T11:30:00Z)."

168 ) from exc

169 return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC)

170

171

172# ---------------------------------------------------------------------------

173# Instance-type normalization

174# ---------------------------------------------------------------------------

175

176#: Friendly / shorthand names mapped to their canonical EC2 instance type. Lets a

177#: caller say ``p6-b200`` and get the full ``p6-b200.48xlarge``.

178INSTANCE_TYPE_ALIASES: dict[str, str] = {

179 "p6-b200": "p6-b200.48xlarge",

180 "p6b200": "p6-b200.48xlarge",

181 "b200": "p6-b200.48xlarge",

182 "p6-b300": "p6-b300.48xlarge",

183 "p6b300": "p6-b300.48xlarge",

184 "b300": "p6-b300.48xlarge",

185 "p5": "p5.48xlarge",

186 "p5e": "p5e.48xlarge",

187 "p5en": "p5en.48xlarge",

188 "p4d": "p4d.24xlarge",

189 "p4de": "p4de.24xlarge",

190 "p3dn": "p3dn.24xlarge",

191}

192

193#: Accelerator families that are NOT sold as standalone EC2 instance types — the

194#: Grace-Blackwell *superchips* (GB200/GB300) ship only as P6e-GB UltraServers, so

195#: ``DescribeCapacityBlockOfferings(InstanceType=…)`` never returns them. The note

196#: steers callers to the UltraServer search flow and to the standalone B200/B300

197#: EC2 types where one exists. NB: the discrete ``b200``/``b300`` GPUs *are* sold

198#: as standalone ``p6-b200.48xlarge`` / ``p6-b300.48xlarge`` instances (see

199#: ``INSTANCE_TYPE_ALIASES``) — only the GB NVL72 superchips are UltraServer-only.

200NON_STANDALONE_INSTANCE_NOTES: dict[str, str] = {

201 "p6e-gb300": (

202 "P6e-GB300 is an UltraServer family, not a standalone EC2 instance type. "

203 "Search Capacity Blocks with UltraserverType=u-p6e-gb300x... instead of "

204 "InstanceType. For a standalone Blackwell Ultra EC2 type use p6-b300.48xlarge."

205 ),

206 "gb300": (

207 "GB300 (Grace Blackwell Ultra NVL72) ships as P6e-GB300 UltraServers, not "

208 "a standalone EC2 instance type. For standalone Blackwell Ultra EC2 "

209 "capacity use p6-b300.48xlarge instead."

210 ),

211 "p6e-gb200": (

212 "P6e-GB200 is an UltraServer family, not a standalone EC2 instance type. "

213 "Search Capacity Blocks with UltraserverType=u-p6e-gb200x... instead of "

214 "InstanceType. For a standalone Blackwell EC2 type use p6-b200.48xlarge."

215 ),

216 "gb200": (

217 "GB200 (Grace Blackwell NVL72) ships as P6e-GB200 UltraServers. For "

218 "standalone Blackwell EC2 capacity use p6-b200.48xlarge instead."

219 ),

220}

221

222

223def normalize_instance_type(name: str) -> tuple[str, str | None]:

224 """Normalize an instance-type string to its canonical form.

225

226 Returns ``(canonical_type, note)`` where ``note`` is a human-readable string

227 when the input was expanded from an alias or names a not-standalone

228 UltraServer family (otherwise ``None``). The canonical type is returned

229 unchanged for an already-canonical or unknown input so the caller can still

230 let the EC2 API be the source of truth on validity.

231 """

232 raw = (name or "").strip()

233 lowered = raw.lower()

234

235 note = NON_STANDALONE_INSTANCE_NOTES.get(lowered)

236 if note:

237 return raw, note

238

239 canonical = INSTANCE_TYPE_ALIASES.get(lowered)

240 if canonical and canonical != raw:

241 return canonical, f"Interpreted '{raw}' as '{canonical}'."

242

243 return raw, None

244

245

246# ---------------------------------------------------------------------------

247# Offering pricing

248# ---------------------------------------------------------------------------

249

250

251def parse_upfront_fee(value: Any) -> float | None:

252 """Parse an upfront-fee value (the API returns it as a string) to a float.

253

254 Returns ``None`` for missing/unparseable values rather than raising, so a

255 malformed fee never breaks the whole search.

256 """

257 if value is None:

258 return None

259 if isinstance(value, bool): # guard: bool is a subclass of int

260 return None

261 if isinstance(value, int | float):

262 return float(value)

263 try:

264 return float(str(value).strip())

265 except TypeError, ValueError:

266 return None

267

268

269def compute_offering_pricing(

270 upfront_fee: Any,

271 duration_hours: int | None,

272 instance_count: int | None,

273 gpus_per_instance: int | None,

274) -> dict[str, float | None]:

275 """Derive per-hour and per-GPU-hour pricing from a block's upfront fee.

276

277 ``upfront_fee`` is the total one-time charge (USD) for the whole block — all

278 instances for the whole duration. The derived figures:

279

280 * ``upfront_fee_usd`` — the parsed total upfront fee.

281 * ``price_per_hour`` — whole-block amortized cost per hour.

282 * ``price_per_instance_hour`` — per-instance amortized cost per hour.

283 * ``price_per_gpu_hour`` — per-GPU amortized cost per hour (needs GPU count).

284

285 Any figure that can't be computed (missing fee, zero duration/count, unknown

286 GPU count) is ``None`` rather than a misleading zero.

287 """

288 fee = parse_upfront_fee(upfront_fee)

289 result: dict[str, float | None] = {

290 "upfront_fee_usd": fee,

291 "price_per_hour": None,

292 "price_per_instance_hour": None,

293 "price_per_gpu_hour": None,

294 }

295 if fee is None or not duration_hours or duration_hours <= 0:

296 return result

297

298 result["price_per_hour"] = round(fee / duration_hours, 4)

299

300 if instance_count and instance_count > 0: 300 ↛ 306line 300 didn't jump to line 306 because the condition on line 300 was always true

301 per_instance_hour = fee / (duration_hours * instance_count)

302 result["price_per_instance_hour"] = round(per_instance_hour, 4)

303 if gpus_per_instance and gpus_per_instance > 0:

304 result["price_per_gpu_hour"] = round(per_instance_hour / gpus_per_instance, 4)

305

306 return result

307

308

309def compute_reservation_pricing(

310 on_demand_price_per_hour: float | None,

311 instance_count: int | None,

312 gpus_per_instance: int | None,

313) -> dict[str, float | None]:

314 """Derive per-hour / per-instance-hour / per-GPU-hour pricing for an ODCR.

315

316 On-Demand Capacity Reservations bill at the standard On-Demand rate for the

317 reserved instance type for as long as the reservation is held, whether or not

318 the capacity is actually in use. So the reservation's running cost is the

319 On-Demand hourly price times the number of instances reserved. The derived

320 figures mirror :func:`compute_offering_pricing` so an ODCR and a Capacity

321 Block can be compared on the same per-instance-hour / per-GPU-hour footing:

322

323 * ``price_per_instance_hour`` — the On-Demand hourly rate for one instance.

324 * ``price_per_hour`` — whole-reservation hourly cost (rate x instance count).

325 * ``price_per_gpu_hour`` — per-GPU hourly rate (needs a GPU count).

326

327 Any figure that can't be computed (missing rate, zero/unknown count) is

328 ``None`` rather than a misleading zero.

329 """

330 result: dict[str, float | None] = {

331 "price_per_instance_hour": None,

332 "price_per_hour": None,

333 "price_per_gpu_hour": None,

334 }

335 if on_demand_price_per_hour is None:

336 return result

337 try:

338 rate = float(on_demand_price_per_hour)

339 except TypeError, ValueError:

340 return result

341 if rate < 0: 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true

342 return result

343

344 result["price_per_instance_hour"] = round(rate, 4)

345

346 if instance_count and instance_count > 0: 346 ↛ 349line 346 didn't jump to line 349 because the condition on line 346 was always true

347 result["price_per_hour"] = round(rate * instance_count, 4)

348

349 if gpus_per_instance and gpus_per_instance > 0:

350 result["price_per_gpu_hour"] = round(rate / gpus_per_instance, 4)

351

352 return result

353

354

355# ---------------------------------------------------------------------------

356# De-dup / sort

357# ---------------------------------------------------------------------------

358

359

360def offering_identity(offering: dict[str, Any]) -> tuple[Any, ...]:

361 """Stable identity for an offering used to de-dup across duration probes.

362

363 Prefers the AWS offering id; falls back to the (region, AZ, start, duration)

364 tuple when an id is absent (e.g. partially-mocked test data).

365 """

366 offering_id = offering.get("offering_id")

367 if offering_id:

368 return ("id", offering_id)

369 return (

370 "tuple",

371 offering.get("region"),

372 offering.get("availability_zone"),

373 offering.get("start_date"),

374 offering.get("duration_hours"),

375 )

376

377

378def dedupe_offerings(offerings: list[dict[str, Any]]) -> list[dict[str, Any]]:

379 """Drop duplicate offerings (same identity) while preserving first-seen order.

380

381 Probing adjacent durations against the same date window routinely returns the

382 same physical block more than once; this collapses those to a single entry.

383 """

384 seen: set[tuple[Any, ...]] = set()

385 unique: list[dict[str, Any]] = []

386 for offering in offerings:

387 identity = offering_identity(offering)

388 if identity in seen:

389 continue

390 seen.add(identity)

391 unique.append(offering)

392 return unique

393

394

395def sort_offerings(offerings: list[dict[str, Any]]) -> list[dict[str, Any]]:

396 """Sort offerings by region, then AZ, then start date — a stable browse order."""

397 return sorted(

398 offerings,

399 key=lambda o: (

400 str(o.get("region") or ""),

401 str(o.get("availability_zone") or ""),

402 str(o.get("start_date") or ""),

403 ),

404 )

405

406

407def _rank_key(offering: dict[str, Any]) -> tuple[float, float, str]:

408 gpu_hour = offering.get("price_per_gpu_hour")

409 per_hour = offering.get("price_per_hour")

410 return (

411 float(gpu_hour) if gpu_hour is not None else float("inf"),

412 float(per_hour) if per_hour is not None else float("inf"),

413 str(offering.get("start_date") or ""),

414 )

415

416

417def rank_offerings(offerings: list[dict[str, Any]]) -> list[dict[str, Any]]:

418 """Rank offerings cheapest-first by per-GPU-hour, then per-hour, then start date.

419

420 Offerings missing a price sort last. Used to surface the single best block in

421 a consolidated multi-region report.

422 """

423 return sorted(offerings, key=_rank_key)

424

425

426def longest_offering(offerings: list[dict[str, Any]]) -> dict[str, Any] | None:

427 """Return the offering with the greatest actual duration, or ``None`` if empty.

428

429 Ties break toward the cheaper per-GPU-hour block so "find longest" still

430 favours value when two blocks share the top duration.

431 """

432 if not offerings:

433 return None

434 return max(

435 offerings,

436 key=lambda o: (

437 int(o.get("duration_hours") or 0),

438 -_rank_key(o)[0],

439 ),

440 )

441

442

443# ---------------------------------------------------------------------------

444# Reservation (ODCR) sort / rank

445# ---------------------------------------------------------------------------

446

447

448def sort_reservations(reservations: list[dict[str, Any]]) -> list[dict[str, Any]]:

449 """Sort reservations by region, then AZ, then instance type — stable browse order.

450

451 The ODCR counterpart to :func:`sort_offerings`, used to present a consolidated

452 multi-region reservation report in a predictable order.

453 """

454 return sorted(

455 reservations,

456 key=lambda r: (

457 str(r.get("region") or ""),

458 str(r.get("availability_zone") or ""),

459 str(r.get("instance_type") or ""),

460 ),

461 )

462

463

464def _reservation_rank_key(reservation: dict[str, Any]) -> tuple[float, float, str]:

465 """Rank key: most available capacity first, then cheapest per-GPU-hour.

466

467 Available instances sort descending (negated), so a reservation with more

468 free capacity ranks ahead of a fuller one; ties break toward the cheaper

469 per-GPU-hour rate, then a stable region string. This mirrors the intent of

470 :func:`_rank_key` for offerings (surface the single best pick first) while

471 optimizing for what matters for an existing reservation: free capacity.

472 """

473 available = reservation.get("available_instances")

474 gpu_hour = reservation.get("price_per_gpu_hour")

475 return (

476 -float(available) if available is not None else 0.0,

477 float(gpu_hour) if gpu_hour is not None else float("inf"),

478 str(reservation.get("region") or ""),

479 )

480

481

482def rank_reservations(reservations: list[dict[str, Any]]) -> list[dict[str, Any]]:

483 """Rank reservations by most-available-first, then cheapest per-GPU-hour.

484

485 The ODCR counterpart to :func:`rank_offerings`. Used to surface the single

486 best existing reservation (most free capacity, cheapest) in a consolidated

487 multi-region search report.

488 """

489 return sorted(reservations, key=_reservation_rank_key)