From 549a36cb6d977a67e09298534aca3021122623f5 Mon Sep 17 00:00:00 2001 From: Ciprian Date: Tue, 1 Sep 2026 12:09:41 +0400 Subject: [PATCH 1/2] fix(diarizer): propagate offline cancellation --- .../Offline/Core/OfflineDiarizerManager.swift | 133 +++++++++++------- .../OfflineDiarizerCancellationTests.swift | 100 +++++++++++++ 2 files changed, 184 insertions(+), 49 deletions(-) create mode 100644 Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift diff --git a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift index 17f04669b..e618f980f 100644 --- a/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift +++ b/Sources/FluidAudio/Diarizer/Offline/Core/OfflineDiarizerManager.swift @@ -3,6 +3,12 @@ import Accelerate import Foundation import OSLog +@available(macOS 14.0, iOS 17.0, *) +private enum OfflinePreparationWorkerResult: Sendable { + case segmentation(SegmentationOutput, TimeInterval) + case embeddings([TimedEmbedding], TimeInterval) +} + @available(macOS 14.0, iOS 17.0, *) public final class OfflineDiarizerManager { private let logger = AppLogger(category: "OfflineDiarizer") @@ -187,64 +193,93 @@ public final class OfflineDiarizerManager { let capturedModels = models let capturedConfig = config - let segmentationTask = Task.detached(priority: .userInitiated) { - [capturedModels, capturedConfig] () throws -> (SegmentationOutput, TimeInterval) in - let processor = OfflineSegmentationProcessor() - let start = Date() - do { - let segmentation = try await processor.process( - audioSource: audioSource, - segmentationModel: capturedModels.segmentationModel, - config: capturedConfig, - chunkHandler: { chunk in - progressCallback?(chunk.chunkIndex + 1, totalChunks) - switch chunkContinuation.yield(chunk) { - case .enqueued, .dropped: - return .continue - case .terminated: - return .stop - @unknown default: - return .stop + let results = try await withThrowingTaskGroup( + of: OfflinePreparationWorkerResult.self, + returning: ( + segmentation: (SegmentationOutput, TimeInterval), + embeddings: ([TimedEmbedding], TimeInterval) + ).self + ) { group in + group.addTask(priority: .userInitiated) { [capturedModels, capturedConfig] in + let processor = OfflineSegmentationProcessor() + let start = Date() + do { + let segmentation = try await processor.process( + audioSource: audioSource, + segmentationModel: capturedModels.segmentationModel, + config: capturedConfig, + chunkHandler: { chunk in + progressCallback?(chunk.chunkIndex + 1, totalChunks) + switch chunkContinuation.yield(chunk) { + case .enqueued, .dropped: + return .continue + case .terminated: + return .stop + @unknown default: + return .stop + } } - } + ) + chunkContinuation.finish() + return .segmentation( + segmentation, + Date().timeIntervalSince(start) + ) + } catch { + chunkContinuation.finish(throwing: error) + throw error + } + } + + group.addTask(priority: .userInitiated) { [capturedModels, capturedConfig] in + let extractor = OfflineEmbeddingExtractor( + fbankModel: capturedModels.fbankModel, + embeddingModel: capturedModels.embeddingModel, + pldaTransform: PLDATransform( + pldaRhoModel: capturedModels.pldaRhoModel, + psi: capturedModels.pldaPsi + ), + config: capturedConfig + ) + let start = Date() + let embeddings = try await extractor.extractEmbeddings( + audioSource: audioSource, + segmentationStream: chunkStream + ) + return .embeddings( + embeddings, + Date().timeIntervalSince(start) ) - chunkContinuation.finish() - return (segmentation, Date().timeIntervalSince(start)) + } + + var segmentationResult: (SegmentationOutput, TimeInterval)? + var embeddingResult: ([TimedEmbedding], TimeInterval)? + + do { + while let result = try await group.next() { + switch result { + case .segmentation(let segmentation, let duration): + segmentationResult = (segmentation, duration) + case .embeddings(let embeddings, let duration): + embeddingResult = (embeddings, duration) + } + } } catch { + group.cancelAll() chunkContinuation.finish(throwing: error) throw error } - } - let embeddingTask = Task.detached(priority: .userInitiated) { - [capturedModels, capturedConfig] () throws -> ([TimedEmbedding], TimeInterval) in - let extractor = OfflineEmbeddingExtractor( - fbankModel: capturedModels.fbankModel, - embeddingModel: capturedModels.embeddingModel, - pldaTransform: PLDATransform(pldaRhoModel: capturedModels.pldaRhoModel, psi: capturedModels.pldaPsi), - config: capturedConfig - ) - let start = Date() - let embeddings = try await extractor.extractEmbeddings( - audioSource: audioSource, - segmentationStream: chunkStream - ) - return (embeddings, Date().timeIntervalSince(start)) + guard let segmentationResult, let embeddingResult else { + throw OfflineDiarizationError.processingFailed( + "Offline preparation workers ended without complete results" + ) + } + return (segmentationResult, embeddingResult) } - let segmentationResult: (SegmentationOutput, TimeInterval) - let embeddingResult: ([TimedEmbedding], TimeInterval) - do { - async let awaitedSegmentation = segmentationTask.value - async let awaitedEmbeddings = embeddingTask.value - segmentationResult = try await awaitedSegmentation - embeddingResult = try await awaitedEmbeddings - } catch { - segmentationTask.cancel() - embeddingTask.cancel() - chunkContinuation.finish(throwing: error) - throw error - } + let segmentationResult = results.segmentation + let embeddingResult = results.embeddings let (segmentation, segmentationTime) = segmentationResult logger.debug("Segmentation completed in \(segmentationTime)s (async)") diff --git a/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift new file mode 100644 index 000000000..72f3054f1 --- /dev/null +++ b/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift @@ -0,0 +1,100 @@ +import AVFoundation +import XCTest + +@testable import FluidAudio + +@available(macOS 14.0, iOS 17.0, *) +final class OfflineDiarizerCancellationTests: XCTestCase { + func testCancellingFileProcessingStopsAllInferenceWorkers() async throws { + try requireOfflineDiarizerModels() + + let fixture = try DiarizationTestFixtures.fixtureAudio(sampleRate: 16_000) + var longMeeting: [Float] = [] + longMeeting.reserveCapacity(fixture.count * 60) + for _ in 0..<60 { + longMeeting.append(contentsOf: fixture) + } + + let audioURL = try writeWAV(samples: longMeeting, sampleRate: 16_000) + defer { try? FileManager.default.removeItem(at: audioURL) } + + let manager = OfflineDiarizerManager() + let progress = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) + var progressIterator = progress.stream.makeAsyncIterator() + + let processingTask = Task { + try await manager.process(audioURL) { _, _ in + progress.continuation.yield(()) + } + } + + guard await progressIterator.next() != nil else { + processingTask.cancel() + XCTFail("Diarization ended before inference began") + return + } + + let cancellationStart = ContinuousClock.now + processingTask.cancel() + + do { + _ = try await processingTask.value + XCTFail("Cancelled diarization unexpectedly completed") + } catch is CancellationError { + // Expected. Returning from the structured task group also proves both + // inference workers have stopped rather than continuing in the background. + } catch { + XCTFail("Expected CancellationError, received \(error)") + } + + let cancellationDuration = cancellationStart.duration(to: .now) + XCTAssertLessThan( + cancellationDuration, + .seconds(5), + "Cancellation should stop segmentation and embedding between model calls" + ) + } + + private func requireOfflineDiarizerModels() throws { + let repoDirectory = OfflineDiarizerModels.defaultModelsDirectory() + .appendingPathComponent(Repo.diarizer.folderName, isDirectory: true) + let allPresent = ModelNames.OfflineDiarizer.requiredModels.allSatisfy { + FileManager.default.fileExists(atPath: repoDirectory.appendingPathComponent($0).path) + } + guard allPresent else { + throw XCTSkip("Offline diarizer models not available") + } + } + + private func writeWAV(samples: [Float], sampleRate: Double) throws -> URL { + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("offline-diarizer-cancellation-\(UUID().uuidString)") + .appendingPathExtension("wav") + let format = AVAudioFormat( + commonFormat: .pcmFormatFloat32, + sampleRate: sampleRate, + channels: 1, + interleaved: false + )! + let buffer = AVAudioPCMBuffer( + pcmFormat: format, + frameCapacity: AVAudioFrameCount(samples.count) + )! + buffer.frameLength = buffer.frameCapacity + samples.withUnsafeBufferPointer { source in + buffer.floatChannelData![0].update( + from: source.baseAddress!, + count: samples.count + ) + } + + let file = try AVAudioFile( + forWriting: url, + settings: format.settings, + commonFormat: .pcmFormatFloat32, + interleaved: false + ) + try file.write(from: buffer) + return url + } +} From d9280d0a3ff965d581c1d7d6f3da366025270d95 Mon Sep 17 00:00:00 2001 From: Ciprian Date: Tue, 1 Sep 2026 12:12:40 +0400 Subject: [PATCH 2/2] test(diarizer): cancel real file inference --- .../OfflineDiarizerCancellationTests.swift | 64 +++++++------------ 1 file changed, 22 insertions(+), 42 deletions(-) diff --git a/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift b/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift index 72f3054f1..8491ce7bf 100644 --- a/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift +++ b/Tests/FluidAudioTests/Diarizer/Offline/OfflineDiarizerCancellationTests.swift @@ -1,4 +1,4 @@ -import AVFoundation +import os import XCTest @testable import FluidAudio @@ -8,22 +8,26 @@ final class OfflineDiarizerCancellationTests: XCTestCase { func testCancellingFileProcessingStopsAllInferenceWorkers() async throws { try requireOfflineDiarizerModels() - let fixture = try DiarizationTestFixtures.fixtureAudio(sampleRate: 16_000) - var longMeeting: [Float] = [] - longMeeting.reserveCapacity(fixture.count * 60) - for _ in 0..<60 { - longMeeting.append(contentsOf: fixture) + guard + let audioPath = ProcessInfo.processInfo.environment[ + "FLUIDAUDIO_CANCELLATION_TEST_AUDIO" + ], + FileManager.default.fileExists(atPath: audioPath) + else { + throw XCTSkip( + "Set FLUIDAUDIO_CANCELLATION_TEST_AUDIO to a real meeting recording" + ) } - - let audioURL = try writeWAV(samples: longMeeting, sampleRate: 16_000) - defer { try? FileManager.default.removeItem(at: audioURL) } + let audioURL = URL(fileURLWithPath: audioPath) let manager = OfflineDiarizerManager() let progress = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) + let progressCount = OSAllocatedUnfairLock(initialState: 0) var progressIterator = progress.stream.makeAsyncIterator() let processingTask = Task { try await manager.process(audioURL) { _, _ in + progressCount.withLock { $0 += 1 } progress.continuation.yield(()) } } @@ -50,9 +54,17 @@ final class OfflineDiarizerCancellationTests: XCTestCase { let cancellationDuration = cancellationStart.duration(to: .now) XCTAssertLessThan( cancellationDuration, - .seconds(5), + .seconds(2), "Cancellation should stop segmentation and embedding between model calls" ) + + let countWhenCancelledCallReturned = progressCount.withLock { $0 } + try await Task.sleep(for: .milliseconds(250)) + XCTAssertEqual( + progressCount.withLock { $0 }, + countWhenCancelledCallReturned, + "No inference progress may continue after process(URL) returns" + ) } private func requireOfflineDiarizerModels() throws { @@ -65,36 +77,4 @@ final class OfflineDiarizerCancellationTests: XCTestCase { throw XCTSkip("Offline diarizer models not available") } } - - private func writeWAV(samples: [Float], sampleRate: Double) throws -> URL { - let url = FileManager.default.temporaryDirectory - .appendingPathComponent("offline-diarizer-cancellation-\(UUID().uuidString)") - .appendingPathExtension("wav") - let format = AVAudioFormat( - commonFormat: .pcmFormatFloat32, - sampleRate: sampleRate, - channels: 1, - interleaved: false - )! - let buffer = AVAudioPCMBuffer( - pcmFormat: format, - frameCapacity: AVAudioFrameCount(samples.count) - )! - buffer.frameLength = buffer.frameCapacity - samples.withUnsafeBufferPointer { source in - buffer.floatChannelData![0].update( - from: source.baseAddress!, - count: samples.count - ) - } - - let file = try AVAudioFile( - forWriting: url, - settings: format.settings, - commonFormat: .pcmFormatFloat32, - interleaved: false - ) - try file.write(from: buffer) - return url - } }