spb/forge-studio

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.8import Foundation910@Observable11@MainActor12final class RunSupervisor {13    let store: RunStore14    let metrics = MetricsStore()1516    private(set) var activeRunID: UUID?17    private(set) var consoleLines: [String] = []18    private var runner: ProcessRunner?19    private var csvPoller: Task<Void, Never>?2021    init(store: RunStore) {22        self.store = store23    }2425    enum SupervisorError: LocalizedError {26        case busy, invalidConfig([String]), forgeNotConfigured27        var errorDescription: String? {28            switch self {29            case .busy:30                return "Un entraînement est déjà en cours — Forge Studio les exécute séquentiellement."31            case .invalidConfig(let errs):32                return "Config invalide : " + errs.joined(separator: " · ")33            case .forgeNotConfigured:34                return "Binaire forge non configuré (Réglages)."35            }36        }37    }3839    private func transition(_ run: inout Run, to next: RunState) {40        guard RunState.transitions[run.state]?.contains(next) == true else { return }41        run.state = next42        store.upsert(run)43    }4445    func start(config: ForgeConfig, datasetPath: String, name: String,46               resumeFrom: String? = nil) async throws {47        guard activeRunID == nil else { throw SupervisorError.busy }48        let errors = config.validationErrors49        guard errors.isEmpty else { throw SupervisorError.invalidConfig(errors) }50        guard let forge = ForgeBinaryLocator.savedBinaryURL else {51            throw SupervisorError.forgeNotConfigured52        }5354        let stamp = ISO8601DateFormatter().string(from: .now)55            .replacingOccurrences(of: ":", with: "-")56        let outDir = store.workspaceURL57            .appendingPathComponent("runs/\(name)-\(stamp)")58        try FileManager.default.createDirectory(at: outDir,59                                                withIntermediateDirectories: true)60        var run = Run(name: name, createdAt: .now, config: config,61                      datasetPath: datasetPath, outDirectory: outDir.path,62                      resumeCheckpoint: resumeFrom)63        try config.exportJSON().write(to: URL(fileURLWithPath: run.configPath))64        store.upsert(run)6566        var args = ["train", "--config", run.configPath, "--data", datasetPath,67                    "--out", run.outDirectory]68        if let resume = resumeFrom { args += ["--resume", resume] }6970        transition(&run, to: .launching)71        activeRunID = run.id72        consoleLines.removeAll()73        await metrics.reset()7475        let processRunner = ProcessRunner()76        runner = processRunner77        let (lines, exit) = try await processRunner.launch(78            executable: forge, arguments: args,79            currentDirectory: forge.deletingLastPathComponent())80        run.pid = await processRunner.pid81        transition(&run, to: .running)8283        // Console + run.log + stdout events.84        let logURL = URL(fileURLWithPath: run.runLogPath)85        FileManager.default.createFile(atPath: logURL.path, contents: nil)86        let logHandle = try? FileHandle(forWritingTo: logURL)87        let runID = run.id88        Task { [weak self] in89            for await (line, isErr) in lines {90                logHandle?.write(Data((line + "\n").utf8))91                await MainActor.run {92                    guard let self else { return }93                    self.consoleLines.append(isErr ? "⚠︎ " + line : line)94                    if self.consoleLines.count > 5000 {95                        self.consoleLines.removeFirst(1000)96                    }97                }98            }99            try? logHandle?.close()100        }101102        // CSV poller: 1 Hz incremental tail, updates run summary fields.103        let csvURL = URL(fileURLWithPath: run.logCSVPath)104        csvPoller = Task { [weak self] in105            while !Task.isCancelled {106                guard let self else { return }107                await self.metrics.ingestCSV(at: csvURL)108                let snap = await self.metrics.snapshot(maxPoints: 4, smoothing: 0)109                await MainActor.run {110                    guard var r = self.store.runs.first(where: { $0.id == runID })111                    else { return }112                    if let last = snap.lastPoint {113                        r.lastStep = last.step114                        r.lastTrainLoss = last.trainLoss115                    }116                    if let best = snap.bestVal {117                        r.bestValLoss = best.loss118                        r.bestValStep = best.step119                    }120                    self.store.upsert(r)121                }122                try? await Task.sleep(for: .seconds(1))123            }124        }125126        // Exit watcher.127        Task { [weak self] in128            let status = await exit.value129            await MainActor.run {130                guard let self,131                      var r = self.store.runs.first(where: { $0.id == runID })132                else { return }133                self.csvPoller?.cancel()134                r.pid = nil135                r.exitCode = status.code136                if status.code == 0 {137                    self.transition(&r, to: .finished)138                } else if r.state == .finishing || status.wasSignaled {139                    self.transition(&r, to: .stopped)140                } else {141                    r.failureReason = "forge s'est terminé avec le code \(status.code) — voir run.log"142                    self.transition(&r, to: .failed)143                }144                self.activeRunID = nil145                self.runner = nil146            }147        }148    }149150    /// SIGTERM. Forge has no signal handler: the process dies immediately and151    /// the last ckpt_latest.bin is the recovery point (UI says so).152    func stop() async {153        guard let id = activeRunID,154              var run = store.runs.first(where: { $0.id == id }) else { return }155        transition(&run, to: .finishing)156        try? await runner?.terminate()157    }158}159