spb/zyquo-mlx Public MIT

1//2//  TrainingService.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011enum TrainingServiceError: LocalizedError {12    case memoryGate(verdict: MemoryVerdict, suggestions: [String])13    case qloraNeedsQuantizedBase14    case alreadyRunning(UUID)15    case nothingToResume1617    var errorDescription: String? {18        switch self {19        case .memoryGate(_, let suggestions):20            "This configuration won't fit in memory on this Mac. Try: "21                + suggestions.joined(separator: " · ")22        case .qloraNeedsQuantizedBase:23            "QLoRA requires a quantized base model. Pick a quantized base (e.g. a -4bit model) or quantize it first in Convert."24        case .alreadyRunning(let id):25            "Run \(id.uuidString.prefix(8)) is already active."26        case .nothingToResume:27            "No saved adapter weights found to resume from."28        }29    }30}3132/// Orchestrates fine-tuning runs (charter 3.B): builds the mlx-lm YAML33/// config, drives `zyquo_train.py` through `PythonRunner`, persists state and34/// metrics via `RunStore`, and exposes a typed event stream. Runs are35/// cancellable; resume is an honest warm start (docs/TRAINING-RESEARCH.md §5.3).36actor TrainingService {3738    static let shared = TrainingService()3940    private var activeRunID: UUID?41    private var activeStreamTask: Task<Void, Never>?4243    // MARK: - Configure4445    /// Create (and persist) a new run after memory-gating it.46    func createRun(47        name: String,48        baseModel: LocalModel,49        dataset: Dataset,50        method: FineTuneMethod,51        hyperParams: HyperParams52    ) async throws -> TrainingRun {53        // QLoRA = LoRA on a quantized base; LoRA on a quantized base IS QLoRA —54        // keep the user's mental model consistent with reality.55        if method == .qlora && baseModel.quantization == nil {56            throw TrainingServiceError.qloraNeedsQuantizedBase57        }5859        let verdict = MemoryAdvisor.trainingVerdict(60            for: baseModel, method: method, params: hyperParams)61        if verdict == .wontFit {62            throw TrainingServiceError.memoryGate(63                verdict: verdict,64                suggestions: MemoryAdvisor.suggestions(65                    for: baseModel, method: method, params: hyperParams))66        }6768        let id = UUID()69        let run = TrainingRun(70            id: id,71            name: name,72            baseModelID: baseModel.id,73            datasetID: dataset.id,74            method: method,75            hyperParams: hyperParams,76            state: .configured,77            directory: await RunStore.shared.directory(for: id),78            createdAt: .now,79            startedAt: nil,80            finishedAt: nil,81            completedIterations: 0,82            checkpoints: [],83            lastError: nil84        )85        try await RunStore.shared.create(run)86        try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: nil)87        return run88    }8990    // MARK: - Start / resume9192    /// Start (or warm-start-resume) a run, streaming typed events.93    func start(run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resume: Bool = false)94        async throws -> AsyncStream<TrainingEvent>95    {96        if let activeRunID { throw TrainingServiceError.alreadyRunning(activeRunID) }9798        var run = run99        var resumeAdapter: URL?100        if resume {101            guard let adapter = await RunStore.shared.latestAdapter(for: run) else {102                throw TrainingServiceError.nothingToResume103            }104            resumeAdapter = adapter105        }106        try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: resumeAdapter)107108        run.state = .running109        run.startedAt = run.startedAt ?? .now110        run.lastError = nil111        try await RunStore.shared.save(run)112        activeRunID = run.id113114        let configPath = run.directory.appendingPathComponent("config.yaml").path115        let pythonStream = try await PythonRunner.shared.stream(116            script: "zyquo_train", arguments: ["--config", configPath])117118        let (stream, continuation) = AsyncStream.makeStream(of: TrainingEvent.self)119        let runSnapshot = run120121        activeStreamTask = Task {122            var current = runSnapshot123            do {124                for try await pythonEvent in pythonStream {125                    guard let event = MetricsStream.decode(pythonEvent) else { continue }126                    switch event {127                    case .metric(let metric):128                        await RunStore.shared.append(metric: metric, to: current)129                        current.completedIterations = max(130                            current.completedIterations, metric.iteration)131                    case .failed(let message):132                        current.lastError = message133                    default:134                        break135                    }136                    continuation.yield(event)137                }138                // A cancelled Task ends the stream without throwing — check139                // cancellation before declaring completion.140                if Task.isCancelled {141                    current.state = .cancelled142                } else {143                    current.state = current.lastError == nil ? .completed : .failed144                }145            } catch is CancellationError {146                current.state = .cancelled147            } catch {148                current.state = .failed149                current.lastError = error.localizedDescription150                continuation.yield(.failed(message: error.localizedDescription))151            }152            current.finishedAt = .now153            current.checkpoints = await RunStore.shared.checkpoints(for: current)154            try? await RunStore.shared.save(current)155            self.clearActive()156            continuation.finish()157        }158159        return stream160    }161162    /// Cancel the active run (terminates the Python process via the stream's163    /// onTermination; state persists as cancelled with checkpoints intact for164    /// warm-start resume).165    func cancel() {166        activeStreamTask?.cancel()167    }168169170    private func clearActive() {171        activeRunID = nil172        activeStreamTask = nil173    }174175    // MARK: - Config generation176177    /// Serialize the run into mlx-lm's YAML schema178    /// (exact keys/defaults per docs/TRAINING-RESEARCH.md §1.2; the179    /// config-only keys `lora_parameters` make YAML mandatory).180    private func writeConfig(181        run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resumeFrom: URL?182    ) throws {183        let hp = run.hyperParams184        let adapterPath = run.directory.appendingPathComponent("adapters").path185186        // .qlora maps to upstream "lora" — QLoRA is implied by the quantized base.187        let fineTuneType = run.method == .qlora ? "lora" : run.method.rawValue188189        var lines: [String] = [190            "# Generated by Zyquo MLX — run \(run.id.uuidString)",191            "model: \(yamlQuote(baseModel.directory.path))",192            "train: true",193            "data: \(yamlQuote(dataset.directory.path))",194            "fine_tune_type: \(fineTuneType)",195            "optimizer: \(hp.optimizer)",196            "mask_prompt: \(hp.maskPrompt)",197            "num_layers: \(hp.numLayers)",198            "batch_size: \(hp.batchSize)",199            "iters: \(hp.iterations)",200            "val_batches: \(hp.valBatches)",201            "learning_rate: \(hp.learningRate)",202            "steps_per_report: \(hp.stepsPerReport)",203            "steps_per_eval: \(hp.stepsPerEval)",204            "grad_accumulation_steps: \(hp.gradAccumulationSteps)",205            "save_every: \(hp.saveEvery)",206            "adapter_path: \(yamlQuote(adapterPath))",207            "max_seq_length: \(hp.maxSeqLength)",208            "grad_checkpoint: \(hp.gradCheckpoint)",209            "seed: \(hp.seed)",210        ]211212        if let resumeFrom {213            lines.append("resume_adapter_file: \(yamlQuote(resumeFrom.path))")214        }215216        if run.method != .full {217            lines.append("lora_parameters:")218            lines.append("  rank: \(hp.rank)")219            lines.append("  scale: \(hp.scale)")220            lines.append("  dropout: \(hp.dropout)")221            if let keys = hp.keys, !keys.isEmpty {222                lines.append("  keys: [\(keys.map(yamlQuote).joined(separator: ", "))]")223            }224        }225226        try (lines.joined(separator: "\n") + "\n")227            .write(228                to: run.directory.appendingPathComponent("config.yaml"),229                atomically: true, encoding: .utf8)230    }231232    private func yamlQuote(_ value: String) -> String {233        "\"" + value.replacingOccurrences(of: "\"", with: "\\\"") + "\""234    }235}236