In a distributed ML inference system with a multi-backend serving layer, naive round-robin load balancing fails catastrophically. Model inference has fundamentally different performance characteristics than typical web requests — latency varies by orders of magnitude depending on model size, batch accumulation, GPU memory pressure, and whether the model weights are cached in device memory. A request routed to a node that just evicted the target model from GPU memory can take 10× longer than one routed to a node with the model already resident.
Model-aware load balancing solves this by incorporating model-specific metadata into routing decisions. For read-heavy inference workloads, consistent hashing pins a model’s requests to a stable subset of nodes, maximizing GPU cache hit rates. For write-heavy workloads like online learning or embedding index updates, least-loaded routing distributes work evenly to prevent hotspots. A well-designed inference gateway combines both strategies with circuit breaker patterns that prevent cascading failures when individual serving nodes degrade.
The circuit breaker pattern — borrowed from electrical engineering — adds a state machine (closed → open → half-open) around each backend connection. When a node’s error rate or latency exceeds a threshold, the breaker trips open, immediately failing fast rather than queuing requests behind a degraded node. This is critical in ML serving where a single slow node can exhaust connection pools and propagate backpressure across the entire inference fleet.
Consistent hashing maps model IDs onto a hash ring, assigning each model to a stable set of k replica nodes. When a node joins or leaves, only 1/n of keys remap rather than reshuffling the entire assignment. For ML inference, this means model weights stay resident in GPU memory on the same nodes across requests, avoiding the 2–30 second penalty of loading weights from disk or network storage. For instance, using virtual nodes (typically 150–200 per physical node) ensures even distribution across heterogeneous GPU hardware.
Write-heavy operations — embedding index updates, model fine-tuning, gradient aggregation — benefit from least-loaded (also called “power of two choices”) routing. The balancer samples two or more candidate nodes and routes to the one with the lowest in-flight request count or queue depth. This prevents convoy effects where a temporarily slow node accumulates an unbounded backlog. A production system tracks per-node GPU utilization, pending batch sizes, and memory pressure to compute a composite load score.
The circuit breaker operates in three states: Closed (normal operation, requests flow through), Open (node is considered failed, requests immediately return an error or are rerouted), and Half-Open (a probe request is sent to test recovery). Transitions are governed by configurable thresholds: consecutive failures, error rate over a sliding window, or p99 latency exceeding an SLO. The half-open state uses exponential backoff with jitter to avoid thundering herd recovery storms.
Rather than a binary healthy/unhealthy signal, a production load balancer maintains a continuous health score per backend node ranging from 0.0 (dead) to 1.0 (fully healthy). The score is a weighted combination of recent error rate (40%), p99 latency relative to baseline (30%), GPU memory availability (20%), and active connection count (10%). Routing weight is proportional to this score, allowing graceful degradation — a node at 0.6 health still receives 60% of its normal traffic share rather than being fully removed.
For latency-critical inference paths, hedged requests work as follows: after a configurable delay (typically the p95 latency), a duplicate request is sent to an alternate node. The first response wins and the other is cancelled. This trades a small amount of extra compute (typically 2–5% overhead) for dramatically tighter tail latency. Hedging is only enabled when the circuit breaker is closed and the system-wide load is below a configurable threshold to prevent amplification during overload.
The core abstraction defines a LoadBalancer interface that all routing strategies implement. Each strategy receives a request context containing the model ID, request type (inference vs. update), and urgency level, and returns the best backend node.
package mlbalancer
import (
"context"
"time"
)
// Backend represents a single ML serving node.
type Backend struct {
ID string
Address string
GPUMemoryMB int64
LoadedModels map[string]bool
HealthScore float64
InFlight int64
}
// RequestContext carries model-aware routing metadata.
type RequestContext struct {
ModelID string
RequestType RequestType // Inference or Update
Urgency Urgency // Normal, High, Critical
PayloadSize int64
}
type RequestType int
const (
Inference RequestType = iota
Update
)
type Urgency int
const (
Normal Urgency = iota
High
Critical
)
// LoadBalancer selects a backend for a given request.
type LoadBalancer interface {
// Pick returns the best backend for the request context.
Pick(ctx context.Context, req *RequestContext) (*Backend, error)
// UpdateBackends refreshes the set of available backends.
UpdateBackends(backends []*Backend)
// ReportResult feeds request outcome back for adaptive routing.
ReportResult(backendID string, latency time.Duration, err error)
}The consistent hash ring uses xxhash for fast, non-cryptographic hashing and virtual nodes for even distribution. When a model ID is hashed, the ring returns the k nearest distinct physical nodes, enabling replica-aware routing.
package mlbalancer
import (
"fmt"
"sort"
"sync"
"github.com/cespare/xxhash/v2"
)
const defaultVirtualNodes = 150
// HashRing implements consistent hashing with virtual nodes.
type HashRing struct {
mu sync.RWMutex
vnodes int
ring []uint64 // sorted hash values
hashToNode map[uint64]string // hash → physical node ID
nodes map[string]*Backend
}
func NewHashRing(vnodes int) *HashRing {
if vnodes <= 0 {
vnodes = defaultVirtualNodes
}
return &HashRing{
vnodes: vnodes,
hashToNode: make(map[uint64]string),
nodes: make(map[string]*Backend),
}
}
// AddNode inserts a backend with its virtual nodes onto the ring.
func (h *HashRing) AddNode(b *Backend) {
h.mu.Lock()
defer h.mu.Unlock()
h.nodes[b.ID] = b
for i := 0; i < h.vnodes; i++ {
key := fmt.Sprintf("%s#%d", b.ID, i)
hash := xxhash.Sum64String(key)
h.ring = append(h.ring, hash)
h.hashToNode[hash] = b.ID
}
sort.Slice(h.ring, func(i, j int) bool {
return h.ring[i] < h.ring[j]
})
}
// GetNodes returns the k closest distinct physical nodes for a key.
func (h *HashRing) GetNodes(key string, k int) []*Backend {
h.mu.RLock()
defer h.mu.RUnlock()
if len(h.ring) == 0 {
return nil
}
hash := xxhash.Sum64String(key)
idx := sort.Search(len(h.ring), func(i int) bool {
return h.ring[i] >= hash
})
if idx == len(h.ring) {
idx = 0 // wrap around
}
seen := make(map[string]bool)
var result []*Backend
for i := 0; i < len(h.ring) && len(result) < k; i++ {
pos := (idx + i) % len(h.ring)
nodeID := h.hashToNode[h.ring[pos]]
if !seen[nodeID] {
seen[nodeID] = true
if b, ok := h.nodes[nodeID]; ok {
result = append(result, b)
}
}
}
return result
}The circuit breaker tracks failures within a sliding window and transitions between states atomically. The half-open state permits a single probe request to test whether the backend has recovered before fully closing the breaker.
package mlbalancer
import (
"sync"
"sync/atomic"
"time"
)
type BreakerState int32
const (
StateClosed BreakerState = iota // normal operation
StateOpen // failing fast
StateHalfOpen // probing recovery
)
// CircuitBreaker wraps a single backend connection.
type CircuitBreaker struct {
mu sync.Mutex
state atomic.Int32
failureCount int64
successCount int64
failureThreshold int64
successThreshold int64 // successes needed in half-open to close
openTimeout time.Duration // how long to stay open before half-open
lastFailure time.Time
slidingWindow []time.Time // timestamps of recent failures
windowSize time.Duration
halfOpenPermit atomic.Bool // only one probe at a time
}
// NewCircuitBreaker creates a breaker with sensible ML serving defaults.
func NewCircuitBreaker(failThreshold, successThreshold int64, openTimeout time.Duration) *CircuitBreaker {
cb := &CircuitBreaker{
failureThreshold: failThreshold,
successThreshold: successThreshold,
openTimeout: openTimeout,
windowSize: 30 * time.Second,
}
cb.state.Store(int32(StateClosed))
return cb
}
// Allow checks whether a request should be permitted through.
func (cb *CircuitBreaker) Allow() bool {
state := BreakerState(cb.state.Load())
switch state {
case StateClosed:
return true
case StateOpen:
cb.mu.Lock()
defer cb.mu.Unlock()
if time.Since(cb.lastFailure) > cb.openTimeout {
cb.state.Store(int32(StateHalfOpen))
cb.halfOpenPermit.Store(true)
return true
}
return false
case StateHalfOpen:
// Only allow one probe request at a time.
return cb.halfOpenPermit.CompareAndSwap(true, false)
}
return false
}
// RecordSuccess records a successful request.
func (cb *CircuitBreaker) RecordSuccess() {
state := BreakerState(cb.state.Load())
if state == StateHalfOpen {
cb.mu.Lock()
defer cb.mu.Unlock()
cb.successCount++
if cb.successCount >= cb.successThreshold {
cb.state.Store(int32(StateClosed))
cb.failureCount = 0
cb.successCount = 0
cb.slidingWindow = nil
}
cb.halfOpenPermit.Store(true) // allow next probe
return
}
// In closed state, reset consecutive failure tracking.
cb.mu.Lock()
cb.failureCount = 0
cb.mu.Unlock()
}
// RecordFailure records a failed request and may trip the breaker.
func (cb *CircuitBreaker) RecordFailure() {
cb.mu.Lock()
defer cb.mu.Unlock()
now := time.Now()
cb.lastFailure = now
state := BreakerState(cb.state.Load())
if state == StateHalfOpen {
// Probe failed — reopen immediately.
cb.state.Store(int32(StateOpen))
cb.successCount = 0
return
}
// Sliding window: prune old failures, add new one.
cutoff := now.Add(-cb.windowSize)
pruned := cb.slidingWindow[:0]
for _, t := range cb.slidingWindow {
if t.After(cutoff) {
pruned = append(pruned, t)
}
}
pruned = append(pruned, now)
cb.slidingWindow = pruned
if int64(len(cb.slidingWindow)) >= cb.failureThreshold {
cb.state.Store(int32(StateOpen))
}
}The health checker runs as a background goroutine, periodically probing each backend and updating its composite health score. Scores feed directly into routing weight calculations.
package mlbalancer
import (
"context"
"math"
"sync"
"time"
)
// HealthChecker periodically evaluates backend health.
type HealthChecker struct {
mu sync.RWMutex
backends map[string]*BackendHealth
interval time.Duration
prober Prober
}
type BackendHealth struct {
Backend *Backend
ErrorRate float64 // 0.0–1.0 over sliding window
P99Latency time.Duration
BaseLatency time.Duration // established baseline p99
GPUMemFree float64 // fraction of GPU memory available
ActiveConns int64
Score float64 // composite 0.0–1.0
LastProbe time.Time
}
// Prober sends a lightweight health probe to a backend.
type Prober interface {
Probe(ctx context.Context, addr string) (latency time.Duration, gpuMemFree float64, err error)
}
// ComputeScore calculates the composite health score.
func (bh *BackendHealth) ComputeScore() float64 {
// Error rate component (40% weight): 1.0 when no errors, 0.0 at 100%.
errScore := 1.0 - bh.ErrorRate
// Latency component (30% weight): ratio of baseline to current p99.
latScore := 1.0
if bh.P99Latency > 0 && bh.BaseLatency > 0 {
ratio := float64(bh.BaseLatency) / float64(bh.P99Latency)
latScore = math.Min(ratio, 1.0) // cap at 1.0
}
// GPU memory component (20% weight).
gpuScore := bh.GPUMemFree
// Connection pressure component (10% weight).
connScore := 1.0
if bh.ActiveConns > 100 {
connScore = math.Max(0.0, 1.0-float64(bh.ActiveConns-100)/500.0)
}
bh.Score = 0.4*errScore + 0.3*latScore + 0.2*gpuScore + 0.1*connScore
return bh.Score
}
// Run starts the periodic health check loop.
func (hc *HealthChecker) Run(ctx context.Context) {
ticker := time.NewTicker(hc.interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
hc.mu.RLock()
backends := make([]*BackendHealth, 0, len(hc.backends))
for _, bh := range hc.backends {
backends = append(backends, bh)
}
hc.mu.RUnlock()
for _, bh := range backends {
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
lat, gpuMem, err := hc.prober.Probe(probeCtx, bh.Backend.Address)
cancel()
hc.mu.Lock()
if err != nil {
bh.ErrorRate = bh.ErrorRate*0.8 + 0.2*1.0 // EWMA
} else {
bh.ErrorRate = bh.ErrorRate*0.8 + 0.2*0.0
bh.P99Latency = lat
bh.GPUMemFree = gpuMem
}
bh.LastProbe = time.Now()
bh.ComputeScore()
bh.Backend.HealthScore = bh.Score
hc.mu.Unlock()
}
}
}
}A production load balancer selects a routing strategy based on request type. Inference requests use consistent hashing for cache affinity, while update requests use least-loaded routing. Both strategies are wrapped in circuit breakers and respect health scores.
package mlbalancer
import (
"context"
"errors"
"math/rand"
"sync"
"time"
)
var ErrNoHealthyBackend = errors.New("no healthy backend available")
// CompositeBalancer selects strategy based on request type.
type CompositeBalancer struct {
mu sync.RWMutex
ring *HashRing
backends []*Backend
breakers map[string]*CircuitBreaker
replicas int // number of hash ring replicas to consider
}
func NewCompositeBalancer(replicas int) *CompositeBalancer {
return &CompositeBalancer{
ring: NewHashRing(defaultVirtualNodes),
breakers: make(map[string]*CircuitBreaker),
replicas: replicas,
}
}
func (cb *CompositeBalancer) Pick(ctx context.Context, req *RequestContext) (*Backend, error) {
switch req.RequestType {
case Inference:
return cb.pickConsistentHash(req)
case Update:
return cb.pickLeastLoaded(req)
default:
return cb.pickLeastLoaded(req)
}
}
// pickConsistentHash routes inference to consistent hash targets,
// falling back through replicas if circuit breakers are open.
func (cb *CompositeBalancer) pickConsistentHash(req *RequestContext) (*Backend, error) {
cb.mu.RLock()
defer cb.mu.RUnlock()
candidates := cb.ring.GetNodes(req.ModelID, cb.replicas)
for _, backend := range candidates {
breaker, ok := cb.breakers[backend.ID]
if !ok {
continue
}
if breaker.Allow() && backend.HealthScore > 0.2 {
return backend, nil
}
}
return nil, ErrNoHealthyBackend
}
// pickLeastLoaded implements power-of-two-choices for write workloads.
func (cb *CompositeBalancer) pickLeastLoaded(req *RequestContext) (*Backend, error) {
cb.mu.RLock()
defer cb.mu.RUnlock()
healthy := make([]*Backend, 0)
for _, b := range cb.backends {
breaker, ok := cb.breakers[b.ID]
if ok && breaker.Allow() && b.HealthScore > 0.2 {
healthy = append(healthy, b)
}
}
if len(healthy) == 0 {
return nil, ErrNoHealthyBackend
}
if len(healthy) == 1 {
return healthy[0], nil
}
// Power of two choices: pick two random candidates, return least loaded.
i := rand.Intn(len(healthy))
j := rand.Intn(len(healthy) - 1)
if j >= i {
j++
}
a, b := healthy[i], healthy[j]
// Composite load: in-flight requests weighted by inverse health score.
loadA := float64(a.InFlight) / a.HealthScore
loadB := float64(b.InFlight) / b.HealthScore
if loadA <= loadB {
return a, nil
}
return b, nil
}
func (cb *CompositeBalancer) UpdateBackends(backends []*Backend) {
cb.mu.Lock()
defer cb.mu.Unlock()
cb.backends = backends
cb.ring = NewHashRing(defaultVirtualNodes)
for _, b := range backends {
cb.ring.AddNode(b)
if _, ok := cb.breakers[b.ID]; !ok {
cb.breakers[b.ID] = NewCircuitBreaker(5, 3, 30*time.Second)
}
}
}
func (cb *CompositeBalancer) ReportResult(backendID string, latency time.Duration, err error) {
cb.mu.RLock()
breaker, ok := cb.breakers[backendID]
cb.mu.RUnlock()
if !ok {
return
}
if err != nil {
breaker.RecordFailure()
} else {
breaker.RecordSuccess()
}
}For latency-sensitive inference, a hedged request fires a backup after the p95 latency threshold. The first response wins and the loser’s context is cancelled.
package mlbalancer
import (
"context"
"time"
)
// HedgedRequest sends a primary request and, if it hasn’t completed
// after hedgeDelay, fires a backup to an alternate backend.
func HedgedRequest(
ctx context.Context,
balancer *CompositeBalancer,
req *RequestContext,
hedgeDelay time.Duration,
doInference func(ctx context.Context, backend *Backend) ([]byte, error),
) ([]byte, error) {
primary, err := balancer.Pick(ctx, req)
if err != nil {
return nil, err
}
type result struct {
data []byte
err error
}
ch := make(chan result, 2)
primaryCtx, primaryCancel := context.WithCancel(ctx)
defer primaryCancel()
// Launch primary request.
go func() {
data, err := doInference(primaryCtx, primary)
ch <- result{data, err}
}()
// Wait for hedgeDelay, then launch backup if primary hasn’t returned.
timer := time.NewTimer(hedgeDelay)
defer timer.Stop()
select {
case r := <-ch:
return r.data, r.err
case <-timer.C:
// Primary is slow — hedge with a backup.
backup, err := balancer.Pick(ctx, req)
if err != nil || backup.ID == primary.ID {
// No alternate available; wait for primary.
r := <-ch
return r.data, r.err
}
hedgeCtx, hedgeCancel := context.WithCancel(ctx)
defer hedgeCancel()
go func() {
data, err := doInference(hedgeCtx, backup)
ch <- result{data, err}
}()
// First response wins.
r := <-ch
return r.data, r.err
case <-ctx.Done():
return nil, ctx.Err()
}
}