Distributed Inference and Training with MLX

import Foundation

// MARK: - Cluster Node Configuration
struct MLXClusterNode: Codable {
    let ssh: String
    let ips: [String]
    let rdma: [String]
    var env: [String: String]?
}

// MARK: - Hostfile Generator
struct MLXHostfileGenerator {
    static func buildHostfile(nodes: [MLXClusterNode]) throws -> URL {
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.prettyPrinted]
        let data = try encoder.encode(nodes)
        let url = FileManager.default
            .temporaryDirectory
            .appendingPathComponent("mlx_cluster.json")
        try data.write(to: url)
        return url
    }
}

// MARK: - Distributed Job Launcher
struct MLXDistributedLauncher {
    let hostfileURL: URL

    /// Launches a distributed MLX LM chat session across the cluster.
    /// Equivalent to: mlx.launch --hostfile cluster.json -- /path/to/mlx_lm.chat --model Qwen3-6 --max-tokens 2048
    func launchDistributedChat(model: String, maxTokens: Int) throws -> Process {
        let process = Process()
        process.executableURL = URL(fileURLWithPath: "/usr/local/bin/python3")
        process.arguments = [
            "-m", "mlx.launch",
            "--hostfile", hostfileURL.path,
            "--",
            "/usr/local/bin/mlx_lm.chat",
            "--model", model,
            "--max-tokens", String(maxTokens)
        ]
        process.environment = ProcessInfo.processInfo.environment.merging(
            ["MLX_METAL_FAST_SYNCH": "1"],
            uniquingKeysWith: { _, new in new }
        )
        try process.run()
        return process
    }
}

// MARK: - Example Usage
func runDistributedInferenceExample() throws {
    // Define 4-node M3 Ultra cluster in a full mesh topology
    let nodes: [MLXClusterNode] = [
        MLXClusterNode(
            ssh: "mac-ultra-1.local",
            ips: ["192.168.1.101"],
            rdma: ["rdma0", "rdma1", "rdma2"]
        ),
        MLXClusterNode(
            ssh: "mac-ultra-2.local",
            ips: ["192.168.1.102"],
            rdma: ["rdma0", "rdma1", "rdma2"]
        ),
        MLXClusterNode(
            ssh: "mac-ultra-3.local",
            ips: ["192.168.1.103"],
            rdma: ["rdma0", "rdma1", "rdma2"]
        ),
        MLXClusterNode(
            ssh: "mac-ultra-4.local",
            ips: ["192.168.1.104"],
            rdma: ["rdma0", "rdma1", "rdma2"]
        )
    ]

    // Write hostfile to disk
    let hostfileURL = try MLXHostfileGenerator.buildHostfile(nodes: nodes)
    print("Hostfile written to: \(hostfileURL.path)")

    // Launch distributed inference job
    let launcher = MLXDistributedLauncher(hostfileURL: hostfileURL)
    let job = try launcher.launchDistributedChat(
        model: "mlx-community/Qwen3-6-8bit",
        maxTokens: 2048
    )
    print("Distributed inference job started with PID: \(job.processIdentifier)")
    job.waitUntilExit()
    print("Job finished with status: \(job.terminationStatus)")
}

// Entry point
try runDistributedInferenceExample()