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