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

1"""Cost visibility for GCO workloads. 

2 

3Uses AWS Cost Explorer for historical spend and the Pricing API 

4for real-time cost estimates on running workloads. 

5""" 

6 

7from __future__ import annotations 

8 

9import logging 

10from collections.abc import Sequence 

11from dataclasses import dataclass, field 

12from datetime import UTC, datetime, timedelta 

13from typing import Any 

14 

15import boto3 

16 

17from .config import GCOConfig, get_config 

18 

19logger = logging.getLogger(__name__) 

20 

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") 

29 

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 

33 

34 

35@dataclass 

36class ResourceCost: 

37 """Cost for a single resource or service.""" 

38 

39 service: str 

40 amount: float 

41 currency: str = "USD" 

42 region: str | None = None 

43 detail: str | None = None 

44 

45 

46@dataclass 

47class CostSummary: 

48 """Aggregated cost summary.""" 

49 

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) 

56 

57 

58@dataclass 

59class WorkloadCost: 

60 """Estimated cost for a running workload.""" 

61 

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 

71 

72 

73class CostTracker: 

74 """Track and estimate costs for GCO resources.""" 

75 

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] = {} 

80 

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") 

89 

90 end = datetime.now(UTC).date() 

91 start = end - timedelta(days=days) 

92 

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 } 

104 

105 if not unfiltered: 

106 kwargs["Filter"] = { 

107 "Tags": { 

108 "Key": "Project", 

109 "Values": ["GCO"], 

110 } 

111 } 

112 

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 

117 

118 summary = CostSummary( 

119 total=0.0, 

120 period_start=start.isoformat(), 

121 period_end=end.isoformat(), 

122 ) 

123 

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 

131 

132 # Sort by cost descending 

133 summary.by_service.sort(key=lambda x: x.amount, reverse=True) 

134 

135 return summary 

136 

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") 

140 

141 end = datetime.now(UTC).date() 

142 start = end - timedelta(days=days) 

143 

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 

164 

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 

172 

173 return dict(sorted(by_region.items(), key=lambda x: x[1], reverse=True)) 

174 

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") 

178 

179 end = datetime.now(UTC).date() 

180 start = end - timedelta(days=days) 

181 

182 kwargs: dict[str, Any] = { 

183 "TimePeriod": { 

184 "Start": start.isoformat(), 

185 "End": end.isoformat(), 

186 }, 

187 "Granularity": "DAILY", 

188 "Metrics": ["UnblendedCost"], 

189 } 

190 

191 if not unfiltered: 

192 kwargs["Filter"] = { 

193 "Tags": { 

194 "Key": "Project", 

195 "Values": ["GCO"], 

196 } 

197 } 

198 

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 

203 

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}) 

209 

210 return trend 

211 

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 [] 

218 

219 checker = get_capacity_checker(self._config) 

220 estimates: list[WorkloadCost] = [] 

221 

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}" 

226 

227 from .kubectl_helpers import update_kubeconfig 

228 

229 update_kubeconfig(cluster_name, region) 

230 

231 from kubernetes import client as k8s_client 

232 from kubernetes import config as k8s_config 

233 

234 k8s_config.load_kube_config() 

235 v1 = k8s_client.CoreV1Api() 

236 

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 

244 

245 for pod in pods.items: 

246 if pod.status.phase not in ("Running", "Pending"): 

247 continue 

248 

249 name = pod.metadata.name 

250 gpu_count = 0 

251 instance_type = "unknown" 

252 

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) 

260 

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 ) 

272 

273 # Calculate cost 

274 hourly_rate = checker.get_on_demand_price(instance_type, region) or 0.0 

275 

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 

283 

284 workload_type = "inference" if ns == "gco-inference" else "job" 

285 

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 ) 

299 

300 except Exception as e: 

301 logger.debug("Failed to estimate workload costs: %s", e) 

302 

303 return estimates 

304 

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") 

308 

309 start = datetime.now(UTC).date() 

310 end = start + timedelta(days=days_ahead) 

311 

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 ) 

327 

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)} 

335 

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 # ------------------------------------------------------------------ 

349 

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. 

354 

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 ] 

396 

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. 

401 

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} 

431 

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``. 

434 

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 } 

448 

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 

472 

473 

474def get_cost_tracker(config: GCOConfig | None = None) -> CostTracker: 

475 """Factory function for CostTracker.""" 

476 return CostTracker(config=config)