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
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Bedrock-powered AI capacity advisor."""
3from __future__ import annotations
5import json
6import logging
7from dataclasses import dataclass, field
8from datetime import UTC, datetime, timedelta
9from typing import Any
11import boto3
12from botocore.config import Config
13from botocore.exceptions import ClientError
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)
25from .checker import CapacityChecker
26from .multi_region import MultiRegionCapacityChecker, compute_price_trend
28logger = logging.getLogger(__name__)
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 "")
37@dataclass
38class BedrockCapacityRecommendation:
39 """AI-generated capacity recommendation from Bedrock."""
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 = ""
52@dataclass
53class CapacityPredictionResult:
54 """Bedrock prediction of the best time(s) to acquire capacity."""
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 = ""
65class _CapacityAdvisorModelDefault:
66 """Lazily expose the advisor's canonical model default as a string."""
68 def __get__(self, instance: object, owner: type[Any] | None = None) -> str:
69 return get_default_capacity_advisor_model_id()
72class BedrockCapacityAdvisor:
73 """
74 AI-powered capacity advisor using Amazon Bedrock.
76 Gathers comprehensive capacity data and uses an LLM to provide
77 intelligent recommendations for workload placement.
79 DISCLAIMER: Recommendations are AI-generated and should be validated
80 before making production decisions.
81 """
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()
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
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 )
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.
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)
118 Returns:
119 Dictionary containing all gathered capacity data
120 """
121 from cli.aws_client import get_aws_client
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 ]
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]
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 }
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)
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"] = []
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 )
205 for instance_type in instance_types:
206 data["spot_data"][instance_type] = {}
207 data["on_demand_data"][instance_type] = {}
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
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 )
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
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 )
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 )
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 )
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)
367 return data
369 def _gather_historical_context(self, capacity_data: dict[str, Any]) -> dict[str, Any]:
370 """Best-effort historical enrichment for the Bedrock prompt.
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
381 store = get_capacity_history_store()
382 except Exception as e:
383 logger.debug("Capacity history store unavailable: %s", e)
384 return {}
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
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 {}
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.
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
437"""
439 if workload_description:
440 prompt += f"WORKLOAD DESCRIPTION:\n{workload_description}\n\n"
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"
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 )
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"
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"
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"
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"
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"
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"
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"
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"
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"
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}
643Respond ONLY with the JSON object, no additional text."""
645 return prompt
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.
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.)
663 Returns:
664 BedrockCapacityRecommendation with the AI's recommendation
665 """
666 # Gather capacity data
667 capacity_data = self.gather_capacity_data(instance_types, regions)
669 # Gather best-effort historical context (skipped when unavailable)
670 historical_context = self._gather_historical_context(capacity_data)
672 # Build prompt
673 prompt = self._build_prompt(
674 capacity_data, workload_description, requirements, historical_context
675 )
677 # Call Bedrock
678 bedrock = self._get_bedrock_client()
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 )
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)
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 )
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 )
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
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)
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.
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
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)
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)
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 )
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.
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
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
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)