Coverage for gco_mcp/tools/capacity.py: 100.00%
181 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Capacity checking and recommendation MCP tools."""
3import cli_runner
4from audit import audit_logged
5from feature_flags import FLAG_CAPACITY_PURCHASE, FLAG_DESTRUCTIVE_OPERATIONS, is_enabled
6from server import mcp
9@mcp.tool(tags={"safe", "capacity"})
10@audit_logged
11def check_capacity(instance_type: str, region: str) -> str:
12 """Check spot and on-demand capacity for a specific instance type.
14 Args:
15 instance_type: EC2 instance type (e.g. g4dn.xlarge, g5.2xlarge, p4d.24xlarge).
16 region: AWS region to check.
17 """
18 return cli_runner._run_cli("capacity", "check", "-i", instance_type, "-r", region)
21@mcp.tool(tags={"safe", "capacity"})
22@audit_logged
23def instance_info(instance_type: str) -> str:
24 """Get hardware and pricing metadata for an EC2 instance type.
26 Args:
27 instance_type: EC2 instance type (for example, g5.2xlarge or p5.48xlarge).
28 """
29 return cli_runner._run_cli("capacity", "instance-info", instance_type)
32@mcp.tool(tags={"safe", "capacity"})
33@audit_logged
34def recommend_capacity(
35 instance_type: str,
36 region: str,
37 fault_tolerance: str = "medium",
38) -> str:
39 """Recommend spot or on-demand capacity for a workload.
41 Args:
42 instance_type: EC2 instance type to evaluate.
43 region: AWS region in which the workload will run.
44 fault_tolerance: Interruption tolerance: ``high``, ``medium``, or ``low``.
45 """
46 return cli_runner._run_cli(
47 "capacity",
48 "recommend",
49 "-i",
50 instance_type,
51 "-r",
52 region,
53 "-f",
54 fault_tolerance,
55 )
58@mcp.tool(tags={"safe", "capacity"})
59@audit_logged
60def capacity_status(region: str | None = None) -> str:
61 """View capacity status across all deployed regions.
63 Args:
64 region: Specific region, or omit for all regions.
65 """
66 args = ["capacity", "status"]
67 if region:
68 args += ["-r", region]
69 return cli_runner._run_cli(*args)
72@mcp.tool(tags={"safe", "capacity"})
73@audit_logged
74def recommend_region(
75 gpu: bool = False, instance_type: str | None = None, gpu_count: int = 0
76) -> str:
77 """Get optimal region recommendation based on capacity.
79 Args:
80 gpu: Whether the workload requires GPUs.
81 instance_type: Specific instance type to check. When provided, uses weighted
82 multi-signal scoring (spot placement scores, pricing, queue depth, etc.).
83 gpu_count: Number of GPUs required for the workload.
84 """
85 args = ["capacity", "recommend-region"]
86 if gpu:
87 args.append("--gpu")
88 if instance_type:
89 args += ["-i", instance_type]
90 if gpu_count:
91 args += ["--gpu-count", str(gpu_count)]
92 return cli_runner._run_cli(*args)
95@mcp.tool(tags={"safe", "capacity"})
96@audit_logged
97def spot_prices(instance_type: str, region: str) -> str:
98 """Get current spot prices for an instance type.
100 Args:
101 instance_type: EC2 instance type.
102 region: AWS region.
103 """
104 return cli_runner._run_cli("capacity", "spot-prices", "-i", instance_type, "-r", region)
107@mcp.tool(tags={"safe", "capacity"})
108@audit_logged
109def ai_recommend(
110 workload: str,
111 instance_type: str | None = None,
112 region: str | None = None,
113 gpu: bool = False,
114 min_gpus: int = 0,
115 min_memory_gb: int = 0,
116 fault_tolerance: str = "low",
117 max_cost: float | None = None,
118 model: str | None = None,
119) -> str:
120 """Get AI-powered capacity recommendation using Amazon Bedrock.
122 Gathers comprehensive capacity data (spot scores, pricing, cluster
123 utilization, queue depth) and sends it to an LLM for analysis.
124 Returns a recommended region, instance type, capacity type, and reasoning.
126 Requires AWS credentials with bedrock:InvokeModel permission and the
127 specified model enabled in your account.
129 Args:
130 workload: Description of the workload (e.g. "Fine-tuning a 20B parameter LLM").
131 instance_type: Specific instance type(s) to consider (e.g. "p4d.24xlarge").
132 region: Specific region(s) to consider (e.g. "us-east-1").
133 gpu: Whether the workload requires GPUs.
134 min_gpus: Minimum number of GPUs required.
135 min_memory_gb: Minimum GPU memory in GB.
136 fault_tolerance: Tolerance for interruptions ("low", "medium", "high").
137 max_cost: Maximum acceptable cost per hour in USD.
138 model: Bedrock model ID to use for analysis. Omit to use
139 `cdk.json` `context.bedrock.default_model_id`. Mission and the
140 capacity advisor consume that one shared configuration value.
141 """
142 args = ["capacity", "ai-recommend", "-w", workload]
143 if instance_type:
144 args += ["-i", instance_type]
145 if region:
146 args += ["-r", region]
147 if gpu:
148 args.append("--gpu")
149 if min_gpus > 0:
150 args += ["--min-gpus", str(min_gpus)]
151 if min_memory_gb > 0:
152 args += ["--min-memory-gb", str(min_memory_gb)]
153 if fault_tolerance != "low":
154 args += ["--fault-tolerance", fault_tolerance]
155 if max_cost is not None:
156 args += ["--max-cost", str(max_cost)]
157 if model:
158 args += ["--model", model]
159 return cli_runner._run_cli(*args)
162@mcp.tool(tags={"safe", "capacity"})
163@audit_logged
164def list_reservations(
165 instance_type: str | None = None,
166 region: str | None = None,
167) -> str:
168 """List On-Demand Capacity Reservations (ODCRs) across regions.
170 Shows all active capacity reservations with utilization details.
172 Args:
173 instance_type: Filter by instance type (e.g. p5.48xlarge).
174 region: Filter by specific region.
175 """
176 args = ["capacity", "reservations"]
177 if instance_type:
178 args += ["-i", instance_type]
179 if region:
180 args += ["-r", region]
181 return cli_runner._run_cli(*args)
184@mcp.tool(tags={"safe", "capacity"})
185@audit_logged
186def reservation_check(
187 instance_type: str,
188 regions: list[str] | None = None,
189 count: int = 1,
190 include_blocks: bool = True,
191 block_duration: int = 24,
192 block_duration_days: int | None = None,
193 earliest_start: str | None = None,
194 latest_start: str | None = None,
195) -> str:
196 """Check ODCR and Capacity Block availability for an instance type.
198 Checks existing On-Demand Capacity Reservations and purchasable Capacity
199 Blocks for ML. By default it searches a 24h block soonest-available, but you
200 can widen the search: pass several regions to fan out in parallel, set a
201 start-date window (earliest_start / latest_start) to ask for blocks starting
202 near a date, and set the block duration in hours or days. For a full
203 duration-range sweep that returns one ranked, de-duplicated report, use
204 find_capacity_blocks instead.
206 Args:
207 instance_type: GPU instance type (e.g. p4d.24xlarge, p5.48xlarge, p6-b200).
208 regions: Regions to check in parallel (any regions, not just deployed);
209 omit to check all deployed regions.
210 count: Minimum number of instances needed.
211 include_blocks: Whether to include Capacity Block offerings.
212 block_duration: Capacity Block duration in hours (default 24).
213 block_duration_days: Capacity Block duration in days (overrides hours).
214 earliest_start: Earliest block start date (YYYY-MM-DD or ISO datetime).
215 latest_start: Latest block start date (YYYY-MM-DD or ISO datetime).
216 """
217 args = ["capacity", "reservation-check", "-i", instance_type, "-c", str(count)]
218 for r in regions or []:
219 args += ["-r", r]
220 if not include_blocks:
221 args.append("--no-blocks")
222 if block_duration != 24:
223 args += ["--block-duration", str(block_duration)]
224 if block_duration_days is not None:
225 args += ["--block-duration-days", str(block_duration_days)]
226 if earliest_start:
227 args += ["--earliest-start", earliest_start]
228 if latest_start:
229 args += ["--latest-start", latest_start]
230 return cli_runner._run_cli(*args)
233@mcp.tool(tags={"safe", "capacity"})
234@audit_logged
235def find_capacity_blocks(
236 instance_type: str,
237 regions: list[str] | None = None,
238 count: int = 1,
239 duration_days: int | None = None,
240 duration_hours: int | None = None,
241 min_duration_days: int | None = None,
242 max_duration_days: int | None = None,
243 min_duration_hours: int | None = None,
244 max_duration_hours: int | None = None,
245 earliest_start: str | None = None,
246 latest_start: str | None = None,
247 find_longest: bool = False,
248) -> str:
249 """Find EC2 Capacity Blocks across regions x durations x a start-date window.
251 This is the one-call sweep for "where and when can I get N of this GPU
252 instance for D days?". It searches every requested region and every valid
253 Capacity Block duration in the range, in parallel, then returns a single
254 consolidated, de-duplicated, ranked report (cheapest per-GPU-hour first) with
255 per-hour and per-GPU-hour pricing and the longest available block.
257 AWS allows durations in 1-day increments up to 14 days, then 7-day increments
258 up to 182 days; a duration range is expanded to those discrete values
259 automatically. Friendly names are normalized (p6-b200 -> p6-b200.48xlarge,
260 p6-b300 -> p6-b300.48xlarge), and UltraServer-only families (the Grace-
261 Blackwell GB200/GB300 superchips / P6e-GB UltraServers) are flagged rather
262 than silently returning nothing.
264 Args:
265 instance_type: GPU instance type or alias (e.g. p6-b200, p5.48xlarge).
266 regions: Regions to search (any regions; defaults to deployed regions).
267 count: Instances per block.
268 duration_days / duration_hours: A single target duration.
269 min_duration_days / max_duration_days: Duration range bounds in days.
270 min_duration_hours / max_duration_hours: Duration range bounds in hours.
271 earliest_start: Earliest block start (YYYY-MM-DD or ISO datetime).
272 latest_start: Latest block start (YYYY-MM-DD or ISO datetime).
273 find_longest: Sweep the duration ladder and surface the longest block.
274 """
275 args = ["capacity", "find-blocks", "-i", instance_type]
276 for r in regions or []:
277 args += ["-r", r]
278 if count != 1:
279 args += ["-c", str(count)]
280 if duration_days is not None:
281 args += ["--duration-days", str(duration_days)]
282 if duration_hours is not None:
283 args += ["--duration-hours", str(duration_hours)]
284 if min_duration_days is not None:
285 args += ["--min-duration-days", str(min_duration_days)]
286 if max_duration_days is not None:
287 args += ["--max-duration-days", str(max_duration_days)]
288 if min_duration_hours is not None:
289 args += ["--min-duration-hours", str(min_duration_hours)]
290 if max_duration_hours is not None:
291 args += ["--max-duration-hours", str(max_duration_hours)]
292 if earliest_start:
293 args += ["--earliest-start", earliest_start]
294 if latest_start:
295 args += ["--latest-start", latest_start]
296 if find_longest:
297 args.append("--find-longest")
298 return cli_runner._run_cli(*args)
301@mcp.tool(tags={"safe", "capacity"})
302@audit_logged
303def capacity_history_show(instance_type: str, region: str, hours: int = 168) -> str:
304 """Show the recorded capacity time-series for an instance type in a region.
306 Requires the historical capacity surface (an optional add-on to the global
307 stack, enabled by default). Returns spot score, spot price, AZ coverage,
308 queue depth, and capacity-block availability over the window.
310 Args:
311 instance_type: EC2 instance type (e.g. g5.xlarge, p5.48xlarge).
312 region: AWS region.
313 hours: Hours of history to show (default 168 = 7 days).
314 """
315 return cli_runner._run_cli(
316 "capacity", "history", "show", "-i", instance_type, "-r", region, "-H", str(hours)
317 )
320@mcp.tool(tags={"safe", "capacity"})
321@audit_logged
322def capacity_history_stats(instance_type: str, region: str, hours: int = 168) -> str:
323 """Show p25/p50/p75/min/max/stddev per capacity metric over a time window.
325 Args:
326 instance_type: EC2 instance type.
327 region: AWS region.
328 hours: Hours of history to summarize (default 168 = 7 days).
329 """
330 return cli_runner._run_cli(
331 "capacity", "history", "stats", "-i", instance_type, "-r", region, "-H", str(hours)
332 )
335@mcp.tool(tags={"safe", "capacity"})
336@audit_logged
337def capacity_history_patterns(instance_type: str, region: str, hours: int = 168) -> str:
338 """Show a day-of-week by hour heatmap of average spot placement scores.
340 Args:
341 instance_type: EC2 instance type.
342 region: AWS region.
343 hours: Hours of history to analyze (default 168 = 7 days).
344 """
345 return cli_runner._run_cli(
346 "capacity", "history", "patterns", "-i", instance_type, "-r", region, "-H", str(hours)
347 )
350@mcp.tool(tags={"safe", "capacity"})
351@audit_logged
352def capacity_predict(
353 instance_type: str,
354 region: str | None = None,
355 hours: int = 168,
356 all_regions: bool = False,
357) -> str:
358 """Predict the best time to acquire capacity from historical patterns (Bedrock).
360 Uses the historical capacity surface plus Amazon Bedrock to recommend the
361 day/hour windows with the best spot availability and pricing.
363 Args:
364 instance_type: EC2 instance type.
365 region: AWS region. Omit when all_regions is true.
366 hours: Hours of history to analyze (default 168 = 7 days).
367 all_regions: Predict across every region that has data for the type.
368 """
369 args = ["capacity", "predict", "-i", instance_type, "-H", str(hours)]
370 if all_regions:
371 args.append("--all-regions")
372 elif region:
373 args += ["-r", region]
374 return cli_runner._run_cli(*args)
377@mcp.tool(tags={"safe", "capacity"})
378@audit_logged
379def find_capacity_reservations(
380 instance_type: str | None = None,
381 regions: list[str] | None = None,
382 count: int = 1,
383 state: str = "active",
384 pricing: bool = True,
385) -> str:
386 """Find existing ODCRs across regions in one parallel, ranked report.
388 The ODCR counterpart to find_capacity_blocks: it searches every requested
389 region in parallel, normalizes friendly instance-type aliases (p6-b200 ->
390 p6-b200.48xlarge), enriches each reservation with On-Demand pricing, and
391 ranks them most-available-first (then cheapest per-GPU-hour). Use this to
392 answer "where do I already have free reserved capacity?" rather than
393 list_reservations' plain per-region aggregation.
395 Args:
396 instance_type: Instance type or alias to filter by (e.g. p6-b200); omit
397 to return every reservation.
398 regions: Regions to search in parallel (any regions; defaults to deployed).
399 count: Minimum available instances to consider the search satisfied.
400 state: Reservation state filter ("active" default; "all" for any state).
401 pricing: Enrich reservations with On-Demand pricing.
402 """
403 args = ["capacity", "find-reservations"]
404 if instance_type:
405 args += ["-i", instance_type]
406 for r in regions or []:
407 args += ["-r", r]
408 if count != 1:
409 args += ["-c", str(count)]
410 if state != "active":
411 args += ["--state", state]
412 if not pricing:
413 args.append("--no-pricing")
414 return cli_runner._run_cli(*args)
417# Capacity purchasing / creation — disabled by default.
418# Set GCO_ENABLE_CAPACITY_PURCHASE=true to enable.
419if is_enabled(FLAG_CAPACITY_PURCHASE):
421 @mcp.tool(tags={"cost-incurring", "capacity"})
422 @audit_logged
423 def reserve_capacity(
424 offering_id: str,
425 region: str,
426 dry_run: bool = False,
427 ) -> str:
428 """Purchase a Capacity Block offering by its ID.
430 Use reservation_check first to find available offerings and their IDs,
431 then purchase with this tool. Use dry_run=True to validate without purchasing.
433 Args:
434 offering_id: Capacity Block offering ID (cb-xxx) from reservation_check.
435 region: AWS region where the offering exists.
436 dry_run: If True, validate the offering without purchasing (no cost).
437 """
438 args = ["capacity", "reserve", "-o", offering_id, "-r", region]
439 if dry_run:
440 args.append("--dry-run")
441 return cli_runner._run_cli(*args)
443 @mcp.tool(tags={"cost-incurring", "capacity"})
444 @audit_logged
445 def create_reservation(
446 instance_type: str,
447 region: str,
448 availability_zone: str,
449 count: int = 1,
450 platform: str = "Linux/UNIX",
451 tenancy: str = "default",
452 match_criteria: str = "open",
453 end_date: str | None = None,
454 ebs_optimized: bool = False,
455 dry_run: bool = False,
456 ) -> str:
457 """Create a new On-Demand Capacity Reservation (ODCR).
459 The ODCR counterpart to reserve_capacity. Reserves On-Demand capacity for
460 an instance type in a specific AZ. Charges accrue for the reserved
461 capacity until it is cancelled — use dry_run=True to validate first.
463 Args:
464 instance_type: EC2 instance type or alias (e.g. p5.48xlarge, p6-b200).
465 region: AWS region.
466 availability_zone: Target AZ (e.g. us-east-1a).
467 count: Number of instances to reserve.
468 platform: Instance platform/OS (default "Linux/UNIX").
469 tenancy: "default" or "dedicated".
470 match_criteria: "open" or "targeted".
471 end_date: Optional end date (YYYY-MM-DD or ISO); omit for unlimited.
472 ebs_optimized: Reserve EBS-optimized capacity.
473 dry_run: If True, validate without creating (no cost).
474 """
475 args = [
476 "capacity",
477 "create-reservation",
478 "-i",
479 instance_type,
480 "-r",
481 region,
482 "-z",
483 availability_zone,
484 "-c",
485 str(count),
486 ]
487 if platform != "Linux/UNIX":
488 args += ["--platform", platform]
489 if tenancy != "default":
490 args += ["--tenancy", tenancy]
491 if match_criteria != "open":
492 args += ["--match-criteria", match_criteria]
493 if end_date:
494 args += ["--end-date", end_date]
495 if ebs_optimized:
496 args.append("--ebs-optimized")
497 if dry_run:
498 args.append("--dry-run")
499 return cli_runner._run_cli(*args)
502# Capacity reservation cancellation — destructive, disabled by default.
503# Set GCO_ENABLE_DESTRUCTIVE_OPERATIONS=true to enable.
504if is_enabled(FLAG_DESTRUCTIVE_OPERATIONS):
506 @mcp.tool(tags={"destructive", "capacity"})
507 @audit_logged
508 def cancel_reservation(
509 reservation_id: str,
510 region: str,
511 dry_run: bool = False,
512 ) -> str:
513 """[gated by GCO_ENABLE_DESTRUCTIVE_OPERATIONS] destructive.
515 Cancel an On-Demand Capacity Reservation, releasing its capacity.
517 Stops On-Demand charges for the reserved capacity. Only ODCRs can be
518 cancelled; a Capacity Block runs for its fixed term. Instances already
519 running against the reservation are not terminated.
521 Args:
522 reservation_id: Capacity Reservation ID (cr-xxx) to cancel.
523 region: AWS region where the reservation exists.
524 dry_run: If True, validate the cancellation without cancelling.
525 """
526 args = ["capacity", "cancel-reservation", "-o", reservation_id, "-r", region, "-y"]
527 if dry_run:
528 args.append("--dry-run")
529 return cli_runner._run_cli(*args)