Coverage for cli/capacity/multi_region.py: 93.37%
246 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Multi-region capacity checking and weighted scoring."""
3from __future__ import annotations
5import logging
6import statistics
7from dataclasses import dataclass, field
8from datetime import UTC, datetime, timedelta
9from typing import Any
11import boto3
12from botocore.exceptions import ClientError
14from cli.config import GCOConfig, get_config
16from .checker import CapacityChecker
18logger = logging.getLogger(__name__)
20_TELEMETRY_MISSING_SIGNAL_PENALTY = 1000.0
21_SCORED_TELEMETRY_SIGNALS = frozenset({"queue", "gpu"})
24def _missing_scored_signal_count(capacity: RegionCapacity) -> int:
25 """Return the number of unavailable signals used for placement scoring."""
26 return len(_SCORED_TELEMETRY_SIGNALS.intersection(capacity.unavailable_signals))
29@dataclass
30class RegionCapacity:
31 """Capacity information for a region."""
33 region: str
34 queue_depth: int = 0
35 pending_jobs: int = 0
36 running_jobs: int = 0
37 gpu_utilization: float = 0.0
38 cpu_utilization: float = 0.0
39 available_gpus: int = 0
40 total_gpus: int = 0
41 avg_wait_time_seconds: int = 0
42 recommendation_score: float = 0.0
43 telemetry_status: str = "unknown"
44 unavailable_signals: list[str] = field(default_factory=list)
45 telemetry_errors: list[str] = field(default_factory=list)
48class MultiRegionCapacityChecker:
49 """
50 Checks capacity across multiple GCO regions.
52 Provides:
53 - Multi-region capacity overview
54 - Intelligent region recommendation
55 - Queue depth analysis
56 - Resource utilization metrics
57 """
59 def __init__(self, config: GCOConfig | None = None):
60 self.config = config or get_config()
61 self._session = boto3.Session()
62 # Errors from the most recent get_all_regions_capacity() sweep, so an
63 # empty result can be distinguished as "checks failed" vs "no regions".
64 self._last_region_errors: list[str] = []
66 def get_region_capacity(self, region: str) -> RegionCapacity:
67 """Get capacity information while retaining telemetry uncertainty."""
68 from cli.aws_client import get_aws_client
70 aws_client = get_aws_client(self.config)
71 stack = aws_client.get_regional_stack(region)
72 capacity = RegionCapacity(region=region)
74 if not stack:
75 capacity.telemetry_status = "unavailable"
76 capacity.unavailable_signals = ["queue", "gpu", "cpu"]
77 capacity.telemetry_errors = ["Regional stack was not found"]
78 capacity.recommendation_score = (
79 len(_SCORED_TELEMETRY_SIGNALS) * _TELEMETRY_MISSING_SIGNAL_PENALTY
80 )
81 return capacity
83 available = {"queue": False, "gpu": False, "cpu": False}
85 # Get queue depth from SQS.
86 try:
87 cfn = self._session.client("cloudformation", region_name=region)
88 response = cfn.describe_stacks(StackName=stack.stack_name)
89 outputs = {
90 output["OutputKey"]: output["OutputValue"]
91 for output in response["Stacks"][0].get("Outputs", [])
92 }
94 queue_url = outputs.get("JobQueueUrl")
95 if queue_url:
96 sqs = self._session.client("sqs", region_name=region)
97 attrs = sqs.get_queue_attributes(
98 QueueUrl=queue_url,
99 AttributeNames=[
100 "ApproximateNumberOfMessages",
101 "ApproximateNumberOfMessagesNotVisible",
102 ],
103 )["Attributes"]
104 capacity.queue_depth = int(attrs.get("ApproximateNumberOfMessages", 0))
105 capacity.running_jobs = int(attrs.get("ApproximateNumberOfMessagesNotVisible", 0))
106 available["queue"] = True
107 else:
108 capacity.telemetry_errors.append("Job queue URL was not present in stack outputs")
109 except ClientError as e:
110 logger.debug("Failed to get queue metrics for %s: %s", region, e)
111 capacity.telemetry_errors.append(f"Queue telemetry failed: {e}")
112 except Exception as e:
113 logger.warning("Unexpected error getting queue metrics for %s: %s", region, e)
114 capacity.telemetry_errors.append(f"Queue telemetry failed: {e}")
116 # Query GPU and CPU independently so one failed metric does not erase
117 # the other useful signal.
118 try:
119 cloudwatch = self._session.client("cloudwatch", region_name=region)
120 except Exception as e:
121 logger.warning("Failed to create CloudWatch client for %s: %s", region, e)
122 capacity.telemetry_errors.append(f"CloudWatch telemetry failed: {e}")
123 else:
124 metric_specs = (
125 ("gpu", "node_gpu_utilization", "gpu_utilization"),
126 ("cpu", "node_cpu_utilization", "cpu_utilization"),
127 )
128 for signal, metric_name, attribute in metric_specs:
129 try:
130 response = cloudwatch.get_metric_statistics(
131 Namespace="ContainerInsights",
132 MetricName=metric_name,
133 Dimensions=[{"Name": "ClusterName", "Value": stack.cluster_name}],
134 StartTime=datetime.now(UTC) - timedelta(minutes=5),
135 EndTime=datetime.now(UTC),
136 Period=300,
137 Statistics=["Average"],
138 )
139 datapoints = response.get("Datapoints", [])
140 if datapoints:
141 setattr(capacity, attribute, datapoints[0]["Average"])
142 available[signal] = True
143 else:
144 capacity.telemetry_errors.append(
145 f"{signal.upper()} telemetry returned no datapoints"
146 )
147 except ClientError as e:
148 logger.debug("Failed to get %s metrics for %s: %s", signal, region, e)
149 capacity.telemetry_errors.append(f"{signal.upper()} telemetry failed: {e}")
150 except Exception as e:
151 logger.warning(
152 "Unexpected error getting %s metrics for %s: %s", signal, region, e
153 )
154 capacity.telemetry_errors.append(f"{signal.upper()} telemetry failed: {e}")
156 capacity.unavailable_signals = [
157 signal for signal, is_available in available.items() if not is_available
158 ]
159 if not capacity.unavailable_signals:
160 capacity.telemetry_status = "complete"
161 elif len(capacity.unavailable_signals) == len(available):
162 capacity.telemetry_status = "unavailable"
163 else:
164 capacity.telemetry_status = "partial"
166 base_score = (
167 capacity.queue_depth * 10 + capacity.gpu_utilization + capacity.running_jobs * 5
168 )
169 missing_scored_signals = _SCORED_TELEMETRY_SIGNALS.intersection(
170 capacity.unavailable_signals
171 )
172 capacity.recommendation_score = base_score + (
173 len(missing_scored_signals) * _TELEMETRY_MISSING_SIGNAL_PENALTY
174 )
175 return capacity
177 def get_all_regions_capacity(self) -> list[RegionCapacity]:
178 """Get capacity information for all deployed regions."""
179 from cli.aws_client import get_aws_client
181 self._last_region_errors = []
182 try:
183 aws_client = get_aws_client(self.config)
184 stacks = aws_client.discover_regional_stacks()
185 except Exception as e:
186 logger.warning("Failed to discover regional stacks: %s", e)
187 self._last_region_errors.append(f"region discovery: {e}")
188 return []
190 capacities = []
191 for region in stacks:
192 try:
193 capacity = self.get_region_capacity(region)
194 capacities.append(capacity)
195 except Exception as e:
196 logger.warning("Failed to get capacity for region %s: %s", region, e)
197 self._last_region_errors.append(f"{region}: {e}")
198 continue
200 return capacities
202 def recommend_region_for_job(
203 self,
204 gpu_required: bool = False,
205 min_gpus: int = 0,
206 instance_type: str | None = None,
207 gpu_count: int = 0,
208 ) -> dict[str, Any]:
209 """
210 Recommend the optimal region for job placement.
212 When instance_type is provided, uses weighted multi-signal scoring that
213 combines spot placement scores, pricing, queue depth, GPU utilization,
214 and running job counts. Falls back to simple scoring when instance_type
215 is not specified.
217 Args:
218 gpu_required: Whether the job requires GPUs
219 min_gpus: Minimum number of GPUs required
220 instance_type: Specific instance type for workload-aware scoring
221 gpu_count: Number of GPUs required
223 Returns:
224 Dictionary with recommended region and justification
225 """
226 capacities = self.get_all_regions_capacity()
228 if not capacities:
229 conservative_score = len(_SCORED_TELEMETRY_SIGNALS) * _TELEMETRY_MISSING_SIGNAL_PENALTY
230 if self._last_region_errors:
231 # The emptiness is due to underlying AWS failures, not an absence
232 # of configured regions — surface the real error instead of a
233 # benign "no capacity" message that masks it.
234 details = "; ".join(self._last_region_errors)
235 return {
236 "region": self.config.default_region,
237 "reason": f"Capacity checks failed for all regions: {details}",
238 "score": conservative_score,
239 "error": details,
240 "telemetry_status": "unavailable",
241 }
242 return {
243 "region": self.config.default_region,
244 "reason": "No capacity data available, using default region",
245 "score": conservative_score,
246 "telemetry_status": "unavailable",
247 }
249 # If every region lacks cluster telemetry, prefer the configured
250 # default as a deterministic fallback rather than treating zero-valued
251 # failed measurements as evidence of ideal capacity.
252 if all(cap.telemetry_status == "unavailable" for cap in capacities): 252 ↛ 253line 252 didn't jump to line 253 because the condition on line 252 was never true
253 capacities = sorted(
254 capacities,
255 key=lambda cap: (
256 cap.region != self.config.default_region,
257 cap.recommendation_score,
258 ),
259 )
261 # When instance_type is provided, use weighted scoring with capacity data
262 if instance_type:
263 return self._weighted_recommend(capacities, instance_type, gpu_count or min_gpus)
265 # Fallback: simple scoring (existing behavior)
266 return self._simple_recommend(capacities)
268 def _simple_recommend(self, capacities: list[RegionCapacity]) -> dict[str, Any]:
269 """Recommend observed capacity before comparing numeric load scores."""
270 all_unavailable = all(cap.telemetry_status == "unavailable" for cap in capacities)
271 sorted_capacities = sorted(
272 capacities,
273 key=lambda cap: (
274 cap.region != self.config.default_region if all_unavailable else False,
275 _missing_scored_signal_count(cap),
276 cap.recommendation_score,
277 cap.region,
278 ),
279 )
280 best = sorted_capacities[0]
282 reasons = []
283 if best.telemetry_errors: 283 ↛ 284line 283 didn't jump to line 284 because the condition on line 283 was never true
284 if best.telemetry_status == "unavailable":
285 reasons.append("capacity telemetry unavailable; using conservative fallback")
286 else:
287 reasons.append("capacity telemetry is partial")
289 if "queue" not in best.unavailable_signals: 289 ↛ 295line 289 didn't jump to line 295 because the condition on line 289 was always true
290 if best.queue_depth == 0:
291 reasons.append("empty queue")
292 elif best.queue_depth < 5:
293 reasons.append(f"low queue depth ({best.queue_depth})")
295 if "gpu" not in best.unavailable_signals: 295 ↛ 301line 295 didn't jump to line 301 because the condition on line 295 was always true
296 if best.gpu_utilization < 50:
297 reasons.append(f"{100 - best.gpu_utilization:.0f}% GPU available")
298 elif best.gpu_utilization < 80:
299 reasons.append(f"moderate GPU utilization ({best.gpu_utilization:.0f}%)")
301 if "queue" not in best.unavailable_signals: 301 ↛ 307line 301 didn't jump to line 307 because the condition on line 301 was always true
302 if best.running_jobs == 0:
303 reasons.append("no running jobs")
304 elif best.running_jobs < 5:
305 reasons.append(f"few running jobs ({best.running_jobs})")
307 reason = ", ".join(reasons) if reasons else "best overall capacity"
309 return {
310 "region": best.region,
311 "reason": reason,
312 "score": best.recommendation_score,
313 "queue_depth": best.queue_depth,
314 "gpu_utilization": best.gpu_utilization,
315 "running_jobs": best.running_jobs,
316 "telemetry_status": best.telemetry_status,
317 "telemetry_errors": best.telemetry_errors or None,
318 "all_regions": [
319 {
320 "region": c.region,
321 "score": c.recommendation_score,
322 "queue_depth": c.queue_depth,
323 "gpu_utilization": c.gpu_utilization,
324 "telemetry_status": c.telemetry_status,
325 }
326 for c in sorted_capacities
327 ],
328 }
330 def _weighted_recommend(
331 self,
332 capacities: list[RegionCapacity],
333 instance_type: str,
334 gpu_count: int = 0,
335 ) -> dict[str, Any]:
336 """
337 Workload-aware recommendation using weighted multi-signal scoring.
339 Gathers per-region capacity data for the specific instance type and
340 combines it with cluster metrics using weighted scoring.
341 """
342 capacity_checker = CapacityChecker(self.config)
344 scored_regions: list[dict[str, Any]] = []
346 for cap in capacities:
347 region = cap.region
349 # Gather instance-specific signals for this region
350 spot_score = 0.0
351 spot_price_ratio = 1.0 # spot/on-demand ratio (lower = better savings)
353 try:
354 placement_scores = capacity_checker.get_spot_placement_score(
355 instance_type, region, target_capacity=max(1, gpu_count)
356 )
357 if placement_scores:
358 spot_score = placement_scores.get("regional", 0) / 10.0 # Normalize to 0-1
359 except Exception as e:
360 logger.debug(
361 "Failed to get spot placement score for %s in %s: %s", instance_type, region, e
362 )
364 try:
365 spot_prices = capacity_checker.get_spot_price_history(instance_type, region)
366 on_demand_price = capacity_checker.get_on_demand_price(instance_type, region)
368 if spot_prices and on_demand_price and on_demand_price > 0:
369 avg_spot = statistics.mean(sp.current_price for sp in spot_prices)
370 spot_price_ratio = avg_spot / on_demand_price
371 except Exception as e:
372 logger.debug(
373 "Failed to get spot pricing for %s in %s: %s", instance_type, region, e
374 )
376 # Capacity Block trend — compares near-term vs far-term offering
377 # density to detect whether AWS is adding or consuming capacity
378 # in this region for the requested instance type.
379 try:
380 cb_trend = capacity_checker.get_capacity_block_trend(instance_type, region)
381 except Exception:
382 cb_trend = 0.0
384 weighted_score = compute_weighted_score(
385 spot_placement_score=spot_score,
386 spot_price_ratio=spot_price_ratio,
387 queue_depth=cap.queue_depth,
388 gpu_utilization=cap.gpu_utilization,
389 running_jobs=cap.running_jobs,
390 capacity_block_trend=cb_trend,
391 )
392 missing_scored_signals = _missing_scored_signal_count(cap)
394 scored_regions.append(
395 {
396 "region": region,
397 "score": weighted_score,
398 "queue_depth": cap.queue_depth,
399 "gpu_utilization": cap.gpu_utilization,
400 "running_jobs": cap.running_jobs,
401 "spot_placement_score": spot_score,
402 "spot_price_ratio": spot_price_ratio,
403 "capacity_block_trend": cb_trend,
404 "telemetry_status": cap.telemetry_status,
405 "telemetry_errors": cap.telemetry_errors or None,
406 "unavailable_scored_signals": missing_scored_signals,
407 }
408 )
410 scored_regions.sort(
411 key=lambda item: (
412 item["unavailable_scored_signals"],
413 item["score"],
414 item["region"],
415 )
416 )
417 best = scored_regions[0]
419 # Build justification from the signals
420 reasons = []
421 if best.get("telemetry_errors"): 421 ↛ 422line 421 didn't jump to line 422 because the condition on line 421 was never true
422 if best["telemetry_status"] == "unavailable":
423 reasons.append("cluster telemetry unavailable")
424 else:
425 reasons.append("cluster telemetry is partial")
427 if best["spot_placement_score"] >= 0.7:
428 reasons.append(f"high spot availability ({best['spot_placement_score']:.0%})")
429 elif best["spot_placement_score"] >= 0.4:
430 reasons.append(f"moderate spot availability ({best['spot_placement_score']:.0%})")
432 if best["spot_price_ratio"] < 0.5:
433 reasons.append(f"good spot savings ({1 - best['spot_price_ratio']:.0%} off on-demand)")
435 if "queue" not in next( 435 ↛ 443line 435 didn't jump to line 443 because the condition on line 435 was always true
436 cap.unavailable_signals for cap in capacities if cap.region == best["region"]
437 ):
438 if best["queue_depth"] == 0:
439 reasons.append("empty queue")
440 elif best["queue_depth"] < 5:
441 reasons.append(f"low queue depth ({best['queue_depth']})")
443 if (
444 "gpu"
445 not in next(
446 cap.unavailable_signals for cap in capacities if cap.region == best["region"]
447 )
448 and best["gpu_utilization"] < 50
449 ):
450 reasons.append(f"{100 - best['gpu_utilization']:.0f}% GPU available")
452 if "queue" not in next( 452 ↛ 460line 452 didn't jump to line 460 because the condition on line 452 was always true
453 cap.unavailable_signals for cap in capacities if cap.region == best["region"]
454 ):
455 if best["running_jobs"] == 0:
456 reasons.append("no running jobs")
457 elif best["running_jobs"] < 5:
458 reasons.append(f"few running jobs ({best['running_jobs']})")
460 if best.get("capacity_block_trend", 0) > 0.2:
461 reasons.append("capacity block availability trending up")
462 elif best.get("capacity_block_trend", 0) < -0.2:
463 reasons.append("capacity block availability trending down")
465 reason = ", ".join(reasons) if reasons else f"best weighted score for {instance_type}"
467 return {
468 "region": best["region"],
469 "reason": reason,
470 "score": best["score"],
471 "queue_depth": best["queue_depth"],
472 "gpu_utilization": best["gpu_utilization"],
473 "running_jobs": best["running_jobs"],
474 "instance_type": instance_type,
475 "scoring_method": "weighted",
476 "telemetry_status": best["telemetry_status"],
477 "telemetry_errors": best["telemetry_errors"],
478 "all_regions": scored_regions,
479 }
482def compute_price_trend(prices: list[float]) -> dict[str, Any]:
483 """
484 Compute a linear regression trend over a price time series.
486 Prices are assumed to be ordered most-recent-first (as returned by
487 the EC2 spot price history API). The series is reversed internally
488 so the slope represents change over time (positive = prices rising).
490 Args:
491 prices: List of price points, most recent first.
493 Returns:
494 Dict with:
495 slope: price change per sample period (positive = rising)
496 normalized_slope: slope / mean_price (scale-independent, -1 to 1 clamped)
497 price_changes: number of distinct price transitions (proxy for volatility)
498 direction: "rising", "falling", or "stable"
499 """
500 if len(prices) < 2:
501 return {
502 "slope": 0.0,
503 "normalized_slope": 0.0,
504 "price_changes": 0,
505 "direction": "stable",
506 }
508 # Reverse so index 0 = oldest, index N = newest
509 series = list(reversed(prices))
510 n = len(series)
512 # Count distinct price transitions (proxy for interruption frequency)
513 price_changes = sum(1 for i in range(1, n) if series[i] != series[i - 1])
515 # Linear regression: slope of price over time
516 x_mean = (n - 1) / 2.0
517 y_mean = statistics.mean(series)
519 numerator = sum((i - x_mean) * (series[i] - y_mean) for i in range(n))
520 denominator = sum((i - x_mean) ** 2 for i in range(n))
522 if denominator == 0 or y_mean == 0:
523 return {
524 "slope": 0.0,
525 "normalized_slope": 0.0,
526 "price_changes": price_changes,
527 "direction": "stable",
528 }
530 slope = numerator / denominator
531 normalized = max(-1.0, min(1.0, slope / y_mean))
533 if normalized > 0.05:
534 direction = "rising"
535 elif normalized < -0.05:
536 direction = "falling"
537 else:
538 direction = "stable"
540 return {
541 "slope": round(slope, 6),
542 "normalized_slope": round(normalized, 4),
543 "price_changes": price_changes,
544 "direction": direction,
545 }
548def compute_weighted_score(
549 spot_placement_score: float = 0.0,
550 spot_price_ratio: float = 1.0,
551 queue_depth: int = 0,
552 gpu_utilization: float = 0.0,
553 running_jobs: int = 0,
554 capacity_block_trend: float = 0.0,
555 *,
556 weights: dict[str, float] | None = None,
557) -> float:
558 """
559 Compute a weighted recommendation score for a region. Lower is better.
561 Combines multiple capacity signals into a single score using configurable
562 weights. Each signal is normalized to 0-1 where 0 is best, then weighted.
564 Args:
565 spot_placement_score: Normalized spot placement score (0-1, higher = better availability)
566 spot_price_ratio: Spot price / on-demand price (0-1, lower = better savings)
567 queue_depth: Number of pending jobs in the region's queue
568 gpu_utilization: GPU utilization percentage (0-100)
569 running_jobs: Number of currently running jobs
570 capacity_block_trend: Trend of capacity block offerings over a 26-week window
571 (-1 to 1). Positive means capacity is growing (regression slope positive),
572 negative means shrinking. Derived from linear regression over weekly
573 offering counts from the describe-capacity-block-offerings API.
574 weights: Optional custom weights dict. Keys: spot_placement, spot_price,
575 queue_depth, gpu_utilization, running_jobs, capacity_blocks.
576 Values should sum to 1.0.
578 Returns:
579 Weighted score (lower is better, range roughly 0-1)
580 """
581 w = weights or {
582 "spot_placement": 0.25,
583 "spot_price": 0.20,
584 "queue_depth": 0.20,
585 "gpu_utilization": 0.15,
586 "running_jobs": 0.10,
587 "capacity_blocks": 0.10,
588 }
590 # Normalize each signal to 0-1 where 0 is best
591 # Spot placement: invert (high score = good, so 1 - score = low = good)
592 norm_spot = 1.0 - min(max(spot_placement_score, 0.0), 1.0)
594 # Spot price ratio: already 0-1 where lower is better
595 norm_price = min(max(spot_price_ratio, 0.0), 1.0)
597 # Queue depth: normalize with diminishing returns (0 = best)
598 # Use tanh-like curve: depth / (depth + k) where k controls sensitivity
599 norm_queue = queue_depth / (queue_depth + 10.0) if queue_depth >= 0 else 0.0
601 # GPU utilization: normalize 0-100 to 0-1
602 norm_gpu = min(max(gpu_utilization, 0.0), 100.0) / 100.0
604 # Running jobs: normalize with diminishing returns
605 norm_jobs = running_jobs / (running_jobs + 20.0) if running_jobs >= 0 else 0.0
607 # Capacity block trend: invert and shift from [-1,1] to [0,1]
608 # trend +1 (growing) → 0.0 (best), trend -1 (shrinking) → 1.0 (worst)
609 clamped_trend = min(max(capacity_block_trend, -1.0), 1.0)
610 norm_blocks = (1.0 - clamped_trend) / 2.0
612 score = (
613 w["spot_placement"] * norm_spot
614 + w["spot_price"] * norm_price
615 + w["queue_depth"] * norm_queue
616 + w["gpu_utilization"] * norm_gpu
617 + w["running_jobs"] * norm_jobs
618 + w.get("capacity_blocks", 0) * norm_blocks
619 )
621 return round(score, 4)
624def get_multi_region_capacity_checker(
625 config: GCOConfig | None = None,
626) -> MultiRegionCapacityChecker:
627 """Get a configured multi-region capacity checker instance."""
628 return MultiRegionCapacityChecker(config)