Files
IQ.Pilot/tinygrad_repo/extra/usbgpu/tbgpu/installer/Shared/TinyGPUCLIRunner.swift
2026-08-22 23:42:42 -05:00

123 lines
4.9 KiB
Swift

import Foundation
import SystemExtensions
enum TinyGPUCLIExit: Int32 { case ok = 0, usage = 2, failed = 3, needsApproval = 4 }
enum DextState { case unloaded, activating, needsApproval, activated }
final class TinyGPUCLIRunner: NSObject, OSSystemExtensionRequestDelegate {
private let dextID: String
private var done: ((TinyGPUCLIExit) -> Void)?
private var isInstall = true
init(_ dextID: String) { self.dextID = dextID }
static func queryDextState(_ bundleID: String) -> DextState {
let p = Process()
p.executableURL = URL(fileURLWithPath: "/usr/bin/systemextensionsctl")
p.arguments = ["list"]
let pipe = Pipe()
p.standardOutput = pipe
p.standardError = Pipe()
guard (try? p.run()) != nil else { return .unloaded }
p.waitUntilExit()
guard let output = String(data: pipe.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8),
let line = output.split(separator: "\n").first(where: { $0.contains(bundleID) }) else { return .unloaded }
if line.contains("[activated enabled]") { return .activated }
if line.contains("[activated waiting for user]") { return .needsApproval }
return line.contains("terminated waiting to uninstall") ? .unloaded : .activating
}
private static let approvalHelp = """
Please go to System Settings > Privacy & Security and allow the extension.
If previously disabled: System Settings > General > Login Items & Extensions > Driver Extensions > Toggle TinyGPU ON
"""
static func statusText(_ state: DextState) -> String {
switch state {
case .unloaded: return "Driver extension not installed.\n\n"
case .activating: return "Extension is activating...\n\n"
case .needsApproval: return "Extension awaiting approval.\n\n" + approvalHelp
case .activated: return "Extension is ready! Run tinygrad to use your eGPU.\n\n"
}
}
func run(args: [String], done: @escaping (TinyGPUCLIExit) -> Void) {
self.done = done
guard args.count > 1 else { return usage() }
switch args[1] {
case "status":
print(Self.statusText(Self.queryDextState(dextID)))
done(.ok)
case "install":
if Self.queryDextState(dextID) == .needsApproval { print(Self.statusText(.needsApproval)); return done(.needsApproval) }
print("Installing TinyGPU driver extension...\nYou may need to approve in System Settings.\n")
submitRequest(activate: true)
case "uninstall":
guard Self.queryDextState(dextID) != .unloaded else { print("Not installed.\n"); return done(.ok) }
print("Uninstalling TinyGPU driver extension...\n")
isInstall = false
submitRequest(activate: false)
case "server":
guard args.count > 2 else { print("Error: server requires socket path\n"); return usage() }
done(run_server(args[2]) == 0 ? .ok : .failed)
case "help", "-h", "--help":
usage(); done(.ok)
default:
print("Unknown command: \(args[1])\n"); usage()
}
}
private func usage() {
print("""
Usage: TinyGPU <command>
status Show extension status
install Install the driver extension
uninstall Remove the driver extension
server <path> Start server on Unix socket
""")
done?(.usage)
}
private func submitRequest(activate: Bool) {
let req = activate
? OSSystemExtensionRequest.activationRequest(forExtensionWithIdentifier: dextID, queue: .main)
: OSSystemExtensionRequest.deactivationRequest(forExtensionWithIdentifier: dextID, queue: .main)
req.delegate = self
OSSystemExtensionManager.shared.submitRequest(req)
}
// MARK: - OSSystemExtensionRequestDelegate
func requestNeedsUserApproval(_ request: OSSystemExtensionRequest) {
print("\nUser approval required!\n\n\(Self.approvalHelp)After approval, connect the gpu and use it with tinygrad.\n")
done?(.needsApproval)
}
func request(_ request: OSSystemExtensionRequest, didFinishWithResult result: OSSystemExtensionRequest.Result) {
switch result {
case .completed: print("Driver extension \(isInstall ? "installed" : "uninstalled") successfully!\n")
case .willCompleteAfterReboot: print("Will complete after reboot.\n")
@unknown default: print("Completed: \(result)\n")
}
done?(.ok)
}
func request(_ request: OSSystemExtensionRequest, didFailWithError error: Error) {
print("\nError: \(error.localizedDescription)\n")
let code = (error as NSError).code
if code == 4 { print("Missing entitlements. Rebuild with proper signing.\n") }
else if code == 8 { print("Extension not found in app bundle.\n") }
else if code == 9 { print("Extension disabled by user.\n\n\(Self.approvalHelp)") }
done?(.failed)
}
func request(_ request: OSSystemExtensionRequest, actionForReplacingExtension existing: OSSystemExtensionProperties,
withExtension ext: OSSystemExtensionProperties) -> OSSystemExtensionRequest.ReplacementAction {
print("Updating v\(existing.bundleShortVersion) -> v\(ext.bundleShortVersion)...\n")
return .replace
}
}