The GPU Credit Scheduler Pattern

TL;DR: Model three things: an account (id, balance, reserved), a job (id, account, cost, priority, submit time), and the scheduler's state (a max-heap of pending jobs keyed by priority then age, a map of running jobs, the accounts). Submit reserves the cost against the account's balance so the invariant "balance ≥ reserved ≥ 0" holds at every moment; start pops the best affordable job; complete converts reserved into spent (or refunds on failure). Heap operations are O(log n); the follow-ups (per-account concurrency caps, fairness by spend, cancellation, time-based credit refill) each add one field and one check, and the invariant is what you re-state after every change. The interviewer is scoring your model, your invariants and your tests, not the heap.

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.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:
        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
        heapq.heappush(self.pending, Job((-priority, next(self._seq)), job_id, account_id, cost, priority))
        return True

    def start_next(self) -> Job | None:
        while self.pending:
            job = heapq.heappop(self.pending)
            if job.id in self.cancelled:
                self.cancelled.discard(job.id)
                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)
        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")
        self.cancelled.add(job_id)                   # the heap entry is skipped when popped
        # the reservation is released lazily when the tombstone is popped, or eagerly if we track jobs by id

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. Eager release of the reservation needs a jobs_by_id map, which is a one-line addition.

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), and naming it wins a point.

"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.start_next().id == "j2"

Four tests, each named for an invariant, written in the first fifteen minutes. They are the difference between "it seems to work" and "it works," and they make every follow-up safe to attempt.

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(n) jobs plus O(tombstones), bounded by the number of cancels since the last pop
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

Fifteen minutes to the working core with tests, five per follow-up. Narrate the model before typing, state the invariant after each change, and give the complexity unasked. The follow-up held back is usually "what if the process restarts?", answered with a write-ahead log of submit, start, complete and cancel events that rebuilds the state on replay. The answer that sounds right and fails is charging at submit "because it's simpler": it makes failures unfair and the interviewer's second follow-up is exactly that.

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.