split-inference-bench/tests/test_budget_accounting.py at main · kadubon/split-inference-bench · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from src.budget import BudgetConfig, BudgetLedger, BudgetVector


def test_budget_ledger_enforces_additive_limit_and_local_ceiling() -> None:
    config = BudgetConfig(
        total=BudgetVector(compute=1.0, memory=1.0, communication=1.0, verification=1.0, horizon=1.0),
        price_weights=BudgetVector(compute=1.0, memory=1.0, communication=1.0, verification=1.0, horizon=1.0),
        local_context_ceiling=2.0,
        deadline=1,
    )
    ledger = BudgetLedger(config=config)
    ledger.charge(
        category="worker_compute",
        vector=BudgetVector(compute=0.5, horizon=0.5),
        actor="worker",
        note="test charge",
        bucket="work",
    )
    ledger.record_local_footprint("worker", 2.0)
    ledger.assert_feasible()