From 42882280e96b888fd327ff56e2e46a329ad17c45 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 18 Jul 2026 16:36:34 +0000 Subject: [PATCH] Shrink ORT arena per window and abort separation on memory warnings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Device still jetsams during stem separation with prepacking disabled: the transformer's attention intermediates spike ONNX Runtime's CPU arena to its high-water mark, which is never returned between windows and sits on top of live playback for the whole track. Two levers: - Per-run arena shrinkage (memory.enable_memory_arena_shrinkage) frees the spike after every 7.8 s window; re-growing costs some CPU, which an offline cache-once job can afford. - iOS memory warnings now abort the in-flight separation (partial stems are deleted; the track plays without stems and retries on a later ensureStems pass) and drop the cached session — wired from RootView via PreparationQueue.handleMemoryPressure(). --- App/Continuity/Views/RootView.swift | 6 ++++ .../Ingest/PreparationQueue+Stems.swift | 8 +++++ .../Sources/Ingest/StemSeparator.swift | 33 +++++++++++++++++-- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/App/Continuity/Views/RootView.swift b/App/Continuity/Views/RootView.swift index 564f99f..b144755 100644 --- a/App/Continuity/Views/RootView.swift +++ b/App/Continuity/Views/RootView.swift @@ -93,6 +93,12 @@ struct RootView: View { // per-playlist lastSyncedAt staleness (+ failure backoff) is the real rate limiter. // Main-runloop timers pause while suspended, but guard on scenePhase anyway for the // coalesced fire a resume delivers before the app is active again. + // Memory warning precedes a jetsam kill: abort any in-flight stem separation and + // free the ONNX session. Losing stems for one track beats losing the process. + .onReceive(NotificationCenter.default.publisher( + for: UIApplication.didReceiveMemoryWarningNotification)) { _ in + prepQueue.handleMemoryPressure() + } .onReceive(syncTick) { _ in guard scenePhase == .active else { return } prepQueue.autoSyncIfNeeded(in: modelContext) diff --git a/Packages/ContinuityKit/Sources/Ingest/PreparationQueue+Stems.swift b/Packages/ContinuityKit/Sources/Ingest/PreparationQueue+Stems.swift index c51a8a9..f096c16 100644 --- a/Packages/ContinuityKit/Sources/Ingest/PreparationQueue+Stems.swift +++ b/Packages/ContinuityKit/Sources/Ingest/PreparationQueue+Stems.swift @@ -30,6 +30,14 @@ extension PreparationQueue { } } + /// System memory warning: iOS warns before it jetsams. Abort any in-flight stem separation + /// (it deletes its partial output; the track plays without stems and retries on a later + /// `ensureStems` pass) and drop the cached ORT session — staying alive beats finishing stems. + public func handleMemoryPressure() { + OnnxStemSeparator.requestAbort() + OnnxStemSeparator.releaseSession() + } + /// Brings a track's stem links in line with the disk: links cached stems that exist unlinked /// (e.g. a re-added video), clears links whose files were evicted so `hasStems` tells the /// truth and the track becomes eligible for re-separation. diff --git a/Packages/ContinuityKit/Sources/Ingest/StemSeparator.swift b/Packages/ContinuityKit/Sources/Ingest/StemSeparator.swift index f9eea9d..001f8ca 100644 --- a/Packages/ContinuityKit/Sources/Ingest/StemSeparator.swift +++ b/Packages/ContinuityKit/Sources/Ingest/StemSeparator.swift @@ -84,6 +84,25 @@ final class OnnxStemSeparator: StemSeparating { cachedSession = nil } + /// Set on a system memory warning; the streaming loop checks it before each window and + /// bails, deleting the partial stems. The track keeps playing without stems and re-enters + /// separation on a later `ensureStems` pass. Cleared when the next separation starts. + nonisolated(unsafe) private static var abortRequested = false + private static let abortLock = NSLock() + + static var isAbortRequested: Bool { + abortLock.lock(); defer { abortLock.unlock() } + return abortRequested + } + + static func requestAbort() { + abortLock.lock(); abortRequested = true; abortLock.unlock() + } + + static func clearAbort() { + abortLock.lock(); abortRequested = false; abortLock.unlock() + } + private static func sharedSession(modelURL: URL) throws -> ORTSession { sessionLock.lock() defer { sessionLock.unlock() } @@ -109,6 +128,7 @@ final class OnnxStemSeparator: StemSeparating { } func separate(inputURL: URL, vocalsOut: URL, accompanimentOut: URL) throws -> StemPaths { + Self.clearAbort() // a stale warning must not kill this fresh attempt // ONNX Runtime session, CPU EP everywhere. The CoreML EP jetsammed real devices: // converting/compiling the 85M-param HT-Demucs transformer at session creation ballooned // the app past the per-process limit (~3.4 GB observed on iPhone 17 Pro — JetsamEvent @@ -230,13 +250,22 @@ final class OnnxStemSeparator: StemSeparating { /// returns just the vocals source (index 3) as planar (2, segment). private static func ortInference(session: ORTSession, segment: Int, channels: Int, vocalsIndex: Int) -> WindowInference { - { input in + // Shrink the CPU arena back after every Run: the transformer's attention intermediates + // spike the arena to its high-water mark, and by default ORT never returns that memory + // between windows — it sits on top of live playback for the whole track (jetsam margin). + // Re-growing per window costs a little CPU; an offline cache-once job can afford it. + let runOptions = try? ORTRunOptions() + try? runOptions?.addConfigEntry(withKey: "memory.enable_memory_arena_shrinkage", value: "cpu:0") + return { input in let inputData = input.withUnsafeBytes { NSMutableData(bytes: $0.baseAddress, length: $0.count) } do { + if Self.isAbortRequested { + throw StemSeparationError.inference("aborted under memory pressure") + } let inputValue = try ORTValue(tensorData: inputData, elementType: .float, shape: [1, NSNumber(value: channels), NSNumber(value: segment)]) let outputs = try session.run(withInputs: ["mix": inputValue], - outputNames: ["stems"], runOptions: nil) + outputNames: ["stems"], runOptions: runOptions) guard let out = outputs["stems"] else { throw StemSeparationError.inference("no 'stems' output") } let outData = try out.tensorData() let vocBase = (vocalsIndex * channels) * segment