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

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_bedrock_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 _SharedBedrockModelDefault: 

66 """Lazily expose the historical advisor class attribute as a string.""" 

67 

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

69 return get_default_bedrock_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 # Backward-compatible lazy class alias for callers that inspect the 

84 # advisor default. Resolution occurs only when this Bedrock-specific 

85 # attribute (or an advisor without an explicit model) is used. 

86 DEFAULT_MODEL = _SharedBedrockModelDefault() 

87 

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

89 self.config = config or get_config() 

90 self._session = boto3.Session() 

91 self._capacity_checker = CapacityChecker(config) 

92 self._multi_region_checker = MultiRegionCapacityChecker(config) 

93 self._uses_default_model = model_id is None 

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

95 

96 def _get_bedrock_client(self) -> Any: 

97 """Get Bedrock runtime client.""" 

98 return self._session.client( 

99 "bedrock-runtime", 

100 region_name="us-east-1", 

101 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS), 

102 ) 

103 

104 def gather_capacity_data( 

105 self, 

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

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

108 ) -> dict[str, Any]: 

109 """ 

110 Gather comprehensive capacity data for AI analysis. 

111 

112 Args: 

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

114 representative per current GPU generation, T4 through Blackwell) 

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

116 

117 Returns: 

118 Dictionary containing all gathered capacity data 

119 """ 

120 from cli.aws_client import get_aws_client 

121 

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

123 # budget inference through frontier training, so workload questions 

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

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

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

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

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

129 # standalone Blackwell types represent that generation here. 

130 if not instance_types: 

131 instance_types = [ 

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

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

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

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

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

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

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

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

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

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

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

143 ] 

144 

145 # Get deployed regions if not specified 

146 if not regions: 

147 aws_client = get_aws_client(self.config) 

148 stacks = aws_client.discover_regional_stacks() 

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

150 

151 data: dict[str, Any] = { 

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

153 "regions_analyzed": regions, 

154 "instance_types_analyzed": instance_types, 

155 "regional_capacity": {}, 

156 "spot_data": {}, 

157 "on_demand_data": {}, 

158 "cluster_metrics": [], 

159 "queue_status": {}, 

160 } 

161 

162 # Gather regional cluster metrics 

163 for region in regions: 

164 try: 

165 capacity = self._multi_region_checker.get_region_capacity(region) 

166 data["cluster_metrics"].append( 

167 { 

168 "region": region, 

169 "queue_depth": capacity.queue_depth, 

170 "running_jobs": capacity.running_jobs, 

171 "pending_jobs": capacity.pending_jobs, 

172 "gpu_utilization": capacity.gpu_utilization, 

173 "cpu_utilization": capacity.cpu_utilization, 

174 "recommendation_score": capacity.recommendation_score, 

175 } 

176 ) 

177 except Exception as e: 

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

179 

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

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

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

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

184 data["data_gaps"] = [] 

185 

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

187 code = ( 

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

189 if isinstance(error, ClientError) 

190 else "" 

191 ) or type(error).__name__ 

192 data["data_gaps"].append( 

193 { 

194 "instance_type": instance_type, 

195 "region": region, 

196 "source": source, 

197 "error": code, 

198 } 

199 ) 

200 logger.debug( 

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

202 ) 

203 

204 for instance_type in instance_types: 

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

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

207 

208 for region in regions: 

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

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

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

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

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

214 try: 

215 spot_entry["placement_scores"] = ( 

216 self._capacity_checker.get_spot_placement_score(instance_type, region) 

217 ) 

218 except Exception as e: 

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

220 try: 

221 spot_prices = self._capacity_checker.get_spot_price_history( 

222 instance_type, region, days=7 

223 ) 

224 spot_entry["prices"] = [ 

225 { 

226 "az": p.availability_zone, 

227 "current": p.current_price, 

228 "avg_7d": p.avg_price_7d, 

229 "stability": p.price_stability, 

230 } 

231 for p in spot_prices 

232 ] 

233 except Exception as e: 

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

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

236 

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

238 try: 

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

240 raw_resp = ec2.describe_spot_price_history( 

241 InstanceTypes=[instance_type], 

242 ProductDescriptions=["Linux/UNIX"], 

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

244 EndTime=datetime.now(UTC), 

245 ) 

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

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

248 az = item["AvailabilityZone"] 

249 if az not in az_raw: 

250 az_raw[az] = [] 

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

252 az_trends = { 

253 az: compute_price_trend(prices) 

254 for az, prices in az_raw.items() 

255 if len(prices) >= 2 

256 } 

257 if az_trends: 

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

259 except Exception as e: 

260 logger.debug( 

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

262 ) 

263 

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

265 try: 

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

267 instance_type, region 

268 ) 

269 except Exception as e: 

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

271 try: 

272 od_entry["available"] = ( 

273 self._capacity_checker.check_instance_available_in_region( 

274 instance_type, region 

275 ) 

276 ) 

277 except Exception as e: 

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

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

280 

281 # Gather capacity reservation and block data 

282 data["reservations"] = {} 

283 data["capacity_blocks"] = {} 

284 for instance_type in instance_types: 

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

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

287 for region in regions: 

288 try: 

289 odcrs = self._capacity_checker.list_capacity_reservations( 

290 region, instance_type=instance_type 

291 ) 

292 if odcrs: 

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

294 { 

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

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

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

298 "utilization_pct": r["utilization_pct"], 

299 } 

300 for r in odcrs 

301 ] 

302 except Exception as e: 

303 logger.debug( 

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

305 ) 

306 

307 try: 

308 blocks = self._capacity_checker.list_capacity_block_offerings( 

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

310 ) 

311 if blocks: 

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

313 { 

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

315 "duration_hours": b["duration_hours"], 

316 "start_date": b["start_date"], 

317 "upfront_fee": b["upfront_fee"], 

318 } 

319 for b in blocks 

320 ] 

321 except Exception as e: 

322 logger.debug( 

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

324 ) 

325 

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

327 data["capacity_block_trends"] = {} 

328 for instance_type in instance_types: 

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

330 for region in regions: 

331 try: 

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

333 if trend != 0.0: 

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

335 "trend_score": trend, 

336 "interpretation": ( 

337 "capacity growing" 

338 if trend > 0.2 

339 else "capacity shrinking" 

340 if trend < -0.2 

341 else "stable" 

342 ), 

343 } 

344 except Exception as e: 

345 logger.debug( 

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

347 instance_type, 

348 region, 

349 e, 

350 ) 

351 

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

353 try: 

354 weighted_results = self._multi_region_checker.recommend_region_for_job( 

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

356 ) 

357 data["weighted_recommendation"] = { 

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

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

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

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

362 } 

363 except Exception as e: 

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

365 

366 return data 

367 

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

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

370 

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

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

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

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

375 without it. 

376 """ 

377 try: 

378 from cli.capacity.history import get_capacity_history_store 

379 

380 store = get_capacity_history_store() 

381 except Exception as e: 

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

383 return {} 

384 

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

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

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

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

389 if current is None: 389 ↛ 390line 389 didn't jump to line 390 because the condition on line 389 was never true

390 continue 

391 try: 

392 stats = store.get_statistics(instance_type, region) 

393 except Exception as e: 

394 logger.debug("Historical stats lookup failed: %s", e) 

395 return context 

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

397 if not spot_stats: 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true

398 continue 

399 try: 

400 patterns = store.get_temporal_patterns(instance_type, region) 

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

402 except Exception: 

403 best_windows = [] 

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

405 "instance_type": instance_type, 

406 "region": region, 

407 "current_spot_score": current, 

408 "p25": spot_stats["p25"], 

409 "p50": spot_stats["p50"], 

410 "p75": spot_stats["p75"], 

411 "best_windows": best_windows, 

412 } 

413 return context 

414 

415 def _build_prompt( 

416 self, 

417 capacity_data: dict[str, Any], 

418 workload_description: str | None = None, 

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

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

421 ) -> str: 

422 """Build the prompt for Bedrock.""" 

423 requirements = requirements or {} 

424 

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

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

427 

428IMPORTANT DISCLAIMERS: 

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

430- Capacity availability can change rapidly 

431- Spot instances may be interrupted at any time 

432- Pricing data may not reflect real-time prices 

433 

434""" 

435 

436 if workload_description: 

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

438 

439 if requirements: 

440 prompt += "REQUIREMENTS:\n" 

441 if requirements.get("gpu_required"): 

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

443 if requirements.get("min_gpus"): 

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

445 if requirements.get("min_memory_gb"): 

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

447 if requirements.get("fault_tolerance"): 

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

449 if requirements.get("max_cost_per_hour"): 

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

451 prompt += "\n" 

452 

453 prompt += "CAPACITY DATA:\n" 

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

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

456 prompt += ( 

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

458 ) 

459 

460 # Cluster metrics 

461 if capacity_data.get("cluster_metrics"): 

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

463 for m in capacity_data["cluster_metrics"]: 

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

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

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

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

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

469 prompt += "\n" 

470 

471 # Spot data summary 

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

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

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

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

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

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

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

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

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

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

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

483 if trends: 483 ↛ 484line 483 didn't jump to line 484 because the condition on line 483 was never true

484 rendered = ", ".join( 

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

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

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

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

489 ) 

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

491 prompt += "\n" 

492 

493 # On-demand data summary 

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

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

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

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

498 price = od_info.get("price_per_hour") 

499 available = od_info.get("available") 

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

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

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

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

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

505 prompt += "\n" 

506 

507 # Capacity reservations (ODCRs) 

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

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

510 if has_reservations: 

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

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

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

514 for r in odcrs: 

515 prompt += ( 

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

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

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

519 ) 

520 prompt += "\n" 

521 

522 # Capacity Blocks for ML 

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

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

525 if has_blocks: 

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

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

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

529 for b in offerings: 

530 prompt += ( 

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

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

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

534 ) 

535 prompt += "\n" 

536 

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

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

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

540 if has_block_trends: 540 ↛ 541line 540 didn't jump to line 541 because the condition on line 540 was never true

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

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

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

544 prompt += ( 

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

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

547 ) 

548 prompt += "\n" 

549 

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

551 weighted = capacity_data.get("weighted_recommendation") 

552 if weighted and weighted.get("all_regions"): 552 ↛ 553line 552 didn't jump to line 553 because the condition on line 552 was never true

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

554 scored_for = ( 

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

556 ) 

557 prompt += ( 

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

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

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

561 ) 

562 for entry in weighted["all_regions"]: 

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

564 details = [] 

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

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

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

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

569 if entry.get("capacity_block_trend"): 

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

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

572 gpu_util = entry.get("gpu_utilization") 

573 if gpu_util is not None: 

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

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

576 prompt += "\n" 

577 

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

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

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

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

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

583 if data_gaps: 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true

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

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

586 for gap in data_gaps: 

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

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

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

590 prompt += ( 

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

592 ) 

593 prompt += ( 

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

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

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

597 ) 

598 prompt += "\n" 

599 

600 if historical_context: 

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

602 for ctx in historical_context.values(): 

603 current = ctx["current_spot_score"] 

604 p25 = ctx["p25"] 

605 p50 = ctx["p50"] 

606 p75 = ctx["p75"] 

607 if current < p25: 

608 interpretation = "likely transient contention" 

609 elif current <= p75: 

610 interpretation = "within normal range" 

611 else: 

612 interpretation = "unusually favorable" 

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

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

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

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

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

618 if windows: 618 ↛ 602line 618 didn't jump to line 602 because the condition on line 618 was always true

619 rendered = ", ".join( 

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

621 ) 

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

623 prompt += "\n" 

624 

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

626{ 

627 "recommended_region": "region-name", 

628 "recommended_instance_type": "instance-type", 

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

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

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

632 "cost_estimate": "Estimated hourly cost", 

633 "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.", 

634 "alternative_options": [ 

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

636 ], 

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

638} 

639 

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

641 

642 return prompt 

643 

644 def get_recommendation( 

645 self, 

646 workload_description: str | None = None, 

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

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

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

650 ) -> BedrockCapacityRecommendation: 

651 """ 

652 Get an AI-powered capacity recommendation. 

653 

654 Args: 

655 workload_description: Description of the workload 

656 instance_types: List of instance types to consider 

657 regions: List of regions to consider 

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

659 

660 Returns: 

661 BedrockCapacityRecommendation with the AI's recommendation 

662 """ 

663 # Gather capacity data 

664 capacity_data = self.gather_capacity_data(instance_types, regions) 

665 

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

667 historical_context = self._gather_historical_context(capacity_data) 

668 

669 # Build prompt 

670 prompt = self._build_prompt( 

671 capacity_data, workload_description, requirements, historical_context 

672 ) 

673 

674 # Call Bedrock 

675 bedrock = self._get_bedrock_client() 

676 

677 try: 

678 # Use the Converse API for better compatibility across models 

679 response = bedrock.converse( 

680 modelId=self.model_id, 

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

682 **build_bedrock_converse_options( 

683 self.model_id, 

684 # Deliberately no maxTokens: the Converse default is the 

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

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

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

688 inference_config={"temperature": 0.1}, 

689 apply_default_reasoning=self._uses_default_model, 

690 ), 

691 ) 

692 

693 # Extended reasoning precedes the final answer with a 

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

695 response_text = extract_bedrock_converse_text(response) 

696 

697 # Parse JSON response 

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

699 json_start = response_text.find("{") 

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

701 if json_start >= 0 and json_end > json_start: 

702 json_str = response_text[json_start:json_end] 

703 result = json.loads(json_str) 

704 else: 

705 raise ValueError( 

706 "No JSON object found in the model response " 

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

708 ) 

709 

710 return BedrockCapacityRecommendation( 

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

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

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

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

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

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

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

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

719 raw_response=response_text, 

720 ) 

721 

722 except ClientError as e: 

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

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

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

726 raise_if_bedrock_ftu_form_error(e) 

727 if error_code == "AccessDeniedException": 

728 raise RuntimeError( 

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

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

731 ) from e 

732 if error_code == "ValidationException": 

733 raise RuntimeError( 

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

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

736 ) from e 

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

738 except json.JSONDecodeError as e: 

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

740 # fail after the response text was extracted. 

741 raise RuntimeError( 

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

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

744 ) from e 

745 except BedrockResponseTruncatedError: 

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

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

748 raise 

749 except Exception as e: 

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

751 

752 def _build_predict_prompt( 

753 self, 

754 instance_type: str, 

755 region: str, 

756 stats: dict[str, Any], 

757 patterns: dict[str, Any], 

758 ) -> str: 

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

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

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

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

763 lines = [ 

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

765 "", 

766 ( 

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

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

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

770 ), 

771 "", 

772 ( 

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

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

775 ), 

776 ] 

777 if spot: 

778 lines.append( 

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