spb/zyquo-mlx Public MIT

1//2//  PythonRunner.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// One JSON-lines event from a PyBridge helper script.12/// All scripts emit `{"event": "...", ...}` objects, one per line13/// (protocol defined in `PyBridge/scripts/` — docs/TRAINING-RESEARCH.md §5).14struct PythonEvent: @unchecked Sendable {15    let event: String16    /// JSONSerialization plist values (NSString/NSNumber/…) — value types in17    /// practice, hence the unchecked-Sendable annotation.18    let payload: [String: Any]1920    func double(_ key: String) -> Double? {21        (payload[key] as? Double) ?? (payload[key] as? Int).map(Double.init)22    }23    func int(_ key: String) -> Int? {24        (payload[key] as? Int) ?? (payload[key] as? Double).map(Int.init)25    }26    func string(_ key: String) -> String? { payload[key] as? String }27}2829enum PythonRunnerError: LocalizedError {30    case scriptNotFound(String)31    case environmentNotProvisioned32    case processFailed(exitCode: Int32, stderr: String)3334    var errorDescription: String? {35        switch self {36        case .scriptNotFound(let name):37            "Bundled Python script missing: \(name)"38        case .environmentNotProvisioned:39            "The Python environment is not set up yet. It is provisioned automatically on first use."40        case .processFailed(let code, let stderr):41            "Python pipeline failed (exit \(code)). \(stderr.suffix(500))"42        }43    }44}4546/// All Python interaction goes through here (charter engineering standard):47/// spawns the pinned venv's interpreter on a bundled script and streams the48/// JSON-lines protocol back. Cancellation terminates the process.49actor PythonRunner {5051    static let shared = PythonRunner()5253    /// Resolve a bundled helper script by name.54    static func scriptURL(named name: String) throws -> URL {55        guard56            let url = Bundle.module.url(57                forResource: name, withExtension: "py", subdirectory: "scripts")58        else {59            throw PythonRunnerError.scriptNotFound("\(name).py")60        }61        return url62    }6364    /// Run a script and stream its events. The stream throws on nonzero exit65    /// (unless an `error` event already surfaced the failure) and finishes66    /// after the process exits. Terminating the stream kills the process.67    func stream(script: String, arguments: [String]) async throws68        -> AsyncThrowingStream<PythonEvent, Error>69    {70        let env = PythonEnvironment.shared71        if await !env.isProvisioned {72            try await env.provision()73        }74        let python = await env.pythonURL75        let scriptURL = try Self.scriptURL(named: script)7677        let process = Process()78        process.executableURL = python79        process.arguments = ["-u", scriptURL.path] + arguments80        // Isolated, reproducible interpreter environment. Homebrew paths are81        // included for external tools some pipelines shell out to (ffmpeg82        // for mlx-whisper audio decoding).83        process.environment = [84            "PATH": "/opt/homebrew/bin:/usr/local/bin:/usr/bin:/bin",85            "HOME": NSHomeDirectory(),86            "PYTHONUNBUFFERED": "1",87        ]8889        let stdout = Pipe()90        let stderr = Pipe()91        process.standardOutput = stdout92        process.standardError = stderr9394        return AsyncThrowingStream { continuation in95            let stderrBuffer = LockedBuffer()9697            stderr.fileHandleForReading.readabilityHandler = { handle in98                stderrBuffer.append(handle.availableData)99            }100101            let lineBuffer = LockedBuffer()102            stdout.fileHandleForReading.readabilityHandler = { handle in103                let data = handle.availableData104                guard !data.isEmpty else { return }105                for line in lineBuffer.appendAndExtractLines(data) {106                    if let event = Self.decode(line: line) {107                        continuation.yield(event)108                    }109                }110            }111112            process.terminationHandler = { proc in113                stdout.fileHandleForReading.readabilityHandler = nil114                stderr.fileHandleForReading.readabilityHandler = nil115                // Flush any trailing output.116                let rest = stdout.fileHandleForReading.readDataToEndOfFile()117                for line in lineBuffer.appendAndExtractLines(rest, flush: true) {118                    if let event = Self.decode(line: line) {119                        continuation.yield(event)120                    }121                }122                stderrBuffer.append(stderr.fileHandleForReading.readDataToEndOfFile())123                if proc.terminationStatus == 0 {124                    continuation.finish()125                } else {126                    continuation.finish(127                        throwing: PythonRunnerError.processFailed(128                            exitCode: proc.terminationStatus,129                            stderr: stderrBuffer.string))130                }131            }132133            continuation.onTermination = { termination in134                if case .cancelled = termination, process.isRunning {135                    process.terminate()136                }137            }138139            do {140                try process.run()141            } catch {142                continuation.finish(throwing: error)143            }144        }145    }146147    private static func decode(line: Data) -> PythonEvent? {148        guard149            let object = try? JSONSerialization.jsonObject(with: line) as? [String: Any],150            let event = object["event"] as? String151        else { return nil }152        return PythonEvent(event: event, payload: object)153    }154}155156/// Thread-safe byte buffer for pipe readability handlers (they run on a157/// private DispatchQueue while the stream lives elsewhere).158private final class LockedBuffer: @unchecked Sendable {159    private let lock = NSLock()160    private var data = Data()161162    func append(_ chunk: Data) {163        lock.lock()164        defer { lock.unlock() }165        data.append(chunk)166    }167168    /// Append then split out complete `\n`-terminated lines (keeps the tail).169    func appendAndExtractLines(_ chunk: Data, flush: Bool = false) -> [Data] {170        lock.lock()171        defer { lock.unlock() }172        data.append(chunk)173        var lines: [Data] = []174        while let newline = data.firstIndex(of: UInt8(ascii: "\n")) {175            lines.append(data.subdata(in: data.startIndex..<newline))176            data.removeSubrange(data.startIndex...newline)177        }178        if flush, !data.isEmpty {179            lines.append(data)180            data.removeAll()181        }182        return lines183    }184185    var string: String {186        lock.lock()187        defer { lock.unlock() }188        return String(data: data, encoding: .utf8) ?? ""189    }190}191