spb/forge-studio Public

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Owns a live run end to end: builds the out directory, launches4// `forge train`, streams stdout/stderr to run.log + the console buffer,5// polls log.csv into the MetricsStore, and drives the Run state machine.6// Single-writer: all Run mutations flow through here and persist via7// RunStore immediately. Posts a local notification on finish/fail.8import Foundation9import UserNotifications1011@Observable12@MainActor13final class RunSupervisor {14    let store: RunStore15    let metrics = MetricsStore()1617    private(set) var activeRunID: UUID?18    private(set) var consoleLines: [String] = []19    private var runner: ProcessRunner?20    private var csvPoller: Task<Void, Never>?2122    init(store: RunStore) {23        self.store = store24    }2526    enum SupervisorError: LocalizedError {27        case busy, invalidConfig([String]), forgeNotConfigured28        var errorDescription: String? {29            switch self {30            case .busy:31                return "Un entraînement est déjà en cours — Forge Studio les exécute séquentiellement."32            case .invalidConfig(let errs):33                return "Config invalide : " + errs.joined(separator: " · ")34            case .forgeNotConfigured:35                return "Binaire forge non configuré (Réglages)."36            }37        }38    }3940    private func transition(_ run: inout Run, to next: RunState) {41        guard RunState.transitions[run.state]?.contains(next) == true else { return }42        run.state = next43        store.upsert(run)44    }4546    /// Public entry: creates the run immediately. If a training is live the47    /// run stays `queued` (persisted) and launches automatically when the48    /// active one exits — trainings never overlap.49    func start(config: ForgeConfig, datasetPath: String, name: String,50               resumeFrom: String? = nil) async throws {51        let errors = config.validationErrors52        guard errors.isEmpty else { throw SupervisorError.invalidConfig(errors) }53        guard ForgeBinaryLocator.savedBinaryURL != nil else {54            throw SupervisorError.forgeNotConfigured55        }5657        let stamp = ISO8601DateFormatter().string(from: .now)58            .replacingOccurrences(of: ":", with: "-")59        let outDir = store.workspaceURL60            .appendingPathComponent("runs/\(name)-\(stamp)")61        try FileManager.default.createDirectory(at: outDir,62                                                withIntermediateDirectories: true)63        let run = Run(name: name, createdAt: .now, config: config,64                      datasetPath: datasetPath, outDirectory: outDir.path,65                      resumeCheckpoint: resumeFrom)66        try config.exportJSON().write(to: URL(fileURLWithPath: run.configPath))67        store.upsert(run)68        if activeRunID == nil {69            try await launch(run)70        } // else: stays queued; launchNextQueued() picks it up on exit71    }7273    private func launchNextQueued() {74        guard activeRunID == nil,75              let next = store.runs.filter({ $0.state == .queued })76                  .min(by: { $0.createdAt < $1.createdAt }) else { return }77        Task {78            do {79                try await launch(next)80            } catch {81                var r = next82                r.failureReason = error.localizedDescription83                transition(&r, to: .launching)84                transition(&r, to: .failed)85            }86        }87    }8889    private func launch(_ queued: Run) async throws {90        guard activeRunID == nil else { throw SupervisorError.busy }91        guard let forge = ForgeBinaryLocator.savedBinaryURL else {92            throw SupervisorError.forgeNotConfigured93        }94        var run = queued9596        var args = ["train", "--config", run.configPath, "--data", run.datasetPath,97                    "--out", run.outDirectory]98        if let resume = run.resumeCheckpoint { args += ["--resume", resume] }99100        transition(&run, to: .launching)101        activeRunID = run.id102        consoleLines.removeAll()103        await metrics.reset()104105        let processRunner = ProcessRunner()106        runner = processRunner107        let (lines, exit) = try await processRunner.launch(108            executable: forge, arguments: args,109            currentDirectory: forge.deletingLastPathComponent())110        run.pid = await processRunner.pid111        transition(&run, to: .running)112113        // Console + run.log + stdout events.114        let logURL = URL(fileURLWithPath: run.runLogPath)115        FileManager.default.createFile(atPath: logURL.path, contents: nil)116        let logHandle = try? FileHandle(forWritingTo: logURL)117        let runID = run.id118        Task { [weak self] in119            for await (line, isErr) in lines {120                logHandle?.write(Data((line + "\n").utf8))121                await MainActor.run {122                    guard let self else { return }123                    self.consoleLines.append(isErr ? "⚠︎ " + line : line)124                    if self.consoleLines.count > 5000 {125                        self.consoleLines.removeFirst(1000)126                    }127                }128            }129            try? logHandle?.close()130        }131132        // CSV poller: 1 Hz incremental tail, updates run summary fields.133        let csvURL = URL(fileURLWithPath: run.logCSVPath)134        csvPoller = Task { [weak self] in135            while !Task.isCancelled {136                guard let self else { return }137                await self.metrics.ingestCSV(at: csvURL)138                let snap = await self.metrics.snapshot(maxPoints: 4, smoothing: 0)139                await MainActor.run {140                    guard var r = self.store.runs.first(where: { $0.id == runID })141                    else { return }142                    if let last = snap.lastPoint {143                        r.lastStep = last.step144                        r.lastTrainLoss = last.trainLoss145                    }146                    if let best = snap.bestVal {147                        r.bestValLoss = best.loss148                        r.bestValStep = best.step149                    }150                    self.store.upsert(r)151                }152                try? await Task.sleep(for: .seconds(1))153            }154        }155156        // Exit watcher.157        Task { [weak self] in158            let status = await exit.value159            await MainActor.run {160                guard let self,161                      var r = self.store.runs.first(where: { $0.id == runID })162                else { return }163                self.csvPoller?.cancel()164                r.pid = nil165                r.exitCode = status.code166                if status.code == 0 {167                    self.transition(&r, to: .finished)168                } else if r.state == .finishing || status.wasSignaled {169                    self.transition(&r, to: .stopped)170                } else {171                    r.failureReason = "forge s'est terminé avec le code \(status.code) — voir run.log"172                    self.transition(&r, to: .failed)173                }174                self.activeRunID = nil175                self.runner = nil176                Self.notify(run: r)177                self.launchNextQueued()178            }179        }180    }181182    // Local notification with the outcome; silently skipped when not183    // running from a bundle (bare `swift run`) or when denied.184    static func requestNotificationAuth() {185        guard Bundle.main.bundleIdentifier != nil else { return }186        UNUserNotificationCenter.current()187            .requestAuthorization(options: [.alert, .sound]) { _, _ in }188    }189190    private static func notify(run: Run) {191        guard Bundle.main.bundleIdentifier != nil else { return }192        let content = UNMutableNotificationContent()193        switch run.state {194        case .finished:195            content.title = "Run terminé — \(run.name)"196            var body = ""197            if let loss = run.lastTrainLoss {198                body += String(format: "loss finale %.4f", loss)199            }200            if let best = run.bestValLoss {201                body += String(format: " · best val %.4f", best)202            }203            content.body = body.isEmpty ? "Entraînement achevé." : body204        case .failed:205            content.title = "Run échoué — \(run.name)"206            content.body = run.failureReason ?? "Voir run.log"207        default:208            return209        }210        UNUserNotificationCenter.current().add(211            UNNotificationRequest(identifier: run.id.uuidString,212                                  content: content, trigger: nil))213    }214215    /// SIGTERM. Forge has no signal handler: the process dies immediately and216    /// the last ckpt_latest.bin is the recovery point (UI says so).217    func stop() async {218        guard let id = activeRunID,219              var run = store.runs.first(where: { $0.id == id }) else { return }220        transition(&run, to: .finishing)221        try? await runner?.terminate()222    }223}224