code/self_healing_finetuning/budget/limits.py

from __future__ import annotations


DEFAULT_BUDGET = {
    "max_gpu_hours_per_run": 8,
    "max_gpu_hours_per_candidate": 16,
    "max_cost_usd_per_run": 250,
    "max_cost_usd_per_candidate": 500,
    "max_cost_usd_per_day": 1000,
    "stop_when_budget_exceeded": True,
    "require_approval_to_exceed_budget": True,
}


def check_budget(estimate: dict[str, float], budget: dict[str, float] | None = None) -> list[str]:
    limits = {**DEFAULT_BUDGET, **(budget or {})}
    errors: list[str] = []
    if estimate.get("gpu_hours", 0) > limits["max_gpu_hours_per_run"]:
        errors.append("estimated GPU hours exceed max_gpu_hours_per_run")
    if estimate.get("cost_usd", 0) > limits["max_cost_usd_per_run"]:
        errors.append("estimated cost exceeds max_cost_usd_per_run")
    return errors