Load Balancing a GPU Worker Pool with Redis Keyspace Notifications and Leases

A pool of GPUs is not like a pool of stateless web servers, and load balancing it like one wastes the most expensive hardware you own. A GPU running a heavy inference job is fully occupied. Send it a second job and it does not share, it queues, and now a request waits behind minutes of someone else's work while another GPU sits idle. Round-robin and least-connections, the reflexes from web load balancing, both fail here because they do not know what "busy" means for a GPU. This post builds a busy-aware scheduler from two Redis primitives, a lease that reserves a GPU and a keyspace notification that says the instant one frees up, and covers the failure modes that force the design.

Why ordinary load balancing fails on GPUs

The assumptions under a normal load balancer do not hold for GPU work. A web server handles many concurrent requests, each cheap and short, so spreading requests evenly across servers is close to optimal and a stale view of load barely matters. A GPU handles one heavy job at a time, each expensive and long, so the only thing that matters is whether a given GPU is free right now, and a stale view of that is the difference between an instant start and a multi-minute stall.

Round-robin ignores busyness entirely and will cheerfully assign job N+1 to the GPU still grinding on job N. Least-connections is closer but still wrong, because it counts connections, not GPU occupancy, and a GPU mid-inference may hold zero connections while being completely unavailable. What you need is a scheduler whose unit of truth is "is this specific GPU free," updated the moment it changes, with a hard reservation so two requests cannot both claim the same free GPU. Redis gives you both halves cheaply.

A lease is the reservation

The core state is a per-GPU lease: a Redis key that exists while a GPU is busy and disappears when it is free. Claiming a GPU is taking its lease, and the claim has to be atomic so two requests racing for the same free GPU cannot both win.

// Claim GPU n only if its lease is absent. SET NX is atomic, so exactly one
// caller wins the race. The TTL is a safety net, covered below.
ok, err := rdb.SetNX(ctx, "gpu:lease:"+gpuID, workerID, leaseTTL).Result()
if err != nil {
	return "", err
}
if !ok {
	return "", errBusy // someone else holds this GPU
}

SetNX is compare-and-swap in one command: set the key only if it does not already exist. The GPU that wins gets the lease and starts its job. Everyone else sees ok come back false and moves on to try another GPU. When the job finishes, the worker deletes the lease and the GPU is free again. If it were only this, you would still be polling to find a free GPU, and polling is exactly the latency-and-waste dial we want to avoid. The second primitive removes the polling.

Keyspace notifications turn "free" into an event

Redis can publish an event whenever a key changes, including when a key expires or is deleted. That feature is keyspace notifications, and it is the piece that lets a waiting request learn the instant a GPU frees up instead of asking on a timer. You enable it once on the server for the event classes you care about, generic key events and expiry events, then subscribe to the pattern for lease keys.

// Subscribe to deletions and expirations of GPU lease keys. Each message
// means a specific GPU just became free.
pubsub := rdb.PSubscribe(ctx,
	"__keyevent@0__:del",
	"__keyevent@0__:expired",
)
for msg := range pubsub.Channel() {
	freedGPU := parseGPUID(msg.Payload) // the lease key that just vanished
	scheduler.notifyFree(freedGPU)
}

Now the flow is event-driven end to end. A request tries to claim any free GPU with SetNX. If all leases are held, instead of spinning, it waits. The moment any job anywhere finishes and deletes its lease, or a lease expires, Redis publishes the event, a waiting request wakes up, and it tries to claim the newly free GPU. The expensive hardware goes from busy to reassigned in the time it takes one Redis message to travel, with no polling in the steady state.

The failure mode that TTL exists for

There is a hole in the happy path, and it is the one that matters most because it strands your most expensive resource. Suppose a worker claims a GPU, starts a job, and then crashes, or its network partitions, or its process is killed mid-run. It never reaches the line that deletes the lease. The lease key sits in Redis forever, that GPU is marked busy for eternity, and no keyspace event will ever fire for it because nothing is going to delete a key from a process that no longer exists. One crash permanently removes a GPU from the pool.

This is why the lease carries a TTL, the leaseTTL in the SetNX call above. The TTL is not for the normal case, where the worker deletes the lease on completion well before expiry. It is the recovery mechanism for the abnormal case. If the holder dies, the lease expires on its own, Redis fires the expiry keyspace event, and the GPU rejoins the pool automatically with no human intervention. The lease is really a lease in the same sense as the previous discussion of single-runner jobs: a reservation that a live holder renews and a dead holder forfeits.

That means a long job has to renew its lease before the TTL lapses, exactly like a heartbeat, or it will lose its GPU mid-run to the recovery mechanism. Set the TTL comfortably longer than the longest expected job, or heartbeat the lease periodically for jobs whose length you cannot bound. The tradeoff is the familiar one: a short TTL recovers a crashed GPU quickly but risks yanking a slow job, a long TTL is safe for slow jobs but leaves a crashed GPU stranded longer. Choose it from your job-length distribution, not from a default.

The race between claiming and waiting

A busy-aware scheduler has a subtle ordering bug that is worth calling out because it is easy to write and hard to reproduce. Imagine a request checks all GPUs, finds them busy, and then subscribes to the free-event feed. If a GPU frees up in the tiny gap between the check and the subscribe, the event fires before the subscription exists, the request never hears it, and it waits for the next free event that may be minutes away. The GPU is free, the request is waiting, and neither knows about the other.

The fix is to reverse the order and re-check after subscribing. Subscribe to the free feed first, then attempt to claim a GPU. If the claim succeeds, unsubscribe and run. If it fails because everything is busy, wait on the feed you already subscribed to, and crucially, when an event arrives, attempt the claim again rather than assuming the freed GPU is yours, because another waiter may have taken it first. Subscribe, then check, then wait-and-retry. That ordering closes the gap, and the retry-on-wake handles the multiple-waiters case where the freed GPU is claimed by someone else before you get to it.

Stale state and the reconcile pass

Even with TTLs, distributed state drifts. A lease might be deleted while its worker is briefly partitioned and still thinks it holds the GPU. An expiry event might be missed if a subscriber was reconnecting at the moment it fired. Event-driven systems built only on events accumulate these small divergences over time. The cure is a periodic reconcile: on a slow interval, independently ask each GPU worker what it is actually doing and correct Redis to match reality. The events keep the system fast and current in the common case. The reconcile pass keeps it correct in the long run by catching the drift the events missed. Fast path from events, correctness backstop from reconciliation, is a good default posture for any event-driven scheduler.

When a simpler design is enough

This machinery earns its place when GPUs are your scarce, expensive resource and jobs are long enough that routing one to a busy GPU is a costly mistake. If your jobs are short and cheap, a plain queue that any idle worker pulls from is simpler and just as good, because the cost of a slightly imperfect assignment is negligible when each job takes moments. If you have a managed inference service that handles placement for you, use it and skip all of this, because you are paying for exactly this problem to be solved.

Build the lease-and-notification scheduler when three things are true together: the GPUs are expensive enough that idle time is a real loss, the jobs are long enough that a wrong assignment strands a request for a meaningful time, and you are running the pool yourself rather than renting a managed endpoint. Under those conditions, a Redis lease for the atomic reservation and a keyspace notification for the instant handoff give you a scheduler that keeps expensive hardware busy without polling, and a TTL plus a reconcile pass keep it honest when workers die and events go missing.