main
py 107 lines 3.28 KB
Raw
1 #!/usr/bin/env python3
2 """Budget alerts for GitHub Actions CI.
3
4 Evaluates current run cost and monthly cumulative spend against thresholds,
5 emitting GitHub Actions annotations (::warning:: / ::error::) as appropriate.
6 """
7
8 from __future__ import annotations
9
10 import argparse
11 import json
12 import sys
13 from datetime import UTC, datetime
14 from pathlib import Path
15
16 ROOT = Path(__file__).resolve().parent.parent
17 DEFAULT_METRICS = ROOT / "data" / "metrics" / "token-usage.jsonl"
18
19 # Thresholds
20 SINGLE_RUN_WARNING = 0.50
21 SINGLE_RUN_FAIL = 1.00
22 MONTHLY_WARNING = 5.00
23 MONTHLY_RECOMMEND_SWITCH = 10.00
24
25
26 def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
27 parser = argparse.ArgumentParser(description="Budget alerts for CI cost control.")
28 parser.add_argument(
29 "--run-cost",
30 type=float,
31 default=None,
32 help="Cost of the current run in USD.",
33 )
34 parser.add_argument(
35 "--metrics",
36 type=Path,
37 default=DEFAULT_METRICS,
38 help="Path to token-usage.jsonl ledger.",
39 )
40 return parser.parse_args(argv)
41
42
43 def load_monthly_spend(metrics_path: Path, now: datetime | None = None) -> float:
44 """Sum estimated_cost for entries in the current month."""
45 if not metrics_path.exists():
46 return 0.0
47 now = now or datetime.now(UTC)
48 total = 0.0
49 for line in metrics_path.read_text().splitlines():
50 line = line.strip()
51 if not line:
52 continue
53 try:
54 entry = json.loads(line)
55 except json.JSONDecodeError:
56 continue
57 ts = entry.get("timestamp", "")
58 try:
59 entry_dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
60 except (ValueError, AttributeError):
61 continue
62 if entry_dt.year == now.year and entry_dt.month == now.month:
63 total += entry.get("estimated_cost", 0.0)
64 return total
65
66
67 def evaluate(run_cost: float | None, monthly_spent: float) -> tuple[list[str], int]:
68 """Return (annotations, exit_code)."""
69 annotations: list[str] = []
70 exit_code = 0
71
72 if run_cost is not None:
73 if run_cost > SINGLE_RUN_FAIL:
74 annotations.append(
75 f"::error::Single run cost ${run_cost:.2f} exceeds hard cap ${SINGLE_RUN_FAIL:.2f}"
76 )
77 exit_code = 1
78 elif run_cost > SINGLE_RUN_WARNING:
79 annotations.append(
80 f"::warning::Single run cost ${run_cost:.2f} exceeds warning threshold ${SINGLE_RUN_WARNING:.2f}"
81 )
82
83 if monthly_spent > MONTHLY_RECOMMEND_SWITCH:
84 annotations.append(
85 f"::warning::Monthly spend ${monthly_spent:.2f} exceeds ${MONTHLY_RECOMMEND_SWITCH:.2f} — consider switching to a cheaper model"
86 )
87 elif monthly_spent > MONTHLY_WARNING:
88 annotations.append(
89 f"::warning::Monthly cumulative spend ${monthly_spent:.2f} exceeds ${MONTHLY_WARNING:.2f}"
90 )
91
92 return annotations, exit_code
93
94
95 def main(argv: list[str] | None = None) -> int:
96 args = parse_args(argv)
97 monthly_spent = load_monthly_spend(args.metrics)
98 if args.run_cost is not None:
99 monthly_spent += args.run_cost
100 annotations, exit_code = evaluate(args.run_cost, monthly_spent)
101 for ann in annotations:
102 print(ann)
103 return exit_code
104
105
106 if __name__ == "__main__":
107 sys.exit(main())