Coverage for cli/costs.py: 99.22%
210 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-30 21:22 +0000
1"""Cost visibility for GCO workloads.
3Uses AWS Cost Explorer for historical spend and the Pricing API
4for real-time cost estimates on running workloads.
5"""
7from __future__ import annotations
9import logging
10from collections.abc import Sequence
11from dataclasses import dataclass, field
12from datetime import UTC, datetime, timedelta
13from typing import Any
15import boto3
17from .config import GCOConfig, get_config
19logger = logging.getLogger(__name__)
21# Cost allocation tag keys GCO reporting relies on. ``Project`` is the
22# user-defined tag applied to every CDK-managed resource (cdk.json
23# ``context.tags``) and is the key every Cost Explorer query in this module
24# filters on; ``aws:eks:cluster-name`` is the AWS-generated tag EKS Auto
25# Mode stamps on the EC2 capacity it launches outside CloudFormation, which
26# is what makes per-cluster attribution possible for Karpenter-provisioned
27# GPU spend.
28DEFAULT_COST_ALLOCATION_TAG_KEYS: tuple[str, ...] = ("Project", "aws:eks:cluster-name")
30# UpdateCostAllocationTagsStatus accepts at most 20 entries per call and
31# ListCostAllocationTags at most 20 TagKeys filters; batch both.
32_COST_ALLOCATION_TAG_BATCH_SIZE = 20
35@dataclass
36class ResourceCost:
37 """Cost for a single resource or service."""
39 service: str
40 amount: float
41 currency: str = "USD"
42 region: str | None = None
43 detail: str | None = None
46@dataclass
47class CostSummary:
48 """Aggregated cost summary."""
50 total: float
51 currency: str = "USD"
52 period_start: str = ""
53 period_end: str = ""
54 by_service: list[ResourceCost] = field(default_factory=list)
55 by_region: dict[str, float] = field(default_factory=dict)
58@dataclass
59class WorkloadCost:
60 """Estimated cost for a running workload."""
62 name: str
63 workload_type: str # "job" or "inference"
64 instance_type: str
65 gpu_count: int
66 hourly_rate: float
67 runtime_hours: float
68 estimated_cost: float
69 region: str
70 status: str
73class CostTracker:
74 """Track and estimate costs for GCO resources."""
76 def __init__(self, config: GCOConfig | None = None):
77 self._config = config
78 self._session = boto3.Session()
79 self._pricing_cache: dict[str, float | None] = {}
81 def get_cost_summary(
82 self,
83 days: int = 30,
84 granularity: str = "MONTHLY",
85 unfiltered: bool = False,
86 ) -> CostSummary:
87 """Get cost summary from Cost Explorer filtered by GCO tags."""
88 ce = self._session.client("ce", region_name="us-east-1")
90 end = datetime.now(UTC).date()
91 start = end - timedelta(days=days)
93 kwargs: dict[str, Any] = {
94 "TimePeriod": {
95 "Start": start.isoformat(),
96 "End": end.isoformat(),
97 },
98 "Granularity": granularity,
99 "Metrics": ["UnblendedCost"],
100 "GroupBy": [
101 {"Type": "DIMENSION", "Key": "SERVICE"},
102 ],
103 }
105 if not unfiltered:
106 kwargs["Filter"] = {
107 "Tags": {
108 "Key": "Project",
109 "Values": ["GCO"],
110 }
111 }
113 try:
114 response = ce.get_cost_and_usage(**kwargs)
115 except Exception as e:
116 raise RuntimeError(f"Cost Explorer query failed: {e}") from e
118 summary = CostSummary(
119 total=0.0,
120 period_start=start.isoformat(),
121 period_end=end.isoformat(),
122 )
124 for result in response.get("ResultsByTime", []):
125 for group in result.get("Groups", []):
126 service = group["Keys"][0]
127 amount = float(group["Metrics"]["UnblendedCost"]["Amount"])
128 if amount > 0.001:
129 summary.by_service.append(ResourceCost(service=service, amount=amount))
130 summary.total += amount
132 # Sort by cost descending
133 summary.by_service.sort(key=lambda x: x.amount, reverse=True)
135 return summary
137 def get_cost_by_region(self, days: int = 30) -> dict[str, float]:
138 """Get cost breakdown by region."""
139 ce = self._session.client("ce", region_name="us-east-1")
141 end = datetime.now(UTC).date()
142 start = end - timedelta(days=days)
144 try:
145 response = ce.get_cost_and_usage(
146 TimePeriod={
147 "Start": start.isoformat(),
148 "End": end.isoformat(),
149 },
150 Granularity="MONTHLY",
151 Metrics=["UnblendedCost"],
152 Filter={
153 "Tags": {
154 "Key": "Project",
155 "Values": ["GCO"],
156 }
157 },
158 GroupBy=[
159 {"Type": "DIMENSION", "Key": "REGION"},
160 ],
161 )
162 except Exception as e:
163 raise RuntimeError(f"Cost Explorer query failed: {e}") from e
165 by_region: dict[str, float] = {}
166 for result in response.get("ResultsByTime", []):
167 for group in result.get("Groups", []):
168 region = group["Keys"][0]
169 amount = float(group["Metrics"]["UnblendedCost"]["Amount"])
170 if amount > 0.001:
171 by_region[region] = by_region.get(region, 0) + amount
173 return dict(sorted(by_region.items(), key=lambda x: x[1], reverse=True))
175 def get_daily_trend(self, days: int = 14, unfiltered: bool = False) -> list[dict[str, Any]]:
176 """Get daily cost trend."""
177 ce = self._session.client("ce", region_name="us-east-1")
179 end = datetime.now(UTC).date()
180 start = end - timedelta(days=days)
182 kwargs: dict[str, Any] = {
183 "TimePeriod": {
184 "Start": start.isoformat(),
185 "End": end.isoformat(),
186 },
187 "Granularity": "DAILY",
188 "Metrics": ["UnblendedCost"],
189 }
191 if not unfiltered:
192 kwargs["Filter"] = {
193 "Tags": {
194 "Key": "Project",
195 "Values": ["GCO"],
196 }
197 }
199 try:
200 response = ce.get_cost_and_usage(**kwargs)
201 except Exception as e:
202 raise RuntimeError(f"Cost Explorer query failed: {e}") from e
204 trend = []
205 for result in response.get("ResultsByTime", []):
206 date = result["TimePeriod"]["Start"]
207 amount = float(result["Total"]["UnblendedCost"]["Amount"])
208 trend.append({"date": date, "amount": amount})
210 return trend
212 def estimate_running_workloads(self, region: str) -> list[WorkloadCost]:
213 """Estimate costs for currently running workloads in a region."""
214 try:
215 from .capacity import get_capacity_checker
216 except ImportError:
217 return []
219 checker = get_capacity_checker(self._config)
220 estimates: list[WorkloadCost] = []
222 # Get running pods from EKS
223 try:
224 project_name = (self._config or get_config()).project_name
225 cluster_name = f"{project_name}-{region}"
227 from .kubectl_helpers import update_kubeconfig
229 update_kubeconfig(cluster_name, region)
231 from kubernetes import client as k8s_client
232 from kubernetes import config as k8s_config
234 k8s_config.load_kube_config()
235 v1 = k8s_client.CoreV1Api()
237 # Check inference namespace
238 for ns in ["gco-inference", "gco-jobs"]:
239 try:
240 pods = v1.list_namespaced_pod(namespace=ns)
241 except Exception as e:
242 logger.debug("Failed to list pods in %s: %s", ns, e)
243 continue
245 for pod in pods.items:
246 if pod.status.phase not in ("Running", "Pending"):
247 continue
249 name = pod.metadata.name
250 gpu_count = 0
251 instance_type = "unknown"
253 # Get GPU requests
254 for container in pod.spec.containers or []:
255 requests = container.resources.requests or {}
256 gpu_req = requests.get( # nosec B113 - dict.get(), not HTTP requests
257 "nvidia.com/gpu", "0"
258 )
259 gpu_count += int(gpu_req)
261 # Get node instance type
262 if pod.spec.node_name:
263 try:
264 node = v1.read_node(pod.spec.node_name)
265 instance_type = node.metadata.labels.get(
266 "node.kubernetes.io/instance-type", "unknown"
267 )
268 except Exception as e:
269 logger.debug(
270 "Failed to get node info for %s: %s", pod.spec.node_name, e
271 )
273 # Calculate cost
274 hourly_rate = checker.get_on_demand_price(instance_type, region) or 0.0
276 # Calculate runtime
277 start_time = pod.status.start_time
278 if start_time:
279 runtime = datetime.now(UTC) - start_time
280 runtime_hours = runtime.total_seconds() / 3600
281 else:
282 runtime_hours = 0.0
284 workload_type = "inference" if ns == "gco-inference" else "job"
286 estimates.append(
287 WorkloadCost(
288 name=name,
289 workload_type=workload_type,
290 instance_type=instance_type,
291 gpu_count=gpu_count,
292 hourly_rate=hourly_rate,
293 runtime_hours=round(runtime_hours, 2),
294 estimated_cost=round(hourly_rate * runtime_hours, 4),
295 region=region,
296 status=pod.status.phase,
297 )
298 )
300 except Exception as e:
301 logger.debug("Failed to estimate workload costs: %s", e)
303 return estimates
305 def get_forecast(self, days_ahead: int = 30) -> dict[str, Any]:
306 """Get cost forecast for the next N days."""
307 ce = self._session.client("ce", region_name="us-east-1")
309 start = datetime.now(UTC).date()
310 end = start + timedelta(days=days_ahead)
312 try:
313 response = ce.get_cost_forecast(
314 TimePeriod={
315 "Start": start.isoformat(),
316 "End": end.isoformat(),
317 },
318 Metric="UNBLENDED_COST",
319 Granularity="MONTHLY",
320 Filter={
321 "Tags": {
322 "Key": "Project",
323 "Values": ["GCO"],
324 }
325 },
326 )
328 return {
329 "forecast_total": float(response.get("Total", {}).get("Amount", 0)),
330 "period_start": start.isoformat(),
331 "period_end": end.isoformat(),
332 }
333 except Exception as e:
334 return {"error": str(e)}
336 # ------------------------------------------------------------------
337 # Cost allocation tags — the switches behind every query above.
338 #
339 # Every Cost Explorer query in this module filters on the ``Project``
340 # tag, but that filter only sees spend if the tag key is *activated*
341 # as a cost allocation tag in the billing account. Activation is a
342 # separate account-level switch from tagging the resources, and the
343 # AWS-generated ``aws:eks:cluster-name`` key (stamped by EKS Auto
344 # Mode on the EC2 instances, volumes, and load balancers it launches
345 # outside CloudFormation) adds per-cluster attribution for the
346 # compute that dominates GCO spend. In an AWS Organization only the
347 # management (payer) account may activate keys.
348 # ------------------------------------------------------------------
350 def get_cost_allocation_tag_status(
351 self, tag_keys: Sequence[str] | None = None
352 ) -> list[dict[str, str]]:
353 """Report cost-allocation status for each tag key, in request order.
355 A key that Billing has never seen on billing data is absent from
356 ``ListCostAllocationTags``; it is reported here with status
357 ``NotFound``. New keys appear up to 24 hours after they are first
358 used on a resource that accrues cost.
359 """
360 keys = list(tag_keys) if tag_keys else list(DEFAULT_COST_ALLOCATION_TAG_KEYS)
361 ce = self._session.client("ce", region_name="us-east-1")
362 found: dict[str, dict[str, str]] = {}
363 for start in range(0, len(keys), _COST_ALLOCATION_TAG_BATCH_SIZE):
364 batch = keys[start : start + _COST_ALLOCATION_TAG_BATCH_SIZE]
365 token: str | None = None
366 while True:
367 kwargs: dict[str, Any] = {"TagKeys": batch, "MaxResults": 100}
368 if token:
369 kwargs["NextToken"] = token
370 response = ce.list_cost_allocation_tags(**kwargs)
371 for tag in response.get("CostAllocationTags", []):
372 key = str(tag.get("TagKey", ""))
373 found[key] = {
374 "tag_key": key,
375 "type": str(tag.get("Type", "")),
376 "status": str(tag.get("Status", "")),
377 "last_updated": str(tag.get("LastUpdatedDate", "") or ""),
378 "last_used": str(tag.get("LastUsedDate", "") or ""),
379 }
380 token = response.get("NextToken")
381 if not token:
382 break
383 return [
384 found.get(
385 key,
386 {
387 "tag_key": key,
388 "type": "Unknown",
389 "status": "NotFound",
390 "last_updated": "",
391 "last_used": "",
392 },
393 )
394 for key in keys
395 ]
397 def activate_cost_allocation_tags(
398 self, tag_keys: Sequence[str] | None = None
399 ) -> dict[str, Any]:
400 """Activate tag keys for cost allocation; report per-key outcomes.
402 Activation is idempotent (re-activating an active key succeeds)
403 and only affects billing data from activation onward — use
404 :meth:`start_cost_allocation_tag_backfill` to re-tag past usage.
405 Returns ``{"activated": [keys], "errors": [{tag_key, code,
406 message}]}`` without raising on per-key failures, so one
407 not-yet-discovered key cannot mask the others succeeding.
408 """
409 keys = list(tag_keys) if tag_keys else list(DEFAULT_COST_ALLOCATION_TAG_KEYS)
410 ce = self._session.client("ce", region_name="us-east-1")
411 activated: list[str] = []
412 errors: list[dict[str, str]] = []
413 for start in range(0, len(keys), _COST_ALLOCATION_TAG_BATCH_SIZE):
414 batch = keys[start : start + _COST_ALLOCATION_TAG_BATCH_SIZE]
415 response = ce.update_cost_allocation_tags_status(
416 CostAllocationTagsStatus=[{"TagKey": key, "Status": "Active"} for key in batch]
417 )
418 failed_keys: set[str] = set()
419 for error in response.get("Errors", []):
420 key = str(error.get("TagKey", ""))
421 failed_keys.add(key)
422 errors.append(
423 {
424 "tag_key": key,
425 "code": str(error.get("Code", "")),
426 "message": str(error.get("Message", "")),
427 }
428 )
429 activated.extend(key for key in batch if key not in failed_keys)
430 return {"activated": activated, "errors": errors}
432 def start_cost_allocation_tag_backfill(self, backfill_from: str) -> dict[str, Any]:
433 """Ask Billing to re-tag historical usage back to ``backfill_from``.
435 ``backfill_from`` is an ISO-8601 timestamp (Billing accepts up to
436 12 months back, quarter-start aligned). Returns the request record
437 with its ``BackfillStatus``.
438 """
439 ce = self._session.client("ce", region_name="us-east-1")
440 response = ce.start_cost_allocation_tag_backfill(BackfillFrom=backfill_from)
441 request = response.get("BackfillRequest", {}) or {}
442 return {
443 "backfill_from": str(request.get("BackfillFrom", "")),
444 "requested_at": str(request.get("RequestedAt", "")),
445 "completed_at": str(request.get("CompletedAt", "") or ""),
446 "status": str(request.get("BackfillStatus", "")),
447 }
449 def get_cost_allocation_backfill_history(self) -> list[dict[str, Any]]:
450 """List prior backfill requests, newest first per the service order."""
451 ce = self._session.client("ce", region_name="us-east-1")
452 history: list[dict[str, Any]] = []
453 token: str | None = None
454 while True:
455 kwargs: dict[str, Any] = {}
456 if token:
457 kwargs["NextToken"] = token
458 response = ce.list_cost_allocation_tag_backfill_history(**kwargs)
459 for request in response.get("BackfillRequests", []):
460 history.append(
461 {
462 "backfill_from": str(request.get("BackfillFrom", "")),
463 "requested_at": str(request.get("RequestedAt", "")),
464 "completed_at": str(request.get("CompletedAt", "") or ""),
465 "status": str(request.get("BackfillStatus", "")),
466 }
467 )
468 token = response.get("NextToken")
469 if not token:
470 break
471 return history
474def get_cost_tracker(config: GCOConfig | None = None) -> CostTracker:
475 """Factory function for CostTracker."""
476 return CostTracker(config=config)