Model serving systems that host dozens or hundreds of models simultaneously face a fundamental tension: loading every model into memory at startup ensures low first-request latency but consumes enormous resources, while loading models only when requested (lazy loading) conserves memory but introduces cold-start penalties. Lazy model loading with reference counting resolves this tension by deferring model initialization until the first inference request, then managing the model’s lifecycle through a combination of TTL-based expiration, LRU eviction under a global memory budget, and atomic reference counting that prevents the eviction of models currently serving active requests.
In a system where ML models are embedded directly into the query processing pipeline for tasks such as semantic reranking and embedding generation, the number of distinct models can far exceed available memory. A single serving node might serve reranking models for multiple tenants, each fine-tuned on domain-specific data, alongside shared embedding models of varying dimensionality. Without a disciplined lifecycle management strategy, either memory is exhausted by eagerly loaded models or query latency becomes unpredictable due to on-demand loading stalls.
The architecture described here introduces a ModelHandle abstraction that wraps each loaded model with an atomic reference counter, a TTL expiration timestamp, and LRU tracking metadata. A background reaper goroutine periodically scans for expired handles, but it cannot reclaim a handle whose reference count is non-zero—guaranteeing that in-flight inference requests always complete against a valid model. The global ModelBudget enforces a hard ceiling on total model memory, triggering LRU eviction when a new model load would exceed the budget. All coordination uses lock-free atomic operations on the hot path, falling back to a mutex only during eviction and reaping—operations that occur orders of magnitude less frequently than inference.
Eager loading initializes all registered models at startup, paying the full memory and initialization cost upfront in exchange for guaranteed low latency on every request. Lazy loading defers model initialization until the first inference call for that model, amortizing startup cost and avoiding memory allocation for models that may never be requested. In practice, hybrid strategies are common: frequently accessed models are eagerly loaded (warm-up), while long-tail models use lazy loading. The ModelConfig struct exposes a LoadPolicy field—either LoadEager or LoadLazy—allowing operators to tune per-model behavior based on access patterns and SLA requirements.
Each ModelHandle maintains an atomic.Int32 reference counter that tracks the number of active inference requests using the model. When a goroutine acquires a handle for inference, it atomically increments the refcount via Acquire(); when inference completes, it decrements via Release(). The eviction and reaping subsystems check the refcount before reclaiming a handle—if the count is non-zero, the model is skipped and retried on the next pass. This approach eliminates the need for a read-write lock on the hot inference path, as atomic.Int32.Add is a single CPU instruction on x86 (LOCK XADD) and ARM (LDADDAL), costing roughly 10–20 nanoseconds compared to 50–200 nanoseconds for an uncontended sync.RWMutex lock/unlock pair.
Every loaded model is assigned a time-to-live (TTL) duration, typically configured per model or per category. Each inference request resets the model’s TTL deadline by updating an atomic.Int64 that stores the Unix nanosecond timestamp of the earliest permissible expiration. A background reaper goroutine wakes at a configurable interval (e.g., every 30 seconds), scans all loaded handles, and marks those whose TTL has elapsed and whose refcount is zero as candidates for unloading. TTL-based expiration complements LRU eviction: LRU handles memory pressure spikes, while TTL ensures that models which were accessed once but never again are cleaned up even when overall memory remains within budget.
The ModelBudget defines a hard upper bound on the total bytes consumed by all loaded models. When a lazy load is triggered and the new model’s estimated memory footprint would cause the total to exceed the budget, the loader invokes the LRU eviction policy. The eviction scan walks the doubly-linked LRU list from the tail (least recently used) toward the head, skipping entries whose reference count is non-zero. Eviction continues until sufficient memory is freed or the list is exhausted. If eviction cannot free enough memory—because all loaded models are in active use—the load request returns ErrBudgetExhausted, signaling the caller to retry or degrade gracefully.
The hot path—acquiring a model handle, running inference, and releasing the handle—uses only atomic operations: an atomic load to check the handle’s validity, an atomic increment for the refcount, and an atomic store to update the TTL deadline. No mutexes are held during inference. The cold path—loading a new model, evicting under budget pressure, or reaping expired handles—acquires a sync.Mutex on the model registry. This two-tier locking strategy ensures that the common case (inference on a loaded model) proceeds without contention, while rare structural mutations (load/evict/reap) are serialized for safety. The sync.Mutex is never held while performing I/O (model deserialization, network fetch), further reducing lock contention.
The ModelHandle is the central abstraction. It wraps the underlying model resource with lifecycle metadata—a reference counter, TTL deadline, memory footprint, and LRU linkage. All fields accessed on the inference hot path use atomic types to avoid mutex contention.
// ModelHandle wraps a loaded model with lifecycle management metadata.
type ModelHandle struct {
ModelID string
Model interface{} // the loaded model (backend-specific)
MemoryBytes uint64
// Atomic fields for lock-free hot-path access
refCount atomic.Int32 // active inference count
ttlDeadline atomic.Int64 // Unix nanosecond expiration timestamp
lastAccess atomic.Int64 // Unix nanosecond of last inference call
state atomic.Int32 // 0=unloaded, 1=loading, 2=loaded, 3=evicting
// Protected by registry mutex (cold path only)
lruElement *list.Element
loadedAt time.Time
config ModelConfig
}
// Acquire increments the refcount, preventing eviction.
// Returns false if the handle is in an evicting or unloaded state.
func (h *ModelHandle) Acquire() bool {
for {
st := h.state.Load()
if st != 2 { // not in "loaded" state
return false
}
h.refCount.Add(1)
// Double-check state after incrementing refcount to prevent
// race with concurrent eviction that checked refcount == 0
// before our Add landed.
if h.state.Load() != 2 {
h.refCount.Add(-1)
return false
}
now := time.Now().UnixNano()
h.lastAccess.Store(now)
h.ttlDeadline.Store(now + int64(h.config.TTL))
return true
}
}
// Release decrements the refcount after inference completes.
func (h *ModelHandle) Release() {
new := h.refCount.Add(-1)
if new < 0 {
panic(fmt.Sprintf("ModelHandle %s: negative refcount %d", h.ModelID, new))
}
}
// InUse reports whether the model has active references.
func (h *ModelHandle) InUse() bool {
return h.refCount.Load() > 0
}
// Expired reports whether the model’s TTL has elapsed.
func (h *ModelHandle) Expired() bool {
return time.Now().UnixNano() > h.ttlDeadline.Load()
}The lazy loader intercepts inference requests for models that are not yet loaded. It uses a state machine on the ModelHandle to coordinate concurrent load attempts—only one goroutine performs the actual load, while others wait on a sync.Once-like mechanism. Once loaded, the handle transitions to the loaded state and all waiters proceed.
// LazyLoader manages on-demand model loading with deduplication.
type LazyLoader struct {
registry *ModelRegistry
backend Backend
budget *ModelBudget
loadGroups sync.Map // map[string]*singleflight.Group
}
// GetOrLoad returns a loaded ModelHandle, triggering a lazy load if necessary.
func (ll *LazyLoader) GetOrLoad(ctx context.Context, modelID string) (*ModelHandle, error) {
// Fast path: model already loaded
if handle, ok := ll.registry.Get(modelID); ok {
if handle.Acquire() {
return handle, nil
}
// Handle exists but is being evicted; fall through to reload
}
// Deduplicate concurrent loads for the same model using singleflight
group, _ := ll.loadGroups.LoadOrStore(modelID, &singleflight.Group{})
sf := group.(*singleflight.Group)
result, err, _ := sf.Do(modelID, func() (interface{}, error) {
return ll.loadModel(ctx, modelID)
})
ll.loadGroups.Delete(modelID)
if err != nil {
return nil, fmt.Errorf("lazy load model %q: %w", modelID, err)
}
handle := result.(*ModelHandle)
if !handle.Acquire() {
return nil, fmt.Errorf("model %q loaded but immediately evicted", modelID)
}
return handle, nil
}
// loadModel performs the actual model load, enforcing budget constraints.
func (ll *LazyLoader) loadModel(ctx context.Context, modelID string) (*ModelHandle, error) {
meta, err := ll.registry.GetMeta(modelID)
if err != nil {
return nil, err
}
// Enforce budget before loading
if err := ll.budget.Reserve(meta.EstimatedBytes); err != nil {
// Attempt eviction to free space
freed, evictErr := ll.registry.EvictLRU(meta.EstimatedBytes)
if evictErr != nil {
return nil, fmt.Errorf("budget exhausted, eviction failed: %w", evictErr)
}
log.Printf("evicted %d bytes to make room for model %s", freed, modelID)
}
// Load the model via the backend
model, memUsed, err := ll.backend.LoadModel(ctx, meta)
if err != nil {
ll.budget.Release(meta.EstimatedBytes)
return nil, err
}
handle := &ModelHandle{
ModelID: modelID,
Model: model,
MemoryBytes: memUsed,
loadedAt: time.Now(),
config: meta.Config,
}
handle.state.Store(2) // loaded
now := time.Now().UnixNano()
handle.lastAccess.Store(now)
handle.ttlDeadline.Store(now + int64(meta.Config.TTL))
ll.registry.Put(handle)
return handle, nil
}The reaper runs as a background goroutine that periodically scans for expired model handles. It only unloads handles that are both expired and not in use, ensuring zero interference with active inference requests.
// Reaper periodically unloads expired models with zero active references.
type Reaper struct {
registry *ModelRegistry
backend Backend
budget *ModelBudget
interval time.Duration
stopCh chan struct{}
}
// Run starts the reaper loop. Call Stop() to terminate.
func (r *Reaper) Run() {
ticker := time.NewTicker(r.interval)
defer ticker.Stop()
for {
select {
case <-ticker.C:
r.sweep()
case <-r.stopCh:
return
}
}
}
// sweep scans all loaded handles and unloads expired, unreferenced ones.
func (r *Reaper) sweep() {
candidates := r.registry.ListLoaded()
var reaped int
for _, handle := range candidates {
if !handle.Expired() {
continue
}
if handle.InUse() {
// Model is expired but still serving requests; skip for now.
// It will be reaped on a future sweep after all references are released.
continue
}
// Attempt to transition state from loaded (2) to evicting (3)
if !handle.state.CompareAndSwap(2, 3) {
continue // another goroutine is already evicting this handle
}
// Final refcount check after state transition
if handle.InUse() {
handle.state.Store(2) // revert to loaded
continue
}
if err := r.backend.UnloadModel(handle.ModelID); err != nil {
log.Printf("reaper: failed to unload %s: %v", handle.ModelID, err)
handle.state.Store(2) // revert on failure
continue
}
r.budget.Release(handle.MemoryBytes)
r.registry.Remove(handle.ModelID)
handle.state.Store(0) // unloaded
reaped++
}
if reaped > 0 {
log.Printf("reaper: unloaded %d expired models", reaped)
}
}
// Stop terminates the reaper goroutine.
func (r *Reaper) Stop() {
close(r.stopCh)
}When the lazy loader needs to free memory to stay within the ModelBudget, it walks the LRU list from the tail and evicts unreferenced models. The refcount guard ensures that models actively serving inference are never evicted, even under extreme memory pressure.
// EvictLRU frees at least targetBytes of memory by evicting least recently
// used models that have zero active references.
func (r *ModelRegistry) EvictLRU(targetBytes uint64) (uint64, error) {
r.mu.Lock()
defer r.mu.Unlock()
var freed uint64
var next *list.Element
for elem := r.lruList.Back(); elem != nil && freed < targetBytes; elem = next {
next = elem.Prev() // capture before potential removal
handle := elem.Value.(*ModelHandle)
// Skip models with active references
if handle.InUse() {
continue
}
// Attempt state transition: loaded → evicting
if !handle.state.CompareAndSwap(2, 3) {
continue
}
// Double-check refcount after winning the CAS
if handle.InUse() {
handle.state.Store(2)
continue
}
if err := r.backend.UnloadModel(handle.ModelID); err != nil {
log.Printf("evict: failed to unload %s: %v", handle.ModelID, err)
handle.state.Store(2)
continue
}
freed += handle.MemoryBytes
r.lruList.Remove(elem)
delete(r.handles, handle.ModelID)
handle.state.Store(0)
log.Printf("evicted model %s (freed %d bytes)", handle.ModelID, handle.MemoryBytes)
}
if freed < targetBytes {
return freed, fmt.Errorf(
"eviction incomplete: freed %d of %d bytes requested; all remaining models in use",
freed, targetBytes,
)
}
return freed, nil
}The ModelBudget tracks total memory consumed by loaded models and enforces a hard ceiling. It uses an atomic.Int64 for the current usage counter, allowing lock-free reads on the inference path. Mutations (reserve/release) use a compare-and-swap loop to avoid mutex overhead.
// ModelBudget enforces a global memory ceiling for loaded models.
type ModelBudget struct {
maxBytes uint64
usedBytes atomic.Int64
}
// NewModelBudget creates a budget with the given memory ceiling.
func NewModelBudget(maxBytes uint64) *ModelBudget {
return &ModelBudget{maxBytes: maxBytes}
}
// Reserve attempts to allocate n bytes from the budget.
// Returns ErrBudgetExhausted if the allocation would exceed the ceiling.
func (b *ModelBudget) Reserve(n uint64) error {
for {
current := b.usedBytes.Load()
proposed := current + int64(n)
if proposed > int64(b.maxBytes) {
return ErrBudgetExhausted
}
if b.usedBytes.CompareAndSwap(current, proposed) {
return nil
}
// CAS failed; another goroutine mutated usedBytes. Retry.
runtime.Gosched()
}
}
// Release returns n bytes to the budget after a model is unloaded.
func (b *ModelBudget) Release(n uint64) {
new := b.usedBytes.Add(-int64(n))
if new < 0 {
panic(fmt.Sprintf("ModelBudget underflow: used=%d after releasing %d", new, n))
}
}
// Available returns the number of bytes remaining in the budget.
func (b *ModelBudget) Available() uint64 {
used := b.usedBytes.Load()
if used >= int64(b.maxBytes) {
return 0
}
return b.maxBytes - uint64(used)
}
// Usage returns current utilization as a fraction in [0.0, 1.0].
func (b *ModelBudget) Usage() float64 {
return float64(b.usedBytes.Load()) / float64(b.maxBytes)
}
// ErrBudgetExhausted is returned when a reserve would exceed the memory ceiling.
var ErrBudgetExhausted = fmt.Errorf("model budget exhausted")In production, operators configure critical models (high-QPS reranking, primary embedding) with LoadEager and long-tail models (per-tenant fine-tuned, experimental) with LoadLazy. The model registry processes eager models at startup, ensuring they are loaded and warmed up before the node advertises readiness to the load balancer. Lazy models are loaded on first access, with the singleflight deduplication preventing thundering-herd problems when multiple concurrent queries reference the same unloaded model.
// ModelConfig controls per-model lifecycle behavior.
type ModelConfig struct {
LoadPolicy LoadPolicy // LoadEager or LoadLazy
TTL time.Duration // time-to-live after last access
Priority int // lower = less likely to be evicted
WarmUp bool // run dummy inference after load
}
type LoadPolicy int
const (
LoadLazy LoadPolicy = iota
LoadEager
)
// InitRegistry loads eager models at startup and registers lazy models
// as metadata-only entries that will be loaded on first access.
func InitRegistry(ctx context.Context, configs []ModelConfig, ll *LazyLoader) error {
for _, cfg := range configs {
switch cfg.LoadPolicy {
case LoadEager:
handle, err := ll.GetOrLoad(ctx, cfg.ModelID)
if err != nil {
return fmt.Errorf("eager load %s: %w", cfg.ModelID, err)
}
handle.Release() // drop the reference from the load call
log.Printf("eager-loaded model %s (%d bytes)", cfg.ModelID, handle.MemoryBytes)
case LoadLazy:
ll.registry.RegisterMeta(cfg)
log.Printf("registered lazy model %s (will load on first access)", cfg.ModelID)
}
}
return nil
}To prevent reference leaks, the codebase uses an RAII-style pattern where Acquire returns a scoped handle that automatically calls Release when it goes out of scope. In Go, this is implemented via a deferred closure or a wrapper type with an explicit Close method compatible with resource cleanup patterns.
// ScopedHandle wraps a ModelHandle and ensures Release is called exactly once.
type ScopedHandle struct {
handle *ModelHandle
released atomic.Bool
}
// NewScopedHandle acquires the model and returns a scoped wrapper.
func NewScopedHandle(handle *ModelHandle) (*ScopedHandle, error) {
if !handle.Acquire() {
return nil, fmt.Errorf("model %s is not available for acquisition", handle.ModelID)
}
return &ScopedHandle{handle: handle}, nil
}
// Model returns the underlying model for inference.
func (s *ScopedHandle) Model() interface{} {
return s.handle.Model
}
// Close releases the reference. Safe to call multiple times.
func (s *ScopedHandle) Close() {
if s.released.CompareAndSwap(false, true) {
s.handle.Release()
}
}
// Usage in inference handler:
func (e *Engine) HandleInference(ctx context.Context, modelID string, input Tensor) (Tensor, error) {
handle, err := e.loader.GetOrLoad(ctx, modelID)
if err != nil {
return Tensor{}, err
}
scoped, err := NewScopedHandle(handle)
if err != nil {
return Tensor{}, err
}
defer scoped.Close()
return e.backend.Infer(ctx, scoped.Model(), input)
}Effective model lifecycle management requires visibility into load latency, eviction frequency, refcount highwater marks, and budget utilization. The following pattern exposes these metrics via atomic counters that can be scraped by Prometheus or similar systems without introducing mutex contention on the hot path.
// LifecycleMetrics tracks model loading and eviction statistics.
type LifecycleMetrics struct {
LazyLoads atomic.Int64
EagerLoads atomic.Int64
Evictions atomic.Int64
TTLExpirations atomic.Int64
BudgetExhausted atomic.Int64
LoadLatencyNs atomic.Int64 // cumulative; divide by LazyLoads for average
PeakRefCount atomic.Int32
}
// RecordLoad updates metrics after a successful model load.
func (m *LifecycleMetrics) RecordLoad(lazy bool, latency time.Duration) {
if lazy {
m.LazyLoads.Add(1)
m.LoadLatencyNs.Add(int64(latency))
} else {
m.EagerLoads.Add(1)
}
}
// UpdatePeakRefCount atomically updates the peak if current exceeds it.
func (m *LifecycleMetrics) UpdatePeakRefCount(current int32) {
for {
peak := m.PeakRefCount.Load()
if current <= peak {
return
}
if m.PeakRefCount.CompareAndSwap(peak, current) {
return
}
}
}Acquire() must be paired with a Release(), even on error paths. A leaked reference prevents the model from ever being evicted or reaped, leading to a slow memory leak. Use the ScopedHandle pattern with defer to make reference leaks structurally impossible.Acquire() can land between the read and the CAS. Always transition state first (via CompareAndSwap), then re-check the refcount. If it’s now non-zero, revert the state transition.ErrBudgetExhausted signal. When all loaded models are in use and eviction cannot free enough memory, the budget returns an exhaustion error. Callers must handle this gracefully—for example, by queuing the request and retrying after a short backoff, or by returning a degraded (non-ML) result. Panicking or retrying in a tight loop under budget exhaustion creates cascading failures.sync.RWMutex on the inference hot path. The entire motivation for atomic reference counting is to avoid lock contention during inference. Introducing a read lock “for safety” around the model access negates this benefit. Trust the state machine and refcount protocol—they provide the same safety guarantees with orders-of-magnitude lower overhead under high concurrency.ModelBudget tracks steady-state model memory, but model deserialization can temporarily require 2–3x the final model size (e.g., decompressing weights, building execution graphs). Reserve headroom in the budget or use a separate loading quota to prevent OOM kills during concurrent model loads.