bleam 0.0.10 → 0.0.12

package/templates/native/ios/GenerationService/main.swift

@@ -2,68 +2,130 @@ import Foundation
 import MLX
 
 private final class GenerationRuntime: @unchecked Sendable {
+  private enum OperationKind {
+    case image
+    case text
+  }
+
   private let modelWarmTimeout: Duration = .seconds(60)
+  private let stateLock = NSLock()
+  private var activeID: String?
+  private var activeKind: OperationKind?
   private var activeTask: Task<Void, Never>?
-  private var activeRunner: ImageGenerationRunner?
+  private var activeImageRunner: ImageGenerationRunner?
+  private var activeTextRunner: BonsaiTextGenerationRunner?
   private var trimTask: Task<Void, Never>?
   private var isCancelled = false
+  private var operationEpoch: UInt64 = 0
+  private var isShuttingDown = false
 
   var isRunning: Bool {
-    activeTask != nil
+    stateLock.withLock { activeTask != nil }
   }
 
-  func start(
+  func startImage(
     request: GenerationWorkerRequest,
     emit: @escaping @Sendable (GenerationWorkerEvent) -> Void
   ) -> Bool {
-    guard activeTask == nil else { return false }
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      trimTask?.cancel()
+      trimTask = nil
+      isCancelled = false
+      operationEpoch &+= 1
+      activeID = request.jobID.uuidString
+      activeKind = .image
+      let runner = ImageGenerationRunner(request: request, emit: emit)
+      activeImageRunner = runner
+      activeTask = Task { [weak self] in
+        await self?.runImage(request, runner: runner, emit: emit)
+      }
+      return true
+    }
+  }
+
+  func startText(
+    request: BonsaiTextGenerationRequest,
+    emit: @escaping @Sendable (BonsaiTextGenerationEvent) -> Void
+  ) -> Bool {
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      trimTask?.cancel()
+      trimTask = nil
+      isCancelled = false
+      operationEpoch &+= 1
+      activeID = request.id
+      activeKind = .text
+      let runner = BonsaiTextGenerationRunner(request: request, emit: emit)
+      activeTextRunner = runner
+      activeTask = Task { [weak self] in
+        await self?.runText(request, runner: runner)
+      }
+      return true
+    }
+  }
+
+  func cancelImage() -> Bool {
+    let operation = stateLock.withLock { () -> (Task<Void, Never>?, ImageGenerationRunner?)? in
+      guard activeKind == .image else { return nil }
+      isCancelled = true
+      trimTask?.cancel()
+      trimTask = nil
+      return (activeTask, activeImageRunner)
+    }
+    guard let operation else { return false }
+    operation.0?.cancel()
+    operation.1?.cancel()
+    return true
+  }
 
-    trimTask?.cancel()
-    trimTask = nil
-    isCancelled = false
-    activeTask = Task { [weak self] in
-      await self?.run(request, emit: emit)
+  func cancelText(id: String) -> Bool {
+    let operation = stateLock.withLock { () -> (Task<Void, Never>?, BonsaiTextGenerationRunner?)? in
+      guard activeKind == .text, activeID == id else { return nil }
+      isCancelled = true
+      trimTask?.cancel()
+      trimTask = nil
+      return (activeTask, activeTextRunner)
     }
+    guard let operation else { return false }
+    operation.0?.cancel()
+    operation.1?.cancel()
     return true
   }
 
-  func cancel() {
-    isCancelled = true
-    trimTask?.cancel()
-    trimTask = nil
-    activeTask?.cancel()
-    activeRunner?.cancel()
+  func beginShutdown() -> Bool {
+    stateLock.withLock {
+      guard activeTask == nil, !isShuttingDown else { return false }
+      isShuttingDown = true
+      operationEpoch &+= 1
+      trimTask?.cancel()
+      trimTask = nil
+      return true
+    }
   }
 
   func trimModelsNow() async {
-    trimTask?.cancel()
-    trimTask = nil
-    // We currently reload every stage per generation, but keep the explicit
-    // trim hook from ImageStudio's service architecture. Warm model caching
-    // should be added here later, never back in the app process.
-    MLX.Memory.clearCache()
-    await Task.yield()
-    MLX.Memory.clearCache()
+    stateLock.withLock {
+      trimTask?.cancel()
+      trimTask = nil
+      MLX.Memory.clearCache()
+      MLX.Memory.clearCache()
+    }
     print("[GenerationService] trim active=\(MLX.Memory.activeMemory / 1_048_576)MB cache=\(MLX.Memory.cacheMemory / 1_048_576)MB")
   }
 
-  private func run(
+  private func runImage(
     _ request: GenerationWorkerRequest,
+    runner: ImageGenerationRunner,
     emit: @escaping @Sendable (GenerationWorkerEvent) -> Void
   ) async {
-    let runner = ImageGenerationRunner(request: request, emit: emit)
-    activeRunner = runner
     defer {
-      if activeRunner === runner {
-        activeRunner = nil
-      }
-      activeTask = nil
-      scheduleTrimAfterIdle()
+      finish(id: request.jobID.uuidString, imageRunner: runner)
     }
 
     do {
       try await runner.run(isCancelled: { [weak self] in
-        Task.isCancelled || (self?.isCancelled ?? false)
+        Task.isCancelled || (self?.cancellationRequested ?? false)
       })
     } catch is CancellationError {
       emit(.init(type: .canceled))
@@ -76,17 +138,59 @@ private final class GenerationRuntime: @unchecked Sendable {
     }
   }
 
+  private func runText(
+    _ request: BonsaiTextGenerationRequest,
+    runner: BonsaiTextGenerationRunner
+  ) async {
+    defer { finish(id: request.id, textRunner: runner) }
+    await runner.run(isCancelled: { [weak self] in
+      Task.isCancelled || (self?.cancellationRequested ?? true)
+    })
+  }
+
+  private var cancellationRequested: Bool {
+    stateLock.withLock { isCancelled }
+  }
+
+  private func finish(
+    id: String,
+    imageRunner: ImageGenerationRunner? = nil,
+    textRunner: BonsaiTextGenerationRunner? = nil
+  ) {
+    let finished = stateLock.withLock {
+      guard activeID == id else { return false }
+      if let imageRunner, activeImageRunner !== imageRunner { return false }
+      if let textRunner, activeTextRunner !== textRunner { return false }
+      activeID = nil
+      activeKind = nil
+      activeImageRunner = nil
+      activeTextRunner = nil
+      activeTask = nil
+      isCancelled = false
+      return true
+    }
+    if finished { scheduleTrimAfterIdle() }
+  }
+
   private func scheduleTrimAfterIdle() {
-    trimTask?.cancel()
-    trimTask = Task { [weak self] in
-      guard let self else { return }
-      do {
-        try await Task.sleep(for: modelWarmTimeout)
-      } catch {
-        return
+    stateLock.withLock {
+      trimTask?.cancel()
+      let scheduledEpoch = operationEpoch
+      trimTask = Task { [weak self] in
+        guard let self else { return }
+        do {
+          try await Task.sleep(for: modelWarmTimeout)
+        } catch {
+          return
+        }
+        guard !Task.isCancelled else { return }
+        self.stateLock.withLock {
+          guard self.activeTask == nil, self.operationEpoch == scheduledEpoch else { return }
+          self.trimTask = nil
+          MLX.Memory.clearCache()
+          MLX.Memory.clearCache()
+        }
       }
-      guard !Task.isCancelled else { return }
-      await self.trimModelsNow()
     }
     print("[GenerationService] model trim scheduled 60s")
   }
@@ -94,10 +198,11 @@ private final class GenerationRuntime: @unchecked Sendable {
 
 private final class GenerationService: NSObject, GenerationServiceProtocol {
   private weak var connection: NSXPCConnection?
-  private let runtime = GenerationRuntime()
+  private let runtime: GenerationRuntime
 
-  init(connection: NSXPCConnection) {
+  init(connection: NSXPCConnection, runtime: GenerationRuntime) {
     self.connection = connection
+    self.runtime = runtime
   }
 
   func startGeneration(requestData: Data, reply: @escaping (Bool) -> Void) {
@@ -107,14 +212,27 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
       return
     }
 
-    reply(runtime.start(request: request) { [weak self] event in
+    reply(runtime.startImage(request: request) { [weak self] event in
       self?.emit(event)
     })
   }
 
   func cancelGeneration(reply: @escaping (Bool) -> Void) {
-    runtime.cancel()
-    reply(true)
+    reply(runtime.cancelImage())
+  }
+
+  func startTextGeneration(requestData: Data, reply: @escaping (Bool) -> Void) {
+    guard let request = try? JSONDecoder().decode(BonsaiTextGenerationRequest.self, from: requestData) else {
+      reply(false)
+      return
+    }
+    reply(runtime.startText(request: request) { [weak self] event in
+      self?.emitText(event)
+    })
+  }
+
+  func cancelTextGeneration(id: String, reply: @escaping (Bool) -> Void) {
+    reply(runtime.cancelText(id: id))
   }
 
   func memoryTelemetry(reply: @escaping (Data?) -> Void) {
@@ -126,7 +244,7 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
   }
 
   func shutdown(reply: @escaping (Bool) -> Void) {
-    guard !runtime.isRunning else {
+    guard runtime.beginShutdown() else {
       reply(false)
       return
     }
@@ -142,13 +260,21 @@ private final class GenerationService: NSObject, GenerationServiceProtocol {
     let client = connection?.remoteObjectProxy as? GenerationServiceClientProtocol
     client?.generationWorkerDidEmit(eventData: data)
   }
+
+  private func emitText(_ event: BonsaiTextGenerationEvent) {
+    guard let data = try? JSONEncoder().encode(event) else { return }
+    let client = connection?.remoteObjectProxy as? GenerationServiceClientProtocol
+    client?.textGenerationWorkerDidEmit(eventData: data)
+  }
 }
 
 private final class GenerationServiceDelegate: NSObject, NSXPCListenerDelegate {
+  private let runtime = GenerationRuntime()
+
   func listener(_ listener: NSXPCListener, shouldAcceptNewConnection connection: NSXPCConnection) -> Bool {
     connection.exportedInterface = NSXPCInterface(with: GenerationServiceProtocol.self)
     connection.remoteObjectInterface = NSXPCInterface(with: GenerationServiceClientProtocol.self)
-    connection.exportedObject = GenerationService(connection: connection)
+    connection.exportedObject = GenerationService(connection: connection, runtime: runtime)
     connection.resume()
     return true
   }