Coverage for cli / capacity / advisor.py: 100.00%

391 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-09-14 22:07 +0000

1"""Bedrock-powered AI capacity advisor.""" 

2 

3from __future__ import annotations 

4 

5import json 

6import logging 

7from dataclasses import dataclass, field 

8from datetime import UTC, datetime, timedelta 

9from typing import Any 

10 

11import boto3 

12from botocore.config import Config 

13from botocore.exceptions import ClientError 

14 

15from cli.config import GCOConfig, get_config 

16from gco.bedrock import ( 

17 BEDROCK_READ_TIMEOUT_SECONDS, 

18 BedrockResponseTruncatedError, 

19 build_bedrock_converse_options, 

20 extract_bedrock_converse_text, 

21 get_default_capacity_advisor_model_id, 

22 raise_if_bedrock_ftu_form_error, 

23) 

24 

25from .checker import CapacityChecker 

26from .multi_region import MultiRegionCapacityChecker, compute_price_trend 

27 

28logger = logging.getLogger(__name__) 

29 

30 

31def _snippet(text: str, limit: int = 200) -> str: 

32 """Compact, single-line prefix of ``text`` for parse-failure messages.""" 

33 collapsed = " ".join(text.split()) 

34 return collapsed[:limit] + ("..." if len(collapsed) > limit else "") 

35 

36 

37@dataclass 

38class BedrockCapacityRecommendation: 

39 """AI-generated capacity recommendation from Bedrock.""" 

40 

41 recommended_region: str 

42 recommended_instance_type: str 

43 recommended_capacity_type: str # "spot" or "on-demand" 

44 reasoning: str 

45 confidence: str # "high", "medium", "low" 

46 cost_estimate: str | None = None 

47 alternative_options: list[dict[str, Any]] = field(default_factory=list) 

48 warnings: list[str] = field(default_factory=list) 

49 raw_response: str = "" 

50 

51 

52@dataclass 

53class CapacityPredictionResult: 

54 """Bedrock prediction of the best time(s) to acquire capacity.""" 

55 

56 instance_type: str 

57 region: str 

58 best_windows: list[dict[str, Any]] = field(default_factory=list) 

59 avoid_windows: list[dict[str, Any]] = field(default_factory=list) 

60 reasoning: str = "" 

61 confidence: str = "low" 

62 raw_response: str = "" 

63 

64 

65class _CapacityAdvisorModelDefault: 

66 """Lazily expose the advisor's canonical model default as a string.""" 

67 

68 def __get__(self, instance: object, owner: type[Any] | None = None) -> str: 

69 return get_default_capacity_advisor_model_id() 

70 

71 

72class BedrockCapacityAdvisor: 

73 """ 

74 AI-powered capacity advisor using Amazon Bedrock. 

75 

76 Gathers comprehensive capacity data and uses an LLM to provide 

77 intelligent recommendations for workload placement. 

78 

79 DISCLAIMER: Recommendations are AI-generated and should be validated 

80 before making production decisions. 

81 """ 

82 

83 # Lazy class attribute for callers that inspect the advisor default. 

84 # Resolution occurs only when this Bedrock-specific attribute (or an 

85 # advisor without an explicit model) is used, reading the dedicated 

86 # ``context.bedrock.capacity_advisor_default_model_id`` knob. 

87 DEFAULT_MODEL = _CapacityAdvisorModelDefault() 

88 

89 def __init__(self, config: GCOConfig | None = None, model_id: str | None = None): 

90 self.config = config or get_config() 

91 self._session = boto3.Session() 

92 self._capacity_checker = CapacityChecker(config) 

93 self._multi_region_checker = MultiRegionCapacityChecker(config) 

94 self._uses_default_model = model_id is None 

95 self.model_id: str = self.DEFAULT_MODEL if model_id is None else model_id 

96 

97 def _get_bedrock_client(self) -> Any: 

98 """Get Bedrock runtime client.""" 

99 return self._session.client( 

100 "bedrock-runtime", 

101 region_name="us-east-1", 

102 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS), 

103 ) 

104 

105 def gather_capacity_data( 

106 self, 

107 instance_types: list[str] | None = None, 

108 regions: list[str] | None = None, 

109 ) -> dict[str, Any]: 

110 """ 

111 Gather comprehensive capacity data for AI analysis. 

112 

113 Args: 

114 instance_types: List of instance types to analyze (defaults to one 

115 representative per current GPU generation, T4 through Blackwell) 

116 regions: List of regions to check (defaults to deployed GCO regions) 

117 

118 Returns: 

119 Dictionary containing all gathered capacity data 

120 """ 

121 from cli.aws_client import get_aws_client 

122 

123 # Default to one representative per current GPU generation, spanning 

124 # budget inference through frontier training, so workload questions 

125 # about any generation get real telemetry. Sibling sizes of the same 

126 # GPU (e.g. g5.2xlarge/g5.4xlarge) are deliberately omitted — each 

127 # type costs a full set of AWS API calls per region. GB200/GB300 

128 # NVL72 are UltraServer families, not standalone EC2 instance types 

129 # (see cli/capacity/blocks.py NON_STANDALONE_INSTANCE_NOTES), so the 

130 # standalone Blackwell types represent that generation here. 

131 if not instance_types: 

132 instance_types = [ 

133 "g4dn.xlarge", # T4 — budget inference 

134 "g6.xlarge", # L4 — budget inference 

135 "g5.xlarge", # A10G — mainstream single-GPU 

136 "g6e.xlarge", # L40S — mainstream single-GPU 

137 "g7.2xlarge", # RTX PRO 4500 Blackwell — current-gen budget inference 

138 "g7e.2xlarge", # RTX PRO 6000 Blackwell — current-gen single-GPU inference 

139 "p4d.24xlarge", # 8x A100 — distributed training 

140 "p5.48xlarge", # 8x H100 — large-scale training 

141 "p5en.48xlarge", # 8x H200 — large-scale training 

142 "p6-b200.48xlarge", # 8x B200 (Blackwell) — frontier training 

143 "p6-b300.48xlarge", # 8x B300 (Blackwell Ultra) — frontier training 

144 ] 

145 

146 # Get deployed regions if not specified 

147 if not regions: 

148 aws_client = get_aws_client(self.config) 

149 stacks = aws_client.discover_regional_stacks() 

150 regions = list(stacks.keys()) if stacks else [self.config.default_region] 

151 

152 data: dict[str, Any] = { 

153 "timestamp": datetime.now(UTC).isoformat(), 

154 "regions_analyzed": regions, 

155 "instance_types_analyzed": instance_types, 

156 "regional_capacity": {}, 

157 "spot_data": {}, 

158 "on_demand_data": {}, 

159 "cluster_metrics": [], 

160 "queue_status": {}, 

161 } 

162 

163 # Gather regional cluster metrics 

164 for region in regions: 

165 try: 

166 capacity = self._multi_region_checker.get_region_capacity(region) 

167 data["cluster_metrics"].append( 

168 { 

169 "region": region, 

170 "queue_depth": capacity.queue_depth, 

171 "running_jobs": capacity.running_jobs, 

172 "pending_jobs": capacity.pending_jobs, 

173 "gpu_utilization": capacity.gpu_utilization, 

174 "cpu_utilization": capacity.cpu_utilization, 

175 "recommendation_score": capacity.recommendation_score, 

176 } 

177 ) 

178 except Exception as e: 

179 logger.debug("Failed to get cluster metrics for %s: %s", region, e) 

180 

181 # Failed lookups are recorded here and rendered into the prompt so the 

182 # model reasons about *missing* data instead of inventing a story for 

183 # why a row is absent (e.g. GetSpotPlacementScores' 24-hour 

184 # new-configuration limit must not read as "this type has no spot"). 

185 data["data_gaps"] = [] 

186 

187 def record_gap(instance_type: str, region: str, source: str, error: Exception) -> None: 

188 code = ( 

189 error.response.get("Error", {}).get("Code", "") 

190 if isinstance(error, ClientError) 

191 else "" 

192 ) or type(error).__name__ 

193 data["data_gaps"].append( 

194 { 

195 "instance_type": instance_type, 

196 "region": region, 

197 "source": source, 

198 "error": code, 

199 } 

200 ) 

201 logger.debug( 

202 "Capacity lookup %r failed for %s in %s: %s", source, instance_type, region, error 

203 ) 

204 

205 for instance_type in instance_types: 

206 data["spot_data"][instance_type] = {} 

207 data["on_demand_data"][instance_type] = {} 

208 

209 for region in regions: 

210 # Each lookup is isolated so one failing or throttled API 

211 # cannot discard the other signals for this (type, region) 

212 # pair, which previously erased real on-demand pricing and 

213 # spot history whenever the placement-score call failed. 

214 spot_entry: dict[str, Any] = {"placement_scores": {}, "prices": []} 

215 try: 

216 spot_entry["placement_scores"] = ( 

217 self._capacity_checker.get_spot_placement_score(instance_type, region) 

218 ) 

219 except Exception as e: 

220 record_gap(instance_type, region, "spot placement score", e) 

221 try: 

222 spot_prices = self._capacity_checker.get_spot_price_history( 

223 instance_type, region, days=7 

224 ) 

225 spot_entry["prices"] = [ 

226 { 

227 "az": p.availability_zone, 

228 "current": p.current_price, 

229 "avg_7d": p.avg_price_7d, 

230 "stability": p.price_stability, 

231 } 

232 for p in spot_prices 

233 ] 

234 except Exception as e: 

235 record_gap(instance_type, region, "spot price history", e) 

236 data["spot_data"][instance_type][region] = spot_entry 

237 

238 # Spot price trend analysis per AZ (for AI interpretation) 

239 try: 

240 ec2 = self._session.client("ec2", region_name=region) 

241 raw_resp = ec2.describe_spot_price_history( 

242 InstanceTypes=[instance_type], 

243 ProductDescriptions=["Linux/UNIX"], 

244 StartTime=datetime.now(UTC) - timedelta(days=7), 

245 EndTime=datetime.now(UTC), 

246 ) 

247 az_raw: dict[str, list[float]] = {} 

248 for item in raw_resp.get("SpotPriceHistory", []): 

249 az = item["AvailabilityZone"] 

250 if az not in az_raw: 

251 az_raw[az] = [] 

252 az_raw[az].append(float(item["SpotPrice"])) 

253 az_trends = { 

254 az: compute_price_trend(prices) 

255 for az, prices in az_raw.items() 

256 if len(prices) >= 2 

257 } 

258 if az_trends: 

259 data["spot_data"][instance_type][region]["price_trends"] = az_trends 

260 except Exception as e: 

261 logger.debug( 

262 "Failed to get price trends for %s in %s: %s", instance_type, region, e 

263 ) 

264 

265 od_entry: dict[str, Any] = {"price_per_hour": None, "available": None} 

266 try: 

267 od_entry["price_per_hour"] = self._capacity_checker.get_on_demand_price( 

268 instance_type, region 

269 ) 

270 except Exception as e: 

271 record_gap(instance_type, region, "on-demand price", e) 

272 try: 

273 od_entry["available"] = ( 

274 self._capacity_checker.check_instance_available_in_region( 

275 instance_type, region 

276 ) 

277 ) 

278 except Exception as e: 

279 record_gap(instance_type, region, "region availability", e) 

280 data["on_demand_data"][instance_type][region] = od_entry 

281 

282 # Gather capacity reservation and block data 

283 data["reservations"] = {} 

284 data["capacity_blocks"] = {} 

285 for instance_type in instance_types: 

286 data["reservations"][instance_type] = {} 

287 data["capacity_blocks"][instance_type] = {} 

288 for region in regions: 

289 try: 

290 odcrs = self._capacity_checker.list_capacity_reservations( 

291 region, instance_type=instance_type 

292 ) 

293 if odcrs: 

294 data["reservations"][instance_type][region] = [ 

295 { 

296 "az": r["availability_zone"], 

297 "total": r["total_instances"], 

298 "available": r["available_instances"], 

299 "utilization_pct": r["utilization_pct"], 

300 } 

301 for r in odcrs 

302 ] 

303 except Exception as e: 

304 logger.debug( 

305 "Failed to list reservations for %s in %s: %s", instance_type, region, e 

306 ) 

307 

308 try: 

309 blocks = self._capacity_checker.list_capacity_block_offerings( 

310 region, instance_type=instance_type, instance_count=1, duration_hours=24 

311 ) 

312 if blocks: 

313 data["capacity_blocks"][instance_type][region] = [ 

314 { 

315 "az": b["availability_zone"], 

316 "duration_hours": b["duration_hours"], 

317 "start_date": b["start_date"], 

318 "upfront_fee": b["upfront_fee"], 

319 } 

320 for b in blocks 

321 ] 

322 except Exception as e: 

323 logger.debug( 

324 "Failed to list capacity blocks for %s in %s: %s", instance_type, region, e 

325 ) 

326 

327 # Capacity block availability trends (26-week regression per instance type per region) 

328 data["capacity_block_trends"] = {} 

329 for instance_type in instance_types: 

330 data["capacity_block_trends"][instance_type] = {} 

331 for region in regions: 

332 try: 

333 trend = self._capacity_checker.get_capacity_block_trend(instance_type, region) 

334 if trend != 0.0: 

335 data["capacity_block_trends"][instance_type][region] = { 

336 "trend_score": trend, 

337 "interpretation": ( 

338 "capacity growing" 

339 if trend > 0.2 

340 else "capacity shrinking" 

341 if trend < -0.2 

342 else "stable" 

343 ), 

344 } 

345 except Exception as e: 

346 logger.debug( 

347 "Failed to get capacity block trend for %s in %s: %s", 

348 instance_type, 

349 region, 

350 e, 

351 ) 

352 

353 # Weighted recommendation scores (algorithmic ranking for AI context) 

354 try: 

355 weighted_results = self._multi_region_checker.recommend_region_for_job( 

356 instance_type=instance_types[0] if instance_types else None, 

357 ) 

358 data["weighted_recommendation"] = { 

359 "top_region": weighted_results.get("region"), 

360 "scoring_method": weighted_results.get("scoring_method", "simple"), 

361 "instance_type": weighted_results.get("instance_type"), 

362 "all_regions": weighted_results.get("all_regions", []), 

363 } 

364 except Exception as e: 

365 logger.debug("Failed to compute weighted recommendation: %s", e) 

366 

367 return data 

368 

369 def _gather_historical_context(self, capacity_data: dict[str, Any]) -> dict[str, Any]: 

370 """Best-effort historical enrichment for the Bedrock prompt. 

371 

372 For each (instance_type, region) with a current spot score, look up the 

373 7-day statistics and temporal patterns from the capacity history store. 

374 Returns an empty dict if the history surface is unavailable (table 

375 missing, no access, or feature disabled) so the advisor still works 

376 without it. 

377 """ 

378 try: 

379 from cli.capacity.history import get_capacity_history_store 

380 

381 store = get_capacity_history_store() 

382 except Exception as e: 

383 logger.debug("Capacity history store unavailable: %s", e) 

384 return {} 

385 

386 context: dict[str, Any] = {} 

387 for instance_type, regions_data in capacity_data.get("spot_data", {}).items(): 

388 for region, spot_info in (regions_data or {}).items(): 

389 current = (spot_info.get("placement_scores") or {}).get("regional") 

390 if current is None: 

391 continue 

392 try: 

393 stats = store.get_statistics(instance_type, region) 

394 except Exception as e: 

395 logger.debug( 

396 "Historical stats lookup failed for %s in %s: %s", instance_type, region, e 

397 ) 

398 continue 

399 spot_stats = stats.get("metrics", {}).get("spot_score") 

400 if not spot_stats: 

401 continue 

402 try: 

403 patterns = store.get_temporal_patterns(instance_type, region) 

404 best_windows = patterns.get("best_windows", [])[:3] 

405 except Exception: 

406 best_windows = [] 

407 context[f"{instance_type}#{region}"] = { 

408 "instance_type": instance_type, 

409 "region": region, 

410 "current_spot_score": current, 

411 "p25": spot_stats["p25"], 

412 "p50": spot_stats["p50"], 

413 "p75": spot_stats["p75"], 

414 "best_windows": best_windows, 

415 } 

416 return context 

417 

418 def _build_prompt( 

419 self, 

420 capacity_data: dict[str, Any], 

421 workload_description: str | None = None, 

422 requirements: dict[str, Any] | None = None, 

423 historical_context: dict[str, Any] | None = None, 

424 ) -> str: 

425 """Build the prompt for Bedrock.""" 

426 requirements = requirements or {} 

427 

428 prompt = """You are an expert AWS capacity planning advisor for GPU/ML workloads. 

429Analyze the following capacity data and provide a recommendation for where to place a workload. 

430 

431IMPORTANT DISCLAIMERS: 

432- This is AI-generated advice and should be validated before production use 

433- Capacity availability can change rapidly 

434- Spot instances may be interrupted at any time 

435- Pricing data may not reflect real-time prices 

436 

437""" 

438 

439 if workload_description: 

440 prompt += f"WORKLOAD DESCRIPTION:\n{workload_description}\n\n" 

441 

442 if requirements: 

443 prompt += "REQUIREMENTS:\n" 

444 if requirements.get("gpu_required"): 

445 prompt += "- GPU Required: Yes\n" 

446 if requirements.get("min_gpus"): 

447 prompt += f"- Minimum GPUs: {requirements['min_gpus']}\n" 

448 if requirements.get("min_memory_gb"): 

449 prompt += f"- Minimum Memory: {requirements['min_memory_gb']} GB\n" 

450 if requirements.get("fault_tolerance"): 

451 prompt += f"- Fault Tolerance: {requirements['fault_tolerance']}\n" 

452 if requirements.get("max_cost_per_hour"): 

453 prompt += f"- Max Cost/Hour: ${requirements['max_cost_per_hour']}\n" 

454 prompt += "\n" 

455 

456 prompt += "CAPACITY DATA:\n" 

457 prompt += f"Timestamp: {capacity_data.get('timestamp', 'N/A')}\n" 

458 prompt += f"Regions Analyzed: {', '.join(capacity_data.get('regions_analyzed', []))}\n" 

459 prompt += ( 

460 f"Instance Types: {', '.join(capacity_data.get('instance_types_analyzed', []))}\n\n" 

461 ) 

462 

463 # Cluster metrics 

464 if capacity_data.get("cluster_metrics"): 

465 prompt += "CLUSTER METRICS BY REGION:\n" 

466 for m in capacity_data["cluster_metrics"]: 

467 prompt += f" {m['region']}:\n" 

468 prompt += f" - Queue Depth: {m['queue_depth']}\n" 

469 prompt += f" - Running Jobs: {m['running_jobs']}\n" 

470 prompt += f" - GPU Utilization: {m['gpu_utilization']:.1f}%\n" 

471 prompt += f" - CPU Utilization: {m['cpu_utilization']:.1f}%\n" 

472 prompt += "\n" 

473 

474 # Spot data summary 

475 prompt += "SPOT CAPACITY SUMMARY:\n" 

476 for instance_type, regions_data in capacity_data.get("spot_data", {}).items(): 

477 prompt += f" {instance_type}:\n" 

478 for region, spot_info in regions_data.items(): 

479 scores = spot_info.get("placement_scores", {}) 

480 regional_score = scores.get("regional", "N/A") 

481 prices = spot_info.get("prices", []) 

482 avg_price = sum(p["current"] for p in prices) / len(prices) if prices else "N/A" 

483 prompt += f" {region}: Score={regional_score}/10, " 

484 prompt += f"Avg Price=${avg_price if isinstance(avg_price, str) else f'{avg_price:.4f}'}/hr\n" 

485 trends = spot_info.get("price_trends", {}) 

486 if trends: 

487 rendered = ", ".join( 

488 f"{az} {t['direction']} " 

489 f"(normalized slope {t['normalized_slope']:+.2f}, " 

490 f"{t['price_changes']} price changes)" 

491 for az, t in sorted(trends.items()) 

492 ) 

493 prompt += f" 7-day spot price trend by AZ: {rendered}\n" 

494 prompt += "\n" 

495 

496 # On-demand data summary 

497 prompt += "ON-DEMAND PRICING:\n" 

498 for instance_type, regions_data in capacity_data.get("on_demand_data", {}).items(): 

499 prompt += f" {instance_type}:\n" 

500 for region, od_info in regions_data.items(): 

501 price = od_info.get("price_per_hour") 

502 available = od_info.get("available") 

503 # None means the offerings lookup failed — say "unknown" so the 

504 # model cannot mistake a failed check for "not offered". 

505 availability = "unknown (lookup failed)" if available is None else available 

506 prompt += f" {region}: ${price:.4f}/hr" if price else f" {region}: N/A" 

507 prompt += f" (Available: {availability})\n" 

508 prompt += "\n" 

509 

510 # Capacity reservations (ODCRs) 

511 reservations = capacity_data.get("reservations", {}) 

512 has_reservations = any(bool(regions_data) for regions_data in reservations.values()) 

513 if has_reservations: 

514 prompt += "CAPACITY RESERVATIONS (ODCRs):\n" 

515 for instance_type, regions_data in reservations.items(): 

516 for region, odcrs in regions_data.items(): 

517 for r in odcrs: 

518 prompt += ( 

519 f" {instance_type} in {region} ({r['az']}): " 

520 f"{r['available']}/{r['total']} available " 

521 f"({r['utilization_pct']}% used)\n" 

522 ) 

523 prompt += "\n" 

524 

525 # Capacity Blocks for ML 

526 blocks = capacity_data.get("capacity_blocks", {}) 

527 has_blocks = any(bool(regions_data) for regions_data in blocks.values()) 

528 if has_blocks: 

529 prompt += "CAPACITY BLOCK OFFERINGS (guaranteed GPU blocks):\n" 

530 for instance_type, regions_data in blocks.items(): 

531 for region, offerings in regions_data.items(): 

532 for b in offerings: 

533 prompt += ( 

534 f" {instance_type} in {region} ({b['az']}): " 

535 f"{b['duration_hours']}h starting {b['start_date']}, " 

536 f"${b['upfront_fee']}\n" 

537 ) 

538 prompt += "\n" 

539 

540 # Capacity block availability trends (26-week offering-density regression) 

541 block_trends = capacity_data.get("capacity_block_trends", {}) 

542 has_block_trends = any(bool(regions_data) for regions_data in block_trends.values()) 

543 if has_block_trends: 

544 prompt += "CAPACITY BLOCK AVAILABILITY TRENDS (26-week, near-term vs far-term):\n" 

545 for instance_type, regions_data in block_trends.items(): 

546 for region, trend in regions_data.items(): 

547 prompt += ( 

548 f" {instance_type} in {region}: " 

549 f"{trend['trend_score']:+.2f} ({trend['interpretation']})\n" 

550 ) 

551 prompt += "\n" 

552 

553 # Algorithmic multi-signal ranking (context for the model, not binding) 

554 weighted = capacity_data.get("weighted_recommendation") 

555 if weighted and weighted.get("all_regions"): 

556 scoring_method = weighted.get("scoring_method", "simple") 

557 scored_for = ( 

558 f" for {weighted['instance_type']}" if weighted.get("instance_type") else "" 

559 ) 

560 prompt += ( 

561 f"ALGORITHMIC REGION RANKING ({scoring_method} scoring{scored_for}; " 

562 "lower score = better; advisory pre-computation, weigh it " 

563 "against the raw data above):\n" 

564 ) 

565 for entry in weighted["all_regions"]: 

566 prompt += f" {entry['region']}: score={entry['score']:.1f}" 

567 details = [] 

568 if entry.get("spot_placement_score") is not None: 

569 details.append(f"spot availability {entry['spot_placement_score']:.0%}") 

570 if entry.get("spot_price_ratio") is not None: 

571 details.append(f"spot/on-demand price ratio {entry['spot_price_ratio']:.2f}") 

572 if entry.get("capacity_block_trend") is not None: 

573 details.append(f"block trend {entry['capacity_block_trend']:+.2f}") 

574 details.append(f"queue depth {entry.get('queue_depth', 'N/A')}") 

575 gpu_util = entry.get("gpu_utilization") 

576 if gpu_util is not None: 

577 details.append(f"GPU util {gpu_util:.0f}%") 

578 prompt += f" ({', '.join(details)})\n" 

579 prompt += "\n" 

580 

581 # Failed lookups — spelled out so the model reasons about missing 

582 # data instead of inventing an explanation for absent rows (e.g. the 

583 # placement-score API's 24-hour new-configuration limit must not read 

584 # as "this instance type has no spot pools"). 

585 data_gaps = capacity_data.get("data_gaps") or [] 

586 if data_gaps: 

587 prompt += "DATA GAPS (lookups that FAILED — treat as unknown, not as unavailable):\n" 

588 grouped: dict[tuple[str, str, str], list[str]] = {} 

589 for gap in data_gaps: 

590 key = (gap["source"], gap["error"], gap["region"]) 

591 grouped.setdefault(key, []).append(gap["instance_type"]) 

592 for (source, error, region), types in sorted(grouped.items()): 

593 prompt += ( 

594 f" {source} in {region} failed with {error} for: {', '.join(sorted(types))}\n" 

595 ) 

596 prompt += ( 

597 " Do not draw capacity or availability conclusions from these " 

598 "missing values; rely on the signals that are present and " 

599 "mention the gap in your warnings.\n" 

600 ) 

601 prompt += "\n" 

602 

603 if historical_context: 

604 prompt += "## Historical Context (last 7 days)\n" 

605 for ctx in historical_context.values(): 

606 current = ctx["current_spot_score"] 

607 p25 = ctx["p25"] 

608 p50 = ctx["p50"] 

609 p75 = ctx["p75"] 

610 if current < p25: 

611 interpretation = "likely transient contention" 

612 elif current <= p75: 

613 interpretation = "within normal range" 

614 else: 

615 interpretation = "unusually favorable" 

616 prompt += f" {ctx['instance_type']} in {ctx['region']}:\n" 

617 prompt += f" Current spot score: {current}\n" 

618 prompt += f" Historical p25/p50/p75: {p25}/{p50}/{p75}\n" 

619 prompt += f" Interpretation: {interpretation}\n" 

620 windows = ctx.get("best_windows") or [] 

621 if windows: 

622 rendered = ", ".join( 

623 f"{w['day']} {w['hour']:02d}:00 (avg {w['avg']})" for w in windows 

624 ) 

625 prompt += f" Best historical windows (top 3): {rendered}\n" 

626 prompt += "\n" 

627 

628 prompt += """Based on this data, provide your recommendation in the following JSON format: 

629{ 

630 "recommended_region": "region-name", 

631 "recommended_instance_type": "instance-type", 

632 "recommended_capacity_type": "spot, on-demand, odcr, or capacity-block", 

633 "reasoning": "Detailed explanation of why this is the best choice", 

634 "confidence": "high, medium, or low", 

635 "cost_estimate": "Estimated hourly cost", 

636 "reservation_advice": "If ODCRs or Capacity Blocks are available, explain how to use them. If not, suggest whether the user should consider purchasing a Capacity Block.", 

637 "alternative_options": [ 

638 {"region": "...", "instance_type": "...", "capacity_type": "...", "reason": "..."} 

639 ], 

640 "warnings": ["Any important warnings or caveats"] 

641} 

642 

643Respond ONLY with the JSON object, no additional text.""" 

644 

645 return prompt 

646 

647 def get_recommendation( 

648 self, 

649 workload_description: str | None = None, 

650 instance_types: list[str] | None = None, 

651 regions: list[str] | None = None, 

652 requirements: dict[str, Any] | None = None, 

653 ) -> BedrockCapacityRecommendation: 

654 """ 

655 Get an AI-powered capacity recommendation. 

656 

657 Args: 

658 workload_description: Description of the workload 

659 instance_types: List of instance types to consider 

660 regions: List of regions to consider 

661 requirements: Dictionary of requirements (gpu_required, min_gpus, etc.) 

662 

663 Returns: 

664 BedrockCapacityRecommendation with the AI's recommendation 

665 """ 

666 # Gather capacity data 

667 capacity_data = self.gather_capacity_data(instance_types, regions) 

668 

669 # Gather best-effort historical context (skipped when unavailable) 

670 historical_context = self._gather_historical_context(capacity_data) 

671 

672 # Build prompt 

673 prompt = self._build_prompt( 

674 capacity_data, workload_description, requirements, historical_context 

675 ) 

676 

677 # Call Bedrock 

678 bedrock = self._get_bedrock_client() 

679 

680 try: 

681 # Use the Converse API for better compatibility across models 

682 response = bedrock.converse( 

683 modelId=self.model_id, 

684 messages=[{"role": "user", "content": [{"text": prompt}]}], 

685 **build_bedrock_converse_options( 

686 self.model_id, 

687 # Deliberately no maxTokens: the Converse default is the 

688 # model's own maximum output length, so reasoning plus the 

689 # JSON answer can never hit a GCO-imposed cap. A cap is 

690 # opt-in — pass maxTokens here to restore one. 

691 inference_config={"temperature": 0.1}, 

692 apply_default_reasoning=self._uses_default_model, 

693 ), 

694 ) 

695 

696 # Extended reasoning precedes the final answer with a 

697 # ``reasoningContent`` block; return the first real text block. 

698 response_text = extract_bedrock_converse_text(response) 

699 

700 # Parse JSON response 

701 # Find JSON in response (in case model adds extra text) 

702 json_start = response_text.find("{") 

703 json_end = response_text.rfind("}") + 1 

704 if json_start >= 0 and json_end > json_start: 

705 json_str = response_text[json_start:json_end] 

706 result = json.loads(json_str) 

707 else: 

708 raise ValueError( 

709 "No JSON object found in the model response " 

710 f"(response begins: {_snippet(response_text)!r})" 

711 ) 

712 

713 return BedrockCapacityRecommendation( 

714 recommended_region=result.get("recommended_region", "unknown"), 

715 recommended_instance_type=result.get("recommended_instance_type", "unknown"), 

716 recommended_capacity_type=result.get("recommended_capacity_type", "spot"), 

717 reasoning=result.get("reasoning", ""), 

718 confidence=result.get("confidence", "low"), 

719 cost_estimate=result.get("cost_estimate"), 

720 alternative_options=result.get("alternative_options", []), 

721 warnings=result.get("warnings", []), 

722 raw_response=response_text, 

723 ) 

724 

725 except ClientError as e: 

726 error_code = e.response.get("Error", {}).get("Code", "") 

727 # Raised as a distinct type (still a RuntimeError) so callers can 

728 # tell a fixable account-setup gap from a transient Bedrock fault. 

729 raise_if_bedrock_ftu_form_error(e) 

730 if error_code == "AccessDeniedException": 

731 raise RuntimeError( 

732 "Access denied to Bedrock. Ensure your IAM role has " 

733 "bedrock:InvokeModel permission and the model is enabled in your account." 

734 ) from e 

735 if error_code == "ValidationException": 

736 raise RuntimeError( 

737 f"Model {self.model_id} may not be available. " 

738 "Try a different model with --model option." 

739 ) from e 

740 raise RuntimeError(f"Bedrock API error: {e}") from e 

741 except json.JSONDecodeError as e: 

742 # ``response_text`` is always bound here: the decoder can only 

743 # fail after the response text was extracted. 

744 raise RuntimeError( 

745 f"Failed to parse AI response as JSON: {e} " 

746 f"(response begins: {_snippet(response_text)!r})" 

747 ) from e 

748 except BedrockResponseTruncatedError: 

749 # Already carries its own remediation; wrapping it in the generic 

750 # "Failed to get AI recommendation" message would only bury it. 

751 raise 

752 except Exception as e: 

753 raise RuntimeError(f"Failed to get AI recommendation: {e}") from e 

754 

755 def _build_predict_prompt( 

756 self, 

757 instance_type: str, 

758 region: str, 

759 stats: dict[str, Any], 

760 patterns: dict[str, Any], 

761 ) -> str: 

762 """Build a Bedrock prompt focused on the best time to acquire capacity.""" 

763 metrics = stats.get("metrics", {}) 

764 spot = metrics.get("spot_score", {}) 

765 price = metrics.get("spot_price", {}) 

766 lines = [ 

767 "You are an expert AWS GPU capacity-timing advisor.", 

768 "", 

769 ( 

770 f"Based ONLY on the historical capacity patterns below for " 

771 f"{instance_type} in {region}, recommend the best time window(s) to " 

772 f"acquire this capacity (spot or capacity blocks), and which windows to avoid." 

773 ), 

774 "", 

775 ( 

776 f"## Historical window: last {stats.get('hours_back')} hours, " 

777 f"{stats.get('sample_count')} samples" 

778 ), 

779 ] 

780 if spot: 

781 lines.append( 

782 f"Spot placement score (1-10, higher = better availability): " 

783 f"p25={spot.get('p25')} p50={spot.get('p50')} p75={spot.get('p75')} " 

784 f"min={spot.get('min')} max={spot.get('max')}" 

785 ) 

786 if price: 

787 lines.append( 

788 f"Spot price USD/hr (lower = cheaper): " 

789 f"p25={price.get('p25')} p50={price.get('p50')} p75={price.get('p75')}" 

790 ) 

791 best = patterns.get("best_windows", [])[:10] 

792 if best: 

793 lines.append("") 

794 lines.append( 

795 "Top observed windows by average spot score (day, hour UTC, avg, samples):" 

796 ) 

797 for window in best: 

798 lines.append( 

799 f"- {window['day']} {window['hour']:02d}:00 UTC: " 

800 f"avg {window['avg']} (n={window['count']})" 

801 ) 

802 lines.append("") 

803 lines.append("Respond ONLY with a JSON object of this exact shape:") 

804 lines.append( 

805 '{"best_windows": [{"day": "Monday", "hour_range": "13:00-16:00 UTC", ' 

806 '"why": "..."}], "avoid_windows": [{"day": "...", "hour_range": "...", ' 

807 '"why": "..."}], "reasoning": "...", "confidence": "high|medium|low"}' 

808 ) 

809 return "\n".join(lines) 

810 

811 def predict_capacity_window( 

812 self, 

813 instance_type: str, 

814 region: str, 

815 hours_back: int = 168, 

816 ) -> CapacityPredictionResult: 

817 """Predict the best acquisition window for an instance type in a region. 

818 

819 Reads the historical capacity surface, builds a timing-focused prompt, 

820 and asks Bedrock. Raises ``ValueError`` when there are no samples yet; 

821 propagates the underlying ``ClientError`` (e.g. ResourceNotFoundException) 

822 when the history table does not exist so callers can surface a hint, and 

823 ``BedrockResponseTruncatedError`` when the model's answer was cut off by 

824 an output-token limit. 

825 """ 

826 from cli.capacity.history import get_capacity_history_store 

827 

828 store = get_capacity_history_store() 

829 stats = store.get_statistics(instance_type, region, hours_back) 

830 if stats.get("sample_count", 0) == 0: 

831 raise ValueError( 

832 f"No historical capacity samples for {instance_type} in {region} yet. " 

833 "The poller records one about every 15 minutes once enabled." 

834 ) 

835 patterns = store.get_temporal_patterns(instance_type, region, hours_back) 

836 prompt = self._build_predict_prompt(instance_type, region, stats, patterns) 

837 

838 bedrock = self._get_bedrock_client() 

839 response = bedrock.converse( 

840 modelId=self.model_id, 

841 messages=[{"role": "user", "content": [{"text": prompt}]}], 

842 **build_bedrock_converse_options( 

843 self.model_id, 

844 # No maxTokens by default — see get_recommendation. 

845 inference_config={"temperature": 0.2}, 

846 apply_default_reasoning=self._uses_default_model, 

847 ), 

848 ) 

849 text = extract_bedrock_converse_text(response) 

850 

851 parsed: dict[str, Any] = {} 

852 start = text.find("{") 

853 end = text.rfind("}") + 1 

854 if start >= 0 and end > start: 

855 try: 

856 parsed = json.loads(text[start:end]) 

857 except json.JSONDecodeError: 

858 parsed = {} 

859 return CapacityPredictionResult( 

860 instance_type=instance_type, 

861 region=region, 

862 best_windows=parsed.get("best_windows", []), 

863 avoid_windows=parsed.get("avoid_windows", []), 

864 reasoning=parsed.get("reasoning", ""), 

865 confidence=parsed.get("confidence", "low"), 

866 raw_response=text, 

867 ) 

868 

869 def predict_capacity_windows_all_regions( 

870 self, 

871 instance_type: str, 

872 hours_back: int = 168, 

873 ) -> list[CapacityPredictionResult]: 

874 """Predict acquisition windows for every region that has history. 

875 

876 Discovers the regions with samples for ``instance_type`` via the history 

877 store's ``by-timestamp`` GSI and runs :meth:`predict_capacity_window` 

878 for each. Raises ``ValueError`` when no region has samples yet; 

879 propagates the underlying ``ClientError`` (e.g. ResourceNotFoundException) 

880 when the history table does not exist. 

881 """ 

882 from cli.capacity.history import get_capacity_history_store 

883 

884 store = get_capacity_history_store() 

885 regions = store.get_regions_with_data(instance_type, hours_back) 

886 if not regions: 

887 raise ValueError( 

888 f"No historical capacity samples for {instance_type} in any region yet. " 

889 "The poller records one about every 15 minutes once enabled." 

890 ) 

891 results: list[CapacityPredictionResult] = [] 

892 for region in regions: 

893 try: 

894 results.append(self.predict_capacity_window(instance_type, region, hours_back)) 

895 except ValueError: 

896 continue 

897 return results 

898 

899 

900def get_bedrock_capacity_advisor( 

901 config: GCOConfig | None = None, model_id: str | None = None 

902) -> BedrockCapacityAdvisor: 

903 """Get a configured Bedrock capacity advisor instance.""" 

904 return BedrockCapacityAdvisor(config, model_id)