Load one model at a time

Why one owner

Loading a model is the moment of highest memory use. The runtime reads weights, builds work buffers, and may compile or specialize for the device. If a chat screen and a background summary both start a load at once, peak memory can be close to double. Many teams see the crash only when two features meet.

The fix is simple: one object owns all models. It loads one at a time, keeps track of what is loaded, and unloads on memory pressure.

iOS: an actor that owns the model

A Swift actor runs one call at a time on its own state. But an await inside an actor lets another call start. So we keep a single in-flight Task and make new callers wait for it.

import UIKit

protocol LoadedModel: AnyObject, Sendable {
    var name: String { get }
}

actor ModelOwner {
    private var current: LoadedModel?
    private var inFlight: Task<LoadedModel, Error>?
    private let gate = MemoryGate()

    enum OwnerError: Error { case doesNotFit(String) }

    func model(_ spec: ModelSpec, load: @escaping @Sendable () async throws -> LoadedModel) async throws -> LoadedModel {
        if let current, current.name == spec.name { return current }
        if let inFlight { _ = try? await inFlight.value }   // wait for the other load
        if let current, current.name == spec.name { return current }

        current = nil                                         // free the old model first
        guard gate.decide(for: spec) == .load else { throw OwnerError.doesNotFit(spec.name) }

        let task = Task { try await load() }
        inFlight = task
        defer { inFlight = nil }
        let model = try await task.value
        current = model
        return model
    }

    func unloadAll(reason: String) {
        print("unload:", reason)
        current = nil
    }
}

@MainActor
final class MemoryPressureWatcher {
    private var token: NSObjectProtocol?
    init(owner: ModelOwner) {
        token = NotificationCenter.default.addObserver(
            forName: UIApplication.didReceiveMemoryWarningNotification,
            object: nil, queue: .main
        ) { _ in
            Task { await owner.unloadAll(reason: "memory warning") }
        }
    }
}

didReceiveMemoryWarningNotification is the warning that comes before a jetsam kill. Dropping the model here is often the difference between "the next answer is slower" and "the app is gone".

Android: one owner and onTrimMemory

On Android, use a Mutex from Kotlin coroutines so loads run one at a time. Listen to onTrimMemory through ComponentCallbacks2 and close the model when the level is high.

import android.content.ComponentCallbacks2
import android.content.res.Configuration
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

class ModelOwner<M : AutoCloseable>(
    private val scope: CoroutineScope,
    private val loader: suspend (String) -> M,
) : ComponentCallbacks2 {
    private val mutex = Mutex()
    private var currentName: String? = null
    private var current: M? = null

    suspend fun get(name: String): M = mutex.withLock {
        current?.let { if (currentName == name) return@withLock it }
        current?.close()            // free the old model first
        current = null
        val model = loader(name)
        current = model
        currentName = name
        model
    }

    suspend fun unload(reason: String) = mutex.withLock {
        android.util.Log.i("ModelOwner", "unload: $reason")
        current?.close()
        current = null
        currentName = null
    }

    override fun onTrimMemory(level: Int) {
        if (level >= ComponentCallbacks2.TRIM_MEMORY_BACKGROUND) {
            scope.launch { unload("onTrimMemory $level") }   // never block the main thread
        }
    }
    override fun onConfigurationChanged(newConfig: Configuration) {}
    @Deprecated("Deprecated in Java") override fun onLowMemory() {}
}

Register it with application.registerComponentCallbacks(owner). If you use LiteRT-LM, close the Conversation safely before you close the Engine. That order has its own crash, which the LiteRT-LM course covers.

Fallback ladder

Write down what happens when a load fails or is refused. A clear order keeps the feature working:

1. Full model on GPU or Neural Engine.
2. Smaller model on the same device.
3. Same small model on CPU only.
4. Server or Private Cloud Compute path, if your privacy rules allow it.
5. A clear message: "This feature needs more free memory. Close other apps and try again."

Sources

• Apple: didReceiveMemoryWarningNotification
• Android: ComponentCallbacks2 (onTrimMemory)