Design a GPU credit ledger that supports:
adding credits with start and expiration windows
subtracting usage at a timestamp
querying balance at any timestamp
This is intentionally close to resource accounting problems you can see in OpenAI-like environments (GPU budget grants, usage events, and late-arriving backfills).
Key challenge: operations can arrive out of order. Example: subtract(..., t=30) can be called before add_credit(..., t=20), but get_balance(20) must still return the correct historical result.
Implement a GPUCredit class with the following methods:
class GPUCredit:
def __init__(self):
pass
def add_credit(self, credit_id: str, amount: int, timestamp: int, expiration: int) -> None:
"""
Add a credit valid in [timestamp, timestamp + expiration] (inclusive).
"""
pass
def subtract(self, amount: int, timestamp: int) -> None:
"""
Subtract credit at timestamp, consuming soonest-expiring credit first.
Does not raise if insufficient credit for that timestamp.
"""
pass
def get_balance(self, timestamp: int) -> int | None:
"""
Return balance at timestamp, or None when:
- no credit window is active at that timestamp, or
- reconstructed balance is negative.
"""
pass
Out-of-order operations: must reconstruct by event time, not insertion order.
Inclusive expiration: start=10, expiration=30 means valid during [10, 40].
Usage priority: subtract from credits that expire soonest first.
No exception on subtract: insufficient balance is allowed.
Return None for invalid query result: no active credit window or negative balance.
Single write event type per timestamp: assume caller does not send both add and subtract at the same timestamp.
Repeated credit IDs: credit_id is metadata; two grants with the same credit_id are still separate grants.
Input assumptions for interview: amounts and expiration are non-negative integers (invalid-input validation is optional and not the main signal).
gpu = GPUCredit()
gpu.add_credit('microsoft', 10, 10, 30) # valid [10, 40]
gpu.get_balance(0) # None
gpu.get_balance(10) # 10
gpu.get_balance(40) # 10
gpu.get_balance(41) # None
gpu = GPUCredit()
gpu.add_credit('amazon', 40, 10, 50) # valid [10, 60]
gpu.subtract(30, 30)
gpu.get_balance(10) # 40
gpu.get_balance(30) # 10
gpu.get_balance(60) # 10
gpu.get_balance(61) # None
gpu = GPUCredit()
gpu.subtract(4, 30)
gpu.add_credit('a', 4, 20, 30) # valid [20, 50]
gpu.get_balance(10) # None
gpu.get_balance(20) # 4
gpu.get_balance(30) # 0
gpu.get_balance(50) # 0
gpu.get_balance(51) # None
gpu = GPUCredit()
gpu.add_credit('openai', 10, 10, 30)
gpu.subtract(100, 20)
gpu.get_balance(10) # 10
gpu.get_balance(20) # None (negative)
gpu.get_balance(30) # None (still negative)
gpu = GPUCredit()
gpu.add_credit('openai', 10, 10, 50) # valid [10, 60]
gpu.add_credit('openai', 10, 30, 50) # valid [30, 80]
gpu.subtract(15, 35)
gpu.get_balance(10) # 10
gpu.get_balance(30) # 20
gpu.get_balance(35) # 5
gpu.get_balance(70) # 5
This is the easiest version to explain and debug in an interview.
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
@dataclass(frozen=True)
class CreditGrant:
credit_id: str
amount: int
start: int
end: int
@dataclass(frozen=True)
class SubtractEvent:
amount: int
timestamp: int
class GPUCreditSimple:
def __init__(self) -> None:
self._credits: List[CreditGrant] = []
self._subtracts: List[SubtractEvent] = []
def add_credit(self, credit_id: str, amount: int, timestamp: int, expiration: int) -> None:
if amount < 0 or expiration < 0:
raise ValueError('amount and expiration must be non-negative')
self._credits.append(
CreditGrant(
credit_id=credit_id,
amount=amount,
start=timestamp,
end=timestamp + expiration,
)
)
def subtract(self, amount: int, timestamp: int) -> None:
if amount < 0:
raise ValueError('amount must be non-negative')
self._subtracts.append(SubtractEvent(amount=amount, timestamp=timestamp))
def get_balance(self, timestamp: int) -> Optional[int]:
# Event tuple: (time, priority, seq, kind, payload)
# Priority: expire=0, add=1, subtract=2
events: List[Tuple[int, int, int, str, object]] = []
seq = 0
for grant_key, grant in enumerate(self._credits):
if grant.start <= timestamp:
events.append((grant.start, 1, seq, 'add', (grant_key, grant)))
seq += 1
if grant.end + 1 <= timestamp:
events.append((grant.end + 1, 0, seq, 'expire', (grant_key, grant)))
seq += 1
for sub in self._subtracts:
if sub.timestamp <= timestamp:
events.append((sub.timestamp, 2, seq, 'subtract', sub))
seq += 1
events.sort(key=lambda x: (x[0], x[1], x[2]))
# grant_key -> [remaining_amount, end_timestamp]
active: Dict[int, List[int]] = {}
active_windows = 0
deficit = 0 # outstanding negative balance from underfunded subtracts
for now, _, _, kind, payload in events:
if kind == 'add':
grant_key, grant = payload # type: ignore[misc]
if not isinstance(grant, CreditGrant):
continue
active_windows += 1
amount = grant.amount
# New credits first offset any existing negative balance.
if deficit > 0:
paid = min(deficit, amount)
deficit -= paid
amount -= paid
active[grant_key] = [amount, grant.end]
elif kind == 'expire':
grant_key, grant = payload # type: ignore[misc]
if not isinstance(grant, CreditGrant):
continue
active_windows -= 1
if grant_key in active:
# Remove expired remainder, even if it is 0.
del active[grant_key]
else:
sub = payload # type: ignore[assignment]
if not isinstance(sub, SubtractEvent):
continue
to_use = sub.amount
# Eligible pools at current time, sorted by earliest expiration.
pools: List[Tuple[int, int]] = []
for grant_key, (remaining, end) in active.items():
if remaining > 0 and end >= now:
pools.append((end, grant_key))
pools.sort()
for _, grant_key in pools:
if to_use == 0:
break
remaining = active[grant_key][0]
take = min(remaining, to_use)
active[grant_key][0] -= take
to_use -= take
# Not enough credits at that moment -> go negative.
if to_use > 0:
deficit += to_use
if active_windows == 0:
return None
available = sum(remaining for remaining, _ in active.values())
balance = available - deficit
return balance if balance >= 0 else None
Same replay model, but subtraction uses a min-heap by expiration time.
from dataclasses import dataclass
from heapq import heappop, heappush
from typing import Dict, List, Optional, Tuple
@dataclass(frozen=True)
class CreditGrant:
credit_id: str
amount: int
start: int
end: int
@dataclass(frozen=True)
class SubtractEvent:
amount: int
timestamp: int
class GPUCredit:
def __init__(self) -> None:
self._credits: List[CreditGrant] = []
self._subtracts: List[SubtractEvent] = []
def add_credit(self, credit_id: str, amount: int, timestamp: int, expiration: int) -> None:
if amount < 0 or expiration < 0:
raise ValueError('amount and expiration must be non-negative')
self._credits.append(
CreditGrant(
credit_id=credit_id,
amount=amount,
start=timestamp,
end=timestamp + expiration,
)
)
def subtract(self, amount: int, timestamp: int) -> None:
if amount < 0:
raise ValueError('amount must be non-negative')
self._subtracts.append(SubtractEvent(amount=amount, timestamp=timestamp))
def get_balance(self, timestamp: int) -> Optional[int]:
events: List[Tuple[int, int, int, str, object]] = []
seq = 0
for grant_key, grant in enumerate(self._credits):
if grant.start <= timestamp:
events.append((grant.start, 1, seq, 'add', (grant_key, grant)))
seq += 1
if grant.end + 1 <= timestamp:
events.append((grant.end + 1, 0, seq, 'expire', (grant_key, grant)))
seq += 1
for sub in self._subtracts:
if sub.timestamp <= timestamp:
events.append((sub.timestamp, 2, seq, 'subtract', sub))
seq += 1
events.sort(key=lambda x: (x[0], x[1], x[2]))
remaining: Dict[int, int] = {}
total_available = 0
deficit = 0
active_windows = 0
# Min-heap of (end, grant_key). Stale entries are lazily discarded.
heap: List[Tuple[int, int]] = []
def prune(now: int) -> None:
while heap:
end, grant_key = heap[0]
if end < now or remaining.get(grant_key, 0) == 0:
heappop(heap)
continue
break
for now, _, _, kind, payload in events:
if kind == 'add':
grant_key, grant = payload # type: ignore[misc]
if not isinstance(grant, CreditGrant):
continue
active_windows += 1
amount = grant.amount
if deficit > 0:
paid = min(deficit, amount)
deficit -= paid
amount -= paid
remaining[grant_key] = amount
if amount > 0:
total_available += amount
heappush(heap, (grant.end, grant_key))
elif kind == 'expire':
grant_key, grant = payload # type: ignore[misc]
if not isinstance(grant, CreditGrant):
continue
active_windows -= 1
rem = remaining.get(grant_key, 0)
if rem > 0:
total_available -= rem
remaining[grant_key] = 0
else:
sub = payload # type: ignore[assignment]
if not isinstance(sub, SubtractEvent):
continue
to_use = sub.amount
prune(now)
while to_use > 0:
prune(now)
if not heap:
break
end, grant_key = heappop(heap)
rem = remaining.get(grant_key, 0)
if rem == 0 or end < now:
continue
take = min(rem, to_use)
rem -= take
to_use -= take
total_available -= take
remaining[grant_key] = rem
if rem > 0:
heappush(heap, (end, grant_key))
if to_use > 0:
deficit += to_use
if active_windows == 0:
return None
balance = total_available - deficit
return balance if balance >= 0 else None
C: number of credit grants
S: number of subtract events
K: number of heap push/pop operations while replaying subtracts
def run_smoke_tests(cls):
# Basic
gpu = cls()
gpu.add_credit('a', 4, 20, 30) # [20, 50]
assert gpu.get_balance(10) is None
assert gpu.get_balance(20) == 4
assert gpu.get_balance(50) == 4
assert gpu.get_balance(51) is None
# Earliest-expiration-first
gpu = cls()
gpu.add_credit('a', 4, 20, 40) # [20, 60]
gpu.add_credit('b', 3, 30, 10) # [30, 40]
gpu.subtract(2, 30)
assert gpu.get_balance(30) == 5
assert gpu.get_balance(40) == 5
assert gpu.get_balance(41) == 4
# Out of order
gpu = cls()
gpu.subtract(4, 30)
gpu.add_credit('a', 4, 20, 30)
assert gpu.get_balance(10) is None
assert gpu.get_balance(20) == 4
assert gpu.get_balance(30) == 0
# Negative balance
gpu = cls()
gpu.add_credit('a', 10, 10, 30)
gpu.subtract(100, 20)
assert gpu.get_balance(10) == 10
assert gpu.get_balance(20) is None
# Repeated credit ID
gpu = cls()
gpu.add_credit('openai', 10, 10, 50)
gpu.add_credit('openai', 10, 30, 50)
gpu.subtract(15, 35)
assert gpu.get_balance(10) == 10
assert gpu.get_balance(30) == 20
assert gpu.get_balance(35) == 5
assert gpu.get_balance(70) == 5
run_smoke_tests(GPUCreditSimple)
run_smoke_tests(GPUCredit)