The GPU Credit Scheduler Pattern

TL;DR: Model accounts, jobs and reservations explicitly. Decide when cost is reserved and charged, then preserve the balance invariant through submit, start, complete and cancel. A priority queue chooses eligible work; concurrency and persistence need their own guarantees.

The problem, as it is usually posed

"We have accounts with credit balances. Jobs are submitted with a cost and a priority. Write a scheduler that decides which job to run next: highest priority first, oldest first among equals, and never run a job the account cannot pay for. Then handle completion and failure." Follow-ups arrive one at a time.

Ask two questions before typing: is the cost charged at submit, at start, or at completion (the answer shapes the ledger), and can a job's cost change after submission (usually no, and saying "I'll assume fixed cost and revisit" is fine).

The model

import heapq
import itertools
from dataclasses import dataclass, field

@dataclass
class Account:
    id: str
    balance: int          # credits the account owns
    reserved: int = 0     # credits held for submitted-but-unfinished jobs
    running: int = 0      # jobs currently running (for the concurrency follow-up)

    @property
    def available(self) -> int:
        return self.balance - self.reserved

@dataclass(order=True)
class Job:
    sort_key: tuple = field(compare=True)           # (-priority, seq): highest priority, then oldest
    id: str = field(compare=False)
    account_id: str = field(compare=False)
    cost: int = field(compare=False)
    priority: int = field(compare=False)

class Scheduler:
    def __init__(self):
        self.accounts: dict[str, Account] = {}
        self.pending: list[Job] = []                 # heap
        self.running: dict[str, Job] = {}
        self.jobs: dict[str, Job] = {}                # IDs are unique for this scheduler lifetime
        self.cancelled: set[str] = set()             # lazy deletion from the heap
        self._seq = itertools.count()                # ties broken by submission order

    def submit(self, job_id, account_id, cost, priority) -> bool:
        if job_id in self.jobs or cost < 0:
            raise ValueError("duplicate job ID or negative cost")
        acct = self.accounts[account_id]
        if cost > acct.available:
            return False                             # reject at submit: the account cannot pay
        acct.reserved += cost                        # invariant: balance >= reserved >= 0
        job = Job((-priority, next(self._seq)), job_id, account_id, cost, priority)
        self.jobs[job_id] = job
        heapq.heappush(self.pending, job)
        return True

    def start_next(self) -> Job | None:
        while self.pending:
            job = heapq.heappop(self.pending)
            if job.id in self.cancelled:
                continue                             # skip a tombstone
            self.running[job.id] = job
            self.accounts[job.account_id].running += 1
            return job
        return None

    def complete(self, job_id, success: bool):
        job = self.running.pop(job_id)
        job.sort_key = None                          # terminal state; no longer in the heap
        acct = self.accounts[job.account_id]
        acct.reserved -= job.cost
        acct.running -= 1
        if success:
            acct.balance -= job.cost                 # reserved becomes spent
        # on failure, the reservation is released and nothing is charged

    def cancel(self, job_id):
        if job_id in self.running:
            raise ValueError("cannot cancel a running job")
        job = self.jobs[job_id]
        if job_id in self.cancelled:
            return
        if job.sort_key is None:
            raise ValueError("job already completed")
        self.accounts[job.account_id].reserved -= job.cost
        self.cancelled.add(job_id)                   # release now; skip the heap entry later

The decisions worth saying out loud:

• Reserve at submit. Charging at submit is unfair on failure; charging at completion lets an account overcommit by submitting many jobs. Reserving at submit and settling at completion keeps the invariant balance ≥ reserved ≥ 0 and never lets an account run jobs it cannot pay for, even with many pending.
• Heap key. (-priority, seq) gives highest priority first and oldest first within a priority in O(log n) per push and pop. A sorted list is O(n) per insert; a dict of lists per priority works if priorities are few and is worth mentioning.
• Lazy cancellation. Heaps cannot delete from the middle cheaply; a tombstone set costs O(1) at cancel and O(log n) when the entry is finally popped. The jobs map lets cancellation release its reservation immediately. Retaining terminal IDs prevents accidental reuse; a production store needs an explicit retention policy.

The invariants, stated

for every account:   balance ≥ reserved ≥ 0;   reserved = Σ cost of its jobs in pending ∪ running (minus tombstones)
for the scheduler:   a job is in exactly one of {pending, running, done}; the heap's tombstones are never running
after complete(success):  balance decreases by cost, reserved decreases by cost, available unchanged
after complete(failure):  balance unchanged, reserved decreases by cost, available increases by cost

One account through one job makes the invariant visible, and shows why start is the transition that does nothing to the balance:

Re-stating these after each follow-up is what the interviewer is listening for; a candidate who types the change and then says "and the invariant still holds because..." is doing the job.

The follow-ups, and the one-line answers

"Cap concurrent jobs per account at k." In start_next, skip a job whose account has running ≥ k. Skipping means the popped job must go back: push it to a side list and re-push after the loop, or keep one heap per account and pick among heads (O(accounts) per start, or a second heap of account heads). State the cost: with a single heap and many blocked accounts, a start can pop and re-push O(blocked) entries.

"Be fair across accounts, not just priority-first." Change the key to (-priority, spend_ratio, seq) where spend_ratio = recent_spend / quota computed at push time, or a deficit-round-robin over per-account queues. Say the trade: fairness keyed at push time goes stale; per-account queues make it exact at O(accounts) per pick (Multi-Tenancy, Quotas and Fair Share).

"Credits refill at a rate per hour." Add last_refill and rate to the account; on every access, balance = min(cap, balance + rate × elapsed). That is the token bucket (Rate-Limiting Algorithms), with a clearly defined refill and reservation policy.

"Jobs can be partially refunded on failure." complete(success=False, refund_fraction) charges cost × (1 − refund_fraction); the invariant on reserved is unchanged.

"Make it safe for many threads." One lock around the scheduler's mutations is correct and enough for a scheduler that makes hundreds of decisions per second; say why finer locking is not worth it here, and what you would do at higher rates (shard by account) (Concurrency in Python, Go and C++).

The tests you write before being asked

def test_reject_when_unaffordable():
    s = Scheduler(); s.accounts["a"] = Account("a", balance=10)
    assert s.submit("j1", "a", cost=8, priority=1)
    assert not s.submit("j2", "a", cost=5, priority=9)   # 8 reserved, 2 available

def test_priority_then_age():
    s = Scheduler(); s.accounts["a"] = Account("a", balance=100)
    s.submit("low", "a", 1, priority=1); s.submit("high1", "a", 1, priority=5); s.submit("high2", "a", 1, priority=5)
    assert [s.start_next().id for _ in range(3)] == ["high1", "high2", "low"]

def test_failure_refunds_reservation():
    s = Scheduler(); s.accounts["a"] = Account("a", balance=10)
    s.submit("j", "a", 7, 1); s.start_next(); s.complete("j", success=False)
    assert s.accounts["a"].available == 10 and s.accounts["a"].balance == 10

def test_cancel_pending_is_skipped():
    s = Scheduler(); s.accounts["a"] = Account("a", balance=10)
    s.submit("j1", "a", 3, 1); s.submit("j2", "a", 3, 1); s.cancel("j1")
    assert s.accounts["a"].available == 7
    s.cancel("j1")
    assert s.accounts["a"].available == 7
    assert s.start_next().id == "j2"

These tests check core transitions. Also test negative costs, duplicate IDs, repeated cancellation, completion replay and the available balance after cancellation. This in-memory example rejects repeated completion; a durable API needs a deliberate replay policy.

Complexity, derived

submit: O(log n) heap push;  start_next: O(log n) amortized (tombstones are popped at most once each)
complete: O(1);  cancel: O(1) plus the deferred pop
memory: O(all submitted IDs) because this example retains terminal jobs and cancellation IDs
with per-account concurrency caps on a single heap: start_next is O(b log n) where b is the number of blocked heads skipped;
  per-account heaps with a heap of heads make it O(log a) where a is the number of accounts

Working it in the room

Clarify the charging policy and whether costs can change before writing code. Implement the state transitions and test repeated completion, cancellation and insufficient credit. When extending the scheduler, explain which invariant changes and how concurrent or replayed events remain safe.

What to remember

• Model: Account (balance, reserved), Job (cost, priority, seq), Scheduler (heap, running, tombstones). Reserve at submit, settle at completion.
• Invariant: balance ≥ reserved ≥ 0; a job is in exactly one state; say it after every change.
• Heap key (−priority, seq); lazy cancellation with tombstones; O(log n) per operation.
• Follow-ups add one field and one check each: concurrency caps, fairness by spend ratio, token-bucket refill, partial refunds, one lock.
• Four tests named for invariants, written before the first follow-up.