spb/zyquo-mlx Public MIT

1//2//  RunStore.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// Persists training runs, their metrics history, and checkpoint inventory12/// under `~/Library/Application Support/ZyquoMLX/Runs/<run-id>/`.13///14/// Layout:15///   run.json          — the `TrainingRun` (state, config, our iteration counter)16///   metrics.jsonl     — append-only `TrainingMetric` history (chart replay)17///   adapters/         — mlx-lm adapter dir (adapters.safetensors, NNNNNNN_… checkpoints)18///   config.yaml       — the exact mlx-lm config the run used19actor RunStore {2021    static let shared = RunStore()2223    private let root: URL2425    init(root: URL = PersistenceService.runsDirectory) {26        self.root = root27    }2829    func directory(for id: UUID) -> URL {30        root.appendingPathComponent(id.uuidString, isDirectory: true)31    }3233    func adaptersDirectory(for run: TrainingRun) -> URL {34        run.directory.appendingPathComponent("adapters", isDirectory: true)35    }3637    // MARK: - CRUD3839    func create(_ run: TrainingRun) throws {40        try FileManager.default.createDirectory(41            at: run.directory, withIntermediateDirectories: true)42        try save(run)43    }4445    func save(_ run: TrainingRun) throws {46        try PersistenceService.saveJSON(run, to: run.directory.appendingPathComponent("run.json"))47    }4849    func load(id: UUID) throws -> TrainingRun {50        try PersistenceService.loadJSON(51            TrainingRun.self, from: directory(for: id).appendingPathComponent("run.json"))52    }5354    func scan() throws -> [TrainingRun] {55        let fm = FileManager.default56        guard fm.fileExists(atPath: root.path) else { return [] }57        return try fm.contentsOfDirectory(at: root, includingPropertiesForKeys: nil)58            .compactMap {59                try? PersistenceService.loadJSON(60                    TrainingRun.self, from: $0.appendingPathComponent("run.json"))61            }62            .sorted { $0.createdAt > $1.createdAt }63    }6465    func delete(_ run: TrainingRun) throws {66        try FileManager.default.removeItem(at: run.directory)67    }6869    // MARK: - Metrics history7071    func append(metric: TrainingMetric, to run: TrainingRun) {72        guard let data = try? JSONEncoder().encode(metric),73            let line = String(data: data, encoding: .utf8)74        else { return }75        let url = run.directory.appendingPathComponent("metrics.jsonl")76        if let handle = FileHandle(forWritingAtPath: url.path) {77            handle.seekToEndOfFile()78            handle.write(Data((line + "\n").utf8))79            try? handle.close()80        } else {81            try? (line + "\n").write(to: url, atomically: true, encoding: .utf8)82        }83    }8485    func metricsHistory(for run: TrainingRun) -> [TrainingMetric] {86        guard87            let content = try? String(88                contentsOf: run.directory.appendingPathComponent("metrics.jsonl"),89                encoding: .utf8)90        else { return [] }91        let decoder = JSONDecoder()92        return content.split(separator: "\n").compactMap {93            $0.data(using: .utf8).flatMap { try? decoder.decode(TrainingMetric.self, from: $0) }94        }95    }9697    // MARK: - Checkpoints9899    /// Scan the adapter directory for numbered checkpoints100    /// (`{iter:07d}_adapters.safetensors` — docs/TRAINING-RESEARCH.md §2.1).101    func checkpoints(for run: TrainingRun) -> [Checkpoint] {102        let dir = adaptersDirectory(for: run)103        guard104            let files = try? FileManager.default.contentsOfDirectory(105                at: dir, includingPropertiesForKeys: [.creationDateKey])106        else { return [] }107        return files.compactMap { url -> Checkpoint? in108            let name = url.lastPathComponent109            guard name.hasSuffix("_adapters.safetensors"),110                let iteration = Int(name.prefix(7))111            else { return nil }112            let savedAt =113                (try? url.resourceValues(forKeys: [.creationDateKey]))?.creationDate ?? .now114            return Checkpoint(fileName: name, iteration: iteration, fileURL: url, savedAt: savedAt)115        }116        .sorted { $0.iteration < $1.iteration }117    }118119    /// Latest usable adapter weights for resume/fuse/inference.120    func latestAdapter(for run: TrainingRun) -> URL? {121        let main = adaptersDirectory(for: run).appendingPathComponent("adapters.safetensors")122        if FileManager.default.fileExists(atPath: main.path) { return main }123        return checkpoints(for: run).last?.fileURL124    }125}126