diff --git a/.github/workflows/runner-gpu.yml b/.github/workflows/runner-gpu.yml new file mode 100644 index 0000000..33520a7 --- /dev/null +++ b/.github/workflows/runner-gpu.yml @@ -0,0 +1,49 @@ +# Build the hosted-tier GPU runner image and push to GHCR. Default = fast path from the prebuilt +# binary release asset; set build_from_source=true to compile (see docker/runner-gpu/Dockerfile). +name: runner-gpu +on: + workflow_dispatch: + inputs: + kokkos_arch: + description: 'Kokkos arch (ADA89 | AMPERE80)' + default: ADA89 + type: string + tarball_url: + description: 'Prebuilt lammps tarball URL (release asset)' + default: https://github.com/forcefieldsilicon/mdengine/releases/download/runner-bin-20260905/lammps-kokkos-ADA89.tar.gz + type: string + build_from_source: + description: 'Compile instead of using the tarball' + default: false + type: boolean +jobs: + build: + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v4 + - uses: docker/setup-buildx-action@v3 + - uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Pick Dockerfile + id: pick + run: | + if [ "${{ inputs.build_from_source }}" = "true" ]; then echo "file=docker/runner-gpu/Dockerfile" >> "$GITHUB_OUTPUT"; else echo "file=docker/runner-gpu/Dockerfile.prebuilt" >> "$GITHUB_OUTPUT"; fi + - uses: docker/build-push-action@v6 + with: + context: docker/runner-gpu + file: ${{ steps.pick.outputs.file }} + push: true + build-args: | + KOKKOS_ARCH=${{ inputs.kokkos_arch }} + LMP_TARBALL_URL=${{ inputs.tarball_url }} + tags: | + ghcr.io/forcefieldsilicon/mdengine-runner-gpu:${{ inputs.kokkos_arch }} + ghcr.io/forcefieldsilicon/mdengine-runner-gpu:${{ inputs.kokkos_arch }}-${{ github.sha }} + labels: | + org.opencontainers.image.source=https://github.com/forcefieldsilicon/mdengine diff --git a/.gitignore b/.gitignore index 18f9783..ac5aac2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,13 +1,15 @@ -Packages +__pycache__/ +! .env.example .build -xcuserdata -*.xcodeproj -DerivedData/ .DS_Store -db.sqlite -.swiftpm .env .env.* -! .env.example +.swiftpm .vscode +*.pyc +*.xcodeproj +db.sqlite +DerivedData/ dist/ +Packages +xcuserdata diff --git a/CLAUDE.md b/CLAUDE.md index 9b22674..a6e12c4 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ macOS MD workbench: `MDEngine` (SwiftUI+Metal viewer), `mdengine-cli`, (Gitinama Inc.). ## Commands -- Build: `swift build` · Tests: `swift test` (16 cases, keep green) +- Build: `swift build` · Tests: `swift test` (20 cases, keep green) - Release: `swift build -c release` — **ALWAYS run after changes**: `/opt/homebrew/bin/mdengine` and `mdengine-mcp` are symlinks into `.build/release/`; a debug-only build leaves every other session running @@ -19,6 +19,10 @@ macOS MD workbench: `MDEngine` (SwiftUI+Metal viewer), `mdengine-cli`, `~/.mdengine/hosts.json`; test bed = `localhost-test` host over ssh to this Mac (own key in ~/.ssh/authorized_keys). `lmp` must be an absolute path — no login PATH over ssh. Release packaging: `make_tools.sh` (CLI+MCP tarball). +- Hosted GPU tier (`Sources/LAMMPSCore/HostedClient.swift`, `hosted/CONTRACT.md`): dev + loop = `python3 hosted/mock/mock_endpoint.py --port 8788 --data ` + + `MDENGINE_HOSTED_URL=http://127.0.0.1:8788/v1 MDENGINE_HOSTED_LAUNCH="{lmp} -in {input} -log log.lammps"`, + pod stand-in = `docker/runner-gpu/runner.sh` with `MDE_WORK`/`LMP` (see hosted/README.md). ## Gotchas that already bit - SPM does NOT prune deleted resources from an existing diff --git a/Sources/LAMMPSCore/HostedClient.swift b/Sources/LAMMPSCore/HostedClient.swift new file mode 100644 index 0000000..318a437 --- /dev/null +++ b/Sources/LAMMPSCore/HostedClient.swift @@ -0,0 +1,362 @@ +import Foundation + +// Hosted accelerated runs — the client side of hosted/CONTRACT.md (GJOB-091). +// +// This is a TRANSPORT, not a product: the same job model as the local runner +// and the ssh remote hosts (`~/.mdengine/jobs//` bookkeeping, deck dir +// shipped whole, results pulled back), the difference being that the deck runs +// on a rented GPU behind api.forcefieldsilicon.com and debits a prepaid credit +// balance. One client, three surfaces: `mdengine run --gpu`, MCP +// `submit_lammps host=cloud`, and the app's File ▸ Run Accelerated…. +// +// Credentials: $MDENGINE_API_KEY, else ~/.mdengine/credentials (JSON, mode 0600): +// { "api_key": "mde_…", "endpoint": "http://127.0.0.1:8787/v1" } // endpoint optional +// $MDENGINE_HOSTED_URL overrides the endpoint (how the mock is reached in dev). + +public struct HostedCredentials: Codable, Equatable { + public var apiKey: String + public var endpoint: String? + + public init(apiKey: String, endpoint: String? = nil) { + self.apiKey = apiKey + self.endpoint = endpoint + } + + public static let productionEndpoint = "https://api.forcefieldsilicon.com/v1" + public static let fileURL = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".mdengine/credentials") + + /// Env first (CI, one-off shells), then the credentials file. + public static func load() -> HostedCredentials? { + let env = ProcessInfo.processInfo.environment + var creds: HostedCredentials? + if let data = try? Data(contentsOf: fileURL), + let c = try? JSONDecoder().decode(HostedCredentials.self, from: data) { + creds = c + } + if let k = env["MDENGINE_API_KEY"], !k.isEmpty { + creds = HostedCredentials(apiKey: k, endpoint: creds?.endpoint) + } + if let u = env["MDENGINE_HOSTED_URL"], !u.isEmpty, creds != nil { + creds?.endpoint = u + } + return creds + } + + /// Written 0600: the key is money. + public func save() throws { + let dir = Self.fileURL.deletingLastPathComponent() + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + let enc = JSONEncoder(); enc.outputFormatting = [.prettyPrinted, .sortedKeys] + try enc.encode(self).write(to: Self.fileURL, options: .atomic) + try FileManager.default.setAttributes([.posixPermissions: 0o600], ofItemAtPath: Self.fileURL.path) + } + + public static func clear() { + try? FileManager.default.removeItem(at: fileURL) + } + + public var resolvedEndpoint: String { + ProcessInfo.processInfo.environment["MDENGINE_HOSTED_URL"] ?? endpoint ?? Self.productionEndpoint + } +} + +// MARK: - Wire types (hosted/CONTRACT.md) + +public struct HostedAccount: Codable { + public let balance_usd: Double + public let rate_table: [String: Double] + public let keys_created: String? +} + +public struct HostedJobSpec: Codable { + public var input: String + public var label: String? + public var gpu: String + public var wall_limit_s: Int + public var estimate_s: Int + public var launch: String + + public init(input: String, label: String? = nil, gpu: String = "any", + wallLimitS: Int = 14400, estimateS: Int = 3600, launch: String = "default") { + self.input = input; self.label = label; self.gpu = gpu + self.wall_limit_s = wallLimitS; self.estimate_s = estimateS; self.launch = launch + } +} + +public struct HostedJobStatus: Codable { + public let id: String + public let state: String + public let created: String? + public let started: String? + public let finished: String? + public let gpu: String? + public let rate_usd_per_h: Double? + public let billed_s: Int? + public let cost_usd: Double? + public let thermo_tail: [String]? + public let exitcode: Int? + public let error: String? + public let attempt: Int? + + public var isTerminal: Bool { ["done", "failed", "cancelled"].contains(state) } + + /// One line, the way job_status prints local jobs. + public var summary: String { + var s = "\(id): \(state)" + if let g = gpu, state != "created", state != "uploaded" { s += " on \(g)" } + if let c = cost_usd, c > 0 { s += String(format: " $%.4f", c) } + if let b = billed_s, b > 0 { s += " (\(b) s billed)" } + if let e = error { s += " error: \(e)" } + if let x = exitcode, isTerminal { s += " exit \(x)" } + return s + } +} + +public struct HostedError: LocalizedError { + public let message: String + public let status: Int? + public var errorDescription: String? { message } + init(_ m: String, status: Int? = nil) { message = m; self.status = status } +} + +// MARK: - Client + +public final class HostedClient { + public let creds: HostedCredentials + public let base: URL + private let session: URLSession + + /// Big artifacts a deck directory accumulates; never worth uploading (same list as RemoteHost). + public static let deckExcludes = ["*.traj", "*.lammpstrj", "*.dump", "*.ckpt*", "*.restart*", + "*.log", "*.mp4", "*.gif", "*.xlsx", ".git", ".DS_Store", "results", "results-*"] + + public init(credentials: HostedCredentials) throws { + guard let u = URL(string: credentials.resolvedEndpoint) else { + throw HostedError("bad endpoint URL: \(credentials.resolvedEndpoint)") + } + creds = credentials + base = u + let cfg = URLSessionConfiguration.ephemeral + cfg.timeoutIntervalForRequest = 60 + cfg.timeoutIntervalForResource = 3600 // deck uploads / result downloads can be large + session = URLSession(configuration: cfg) + } + + /// Convenience: fail with a message that says how to log in. + public static func fromSavedCredentials() throws -> HostedClient { + guard let c = HostedCredentials.load() else { + throw HostedError("no API key — run `mdengine login ` (or set $MDENGINE_API_KEY). " + + "Keys come with a credit pack: https://forcefieldsilicon.com/mdengine") + } + return try HostedClient(credentials: c) + } + + // MARK: HTTP (synchronous — CLI and MCP are single-threaded; the app calls off-main) + + private func request(_ method: String, _ path: String, json: Any? = nil, body: Data? = nil, + absolute: URL? = nil, auth: Bool = true) throws -> (Int, Data) { + let url = absolute ?? base.appendingPathComponent(path) + var req = URLRequest(url: url) + req.httpMethod = method + if auth { req.setValue("Bearer \(creds.apiKey)", forHTTPHeaderField: "Authorization") } + if let json { + req.setValue("application/json", forHTTPHeaderField: "Content-Type") + req.httpBody = try JSONSerialization.data(withJSONObject: json) + } else if let body { + req.setValue("application/gzip", forHTTPHeaderField: "Content-Type") + req.httpBody = body + } + let sem = DispatchSemaphore(value: 0) + var out: (Int, Data)? + var failure: Error? + session.dataTask(with: req) { data, resp, error in + if let error { failure = error } + else { out = ((resp as? HTTPURLResponse)?.statusCode ?? 0, data ?? Data()) } + sem.signal() + }.resume() + sem.wait() + if let failure { + throw HostedError("cannot reach \(url.host ?? url.absoluteString): \(failure.localizedDescription)") + } + return out! + } + + private func decode(_ type: T.Type, _ r: (Int, Data), expect: Set = [200, 201, 202]) throws -> T { + guard expect.contains(r.0) else { throw serverError(r) } + do { return try JSONDecoder().decode(type, from: r.1) } + catch { throw HostedError("unexpected response from \(base.host ?? "endpoint") (HTTP \(r.0)): \(String(decoding: r.1.prefix(200), as: UTF8.self))") } + } + + private func serverError(_ r: (Int, Data)) -> HostedError { + let msg = (try? JSONSerialization.jsonObject(with: r.1) as? [String: Any])?["error"] as? String + switch r.0 { + case 401: return HostedError("API key rejected (HTTP 401) — check `mdengine login`", status: 401) + case 402: return HostedError("insufficient credit balance (HTTP 402) — buy credits at https://forcefieldsilicon.com/mdengine", status: 402) + default: return HostedError("endpoint error HTTP \(r.0)\(msg.map { ": \($0)" } ?? "")", status: r.0) + } + } + + // MARK: API + + public func me() throws -> HostedAccount { + try decode(HostedAccount.self, request("GET", "me")) + } + + public func status(_ id: String) throws -> HostedJobStatus { + try decode(HostedJobStatus.self, request("GET", "jobs/\(id)")) + } + + public func list() throws -> [HostedJobStatus] { + struct L: Codable { let jobs: [HostedJobStatus] } + return try decode(L.self, request("GET", "jobs")).jobs + } + + public func cancel(_ id: String) throws -> HostedJobStatus { + try decode(HostedJobStatus.self, request("DELETE", "jobs/\(id)")) + } + + /// Tar the deck's directory (minus trajectories/checkpoints/logs), create the job, + /// upload, start. Returns the endpoint's job id and writes local bookkeeping so + /// job_status / list_jobs / fetch see it like any other job. + public func submit(input: String, spec: HostedJobSpec? = nil) throws -> String { + let inputURL = URL(fileURLWithPath: (input as NSString).expandingTildeInPath).standardizedFileURL + guard FileManager.default.fileExists(atPath: inputURL.path) else { throw HostedError("no such input: \(input)") } + let deckDir = inputURL.deletingLastPathComponent() + var spec = spec ?? HostedJobSpec(input: inputURL.lastPathComponent) + spec.input = inputURL.lastPathComponent + if spec.label == nil { spec.label = inputURL.deletingPathExtension().lastPathComponent } + // Dev only: the mock's "pod" is this Mac's lmp_serial, which has no KOKKOS — let the + // launch line be overridden without touching any surface's code path. + if let l = ProcessInfo.processInfo.environment["MDENGINE_HOSTED_LAUNCH"], !l.isEmpty { spec.launch = l } + + let tarball = try Self.tarDeck(deckDir) + guard tarball.count <= 2_000_000_000 else { + throw HostedError("deck directory tars to \(tarball.count / 1_000_000) MB — the limit is 2 GB; move old results out of it") + } + + struct Created: Codable { let id: String; let upload_url: String } + let specObj = try JSONSerialization.jsonObject(with: JSONEncoder().encode(spec)) + let created = try decode(Created.self, request("POST", "jobs", json: specObj)) + guard let up = URL(string: created.upload_url) else { throw HostedError("bad upload_url from endpoint") } + let put = try request("PUT", "", body: tarball, absolute: up, auth: false) + guard (200..<300).contains(put.0) else { throw HostedError("deck upload failed (HTTP \(put.0))") } + _ = try decode(HostedJobStatus.self, request("POST", "jobs/\(created.id)/start")) + + let jobDir = Self.jobsRoot.appendingPathComponent(created.id) + try FileManager.default.createDirectory(at: jobDir, withIntermediateDirectories: true) + let meta: [String: Any] = [ + "id": created.id, "input": inputURL.path, "cwd": deckDir.path, + "cloud": true, "endpoint": base.absoluteString, "gpu": spec.gpu, "label": spec.label ?? "", + "started": Date().timeIntervalSince1970, "deck_bytes": tarball.count, + ] + try JSONSerialization.data(withJSONObject: meta, options: [.prettyPrinted]) + .write(to: jobDir.appendingPathComponent("job.json")) + return created.id + } + + /// Download the results tarball (work/ + log.lammps + exitcode) into + /// ~/.mdengine/jobs//results/ and mirror log/exitcode into the job dir. + /// Returns the results directory. + @discardableResult + public func fetch(_ id: String) throws -> URL { + struct R: Codable { let download_url: String; let bytes: Int? } + let r = try decode(R.self, request("GET", "jobs/\(id)/results"), expect: [200]) + guard let dl = URL(string: r.download_url) else { throw HostedError("bad download_url from endpoint") } + let got = try request("GET", "", absolute: dl, auth: false) + guard got.0 == 200, !got.1.isEmpty else { throw HostedError("results download failed (HTTP \(got.0))") } + + let jobDir = Self.jobsRoot.appendingPathComponent(id) + let results = jobDir.appendingPathComponent("results") + try FileManager.default.createDirectory(at: results, withIntermediateDirectories: true) + let tmp = jobDir.appendingPathComponent("results.tar.gz") + try got.1.write(to: tmp) + let untar = try Self.run("/usr/bin/tar", ["-xzf", tmp.path, "-C", results.path]) + guard untar.status == 0 else { throw HostedError("untar failed: \(untar.err)") } + try? FileManager.default.removeItem(at: tmp) + // The runner tars `work/ log.lammps exitcode` at the top level; surface the two + // bookkeeping files where local jobs keep them. + for name in ["log.lammps", "exitcode", "stdout.txt"] { + let src = results.appendingPathComponent(name), dst = jobDir.appendingPathComponent(name) + if FileManager.default.fileExists(atPath: src.path) { + try? FileManager.default.removeItem(at: dst) + try? FileManager.default.copyItem(at: src, to: dst) + } + } + return results + } + + /// Poll until terminal, reporting each status change (and new thermo lines) to `progress`. + public func wait(_ id: String, every seconds: Double = 5, + progress: ((HostedJobStatus, [String]) -> Void)? = nil) throws -> HostedJobStatus { + var seenTail: [String] = [] + var lastState = "" + while true { + let s = try status(id) + let tail = s.thermo_tail ?? [] + let fresh = tail.filter { !seenTail.contains($0) } + if s.state != lastState || !fresh.isEmpty { progress?(s, fresh) } + seenTail = tail; lastState = s.state + if s.isTerminal { return s } + Thread.sleep(forTimeInterval: seconds) + } + } + + // MARK: helpers + + public static let jobsRoot = FileManager.default.homeDirectoryForCurrentUser + .appendingPathComponent(".mdengine/jobs") + + /// job.json for a cloud job, if this id is one. + public static func cloudMeta(_ id: String) -> [String: Any]? { + guard let data = try? Data(contentsOf: jobsRoot.appendingPathComponent(id).appendingPathComponent("job.json")), + let obj = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + obj["cloud"] as? Bool == true else { return nil } + return obj + } + + /// The trajectory worth opening from a results dir: newest of the dump-like files. + public static func primaryTrajectory(in results: URL) -> URL? { + let fm = FileManager.default + guard let e = fm.enumerator(at: results, includingPropertiesForKeys: [.contentModificationDateKey, .fileSizeKey]) else { return nil } + let exts: Set = ["traj", "lammpstrj", "xyz", "dump"] + var best: (URL, Int)? + for case let u as URL in e where exts.contains(u.pathExtension.lowercased()) { + let size = (try? u.resourceValues(forKeys: [.fileSizeKey]).fileSize) ?? 0 + if best == nil || size > best!.1 { best = (u, size) } + } + return best?.0 + } + + static func tarDeck(_ dir: URL) throws -> Data { + var args = ["-czf", "-", "-C", dir.path] + for x in deckExcludes { args += ["--exclude", x] } + args.append(".") + let r = try run("/usr/bin/tar", args) + guard r.status == 0 else { throw HostedError("tar of \(dir.path) failed: \(r.err)") } + return r.out + } + + struct Shell { let status: Int32; let out: Data; let err: String } + static func run(_ exe: String, _ args: [String], stdin: Data? = nil) throws -> Shell { + let p = Process() + p.executableURL = URL(fileURLWithPath: exe) + p.arguments = args + let o = Pipe(), e = Pipe() + p.standardOutput = o; p.standardError = e + let i = stdin.map { _ in Pipe() } + if let i { p.standardInput = i } + try p.run() + if let i, let stdin { + DispatchQueue.global().async { // feed concurrently with the drain below: no pipe deadlock either way + i.fileHandleForWriting.write(stdin) + i.fileHandleForWriting.closeFile() + } + } + // Drain stdout before waiting: a multi-MB tarball would fill the pipe and deadlock. + let out = o.fileHandleForReading.readDataToEndOfFile() + let err = String(data: e.fileHandleForReading.readDataToEndOfFile(), encoding: .utf8) ?? "" + p.waitUntilExit() + return Shell(status: p.terminationStatus, out: out, err: err) + } +} diff --git a/Sources/MDEngine/HostedJobsView.swift b/Sources/MDEngine/HostedJobsView.swift new file mode 100644 index 0000000..9e93b38 --- /dev/null +++ b/Sources/MDEngine/HostedJobsView.swift @@ -0,0 +1,306 @@ +import SwiftUI +import AppKit +import LAMMPSCore + +/// The app's side of the hosted GPU tier (GJOB-091): File ▸ Run Accelerated… +/// submits a deck, the Accelerated Runs window shows live thermo, and a finished +/// job's trajectory opens in the viewer by itself. All HostedClient calls are +/// synchronous, so every one of them runs on `queue`, never on main. +@MainActor +final class HostedJobsModel: ObservableObject { + @Published var jobs: [HostedJobStatus] = [] + @Published var balance: Double? + @Published var busy = false + @Published var lastError: String? + @Published var hasCredentials = HostedCredentials.load() != nil + /// Jobs whose results were already downloaded and opened — never re-open on the next poll. + private var opened: Set = [] + private var timer: Timer? + private let queue = DispatchQueue(label: "mdengine.hosted", qos: .userInitiated) + /// Set by the app: how to show a fetched trajectory. + var openTrajectory: ((URL) -> Void)? + + // MARK: polling + + func startPolling() { + refresh() + timer?.invalidate() + timer = Timer.scheduledTimer(withTimeInterval: 5, repeats: true) { [weak self] _ in + Task { @MainActor in self?.refresh() } + } + } + + func stopPolling() { timer?.invalidate(); timer = nil } + + func refresh() { + hasCredentials = HostedCredentials.load() != nil + guard hasCredentials else { return } + queue.async { [weak self] in + do { + let client = try HostedClient.fromSavedCredentials() + let list = try client.list() + let me = try? client.me() + Task { @MainActor in + guard let self else { return } + self.jobs = list + if let me { self.balance = me.balance_usd } + self.lastError = nil + self.autoFetchFinished(list, client) + } + } catch { + Task { @MainActor in self?.lastError = error.localizedDescription } + } + } + } + + /// A job that just reached `done` gets its results pulled and its trajectory shown, once. + private func autoFetchFinished(_ list: [HostedJobStatus], _ client: HostedClient) { + for j in list where j.state == "done" && !opened.contains(j.id) { + // Only jobs this Mac submitted (they have local bookkeeping) — not every job on the key. + guard HostedClient.cloudMeta(j.id) != nil else { continue } + opened.insert(j.id) + let alreadyFetched = FileManager.default.fileExists( + atPath: HostedClient.jobsRoot.appendingPathComponent(j.id).appendingPathComponent("results").path) + if alreadyFetched { continue } + fetch(j.id, client: client) + } + } + + // MARK: actions + + func submitPanel() { + let panel = NSOpenPanel() + panel.message = "Choose the LAMMPS input to run on a hosted GPU. Its whole directory is uploaded (minus trajectories/checkpoints/logs), so the deck must be self-contained." + panel.canChooseDirectories = false + panel.allowsMultipleSelection = false + guard panel.runModal() == .OK, let url = panel.url else { return } + submit(input: url) + } + + func submit(input: URL, gpu: String = "any", wallHours: Double = 4) { + busy = true + queue.async { [weak self] in + do { + let client = try HostedClient.fromSavedCredentials() + let spec = HostedJobSpec(input: input.lastPathComponent, + label: input.deletingPathExtension().lastPathComponent, + gpu: gpu, wallLimitS: Int(wallHours * 3600)) + let id = try client.submit(input: input.path, spec: spec) + Task { @MainActor in + self?.busy = false + self?.lastError = nil + self?.refresh() + _ = id + } + } catch { + Task { @MainActor in + self?.busy = false + self?.lastError = error.localizedDescription + Self.alert("Could not submit \(input.lastPathComponent)", info: error.localizedDescription) + } + } + } + } + + func fetch(_ id: String, client: HostedClient? = nil) { + queue.async { [weak self] in + do { + let c = try client ?? HostedClient.fromSavedCredentials() + let dir = try c.fetch(id) + let traj = HostedClient.primaryTrajectory(in: dir) + Task { @MainActor in + self?.lastError = nil + if let traj { self?.openTrajectory?(traj) } + else { NSWorkspace.shared.open(dir) } // no trajectory in the deck's output: show the folder + } + } catch { + Task { @MainActor in self?.lastError = error.localizedDescription } + } + } + } + + func cancel(_ id: String) { + queue.async { [weak self] in + do { _ = try HostedClient.fromSavedCredentials().cancel(id) } + catch { Task { @MainActor in self?.lastError = error.localizedDescription } } + Task { @MainActor in self?.refresh() } + } + } + + func revealResults(_ id: String) { + let dir = HostedClient.jobsRoot.appendingPathComponent(id).appendingPathComponent("results") + if FileManager.default.fileExists(atPath: dir.path) { NSWorkspace.shared.open(dir) } + else { fetch(id) } + } + + static func alert(_ message: String, info: String) { + let a = NSAlert() + a.messageText = message + a.informativeText = info + a.runModal() + } +} + +/// Window ▸ Accelerated Runs: one row per hosted job, live thermo for the selected one. +struct HostedJobsView: View { + @ObservedObject var model: HostedJobsModel + @State private var selected: String? + + var body: some View { + VStack(spacing: 0) { + if !model.hasCredentials { + ContentUnavailableView { + Label("No API key", systemImage: "key") + } description: { + Text("Accelerated runs use prepaid GPU credits. Add your key in Settings ▸ Accelerated, or get one with a credit pack.") + } actions: { + SettingsLink { Text("Open Settings…") } + Link("Get credits", destination: HostedLinks.credits) + } + } else { + HSplitView { + List(model.jobs, id: \.id, selection: $selected) { j in + HStack { + Circle().fill(color(j.state)).frame(width: 8, height: 8) + VStack(alignment: .leading, spacing: 2) { + Text(j.id).font(.system(.body, design: .monospaced)) + Text(line(j)).font(.caption).foregroundStyle(.secondary) + } + Spacer() + } + .contextMenu { + if j.isTerminal { Button("Open results") { model.revealResults(j.id) } } + else { Button("Cancel job") { model.cancel(j.id) } } + } + .tag(j.id) + } + .frame(minWidth: 320) + detail + .frame(minWidth: 380) + } + } + Divider() + HStack { + Button { model.submitPanel() } label: { Label("Run Accelerated…", systemImage: "bolt.fill") } + .disabled(!model.hasCredentials || model.busy) + if model.busy { ProgressView().controlSize(.small) } + Spacer() + if let e = model.lastError { + Text(e).font(.caption).foregroundStyle(.red).lineLimit(1).help(e) + } + if let b = model.balance { + Text(String(format: "Balance $%.2f", b)).font(.callout).monospacedDigit() + } + Link("Buy credits", destination: HostedLinks.credits).font(.callout) + } + .padding(10) + } + .frame(minWidth: 720, minHeight: 360) + .onAppear { model.startPolling() } + .onDisappear { model.stopPolling() } + } + + @ViewBuilder private var detail: some View { + if let id = selected, let j = model.jobs.first(where: { $0.id == id }) { + VStack(alignment: .leading, spacing: 8) { + Text(j.summary).font(.headline).textSelection(.enabled) + ScrollView { + Text((j.thermo_tail ?? []).isEmpty ? "(no thermo yet)" : (j.thermo_tail ?? []).joined(separator: "\n")) + .font(.system(.caption, design: .monospaced)) + .textSelection(.enabled) + .frame(maxWidth: .infinity, alignment: .leading) + } + HStack { + if j.isTerminal { + Button("Open results") { model.revealResults(j.id) } + } else { + Button("Cancel") { model.cancel(j.id) } + } + Spacer() + } + } + .padding(12) + } else { + ContentUnavailableView("Select a job", systemImage: "list.bullet", + description: Text("Live thermo output shows here. Finished runs open in the viewer automatically.")) + } + } + + private func line(_ j: HostedJobStatus) -> String { + var parts: [String] = [j.state] + if let g = j.gpu, !["created", "uploaded"].contains(j.state) { parts.append(g) } + if let c = j.cost_usd, c > 0 { parts.append(String(format: "$%.3f", c)) } + if let created = j.created { parts.append(created) } + return parts.joined(separator: " · ") + } + + private func color(_ state: String) -> Color { + switch state { + case "done": return .green + case "failed": return .red + case "cancelled": return .gray + case "running", "uploading": return .blue + default: return .orange + } + } +} + +enum HostedLinks { + /// Credit packs (Stripe Payment Links live behind this page — GJOB-096). + static let credits = URL(string: "https://forcefieldsilicon.com/mdengine#credits")! +} + +/// Settings ▸ Accelerated: API key + endpoint, verified against the endpoint before saving. +struct HostedSettingsView: View { + @State private var apiKey = HostedCredentials.load()?.apiKey ?? "" + @State private var endpoint = HostedCredentials.load()?.endpoint ?? "" + @State private var advanced = false + @State private var status: String = "" + @State private var checking = false + + var body: some View { + Form { + SecureField("API key", text: $apiKey, prompt: Text("mde_…")) + .textFieldStyle(.roundedBorder) + DisclosureGroup("Advanced", isExpanded: $advanced) { + TextField("Endpoint", text: $endpoint, prompt: Text(HostedCredentials.productionEndpoint)) + .textFieldStyle(.roundedBorder) + .font(.system(.body, design: .monospaced)) + } + HStack { + Button(checking ? "Checking…" : "Verify & Save") { verifyAndSave() } + .disabled(checking || !apiKey.hasPrefix("mde_")) + Button("Remove key") { + HostedCredentials.clear(); apiKey = ""; status = "key removed" + } + .disabled(HostedCredentials.load() == nil) + Spacer() + Link("Get credits", destination: HostedLinks.credits) + } + if !status.isEmpty { + Text(status).font(.caption).foregroundStyle(status.hasPrefix("✓") ? Color.secondary : Color.red) + } + Text("The key is stored in ~/.mdengine/credentials (readable only by you) and shared with the mdengine CLI and MCP server.") + .font(.caption).foregroundStyle(.secondary) + } + .padding(20) + .frame(width: 420) + } + + private func verifyAndSave() { + checking = true + let creds = HostedCredentials(apiKey: apiKey.trimmingCharacters(in: .whitespaces), + endpoint: endpoint.isEmpty ? nil : endpoint) + DispatchQueue.global(qos: .userInitiated).async { + let result: String + do { + let me = try HostedClient(credentials: creds).me() + try creds.save() + result = String(format: "✓ key saved — balance $%.2f", me.balance_usd) + } catch { + result = error.localizedDescription + } + DispatchQueue.main.async { status = result; checking = false } + } + } +} diff --git a/Sources/MDEngine/MDEngineApp.swift b/Sources/MDEngine/MDEngineApp.swift index 20a52c7..a10ffdc 100644 --- a/Sources/MDEngine/MDEngineApp.swift +++ b/Sources/MDEngine/MDEngineApp.swift @@ -21,6 +21,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate { @main struct MDEngineApp: App { @StateObject private var model = ContentViewModel() + @StateObject private var hosted = HostedJobsModel() @NSApplicationDelegateAdaptor(AppDelegate.self) private var appDelegate init() { @@ -37,21 +38,30 @@ struct MDEngineApp: App { // macOS state restoration was resurrecting confusing blank duplicates. Window("MDEngine", id: "main") { ContentView(model: model) + .modifier(HostedLaunch(hosted: hosted, viewer: model)) } .commands { - AppCommands(model: model) + AppCommands(model: model, hosted: hosted) } + // Hosted GPU runs: submit, watch live thermo, results auto-open (GJOB-091). + Window("Accelerated Runs", id: "hosted") { + HostedJobsView(model: hosted) + } + .defaultSize(width: 820, height: 420) + Settings { SettingsView() } } } -/// Menu bar: MDEngine · File (Load/Export) · Edit (Application Settings…) · Help (User Manual). +/// Menu bar: MDEngine · File (Load/Export/Run Accelerated) · Edit (Application Settings…) · Help (User Manual). struct AppCommands: Commands { @ObservedObject var model: ContentViewModel + @ObservedObject var hosted: HostedJobsModel @AppStorage("orthographicProjection") private var orthographic = false + @Environment(\.openWindow) private var openWindow var body: some Commands { CommandGroup(replacing: .newItem) { @@ -77,6 +87,14 @@ struct AppCommands: Commands { Button("Export Video…") { model.exportVideo(format: .mp4) } .keyboardShortcut("e", modifiers: [.command, .shift]) .disabled(model.frames.count < 2) + Divider() + Button("Run Accelerated…") { + openWindow(id: "hosted") + if hosted.hasCredentials { hosted.submitPanel() } + } + .keyboardShortcut("r", modifiers: [.command, .shift]) + Button("Accelerated Runs") { openWindow(id: "hosted") } + .keyboardShortcut("j", modifiers: [.command, .shift]) } CommandGroup(after: .pasteboard) { Divider() @@ -106,3 +124,23 @@ struct AppCommands: Commands { } } } + +/// Connects the hosted tier to the viewer and honours `MDEngine --run-accelerated ` +/// (used by `mdengine run --gpu --open` and by tests: submit + show the runs window, no panels). +private struct HostedLaunch: ViewModifier { + let hosted: HostedJobsModel + let viewer: ContentViewModel + @Environment(\.openWindow) private var openWindow + + func body(content: Content) -> some View { + content.onAppear { + // A finished hosted run lands in the viewer like any opened file. + hosted.openTrajectory = { url in viewer.load(url: url) } + let args = CommandLine.arguments + if let i = args.firstIndex(of: "--run-accelerated"), i + 1 < args.count { + openWindow(id: "hosted") + hosted.submit(input: URL(fileURLWithPath: args[i + 1])) + } + } + } +} diff --git a/Sources/MDEngine/SettingsView.swift b/Sources/MDEngine/SettingsView.swift index 0dca5a9..054c4d0 100644 --- a/Sources/MDEngine/SettingsView.swift +++ b/Sources/MDEngine/SettingsView.swift @@ -8,6 +8,14 @@ struct SettingsView: View { @AppStorage("backgroundBrightness") private var backgroundBrightness = 0.05 var body: some View { + TabView { + display.tabItem { Label("Display", systemImage: "cube") } + HostedSettingsView().tabItem { Label("Accelerated", systemImage: "bolt.fill") } + } + .frame(width: 460) + } + + private var display: some View { Form { Slider(value: $atomPointSize, in: 4...32) { Text("Atom size") @@ -28,6 +36,5 @@ struct SettingsView: View { } } .padding(20) - .frame(width: 380) } } diff --git a/Sources/MDEngineCLI/main.swift b/Sources/MDEngineCLI/main.swift index eef35d2..c7f82f7 100644 --- a/Sources/MDEngineCLI/main.swift +++ b/Sources/MDEngineCLI/main.swift @@ -20,6 +20,16 @@ USAGE keep every Nth frame (last always kept) mdengine run [--threads N] [--lmp PATH] [--log FILE] run LAMMPS with -sf omp -pk omp N + mdengine run --gpu [--label S] [--gpu-type any|rtx4090|a100] + [--wall-hours H] [--estimate-min M] [--no-wait] + run on a hosted GPU (prepaid credits): + ships the deck's directory, streams + thermo, downloads results when done + mdengine login [--endpoint URL] store the API key (~/.mdengine/credentials) + mdengine account credit balance + rate table + mdengine jobs hosted jobs, newest first + mdengine job [--log|--fetch|--cancel|--wait] + status / thermo tail / download / cancel mdengine gui open MDEngine.app NOTES @@ -31,6 +41,9 @@ NOTES Every subcommand accepts -h/--help. `run` finds LAMMPS via $MDENGINE_LMP, then lmp_mpi / lmp_serial / lmp on PATH. Default --threads = number of performance cores. + --gpu needs an API key (comes with a credit pack: forcefieldsilicon.com/mdengine); + $MDENGINE_API_KEY / $MDENGINE_HOSTED_URL override the stored credentials. + A hosted deck must be self-contained in its directory (data, potentials, molecule files). """ func fail(_ msg: String) -> Never { @@ -127,6 +140,28 @@ func potentialsDir(for lmp: String) -> String? { return candidates.first { FileManager.default.fileExists(atPath: $0) } } +func shellQuoteCLI(_ s: String) -> String { "'" + s.replacingOccurrences(of: "'", with: "'\\''") + "'" } + +/// Stream state changes + new thermo lines until terminal, then download results. Returns the exit code to use. +func hostedWaitAndFetch(_ client: HostedClient, _ id: String) -> Int32 { + do { + let final = try client.wait(id, every: 5) { s, fresh in + for line in fresh { print(" " + line) } + fflush(stdout) + if fresh.isEmpty, !s.isTerminal { print("mdengine: \(s.summary)") } + } + print("mdengine: \(final.summary)") + guard final.state == "done" || final.state == "failed" else { return 1 } + let dir = try client.fetch(id) + print("results → \(dir.path)") + if let t = HostedClient.primaryTrajectory(in: dir) { print("trajectory: \(t.path)") } + return final.state == "done" ? 0 : Int32(final.exitcode ?? 1) + } catch { + FileHandle.standardError.write(Data((error.localizedDescription + "\n").utf8)) + return 1 + } +} + var args = Array(CommandLine.arguments.dropFirst()) guard let command = args.first else { print(usage) @@ -223,6 +258,27 @@ case "decimate": } catch { fail("write failed: \(error.localizedDescription)") } case "run": + if takeFlag("--gpu", &args) { + guard let input = args.first else { fail("usage: mdengine run --gpu [--label S] [--gpu-type T] [--wall-hours H] [--no-wait]") } + args.removeFirst() + let label = takeOption("--label", &args) + let gpu = takeOption("--gpu-type", &args) ?? "any" + let wallH = Double(takeOption("--wall-hours", &args) ?? "") ?? 4 + let estMin = Double(takeOption("--estimate-min", &args) ?? "") ?? 60 + let noWait = takeFlag("--no-wait", &args) + let client: HostedClient + do { client = try HostedClient.fromSavedCredentials() } catch { fail(error.localizedDescription) } + let spec = HostedJobSpec(input: input, label: label, gpu: gpu, + wallLimitS: Int(wallH * 3600), estimateS: Int(estMin * 60)) + let id: String + do { id = try client.submit(input: input, spec: spec) } catch { fail(error.localizedDescription) } + print("mdengine: submitted \(id) → \(client.base.host ?? "endpoint") (gpu \(gpu), wall limit \(wallH) h)") + if noWait { + print("poll: mdengine job \(id)\nfetch: mdengine job \(id) --fetch") + exit(0) + } + exit(hostedWaitAndFetch(client, id)) + } guard let input = args.first else { fail("usage: mdengine run [--threads N] [--lmp PATH]") } args.removeFirst() let threads = Int(takeOption("--threads", &args) ?? "") ?? performanceCores() @@ -251,6 +307,57 @@ case "run": task.waitUntilExit() exit(task.terminationStatus) +case "login": + guard let key = args.first, key.hasPrefix("mde_") else { fail("usage: mdengine login [--endpoint URL]") } + args.removeFirst() + let endpoint = takeOption("--endpoint", &args) + let creds = HostedCredentials(apiKey: key, endpoint: endpoint) + do { + let acct = try HostedClient(credentials: creds).me() // verify before storing + try creds.save() + print("logged in — balance $\(String(format: "%.2f", acct.balance_usd)); key stored in \(HostedCredentials.fileURL.path) (0600)") + } catch { fail(error.localizedDescription) } + +case "account": + do { + let client = try HostedClient.fromSavedCredentials() + let acct = try client.me() + print("endpoint: \(client.base.absoluteString)") + print("balance: $\(String(format: "%.2f", acct.balance_usd))") + let rates = acct.rate_table.sorted { $0.key < $1.key }.map { "\($0.key) $\(String(format: "%.2f", $0.value))/h" } + print("rates: \(rates.joined(separator: ", "))") + print("credits: https://forcefieldsilicon.com/mdengine") + } catch { fail(error.localizedDescription) } + +case "jobs": + do { + let jobs = try HostedClient.fromSavedCredentials().list() + if jobs.isEmpty { print("(no hosted jobs)") } + for j in jobs { print(j.summary + (j.created.map { " \($0)" } ?? "")) } + } catch { fail(error.localizedDescription) } + +case "job": + guard let id = args.first else { fail("usage: mdengine job [--log|--fetch|--cancel|--wait]") } + args.removeFirst() + let client: HostedClient + do { client = try HostedClient.fromSavedCredentials() } catch { fail(error.localizedDescription) } + do { + if takeFlag("--cancel", &args) { + print(try client.cancel(id).summary) + } else if takeFlag("--fetch", &args) { + let dir = try client.fetch(id) + print("results → \(dir.path)") + if let t = HostedClient.primaryTrajectory(in: dir) { print("trajectory: \(t.path)\nopen: mdengine gui && open -a MDEngine \(shellQuoteCLI(t.path))") } + } else if takeFlag("--wait", &args) { + exit(hostedWaitAndFetch(client, id)) + } else { + let s = try client.status(id) + print(s.summary) + _ = takeFlag("--log", &args) // status always includes the thermo tail; --log is accepted for symmetry + for line in (s.thermo_tail ?? []).suffix(12) { print(" " + line) } + } + } catch { fail(error.localizedDescription) } + case "gui": let task = Process() task.executableURL = URL(fileURLWithPath: "/usr/bin/open") diff --git a/Sources/MDEngineMCP/main.swift b/Sources/MDEngineMCP/main.swift index 8bb1bc4..5e0311b 100644 --- a/Sources/MDEngineMCP/main.swift +++ b/Sources/MDEngineMCP/main.swift @@ -302,12 +302,14 @@ let toolDefs: [[String: Any]] = [ "out": ["type": "string", "description": "Output file path"]], "required": ["path", "every", "out"]]], ["name": "submit_lammps", - "description": "Submit a LAMMPS input script as a DETACHED background job — locally (keeps the machine awake, survives this server exiting, records its exit code) or on a configured REMOTE host (~/.mdengine/hosts.json: the deck's directory is rsynced up, minus trajectories/checkpoints/logs, and LAMMPS runs there under nohup with its exit code recorded remotely; a GPU box is just a host whose launch template carries the KOKKOS flags). Runs in the input's own directory so relative data/potential paths work; a remote deck must be self-contained within that directory. Returns a job id — poll with job_status; for remote jobs, fetch_job pulls results back.", + "description": "Submit a LAMMPS input script as a DETACHED background job — locally (keeps the machine awake, survives this server exiting, records its exit code) or on a configured REMOTE host (~/.mdengine/hosts.json: the deck's directory is rsynced up, minus trajectories/checkpoints/logs, and LAMMPS runs there under nohup with its exit code recorded remotely; a GPU box is just a host whose launch template carries the KOKKOS flags). host='cloud' sends the deck to the hosted GPU tier (RTX 4090, KOKKOS; prepaid credits at $2/GPU-h, API key from `mdengine login`) — same job tools, results come back with fetch_job. Runs in the input's own directory so relative data/potential paths work; a remote deck must be self-contained within that directory. Returns a job id — poll with job_status; for remote jobs, fetch_job pulls results back.", "inputSchema": ["type": "object", "properties": ["input": ["type": "string", "description": "Path to the LAMMPS input script"], "threads": ["type": "integer", "description": "OpenMP threads (default: performance-core count locally, or the host's configured threads)"], "label": ["type": "string", "description": "Short slug for the job id"], - "host": ["type": "string", "description": "Remote host name from hosts.json; 'local' forces this machine; omitted = hosts.json default, else local"]], + "host": ["type": "string", "description": "Remote host name from hosts.json; 'cloud' = the hosted GPU tier (prepaid credits, API key via `mdengine login`); 'local' forces this machine; omitted = hosts.json default, else local"], + "gpu": ["type": "string", "description": "host=cloud only: any (cheapest) | rtx4090 | a100"], + "wall_hours": ["type": "number", "description": "host=cloud only: hard wall-clock cap in hours (default 4; billed to the cap if hit)"]], "required": ["input"]]], ["name": "list_hosts", "description": "List the execution hosts configured in ~/.mdengine/hosts.json (ssh target, remote LAMMPS, workdir, launch template) and which is the default.", @@ -584,6 +586,13 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { case "submit_lammps": guard let input = a["input"] as? String else { throw err("invalid arguments: input") } + if (a["host"] as? String) == "cloud" { + let client = try HostedClient.fromSavedCredentials() + let spec = HostedJobSpec(input: input, label: a["label"] as? String, gpu: a["gpu"] as? String ?? "any", + wallLimitS: Int(((a["wall_hours"] as? Double) ?? 4) * 3600)) + let id = try client.submit(input: input, spec: spec) + return "submitted \(id) to the hosted GPU tier (\(client.base.host ?? "endpoint"), gpu \(spec.gpu), wall limit \(spec.wall_limit_s / 3600) h)\nlocal job dir: \(Jobs.dir(id).path)\npoll with job_status (live thermo tail); fetch_job downloads results when done" + } if let host = try RemoteHosts.resolve(a["host"] as? String) { let id = try RemoteJobs.submit(host: host, input: input, threads: a["threads"] as? Int, label: a["label"] as? String) @@ -594,11 +603,23 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { return "submitted \(id)\njob dir: \(Jobs.dir(id).path)\npoll with job_status" case "list_hosts": - return RemoteHosts.describe() + var cloud = "cloud: hosted GPU tier — " + if let c = try? HostedClient.fromSavedCredentials(), let me = try? c.me() { + cloud += "\(c.base.host ?? c.base.absoluteString), balance $\(String(format: "%.2f", me.balance_usd)), rates \(me.rate_table.sorted { $0.key < $1.key }.map { "\($0.key) $\($0.value)/h" }.joined(separator: ", "))" + } else { + cloud += "no API key (mdengine login ; keys come with a credit pack at forcefieldsilicon.com/mdengine)" + } + return RemoteHosts.describe() + "\n" + cloud case "fetch_job": guard let id = a["job_id"] as? String else { throw err("invalid arguments: job_id") } guard Jobs.meta(id) != nil else { throw err("unknown job \(id)") } + if HostedClient.cloudMeta(id) != nil { + let dir = try HostedClient.fromSavedCredentials().fetch(id) + let names = (try? FileManager.default.contentsOfDirectory(atPath: dir.path))?.sorted() ?? [] + let traj = HostedClient.primaryTrajectory(in: dir).map { "\ntrajectory: \($0.path)" } ?? "" + return "fetched \(names.count) files → \(dir.path)\n" + names.prefix(50).map { " " + $0 }.joined(separator: "\n") + traj + } guard let host = Jobs.remoteHost(id) else { return "\(id) ran locally — its files are already in place (see job_files)" } let withTraj = a["include_trajectories"] as? Bool ?? true let excludes = withTraj ? [".git"] : RemoteHost.defaultExcludes @@ -607,6 +628,12 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { case "job_status": guard let id = a["job_id"] as? String else { throw err("invalid arguments: job_id") } guard let meta = Jobs.meta(id) else { throw err("unknown job \(id)") } + if HostedClient.cloudMeta(id) != nil { + let s = try HostedClient.fromSavedCredentials().status(id) + let tail = (s.thermo_tail ?? []).suffix(8).joined(separator: "\n") + return "\(s.summary) · hosted GPU tier\ninput: \(meta["input"] ?? "?")\n" + (tail.isEmpty ? "(no thermo yet)" : tail) + + (s.isTerminal ? "\nfetch_job downloads results" : "") + } let elapsed = (meta["started"] as? Double) .map { String(format: "%.0f s", Date().timeIntervalSince1970 - $0) } ?? "?" let where_ = (meta["host"] as? String).map { " · host \($0)" } ?? "" @@ -619,6 +646,14 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { case "job_log": guard let id = a["job_id"] as? String else { throw err("invalid arguments: job_id") } let n = a["lines"] as? Int ?? 40 + if HostedClient.cloudMeta(id) != nil { + let local = Jobs.dir(id).appendingPathComponent("log.lammps") + if let text = try? String(contentsOf: local, encoding: .utf8) { // fetched already + return text.split(separator: "\n").suffix(n).joined(separator: "\n") + } + let s = try HostedClient.fromSavedCredentials().status(id) + return (s.thermo_tail ?? ["(no thermo yet)"]).joined(separator: "\n") + "\n(live tail from the endpoint; the full log arrives with fetch_job)" + } if let h = Jobs.remoteHost(id) { return RemoteJobs.logTail(id, host: h, lines: n) } let log = Jobs.dir(id).appendingPathComponent("log.lammps") let alt = Jobs.dir(id).appendingPathComponent("stdout.log") @@ -645,6 +680,11 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { } return "\(label): \(dir)\n" + rows.joined(separator: "\n") } + if HostedClient.cloudMeta(id) != nil { + let res = Jobs.dir(id).appendingPathComponent("results").path + return listing(res, label: "fetched results") + "\n" + listing(Jobs.dir(id).path, label: "job bookkeeping") + + (FileManager.default.fileExists(atPath: res) ? "" : "\n(hosted job — fetch_job downloads results when done)") + } if let h = Jobs.remoteHost(id) { return RemoteJobs.files(id, host: h) + "\n" + listing(Jobs.dir(id).path, label: "local job bookkeeping") @@ -656,6 +696,7 @@ func callTool(_ name: String, _ a: [String: Any]) throws -> String { case "cancel_job": guard let id = a["job_id"] as? String else { throw err("invalid arguments: job_id") } + if HostedClient.cloudMeta(id) != nil { return try HostedClient.fromSavedCredentials().cancel(id).summary } return try Jobs.cancel(id) case "run_lammps": diff --git a/Tests/AppTests/HostedClientTests.swift b/Tests/AppTests/HostedClientTests.swift new file mode 100644 index 0000000..8b7bb3d --- /dev/null +++ b/Tests/AppTests/HostedClientTests.swift @@ -0,0 +1,53 @@ +import XCTest +@testable import LAMMPSCore + +/// Offline pieces of the hosted-tier client. The wire protocol itself is exercised +/// against hosted/mock/mock_endpoint.py (see hosted/README.md), not here. +final class HostedClientTests: XCTestCase { + func testDeckTarExcludesArtifacts() throws { + let dir = FileManager.default.temporaryDirectory.appendingPathComponent("mde-deck-\(UUID().uuidString)") + try FileManager.default.createDirectory(at: dir.appendingPathComponent("results"), withIntermediateDirectories: true) + for name in ["in.lmp", "ffield.reax.X", "O2.data", "big.traj", "old.lammpstrj", "run.log", "sim.ckpt.a", "results/traj.lammpstrj"] { + try "x".write(to: dir.appendingPathComponent(name), atomically: true, encoding: .utf8) + } + let tarball = try HostedClient.tarDeck(dir) + let list = try HostedClient.run("/usr/bin/tar", ["-tzf", "-"], stdin: tarball) + let names = Set(String(decoding: list.out, as: UTF8.self).split(separator: "\n").map { $0.replacingOccurrences(of: "./", with: "") }) + XCTAssertTrue(names.isSuperset(of: ["in.lmp", "ffield.reax.X", "O2.data"])) + for excluded in ["big.traj", "old.lammpstrj", "run.log", "sim.ckpt.a", "results/traj.lammpstrj"] { + XCTAssertFalse(names.contains(excluded), "\(excluded) should not be uploaded") + } + try? FileManager.default.removeItem(at: dir) + } + + func testPrimaryTrajectoryPicksLargestDumpLikeFile() throws { + let dir = FileManager.default.temporaryDirectory.appendingPathComponent("mde-res-\(UUID().uuidString)") + try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true) + try String(repeating: "a", count: 10).write(to: dir.appendingPathComponent("small.xyz"), atomically: true, encoding: .utf8) + try String(repeating: "a", count: 1000).write(to: dir.appendingPathComponent("traj.lammpstrj"), atomically: true, encoding: .utf8) + try String(repeating: "a", count: 5000).write(to: dir.appendingPathComponent("log.lammps"), atomically: true, encoding: .utf8) + XCTAssertEqual(HostedClient.primaryTrajectory(in: dir)?.lastPathComponent, "traj.lammpstrj") + try? FileManager.default.removeItem(at: dir) + } + + func testJobStatusSummaryAndTerminal() throws { + let json = """ + {"id":"MDJOB-20260905-K3F9QZ","state":"done","gpu":"rtx4090","rate_usd_per_h":2.0,"billed_s":812,"cost_usd":0.4511,"exitcode":0} + """ + let s = try JSONDecoder().decode(HostedJobStatus.self, from: Data(json.utf8)) + XCTAssertTrue(s.isTerminal) + XCTAssertEqual(s.summary, "MDJOB-20260905-K3F9QZ: done on rtx4090 $0.4511 (812 s billed) exit 0") + let q = try JSONDecoder().decode(HostedJobStatus.self, from: Data(#"{"id":"X","state":"queued","gpu":"any"}"#.utf8)) + XCTAssertFalse(q.isTerminal) + } + + func testCredentialsEnvOverrideEndpoint() { + let c = HostedCredentials(apiKey: "mde_abc", endpoint: nil) + // No env in the test runner → production endpoint. + if ProcessInfo.processInfo.environment["MDENGINE_HOSTED_URL"] == nil { + XCTAssertEqual(c.resolvedEndpoint, HostedCredentials.productionEndpoint) + } + XCTAssertEqual(HostedCredentials(apiKey: "mde_abc", endpoint: "http://127.0.0.1:8787/v1").resolvedEndpoint, + ProcessInfo.processInfo.environment["MDENGINE_HOSTED_URL"] ?? "http://127.0.0.1:8787/v1") + } +} diff --git a/docker/runner-gpu/Dockerfile b/docker/runner-gpu/Dockerfile new file mode 100644 index 0000000..a84d962 --- /dev/null +++ b/docker/runner-gpu/Dockerfile @@ -0,0 +1,55 @@ +# MDEngine GPU runner — LAMMPS + KOKKOS/CUDA for the hosted accelerated tier (GJOB-088). +# +# Deployment model (MVP): this image IS the rented pod (RunPod / Vast / Lambda). It runs sshd, +# so mdengine-mcp's P0.5 transport (ssh + rsync, hosts.json) drives it with no new code: +# hosts.json entry: { "ssh": "root@", "ssh_options": ["-p", ""], "workdir": "/work", +# "lmp": "/usr/local/bin/lmp", "threads": 1, +# "launch": "{lmp} -in {input} -k on g 1 -sf kk -pk kokkos newton on neigh half -log {log}" } +# One pod per job, destroyed after fetch: the pod is the sandbox (it holds nothing of ours but +# this image). Per-job `docker run --network none` hardening (see ../runner) applies only where +# we control the docker host (Lambda/Vast VM) — not inside a RunPod pod. +# +# Build (x86_64 host with nvcc; NO GPU needed to compile — CI or any rented box): +# docker build --build-arg KOKKOS_ARCH=ADA89 -t mdengine-runner-gpu:ada89 docker/runner-gpu +# KOKKOS_ARCH must match the rented GPU (one arch per build): +# ADA89 = RTX 4090 / L4 / L40 AMPERE80 = A100 AMPERE86 = RTX 3090 / A10 HOPPER90 = H100 +# Estimate: ~45-60 min on 4 cores (log it in tools/build_estimates.jsonl, feature-class lammps-kokkos-cuda). + +# ---------- build stage ---------- +FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS build +ARG LAMMPS_TAG=stable_29Aug2024 +ARG KOKKOS_ARCH=ADA89 +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + ca-certificates git cmake g++ make python3 \ + && rm -rf /var/lib/apt/lists/* +RUN git clone --depth 1 --branch ${LAMMPS_TAG} https://github.com/lammps/lammps.git /src +RUN cmake -S /src/cmake -B /build \ + -D CMAKE_BUILD_TYPE=Release \ + -D CMAKE_CXX_COMPILER=/src/lib/kokkos/bin/nvcc_wrapper \ + -D PKG_KOKKOS=yes -D Kokkos_ENABLE_CUDA=yes -D Kokkos_ENABLE_OPENMP=yes \ + -D Kokkos_ARCH_${KOKKOS_ARCH}=yes \ + -D BUILD_OMP=yes -D PKG_OPENMP=yes \ + -D PKG_MANYBODY=yes -D PKG_MOLECULE=yes -D PKG_KSPACE=yes \ + -D PKG_REAXFF=yes -D PKG_QEQ=yes -D PKG_RIGID=yes \ + -D PKG_EXTRA-DUMP=yes -D PKG_MISC=yes -D PKG_EXTRA-FIX=yes \ + && cmake --build /build -j "$(nproc)" \ + && cmake --install /build --prefix /opt/lammps + +# ---------- runtime stage ---------- +FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04 +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + libgomp1 openssh-server rsync ca-certificates curl python3 \ + && rm -rf /var/lib/apt/lists/* \ + && mkdir -p /run/sshd /work /root/.ssh && chmod 700 /root/.ssh +COPY --from=build /opt/lammps/bin/lmp /usr/local/bin/lmp +COPY --from=build /src/potentials /usr/local/share/lammps/potentials +COPY start.sh runner.sh /usr/local/bin/ +RUN chmod +x /usr/local/bin/start.sh /usr/local/bin/runner.sh +ENV LAMMPS_POTENTIALS=/usr/local/share/lammps/potentials \ + OMP_NUM_THREADS=1 +WORKDIR /work +EXPOSE 22 +# Two modes: (a) MDE_JOB_ID set -> pull-runner (production, hosted/CONTRACT.md); (b) otherwise dev mode: +# PUBLIC_KEY env (RunPod convention) -> authorized_keys; then sshd in the foreground. +# Running `lmp` directly (no sshd) also works: docker run --gpus all ... mdengine-runner-gpu lmp -in deck.in -k on g 1 -sf kk +CMD ["/usr/local/bin/start.sh"] diff --git a/docker/runner-gpu/Dockerfile.prebuilt b/docker/runner-gpu/Dockerfile.prebuilt new file mode 100644 index 0000000..eb49c6b --- /dev/null +++ b/docker/runner-gpu/Dockerfile.prebuilt @@ -0,0 +1,22 @@ +# MDEngine GPU runner — FAST PATH: uses the prebuilt LAMMPS KOKKOS/CUDA tarball published as a GitHub +# release asset instead of compiling (~3 min vs ~60). Same runtime contract as Dockerfile. +# docker build -f Dockerfile.prebuilt --build-arg LMP_TARBALL_URL=... -t mdengine-runner-gpu:ada89 . +FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04 +ARG LMP_TARBALL_URL +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ + libgomp1 openssh-server rsync ca-certificates curl python3 \ + && rm -rf /var/lib/apt/lists/* \ + && mkdir -p /run/sshd /work /root/.ssh && chmod 700 /root/.ssh +RUN curl -fsSL "$LMP_TARBALL_URL" -o /tmp/lmp.tgz && tar -xzf /tmp/lmp.tgz -C /opt && rm /tmp/lmp.tgz \ + && ln -s /opt/lammps/bin/lmp /usr/local/bin/lmp \ + && mkdir -p /usr/local/share/lammps \ + && ( [ -d /opt/lammps/share/lammps/potentials ] && ln -s /opt/lammps/share/lammps/potentials /usr/local/share/lammps/potentials || true ) \ + && test -x /opt/lammps/bin/lmp +# NOTE: `lmp -h` cannot run here — libcuda.so.1 comes from the host driver at pod start, not the image. +COPY start.sh runner.sh /usr/local/bin/ +RUN chmod +x /usr/local/bin/start.sh /usr/local/bin/runner.sh +ENV LAMMPS_POTENTIALS=/usr/local/share/lammps/potentials \ + OMP_NUM_THREADS=1 +WORKDIR /work +EXPOSE 22 +CMD ["/usr/local/bin/start.sh"] diff --git a/docker/runner-gpu/README.md b/docker/runner-gpu/README.md new file mode 100644 index 0000000..5dbaa71 --- /dev/null +++ b/docker/runner-gpu/README.md @@ -0,0 +1,29 @@ +# MDEngine GPU runner image (hosted accelerated tier, GJOB-088) + +LAMMPS + KOKKOS/CUDA (REAXFF, QEQ, MANYBODY/EAM, MOLECULE, KSPACE, RIGID) with sshd, so a +rented GPU pod is driven by the existing `mdengine-mcp` remote transport (`submit_lammps host=`, +`fetch_job`) — the pod is just a `hosts.json` host. + +## Per-job flow (MVP, manual → scripted) +1. Start a pod from this image on the provider (RunPod/Vast/Lambda); pass `PUBLIC_KEY`. +2. Add it to `~/.mdengine/hosts.json` (`ssh`, `ssh_options: ["-p", port]`, `workdir: /work`, + `lmp: /usr/local/bin/lmp`, `launch` template with + `-k on g 1 -sf kk -pk kokkos newton on neigh half` — exact entry in the Dockerfile header). +3. `submit_lammps host=` → `job_status` → `fetch_job`. +4. Destroy the pod. GPU-seconds used = the metered quantity (step 4 of the build order). + +## Build +Needs an x86_64 host with nvcc (no GPU required to compile). Do NOT build on rakhsh +(arm64 + QEMU = many hours). Options: GitHub Actions on the public repo → ghcr.io, or a +$0.50 rented box. Pick `KOKKOS_ARCH` to match the GPU you will rent (header of Dockerfile). + +## KOKKOS notes +- ReaxFF and QEq have `/kk` variants (`pair reaxff/kk`, `fix qeq/reaxff/kk`) — `-sf kk` picks them. +- EAM (`eam/alloy/kk`) supported. SMTBQ has NO KOKKOS path — CPU-only forever. +- `OMP_NUM_THREADS=1` in the image: one GPU, one host thread; `-pk kokkos` overrides. +- `-k on g 1` = 1 GPU. Multi-GPU needs MPI (not in this image on purpose — one pod, one GPU, one job). + +## Security model +Decks are programs (`shell`). On a rented pod the sandbox is the pod: ephemeral, one job, +holds no credentials (the operator's public key only), destroyed after fetch. Where we control +the docker host, wrap `lmp` with the `../runner` hardening flags plus `--gpus all`. diff --git a/docker/runner-gpu/build-on-pod.sh b/docker/runner-gpu/build-on-pod.sh new file mode 100755 index 0000000..6da89ee --- /dev/null +++ b/docker/runner-gpu/build-on-pod.sh @@ -0,0 +1,28 @@ +#!/bin/bash +# Build LAMMPS+KOKKOS/CUDA directly on a RunPod pod (same recipe as the Dockerfile; used when +# we build on the rented box instead of in CI). Run as root on a runpod/pytorch *-devel image. +# KOKKOS_ARCH=ADA89 bash build-on-pod.sh -> /opt/lammps, tarball /workspace/lammps-kokkos-.tar.gz +set -euo pipefail +LAMMPS_TAG=${LAMMPS_TAG:-stable_29Aug2024} +KOKKOS_ARCH=${KOKKOS_ARCH:-ADA89} +export DEBIAN_FRONTEND=noninteractive +export PATH=/usr/local/cuda/bin:$PATH # ssh shells on runpod images lack the CUDA PATH +JOBS=${JOBS:-16} # nproc reports the HOST cores (64); the pod gets ~12 +apt-get update -qq && apt-get install -y -qq --no-install-recommends git cmake g++ make python3 rsync > /dev/null +[ -d /src ] || git clone --depth 1 --branch "$LAMMPS_TAG" https://github.com/lammps/lammps.git /src +cmake -S /src/cmake -B /build \ + -D CMAKE_BUILD_TYPE=Release \ + -D CMAKE_CXX_COMPILER=/src/lib/kokkos/bin/nvcc_wrapper \ + -D PKG_KOKKOS=yes -D Kokkos_ENABLE_CUDA=yes -D Kokkos_ENABLE_OPENMP=yes \ + -D Kokkos_ARCH_${KOKKOS_ARCH}=yes \ + -D BUILD_OMP=yes -D PKG_OPENMP=yes \ + -D PKG_MANYBODY=yes -D PKG_MOLECULE=yes -D PKG_KSPACE=yes \ + -D PKG_REAXFF=yes -D PKG_QEQ=yes -D PKG_RIGID=yes \ + -D PKG_EXTRA-DUMP=yes -D PKG_MISC=yes -D PKG_EXTRA-FIX=yes +cmake --build /build -j "$JOBS" +cmake --install /build --prefix /opt/lammps +mkdir -p /opt/lammps/share/lammps && cp -r /src/potentials /opt/lammps/share/lammps/ +/opt/lammps/bin/lmp -h | grep -E "KOKKOS|REAXFF|MANYBODY|OPENMP" | head +tar -C /opt -czf /workspace/lammps-kokkos-${KOKKOS_ARCH}.tar.gz lammps +ls -la /workspace/lammps-kokkos-${KOKKOS_ARCH}.tar.gz +echo BUILD-OK diff --git a/docker/runner-gpu/ci-build.yml b/docker/runner-gpu/ci-build.yml new file mode 100644 index 0000000..4166805 --- /dev/null +++ b/docker/runner-gpu/ci-build.yml @@ -0,0 +1,25 @@ +# Stage for .github/workflows/runner-gpu.yml in the PUBLIC repo (forcefieldsilicon/mdengine). +# Compiles LAMMPS+KOKKOS/CUDA on a plain x86 runner (nvcc needs no GPU) and pushes to ghcr.io. +# ~60-90 min on the 4-core runner; well under the 360-min cap. Manual trigger only. +name: runner-gpu +on: + workflow_dispatch: + inputs: + kokkos_arch: + description: 'Kokkos arch (ADA89 | AMPERE80 | AMPERE86 | HOPPER90)' + default: ADA89 +jobs: + build: + runs-on: ubuntu-latest + permissions: { contents: read, packages: write } + steps: + - uses: actions/checkout@v4 + - uses: docker/setup-buildx-action@v3 + - uses: docker/login-action@v3 + with: { registry: ghcr.io, username: ${{ github.actor }}, password: ${{ secrets.GITHUB_TOKEN }} } + - uses: docker/build-push-action@v6 + with: + context: docker/runner-gpu + push: true + build-args: KOKKOS_ARCH=${{ inputs.kokkos_arch }} + tags: ghcr.io/forcefieldsilicon/mdengine-runner-gpu:${{ inputs.kokkos_arch }} diff --git a/docker/runner-gpu/runner.sh b/docker/runner-gpu/runner.sh new file mode 100755 index 0000000..82d20c7 --- /dev/null +++ b/docker/runner-gpu/runner.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# MDEngine pull-runner (GJOB-093). Implements the pod side of hosted/CONTRACT.md v1. +# Env (set by the launcher on pod creation): MDE_ENDPOINT, MDE_JOB_ID, MDE_JOB_TOKEN. +# Sequence: GET job → download+untar input → run lmp under `timeout` with a heartbeat → tar results +# → PUT to presigned URL → POST done → exit. Never needs inbound network, ssh, or a public IP. +set -uo pipefail +: "${MDE_ENDPOINT:?}" "${MDE_JOB_ID:?}" "${MDE_JOB_TOKEN:?}" +LMP=${LMP:-/usr/local/bin/lmp} +export PATH=/usr/local/cuda/bin:$PATH +API="$MDE_ENDPOINT/internal/jobs/$MDE_JOB_ID" +AUTH="Authorization: Bearer $MDE_JOB_TOKEN" +WORK=${MDE_WORK:-/work}; mkdir -p "$WORK"; cd "$WORK" # MDE_WORK: dev override (mac test) +RES=${MDE_WORK:+$WORK/../results.tar.gz}; RES=${RES:-/results.tar.gz}; IN=${MDE_WORK:+$WORK/../input.tar.gz}; IN=${IN:-/input.tar.gz} +if command -v timeout >/dev/null; then TMO=timeout; elif command -v gtimeout >/dev/null; then TMO=gtimeout; else TMO=""; echo "WARN: no timeout(1); wall limit unenforced (dev only)" >&2; fi +T0=$(date +%s) +elapsed() { echo $(( $(date +%s) - T0 )); } +finish() { # exitcode error + local rc=$1 err=${2:-null} + [ "$err" != null ] && err="\"$err\"" + local bytes=0; [ -f "$RES" ] && bytes=$(wc -c < "$RES" | tr -d " ") + curl -sS -m 30 -X POST -H "$AUTH" -H 'Content-Type: application/json' \ + -d "{\"exitcode\":$rc,\"elapsed_s\":$(elapsed),\"results_bytes\":$bytes,\"error\":$err}" "$API/done" >/dev/null || true + exit "$rc" +} +# 1. job spec +SPEC=$(curl -sS -m 30 -H "$AUTH" "$API") || finish 70 fetch_spec +jq_() { printf '%s' "$SPEC" | python3 -c "import sys,json; d=json.load(sys.stdin); v=d$1; print(v if v is not None else '')"; } +INPUT_URL=$(jq_ "['input_url']"); PUT_URL=$(jq_ "['results_put_url']") +INPUT=$(jq_ "['input']"); WALL=$(jq_ "['wall_limit_s']"); LAUNCH=$(jq_ "['launch']") +[ -n "$INPUT_URL" ] && [ -n "$PUT_URL" ] && [ -n "$INPUT" ] || finish 70 bad_spec +# 2. input +curl -sS -m 600 -o "$IN" "$INPUT_URL" || finish 71 fetch_input +tar -tzf "$IN" | grep -Eq '(^|/)\.\.(/|$)|^/' && finish 72 unsafe_tarball +tar -xzf "$IN" -C "$WORK" || finish 72 untar +[ -f "$WORK/$INPUT" ] || finish 72 input_missing +# 2b. host diagnostics into the results (driver/GPU/CUDA visibility) — makes a bad host explainable +{ echo "== $(date -u +%FT%TZ) pod host diag"; nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader 2>&1 | head -3 + echo "cuda devices: $(ls /dev/nvidia* 2>/dev/null | tr '\n' ' ')"; echo "libcuda: $(ls /usr/lib/x86_64-linux-gnu/libcuda.so.* 2>/dev/null | head -1)" + echo "env: $(env | grep -E '^(NVIDIA|CUDA)' | tr '\n' ' ')"; } > "$WORK/hostdiag.txt" 2>&1 +# 3. run with heartbeat (last 20 thermo-ish lines of the log) +[ -z "$LAUNCH" ] || [ "$LAUNCH" = default ] && LAUNCH='{lmp} -in {input} -k on g 1 -sf kk -pk kokkos newton on neigh half -log log.lammps' +CMD=${LAUNCH//\{lmp\}/$LMP}; CMD=${CMD//\{input\}/$INPUT} +hb() { local tail; tail=$(tail -n 20 "$WORK/log.lammps" 2>/dev/null | python3 -c 'import sys,json; print(json.dumps(sys.stdin.read().splitlines()))') + curl -sS -m 10 -X POST -H "$AUTH" -H 'Content-Type: application/json' -d "{\"thermo_tail\":${tail:-[]},\"elapsed_s\":$(elapsed)}" "$API/heartbeat" >/dev/null || true; } +hb # immediate: state -> running before the first 30 s tick +( while sleep 30; do hb; done ) & HB=$! + +${TMO:+$TMO --signal=TERM --kill-after=30 "${WALL:-86400}"} bash -c "$CMD" > "$WORK/stdout.txt" 2>&1; RC=$? +kill $HB 2>/dev/null; wait $HB 2>/dev/null +echo "$RC" > "$WORK/exitcode" +ERR=null; [ $RC -eq 124 ] && ERR=wall_limit; { [ $RC -ne 0 ] && [ $RC -ne 124 ]; } && ERR=lammps_error +# 4. results (never ship the input tarball back; cap handled by the endpoint's presigned size limit) +tar -czf "$RES" -C "$WORK" . || finish 73 pack +curl -sS -m 1800 -X PUT -H 'Content-Type: application/gzip' --upload-file "$RES" "$PUT_URL" >/dev/null || finish 74 upload +finish "$RC" "$ERR" diff --git a/docker/runner-gpu/start.sh b/docker/runner-gpu/start.sh new file mode 100644 index 0000000..dd492d1 --- /dev/null +++ b/docker/runner-gpu/start.sh @@ -0,0 +1,16 @@ +#!/bin/sh +# Pod entrypoint: install the operator's public key, report the GPU, serve ssh. +set -eu +if [ -n "${MDE_JOB_ID:-}" ]; then + # production: pull one job, exit. Pod-side TTL (CONTRACT "Pod lifecycle" #3): even with the endpoint + # unreachable, this container ends at wall+600 s, so GPU billing is bounded without any outside help. + exec timeout --signal=TERM --kill-after=60 "$(( ${MDE_WALL_LIMIT_S:-86400} + 600 ))" /usr/local/bin/runner.sh +fi +if [ -n "${PUBLIC_KEY:-}" ]; then + printf '%s\n' "$PUBLIC_KEY" >> /root/.ssh/authorized_keys + chmod 600 /root/.ssh/authorized_keys +fi +ssh-keygen -A >/dev/null 2>&1 || true +nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo "no GPU visible" +lmp -h 2>/dev/null | grep -m1 -o "KOKKOS" || echo "WARNING: lmp lacks KOKKOS" +exec /usr/sbin/sshd -D -e -o PasswordAuthentication=no -o PermitRootLogin=prohibit-password diff --git a/hosted/CONTRACT.md b/hosted/CONTRACT.md new file mode 100644 index 0000000..9e04850 --- /dev/null +++ b/hosted/CONTRACT.md @@ -0,0 +1,111 @@ +# MDEngine hosted job contract v1 (GJOB-094; consumed by GJOB-093 runner and GJOB-091 clients) + +_Pinned 2026-09-05. Everything below is what the runner image, the endpoint, and the CLI/MCP/app +client agree on. Change it here first; code follows._ + +## Principles +- Same job model as the local runner (`~/.mdengine/jobs//`: `input`, `work/`, `log`, `exitcode`) + and the P0.5 remote transport — the hosted tier is a **transport**, not a new product. +- Pods **pull**. The endpoint never connects to a pod. A pod holds one job, one one-shot token, + nothing else of ours. A LAMMPS deck is untrusted code: the endpoint never executes anything. +- Files are exchanged via **presigned object-storage URLs**; the endpoint's VPS stores metadata only. + +## Identifiers +- Job id: `MDJOB--<6 base32>` (same shape as local jobs; sortable). +- API key: `mde_` + 32 hex. Sent as `Authorization: Bearer mde_…`. Stored hashed (sha256) at rest. +- Job token (pod-side): `jt_` + 32 hex, single job, expires when the job reaches a terminal state. + +## Client-facing API (base `https://api.forcefieldsilicon.com/v1`) +| Method | Path | Body / notes | Returns | +|---|---|---|---| +| GET | `/me` | — | `{balance_usd, rate_table, keys_created}` | +| POST | `/jobs` | JSON `JobSpec` (below) | `201 {id, upload_url, upload_expires}` — client PUTs the deck tarball to `upload_url` then calls start | +| POST | `/jobs/{id}/start` | — | `202 {id, state:"queued"}` | +| GET | `/jobs/{id}` | — | `JobStatus` (below) | +| GET | `/jobs/{id}/results` | — | `{download_url, expires, bytes}` (tarball of `work/` + `log` + `exitcode`) | +| DELETE | `/jobs/{id}` | — | `202` → state `cancelled`; billed to cancel time | +| GET | `/jobs?limit=&cursor=` | — | list, newest first | + +### JobSpec (client → endpoint) +```json +{ + "input": "in.lmp", // relative path inside the tarball; deck must be self-contained + "label": "Al slab oxidation", // optional, free text ≤120 chars + "gpu": "any", // "any" | "rtx4090" | "a100" — rate differs; "any" = cheapest available + "wall_limit_s": 14400, // hard cap, ≤ 86400; job fails at cap, billed to cap + "estimate_s": 3600, // client's guess; used only for the balance pre-check (≥15 min at rate) + "launch": "default" // "default" = "{lmp} -in {input} -k on g 1 -sf kk -pk kokkos newton on neigh half -log log.lammps" +} +``` +Tarball limits: ≤ 2 GB (matches the MCP 2 GB guard); paths must be relative, no `..`, no symlinks. + +### JobStatus (endpoint → client) +```json +{ + "id": "MDJOB-20260905-K3F9QZ", "state": "running", + "states": "created|uploaded|queued|launching|running|uploading|done|failed|cancelled", + "created": "…Z", "started": "…Z", "finished": null, + "gpu": "rtx4090", "rate_usd_per_h": 2.0, "billed_s": 812, "cost_usd": 0.45, + "thermo_tail": ["Step Temp PotEng …", " 1200 300.1 -20413.7 …"], // last ≤ 20 lines, from heartbeat + "exitcode": null, "error": null, // error: "wall_limit" | "pod_lost" | "lammps_error" | "cancelled" + "attempt": 1 // 2 after one automatic relaunch on pod loss +} +``` + +## Pod-facing API (runner → endpoint; auth = job token as `Authorization: Bearer jt_…`) +| Method | Path | Notes | +|---|---|---| +| GET | `/internal/jobs/{id}` | `{input_url (presigned GET), launch, wall_limit_s, results_put_url (presigned PUT)}` | +| POST | `/internal/jobs/{id}/heartbeat` | every 30 s: `{thermo_tail:[…], elapsed_s}`; 3 missed → `pod_lost` | +| POST | `/internal/jobs/{id}/done` | `{exitcode, elapsed_s, results_bytes}` → endpoint verifies the object exists, sets done/failed, **invalidates token**, terminates pod | + +Runner sequence: boot → GET job → download + untar to `/work` → `timeout wall_limit lmp …` (heartbeat +thread) → tar `work/ log.lammps exitcode` → PUT results → POST done → exit 0. Any failure → POST done +with nonzero exitcode and `error`. Pod never needs inbound network, ssh, or a public IP. + +## Purchase flow — instant credit, no human in the loop (GJOB-096, "the RunPod way", 2026-09-06) +Rule: a buyer is running within minutes of paying, like adding credits on RunPod. Nobody waits for an email. +- **First purchase** (no key yet): Stripe Payment Link → after payment Stripe REDIRECTS to + `GET /welcome?session_id={CHECKOUT_SESSION_ID}` → endpoint fetches the Checkout Session from Stripe + (secret key, server-side), verifies `payment_status=paid`, creates a key, credits `amount_total` at + the pack's GPU-hours, shows the key ONCE with the exact `mdengine login mde_…` line + app instructions. + Idempotent on `session_id` (revisit shows "already issued; check your email/CLI"). +- **Top-up** (has key): CLI `mdengine account --buy` / app "Buy credits" open the Payment Link with + `?client_reference_id=&prefilled_email=` → the same session handler credits THAT key. + The welcome page then says "credited to your existing key" and shows the new balance. +- **Webhook** `POST /v1/stripe/webhook` (`checkout.session.completed`, signature verified) runs the + same idempotent handler — source of truth if the buyer closes the tab before the redirect. +- Ledger table `credits(session_id PK, key_id, usd, gpu_s, created)`; key balance = Σcredits − Σbilled. +- Email of the key = fallback, not the path (no mail infra on day 1; Stripe's receipt goes out anyway). +- Pack → GPU-hours map lives in the endpoint config, keyed by Stripe price id (Starter 12.5 h, Lab 50 h, + Group 275 h). Coupon-discounted payments still grant the full pack hours (price id decides, not amount). + +## Pod lifecycle — no orphan ever bills (GJOB-099) +Invariant: **a pod exists only while a job is in `launching`/`running`.** The `/done → terminate` +path above is the happy path, not the guarantee. Three independent enforcers, any one sufficient: +1. **Endpoint reaper** (cron, every 5 min, idempotent): list every pod on the account; terminate any + whose `pod_id` maps to a terminal job, maps to no job, or has `age > wall_limit_s + 20 min`. + Retries with backoff; a pod that survives 3 reaper passes pages arvand. Also runs at endpoint boot, + so an endpoint outage cannot leave orphans behind it. +2. **`pod_lost` kills the lost pod** before relaunching — the relaunch never adds a second pod. +3. **Pod-side TTL**: `start.sh` runs the runner under `timeout $((wall_limit_s + 600))`; on runner + exit (any code) it POSTs `/done` if not already sent, then **stops its own container** — the pod is + dead-weight from then on and the reaper's job is only to clear the billing shell. Pods never hold an + account API key (a job token can't delete pods by design), so pod-side self-delete is not an option. +Why this section exists: RUN-022 (research, 2026-09-05/06) ran on a hand-launched pod whose only +terminator was a human; the human left for a day and $13 of credit sat one script away from zero. + +## Billing (GJOB-095) +`billed_s` runs from `running` to terminal state. Launch/pull overhead not billed. Rate by `gpu` +from the endpoint's rate table (re-derived after the A100 test, GJOB-092). Refund on `pod_lost` past +the retry, or any endpoint-side failure. Submit refused if `balance < rate × max(estimate_s, 900)`. + +## State machine +created → uploaded (client PUT ok) → queued (start) → launching (pod requested; DC fallback list) → +running (first heartbeat) → uploading (done received) → done | failed | cancelled. +Timeouts: launching > 10 min → next DC/GPU in fallback → after list exhausted `failed:no_capacity` +(not billed). running with 3 missed heartbeats → relaunch once (attempt 2) → then `failed:pod_lost`. + +## Local mock +`hosted/mock/` (to build): same routes, sqlite, files on local disk with `file://`-style URLs, so +runner (docker) and clients develop offline before the VPS exists. diff --git a/hosted/README.md b/hosted/README.md new file mode 100644 index 0000000..c10a5a3 --- /dev/null +++ b/hosted/README.md @@ -0,0 +1,32 @@ +# MDEngine hosted tier (GJOB-088) + +- `CONTRACT.md` — job contract v1: the one spec the runner, endpoint, and CLI/MCP/app clients share. +- `mock/mock_endpoint.py` — stdlib local mock of the endpoint (sqlite + local blob store). Dev only. +- Runner (pod side): `../docker/runner-gpu/runner.sh`, baked into the runner-gpu image; runs when + `MDE_JOB_ID` is set, otherwise the image is an sshd dev box. + +## Offline end-to-end (verified 2026-09-05 on rakhsh, CPU lmp standing in for the pod) +``` +python3 hosted/mock/mock_endpoint.py --port 8787 --data /tmp/mde-mock & # key mde_test, $20 +# client: POST /v1/jobs -> PUT deck tarball to upload_url -> POST /v1/jobs/{id}/start +# pod: set -a; . /tmp/mde-mock/launch.env; MDE_WORK=/tmp/pod/work LMP=/opt/homebrew/bin/lmp_serial docker/runner-gpu/runner.sh +# client: GET /v1/jobs/{id} -> done; GET /v1/jobs/{id}/results -> download_url +``` +## Clients (GJOB-091) — one `HostedClient` in LAMMPSCore, three surfaces +- CLI: `mdengine login mde_… [--endpoint URL]`, `mdengine run --gpu deck.in [--gpu-type any|rtx4090|a100] [--wall-hours H] [--no-wait]`, + `mdengine account`, `mdengine jobs`, `mdengine job [--fetch|--cancel|--wait]`. +- MCP: `submit_lammps host=cloud` (+ `gpu`, `wall_hours`); `job_status` / `job_log` / `job_files` / `cancel_job` / `fetch_job` all branch on the + local `job.json` carrying `"cloud": true`; `list_hosts` shows the balance. +- App: File ▸ Run Accelerated… (⇧⌘R) and File ▸ Accelerated Runs (⇧⌘J) window with live thermo; a finished job's trajectory opens in the + viewer by itself; Settings ▸ Accelerated holds the key. `MDEngine --run-accelerated deck.in` submits from the command line. +- Credentials: `$MDENGINE_API_KEY`, else `~/.mdengine/credentials` (0600). Dev: `$MDENGINE_HOSTED_URL` points at the mock, + `$MDENGINE_HOSTED_LAUNCH="{lmp} -in {input} -log log.lammps"` drops the KOKKOS flags for a CPU stand-in pod. +- Verified 2026-09-05 against the mock from all three surfaces (submit → pod → done → results fetched → balance debited). + +## Production endpoint (GJOB-094) — `endpoint/` +`endpoint/mde_endpoint.py` is the real service, evolved from the mock with the same routes and JSON shapes: hashed API +keys, a credit ledger (balance = credits - billed), `GET /v1/health`, the Stripe purchase flow (`/welcome` redirect + +signed webhook, idempotent on the Checkout Session), and job submission gated behind `MDE_RUNNERS_OPEN=1`. Operator CLI +`endpoint/mde_admin.py`, tests `endpoint/test_endpoint.py`, and `endpoint/deploy/` (Caddy TLS, systemd, Ubuntu 24.04 +bootstrap, deploy script). Ops notes in `endpoint/README.md`. Still to come there: RunPod launcher + reaper and +object-storage presigned URLs. The mock stays as the offline dev loop. diff --git a/hosted/endpoint/README.md b/hosted/endpoint/README.md new file mode 100644 index 0000000..b0c7aed --- /dev/null +++ b/hosted/endpoint/README.md @@ -0,0 +1,142 @@ +# MDEngine hosted endpoint — service + ops notes + +Production implementation of `../CONTRACT.md` v1, evolved from `../mock/mock_endpoint.py` (same routes and +JSON shapes, so the CLI/MCP/app clients verified against the mock work unchanged). Python 3.12 stdlib only, +sqlite, one process behind Caddy. + +``` +mde_endpoint.py the service mde_admin.py operator CLI (same sqlite) +mde_launcher.py RunPod pod launcher test_endpoint.py unittest, fake Stripe, fake launcher +deploy/ Caddyfile, systemd unit, bootstrap, deploy, env example +``` + +## Run locally +``` +cat > /tmp/mde.env <` with signed URLs. | +| `MDE_BIND` | listener, default `127.0.0.1:8080` (Caddy fronts it). | +| `MDE_PUBLIC_URL` | scheme+host clients and pods can reach; embedded in upload/download URLs. | +| `STRIPE_SECRET_KEY` | used server-side to fetch Checkout Sessions (`/welcome` and the webhook). | +| `STRIPE_WEBHOOK_SECRET` | signing secret of the `checkout.session.completed` webhook endpoint. | +| `MDE_PACKS` | `price_id:usd:gpu_hours,...` — the price id decides the hours credited. | +| `MDE_RATES` | `gpu:usd_per_hour,...`; default `any:2,rtx4090:2`. Credited usd = hours x rate("any"). | +| `MDE_RUNNERS_OPEN` | `1` opens job submission; anything else returns 503 `gpu_runners_open_soon`. | +| `MDE_ADMIN_TOKEN` | optional bearer for `GET /v1/admin/stats`; unset disables the route. | +| `RUNPOD_API_KEY` | RunPod account API key (pods read/write). Unset = no launcher: `start` writes `launch.env` for a hand-run pod. | +| `MDE_RUNNER_IMAGE` | pod image; default `ghcr.io/forcefieldsilicon/mdengine-runner-gpu:ADA89` (`docker/runner-gpu`). | +| `MDE_GPU_LADDER` | fallback rungs `CLOUD:gpu id,...` tried in order at launch; default `COMMUNITY:NVIDIA GeForce RTX 4090,SECURE:NVIDIA GeForce RTX 4090`. | +| `MDE_POD_DISK_GB` | pod container disk, default `20`. | +| `MDE_MIN_CUDA` | `gpu.minCudaVersion` on the pod request, default `12.4` (the ADA89 image is built against CUDA 12.4). | +| `MDE_LAUNCH_TIMEOUT_S` | a job still `launching` with no heartbeat after this many seconds -> pod deleted, `failed:no_capacity` (unbilled). Default `600`. | +| `MDE_REAPER_INTERVAL_S` | how often the reaper lists the account's pods, default `300`. It also runs once at boot. | + +Logs are one JSON object per line on stdout (`journalctl -u mde-endpoint -f`). Full API keys, job tokens and +Stripe secrets never appear in logs or in the database (keys are stored as sha256; `key_id` = first 8 hex). + +## Routes +Client (Bearer `mde_...`): `GET /v1/me`, `POST /v1/jobs`, `POST /v1/jobs/{id}/start`, `GET /v1/jobs[?limit=&cursor=]`, +`GET /v1/jobs/{id}`, `GET /v1/jobs/{id}/results`, `DELETE /v1/jobs/{id}` — shapes per `../CONTRACT.md`. +Pod (Bearer `jt_...`, minted at `start`, invalidated at any terminal state): `GET /internal/jobs/{id}`, +`POST /internal/jobs/{id}/heartbeat`, `POST /internal/jobs/{id}/done`. +Blobs: `PUT|GET /blob/.in.tar.gz|.out.tar.gz?exp=&sig=` (HMAC-signed, stands in for presigned object storage; 2 GB cap, streamed to disk). +Public: `GET /v1/health` -> `{"ok":true,"version":"...","runners":"open|closed","launcher":"runpod|none"}`; +`GET /welcome?session_id=`; `POST /v1/stripe/webhook`. + +### Purchase flow +1. Stripe Payment Link redirects to `/welcome?session_id={CHECKOUT_SESSION_ID}`. The endpoint fetches the session + (`expand[]=line_items`), requires `payment_status=paid`, maps the price id through `MDE_PACKS`, and: + - `client_reference_id` equal to an existing `key_id` -> credits that key, page shows the new balance; + - otherwise creates a key, credits it, shows the full key once with the `mdengine login mde_...` line and app steps. +2. The webhook runs the same handler. Both are idempotent on `session_id` (primary key of `credits`). +3. If the webhook lands before the buyer's browser does, the new key is parked in `pending_keys` and revealed on the + first `/welcome` visit, then deleted. Unclaimed entries are purged after 7 days; `mde-admin pending` lists them and + `mde-admin pending --reveal ` prints one for manual delivery (the only path where a plaintext key rests in + the database, and only until it is shown). +4. Coupon-discounted payments grant the full pack hours. An unknown price id is honored at `amount_total` / rate and + logged as `pack.unknown_price`. + +Balance = sum(`credits.usd`) - sum(billed job cost). Jobs with `error` in `pod_lost`/`no_capacity` cost nothing; +cancelled jobs bill to cancel time; running jobs bill live. + +### Pod launcher (`mde_launcher.py`, RunPod REST v2) +With `RUNPOD_API_KEY` set, `POST /v1/jobs/{id}/start` mints the job token, sets `queued`, and hands off to a background +thread so the HTTP response never waits on RunPod: +1. `RunPodLauncher.create` walks `MDE_GPU_LADDER`: one `POST /v2/pods` per rung with + `{name:"mde-", image, cloud, gpu:{id,count:1,minCudaVersion}, disk, env:{MDE_ENDPOINT,MDE_JOB_ID,MDE_JOB_TOKEN}}` + (no `dataCenterIds`: the scheduler picks). Each rung logs `launch.attempt` (status code, error body truncated to 300 + chars; never the request body, which carries the token). First 201 wins: job -> `launching`, `pod_id`, `launched_at`. + Every rung refused -> `job.no_capacity`, job `failed:no_capacity`, token invalidated, nothing billed, no pod. +2. The pod boots `start.sh` -> `runner.sh`, whose first heartbeat flips the job to `running`. +3. Any terminal write (`done`, `failed` incl. `pod_lost`/`no_results`, `cancelled`, launch timeout) is followed by + `DELETE /v2/pods/{pod_id}` in a background thread (`pod.deleted` / `pod.delete_failed`). 204 and 404 both count as + deleted; 429/5xx retry 3x with backoff. A cancel that lands while the create call is in flight deletes the pod the + moment the create returns. +4. Watchdog (every 30 s): `running` with no heartbeat for 120 s -> `failed:pod_lost` + pod deleted; + `launching` past `MDE_LAUNCH_TIMEOUT_S` with no heartbeat -> `job.launch_timeout`, `failed:no_capacity` + pod deleted. +5. Reaper (`reaper_once`, every `MDE_REAPER_INTERVAL_S` and once at boot; CONTRACT "Pod lifecycle", GJOB-099): lists + every pod on the account and deletes any `mde-*` pod whose job is terminal, missing (`pod_id` and name suffix both + unknown), or whose age exceeds the job's `wall_limit_s` + 20 min (`reaper.deleted` with `reason`). Pods not named + `mde-*` are never touched. Consecutive delete failures per pod are counted in memory; the third logs `reaper.stuck` + (page on that). Job rows keep `pod_id` after deletion for the audit trail; `GET /v1/jobs/{id}` echoes it. + +Job tokens end up in the pod's env on RunPod (that is how the runner authenticates), so anyone with the RunPod account +can read them while the pod exists; the endpoint invalidates the token at every terminal state, and a job token cannot +touch pods or keys. The RunPod API key itself lives only in the env file and the `Authorization` header. + +## Deploy (Ubuntu 24.04, Caddy TLS, systemd) +``` +deploy/deploy.sh root@HOST # rsync + bootstrap.sh + restart + https health check +ssh root@HOST 'vi /etc/mde/endpoint.env && systemctl restart mde-endpoint' # first time: fill the env file +ssh root@HOST mde-admin stats +``` +`bootstrap.sh` is idempotent: apt update, unattended-upgrades, ufw 22/80/443, user `mde`, Caddy from its apt repo, +`/opt/mde` (code), `/var/lib/mde` (state, owned by `mde`), `/etc/mde` (env, root:mde 0640), both services enabled. +It copies no secrets; an example env file is placed only if none exists. DNS for the hostname in `deploy/Caddyfile` +must point at the box before Caddy can obtain its certificate. The unit runs with `ProtectSystem=strict`; the only +writable path is `/var/lib/mde`. + +Back up `/var/lib/mde/mde.sqlite` (it is the ledger): `sqlite3 /var/lib/mde/mde.sqlite ".backup /root/mde-$(date +%F).sqlite"` +or continuous replication with litestream. + +## Operations +- **Open the runners**: set `MDE_RUNNERS_OPEN=1` in the env file, `systemctl restart mde-endpoint`, confirm + `curl https://HOST/v1/health` says `"runners":"open"`. Setting it back to `0` closes submission without touching + balances (403/503 happen after auth and balance checks, so `mdengine account` keeps working). +- **Rotate the webhook secret**: in Stripe, add a second webhook endpoint (or roll the secret) for + `checkout.session.completed` -> `https://HOST/v1/stripe/webhook`; put the new `whsec_` in the env file; restart; + send a test event from the Stripe dashboard and check `journalctl` for `credit.*`/`webhook.bad_signature`; delete the old + endpoint. A bad signature returns 400 and Stripe retries, so a short overlap loses nothing. +- **Rotate the Stripe secret key**: replace `STRIPE_SECRET_KEY`, restart. Only session reads are needed + (a restricted key with `checkout_sessions: read` works). +- **Issue a key by hand**: `mde-admin key new --email E --credit 25 [--label L]` (prints the key once); + top up: `mde-admin credit add --key KEYID --usd 50`; inspect: `mde-admin key list`, `mde-admin ledger [--key KEYID]`, `mde-admin stats`. +- **Replace a lost key**: create a new key with `key new --credit 0`, then `credit add` the old balance and note the + old key id in `--label`. Old keys cannot be recovered from their hash. +- **Pod stand-in without a launcher**: with `RUNPOD_API_KEY` unset, `start` writes `/var/lib/mde/launch.env` + (`MDE_ENDPOINT`, `MDE_JOB_ID`, `MDE_JOB_TOKEN`) for a hand-run `docker/runner-gpu/runner.sh` (dev loop). +- **Rotate the RunPod key**: create the new key in RunPod (pods read/write), replace `RUNPOD_API_KEY`, restart, check + `journalctl` for `reaper.boot` (a successful list) rather than `reaper.list_failed`; then revoke the old key. +- **Orphan check by hand**: `journalctl -u mde-endpoint | grep -E 'reaper\.(deleted|stuck|list_failed)|pod\.delete_failed'`. + `reaper.stuck` means three passes failed to delete one pod: delete it in the RunPod console and look at the error text. + +## Not in this skeleton (tracked in `../CONTRACT.md`) +Presigned object-storage (R2) URLs — blobs live on the box behind signed URLs; automatic relaunch on `pod_lost` +(attempt 2) — the watchdog marks the job `failed:pod_lost`, deletes the pod, and does not bill it. The GPU/cloud +fallback ladder is walked at create time (a rung that refuses moves to the next one immediately); a pod that is +accepted but never heartbeats is not moved to the next rung, it times out to `failed:no_capacity`. diff --git a/hosted/endpoint/deploy/Caddyfile b/hosted/endpoint/deploy/Caddyfile new file mode 100644 index 0000000..444d449 --- /dev/null +++ b/hosted/endpoint/deploy/Caddyfile @@ -0,0 +1,29 @@ +# MDEngine hosted endpoint — TLS terminator in front of the Python service. +# Caddy obtains and renews the Let's Encrypt certificate by itself; the DNS A record +# for the hostname must already point at this box before the first start. +api.forcefieldsilicon.com { + encode zstd gzip + # Deck tarballs go through PUT /blob/...; allow the contract's 2 GB cap. + request_body { + max_size 2GB + } + reverse_proxy 127.0.0.1:8080 { + # Long PUT/GET of multi-hundred-MB tarballs on a slow link. + transport http { + response_header_timeout 30m + read_timeout 30m + write_timeout 30m + } + } + header { + Strict-Transport-Security "max-age=31536000" + X-Content-Type-Options "nosniff" + -Server + } + log { + output file /var/log/caddy/api.access.log { + roll_size 50mb + roll_keep 5 + } + } +} diff --git a/hosted/endpoint/deploy/bootstrap.sh b/hosted/endpoint/deploy/bootstrap.sh new file mode 100755 index 0000000..1f49b8d --- /dev/null +++ b/hosted/endpoint/deploy/bootstrap.sh @@ -0,0 +1,66 @@ +#!/usr/bin/env bash +# Bootstrap a fresh Ubuntu 24.04 host for the MDEngine hosted endpoint. Idempotent; run as root. +# Takes NO secrets: /etc/mde/endpoint.env is written separately (see ../README.md). +# Expects mde_endpoint.py, mde_admin.py and deploy/ next to this script (deploy.sh rsyncs them). +set -euo pipefail +[ "$(id -u)" = 0 ] || { echo "run as root" >&2; exit 1; } +HERE=$(cd "$(dirname "$0")" && pwd); SRC=$(dirname "$HERE") +export DEBIAN_FRONTEND=noninteractive + +echo "== packages" +apt-get update -q +apt-get install -y -q python3 sqlite3 ufw unattended-upgrades debian-keyring debian-archive-keyring apt-transport-https curl gnupg rsync +dpkg-reconfigure -f noninteractive unattended-upgrades + +echo "== firewall" +ufw --force default deny incoming >/dev/null +ufw --force default allow outgoing >/dev/null +ufw allow 22/tcp >/dev/null; ufw allow 80/tcp >/dev/null; ufw allow 443/tcp >/dev/null +ufw --force enable >/dev/null +ufw status | head -5 + +echo "== caddy (official apt repo)" +if ! command -v caddy >/dev/null; then + curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/gpg.key' | gpg --dearmor -o /usr/share/keyrings/caddy-stable-archive-keyring.gpg --yes + curl -1sLf 'https://dl.cloudsmith.io/public/caddy/stable/debian.deb.txt' > /etc/apt/sources.list.d/caddy-stable.list + apt-get update -q && apt-get install -y -q caddy +fi + +echo "== user + dirs" +id mde >/dev/null 2>&1 || useradd --system --home /var/lib/mde --shell /usr/sbin/nologin mde +install -d -o root -g root -m 755 /opt/mde +install -d -o mde -g mde -m 750 /var/lib/mde /var/lib/mde/blobs +install -d -o root -g mde -m 750 /etc/mde +install -d -o caddy -g caddy -m 755 /var/log/caddy 2>/dev/null || true + +echo "== files" +install -o root -g root -m 755 "$SRC/mde_endpoint.py" /opt/mde/mde_endpoint.py +install -o root -g root -m 755 "$SRC/mde_admin.py" /opt/mde/mde_admin.py +install -o root -g root -m 644 "$SRC/mde_launcher.py" /opt/mde/mde_launcher.py +ln -sf /opt/mde/mde_admin.py /usr/local/bin/mde-admin +install -o root -g root -m 644 "$HERE/mde-endpoint.service" /etc/systemd/system/mde-endpoint.service +install -o root -g root -m 755 "$HERE/mde-backup.sh" /opt/mde/mde-backup.sh +install -o root -g root -m 644 "$HERE/mde-backup.service" /etc/systemd/system/mde-backup.service +install -o root -g root -m 644 "$HERE/mde-backup.timer" /etc/systemd/system/mde-backup.timer +install -o root -g root -m 644 "$HERE/Caddyfile" /etc/caddy/Caddyfile +install -d -m 755 /etc/systemd/system/caddy.service.d +install -o root -g root -m 644 "$HERE/caddy-override.conf" /etc/systemd/system/caddy.service.d/override.conf +if [ ! -f /etc/mde/endpoint.env ]; then + install -o root -g mde -m 640 "$HERE/endpoint.env.example" /etc/mde/endpoint.env + echo "!! /etc/mde/endpoint.env is the EXAMPLE — fill in the real values, then: systemctl restart mde-endpoint" +fi +python3 -m py_compile /opt/mde/mde_endpoint.py /opt/mde/mde_admin.py /opt/mde/mde_launcher.py + +echo "== services" +systemctl daemon-reload +systemctl enable --now caddy >/dev/null +systemctl enable mde-endpoint >/dev/null +systemctl enable --now mde-backup.timer >/dev/null +systemctl restart mde-endpoint +caddy validate --config /etc/caddy/Caddyfile >/dev/null +chown -R caddy:caddy /var/log/caddy # validate (run as root) may have created the access log root-owned +systemctl restart caddy +sleep 1 +systemctl is-active mde-endpoint caddy +curl -fsS http://127.0.0.1:8080/v1/health && echo +echo "== done" diff --git a/hosted/endpoint/deploy/caddy-override.conf b/hosted/endpoint/deploy/caddy-override.conf new file mode 100644 index 0000000..173810d --- /dev/null +++ b/hosted/endpoint/deploy/caddy-override.conf @@ -0,0 +1,5 @@ +# systemd drop-in for the distro caddy.service: its sandbox denies /var/log/caddy even when the +# directory is caddy-owned; declare it so the access log in Caddyfile can be written. +[Service] +LogsDirectory=caddy +ReadWritePaths=/var/log/caddy diff --git a/hosted/endpoint/deploy/deploy.sh b/hosted/endpoint/deploy/deploy.sh new file mode 100755 index 0000000..ba70228 --- /dev/null +++ b/hosted/endpoint/deploy/deploy.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Deploy the endpoint from a workstation: deploy/deploy.sh root@HOST [PUBLIC_HOSTNAME] +# rsyncs mde_endpoint.py + mde_admin.py + deploy/ to the host, runs bootstrap.sh (idempotent), +# restarts the service, and checks https://PUBLIC_HOSTNAME/v1/health. Copies no secrets. +set -euo pipefail +TARGET=${1:?usage: deploy.sh root@HOST [PUBLIC_HOSTNAME]} +PUBLIC=${2:-api.forcefieldsilicon.com} +HERE=$(cd "$(dirname "$0")" && pwd); SRC=$(dirname "$HERE") +STAGE=/root/mde-deploy + +python3 -m py_compile "$SRC/mde_endpoint.py" "$SRC/mde_admin.py" "$SRC/mde_launcher.py" +( cd "$SRC" && python3 test_endpoint.py -q ) || { echo "tests failed; not deploying" >&2; exit 1; } + +rsync -az --delete --exclude '__pycache__' --exclude '*.pyc' \ + "$SRC/mde_endpoint.py" "$SRC/mde_admin.py" "$SRC/mde_launcher.py" "$SRC/deploy" "$TARGET:$STAGE/" +ssh "$TARGET" "bash $STAGE/deploy/bootstrap.sh && systemctl restart mde-endpoint && sleep 1 && journalctl -u mde-endpoint -n 3 --no-pager" + +echo "== public health" +curl -fsS --max-time 15 "https://$PUBLIC/v1/health" && echo || { + echo "public health check failed: DNS for $PUBLIC not pointing here yet, or Caddy still fetching its certificate (journalctl -u caddy)" >&2; exit 2; } diff --git a/hosted/endpoint/deploy/endpoint.env.example b/hosted/endpoint/deploy/endpoint.env.example new file mode 100644 index 0000000..e05c608 --- /dev/null +++ b/hosted/endpoint/deploy/endpoint.env.example @@ -0,0 +1,37 @@ +# /etc/mde/endpoint.env — read by systemd (EnvironmentFile) and by mde_admin.py. +# chmod 0640, owner root:mde. Never commit a filled-in copy. + +# Storage (inside /var/lib/mde, which the service unit can write). +MDE_DB=/var/lib/mde/mde.sqlite +MDE_BLOBS=/var/lib/mde/blobs + +# Listener (Caddy proxies to it) and the URL clients/pods can reach (used in blob URLs). +MDE_BIND=127.0.0.1:8080 +MDE_PUBLIC_URL=https://api.forcefieldsilicon.com + +# Stripe (Dashboard ▸ Developers). Restricted key with checkout_sessions:read is enough for the secret key. +STRIPE_SECRET_KEY=sk_live_... +STRIPE_WEBHOOK_SECRET=whsec_... + +# Packs: ::, comma separated. The price id decides the hours. +MDE_PACKS=price_xxx:25:12.5,price_yyy:100:50,price_zzz:500:275 + +# Rate table, $/GPU-hour by gpu name. Default when unset: any:2,rtx4090:2 +MDE_RATES=any:2,rtx4090:2 + +# Job submission is refused with 503 gpu_runners_open_soon until this is exactly 1. +MDE_RUNNERS_OPEN=0 + +# Optional: bearer token for GET /v1/admin/stats. Leave empty to disable the route. +MDE_ADMIN_TOKEN= + +# RunPod pod launcher (hosted/endpoint/mde_launcher.py). Leave RUNPOD_API_KEY empty to run without a launcher +# (start writes /var/lib/mde/launch.env for a hand-run pod). Key needs pods read/write; never commit a real one. +RUNPOD_API_KEY= +# Defaults shown; uncomment to override. +#MDE_RUNNER_IMAGE=ghcr.io/forcefieldsilicon/mdengine-runner-gpu:ADA89 +#MDE_GPU_LADDER=COMMUNITY:NVIDIA GeForce RTX 4090,SECURE:NVIDIA GeForce RTX 4090 +#MDE_POD_DISK_GB=20 +#MDE_MIN_CUDA=12.4 +#MDE_LAUNCH_TIMEOUT_S=600 +#MDE_REAPER_INTERVAL_S=300 diff --git a/hosted/endpoint/deploy/mde-backup.service b/hosted/endpoint/deploy/mde-backup.service new file mode 100644 index 0000000..d80e501 --- /dev/null +++ b/hosted/endpoint/deploy/mde-backup.service @@ -0,0 +1,7 @@ +[Unit] +Description=MDEngine ledger sqlite backup +[Service] +Type=oneshot +User=mde +EnvironmentFile=/etc/mde/endpoint.env +ExecStart=/opt/mde/mde-backup.sh diff --git a/hosted/endpoint/deploy/mde-backup.sh b/hosted/endpoint/deploy/mde-backup.sh new file mode 100644 index 0000000..e1cd938 --- /dev/null +++ b/hosted/endpoint/deploy/mde-backup.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# Hourly consistent snapshot of the ledger db (sqlite online backup), 72 h rotation + daily keep 30. +# Installed by bootstrap.sh as /opt/mde/mde-backup.sh with a systemd timer. Off-box copy is pulled +# by the operator workstation (see README "Backups"); Hetzner server snapshots are the second layer. +set -euo pipefail +DB=${MDE_DB:-/var/lib/mde/mde.sqlite}; OUT=/var/lib/mde/backups; mkdir -p "$OUT/hourly" "$OUT/daily" +ts=$(date -u +%Y%m%dT%H%MZ); day=$(date -u +%Y%m%d) +sqlite3 "$DB" ".backup '$OUT/hourly/mde-$ts.sqlite'" +gzip -f "$OUT/hourly/mde-$ts.sqlite" +[ -e "$OUT/daily/mde-$day.sqlite.gz" ] || cp "$OUT/hourly/mde-$ts.sqlite.gz" "$OUT/daily/mde-$day.sqlite.gz" +ls -1t "$OUT/hourly"/*.gz 2>/dev/null | tail -n +73 | xargs -r rm -f +ls -1t "$OUT/daily"/*.gz 2>/dev/null | tail -n +31 | xargs -r rm -f +sha256sum "$OUT/hourly/mde-$ts.sqlite.gz" | cut -c1-16 diff --git a/hosted/endpoint/deploy/mde-backup.timer b/hosted/endpoint/deploy/mde-backup.timer new file mode 100644 index 0000000..70d346b --- /dev/null +++ b/hosted/endpoint/deploy/mde-backup.timer @@ -0,0 +1,8 @@ +[Unit] +Description=Hourly MDEngine ledger backup +[Timer] +OnCalendar=hourly +RandomizedDelaySec=120 +Persistent=true +[Install] +WantedBy=timers.target diff --git a/hosted/endpoint/deploy/mde-endpoint.service b/hosted/endpoint/deploy/mde-endpoint.service new file mode 100644 index 0000000..584ebb7 --- /dev/null +++ b/hosted/endpoint/deploy/mde-endpoint.service @@ -0,0 +1,27 @@ +[Unit] +Description=MDEngine hosted endpoint (api) +After=network-online.target +Wants=network-online.target + +[Service] +Type=simple +User=mde +Group=mde +EnvironmentFile=/etc/mde/endpoint.env +WorkingDirectory=/opt/mde +ExecStart=/usr/bin/python3 /opt/mde/mde_endpoint.py +Restart=always +RestartSec=2 +# Hardening: the service only needs its state dir and loopback. +NoNewPrivileges=yes +ProtectSystem=strict +ProtectHome=yes +PrivateTmp=yes +ReadWritePaths=/var/lib/mde +UMask=0077 +StandardOutput=journal +StandardError=journal +SyslogIdentifier=mde-endpoint + +[Install] +WantedBy=multi-user.target diff --git a/hosted/endpoint/mde_admin.py b/hosted/endpoint/mde_admin.py new file mode 100644 index 0000000..770cd61 --- /dev/null +++ b/hosted/endpoint/mde_admin.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +"""Operator CLI for the MDEngine hosted endpoint. Same sqlite file as mde_endpoint.py; stdlib only. + + mde_admin.py key new --email E --credit USD [--label L] # prints the full key ONCE + mde_admin.py key list + mde_admin.py credit add --key KEYID --usd X [--note TEXT] + mde_admin.py ledger [--key KEYID] + mde_admin.py stats + mde_admin.py pending [--reveal SESSION_ID] # keys bought via webhook, not yet shown + +DB path: --db, else $MDE_DB, else the value in --env-file / /etc/mde/endpoint.env. +""" +import argparse, os, secrets, sys +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import mde_endpoint as E + +OPEN = [] + +def open_db(a): + E.load_env_file(a.env_file) + path = a.db or os.environ.get("MDE_DB") + if not path: sys.exit("no db: pass --db or set MDE_DB (or --env-file)") + db = E.DB(path); OPEN.append(db) + return db, E.Config() + +def rate(cfg): return cfg.rates.get("any", 2.0) + +def cmd_key_new(a): + db, cfg = open_db(a) + full, kid = db.create_key(email=a.email, label=a.label) + if a.credit: + db.add_credit("admin-" + secrets.token_hex(6), kid, a.credit, a.credit / rate(cfg) * 3600, "admin") + print(f"key_id {kid}\nemail {a.email}\nbalance ${db.balance(kid):.2f}\n\n{full}\n\n(shown once; only the hash is stored) mdengine login {full}") + E.log("admin.key_new", key_id=kid, usd=a.credit, has_email=bool(a.email)) + +def cmd_key_list(a): + db, _ = open_db(a) + print(f"{'key_id':<10}{'balance':>10} {'created':<21}{'email':<32}label") + for k in db.q("select * from keys order by created"): + print(f"{k['key_id']:<10}{db.balance(k['key_id']):>10.2f} {k['created']:<21}{(k['email'] or '-'):<32}{k['label'] or ''}") + +def cmd_credit_add(a): + db, cfg = open_db(a) + if not db.key_by_id(a.key): sys.exit(f"no key {a.key}") + sid = "admin-" + secrets.token_hex(6) + db.add_credit(sid, a.key, a.usd, a.usd / rate(cfg) * 3600, a.note or "admin") + print(f"{sid}: +${a.usd:.2f} -> key {a.key} balance ${db.balance(a.key):.2f}") + E.log("admin.credit_add", key_id=a.key, usd=a.usd, session=sid) + +def cmd_ledger(a): + db, _ = open_db(a) + where, args = ("where key_id=?", (a.key,)) if a.key else ("", ()) + print("-- credits") + for c in db.q(f"select * from credits {where} order by created", *args): + print(f"{c['created']} {c['key_id']} +${c['usd']:>8.2f} {c['gpu_s']/3600:>7.2f} h {c['price_id'] or '-':<32} {c['session_id']}") + print("-- billed jobs") + for j in db.q(f"select * from jobs {where} order by created", *args): + if j["state"] in ("created", "uploaded", "queued", "launching"): continue + print(f"{j['created']} {j['key_id']} -${E.job_cost(j):>8.4f} {E.billed_seconds(j):>7d} s {j['state']:<10} {j['error'] or '':<12} {j['id']}") + +def cmd_stats(a): + db, cfg = open_db(a) + keys = db.one("select count(*) n from keys")["n"] + credits = db.one("select coalesce(sum(usd),0) s, count(*) n from credits") + billed = sum(E.job_cost(j) for j in db.q("select * from jobs where state not in ('created','uploaded','queued','launching')")) + print(f"keys {keys}\ncredit rows {credits['n']}\ncredited usd {credits['s']:.2f}\nbilled usd {billed:.4f}\noutstanding {credits['s']-billed:.2f}") + for r in db.q("select state, count(*) n from jobs group by state order by state"): print(f"jobs {r['state']:<10} {r['n']}") + print(f"pending keys {db.one('select count(*) n from pending_keys')['n']}\nrunners {'open' if cfg.runners_open else 'closed'}\nrates {cfg.rates}\npacks {len(cfg.packs)}") + +def cmd_pending(a): + db, _ = open_db(a) + if a.reveal: + p = db.one("select * from pending_keys where session_id=?", a.reveal) + if not p: sys.exit("no pending key for that session") + k = db.key_by_id(p["key_id"]) + print(f"key_id {p['key_id']} email {k['email'] if k else '-'}\n\n{p['full_key']}\n") + if not a.keep: db.x("delete from pending_keys where session_id=?", a.reveal); print("(removed from pending; deliver it now)") + E.log("admin.pending_reveal", key_id=p["key_id"], session=a.reveal) + return + rows = db.q("select p.*, k.email from pending_keys p left join keys k on k.key_id=p.key_id order by p.created") + if not rows: print("no pending keys"); return + for p in rows: print(f"{p['created']} {p['key_id']} {p['email'] or '-':<32} {p['session_id']}") + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--db"); ap.add_argument("--env-file", default=os.environ.get("MDE_ENV_FILE", "/etc/mde/endpoint.env")) + sub = ap.add_subparsers(dest="cmd", required=True) + key = sub.add_parser("key").add_subparsers(dest="sub", required=True) + n = key.add_parser("new"); n.add_argument("--email", required=True); n.add_argument("--credit", type=float, default=0.0); n.add_argument("--label"); n.set_defaults(f=cmd_key_new) + key.add_parser("list").set_defaults(f=cmd_key_list) + cr = sub.add_parser("credit").add_subparsers(dest="sub", required=True) + c = cr.add_parser("add"); c.add_argument("--key", required=True); c.add_argument("--usd", type=float, required=True); c.add_argument("--note"); c.set_defaults(f=cmd_credit_add) + l = sub.add_parser("ledger"); l.add_argument("--key"); l.set_defaults(f=cmd_ledger) + sub.add_parser("stats").set_defaults(f=cmd_stats) + p = sub.add_parser("pending"); p.add_argument("--reveal", metavar="SESSION_ID"); p.add_argument("--keep", action="store_true"); p.set_defaults(f=cmd_pending) + a = ap.parse_args(argv) + try: a.f(a) + finally: + while OPEN: OPEN.pop().c.close() + +if __name__ == "__main__": + main() diff --git a/hosted/endpoint/mde_endpoint.py b/hosted/endpoint/mde_endpoint.py new file mode 100644 index 0000000..52edac4 --- /dev/null +++ b/hosted/endpoint/mde_endpoint.py @@ -0,0 +1,665 @@ +#!/usr/bin/env python3 +"""MDEngine hosted endpoint (hosted/CONTRACT.md v1). Python 3.12 stdlib only. + +Evolved from hosted/mock/mock_endpoint.py: same routes and JSON shapes for the client-facing +(/v1/*), blob (/blob/*), and pod-facing (/internal/*) APIs, plus: + + * API keys stored as sha256 hashes only (table `keys`); balance derived from a credit ledger + (table `credits`) minus billed job cost -- there is no mutable balance column. + * GET /v1/health liveness + whether GPU runners are open + * GET /welcome?session_id=... Stripe Checkout redirect target: issue/credit a key + * POST /v1/stripe/webhook checkout.session.completed backstop (same handler) + * Job submission gated by MDE_RUNNERS_OPEN=1 (auth + balance still checked while closed). + * RunPod pod launcher (mde_launcher.py) when RUNPOD_API_KEY is set: one pod per job, deleted at every + terminal state; watchdog handles launch timeout / pod_lost; reaper enforces "a pod exists only while + a job is launching/running" (CONTRACT.md, GJOB-099). Without the key `start` writes launch.env (dev). + +Runs behind Caddy (TLS) on 127.0.0.1:8080 under systemd; see deploy/ and README.md. + + python3 mde_endpoint.py --env-file /etc/mde/endpoint.env +""" +import argparse, hashlib, hmac, json, os, secrets, sqlite3, sys, threading, time, urllib.error, urllib.parse, urllib.request +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlparse, parse_qs + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from mde_launcher import RunPodLauncher, FakeLauncher, NoCapacity, LauncherError, parse_ladder, pod_age_s + +VERSION = "0.1.0" +STATES = "created uploaded queued launching running uploading done failed cancelled".split() +TERMINAL = ("done", "failed", "cancelled") +NOT_BILLED_ERRORS = ("pod_lost", "no_capacity") +MAX_BLOB = 2_000_000_000 # 2 GB tarball cap (CONTRACT.md) +BLOB_TTL_S = 7 * 86400 # signed blob URLs stay valid for a week (results linger) +HEARTBEAT_LOST_S = 120 # 30 s heartbeat, 3 missed -> pod_lost +PENDING_KEY_TTL_S = 7 * 86400 # unshown keys from webhook-first purchases are purged after this +STRIPE_TOLERANCE_S = 300 +REAPER_GRACE_S = 1200 # reaper kills a pod older than its job's wall_limit_s + this (CONTRACT: +20 min) +POD_PREFIX = "mde-" + +# ----------------------------------------------------------------------------- helpers + +def now(): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) +def parse_ts(s): return time.mktime(time.strptime(s, "%Y-%m-%dT%H:%M:%SZ")) - time.timezone if s else None +def job_id(): return "MDJOB-%s-%s" % (time.strftime("%Y%m%d", time.gmtime()), secrets.token_hex(3).upper()) +def sha256(s): return hashlib.sha256(s.encode() if isinstance(s, str) else s).hexdigest() + +def log(ev, **kw): + """One JSON line per event on stdout (journald). Never pass full keys or secrets.""" + rec = {"ts": now(), "ev": ev}; rec.update(kw) + print(json.dumps(rec, separators=(",", ":"), default=str), flush=True) + +def new_api_key(): + """Returns (full_key, key_id, key_hash). Only key_id/key_hash are ever stored.""" + full = "mde_" + secrets.token_hex(16); h = sha256(full); return full, h[:8], h + +def load_env_file(path): + if not path or not os.path.exists(path): return + with open(path) as fh: + for line in fh: + line = line.strip() + if not line or line.startswith("#") or "=" not in line: continue + k, v = line.split("=", 1); v = v.strip() + if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'": v = v[1:-1] + os.environ.setdefault(k.strip(), v) + +def parse_packs(s): + """MDE_PACKS="price_xxx:25:12.5,price_yyy:100:50" -> {price_id: (usd, gpu_hours)}""" + out = {} + for item in (s or "").split(","): + item = item.strip() + if not item: continue + pid, usd, hours = item.split(":"); out[pid] = (float(usd), float(hours)) + return out + +def parse_rates(s): + if not s: return {"any": 2.0, "rtx4090": 2.0} + return {k.strip(): float(v) for k, v in (kv.split(":") for kv in s.split(",") if kv.strip())} + +class Config: + def __init__(self, env=None): + e = env if env is not None else os.environ + self.db = e.get("MDE_DB", "/var/lib/mde/mde.sqlite") + self.blobs = e.get("MDE_BLOBS", "/var/lib/mde/blobs") + self.bind = e.get("MDE_BIND", "127.0.0.1:8080") + self.public_url = e.get("MDE_PUBLIC_URL", "").rstrip("/") # what clients/pods can reach, e.g. https://api.example.com + self.stripe_secret = e.get("STRIPE_SECRET_KEY", "") + self.webhook_secret = e.get("STRIPE_WEBHOOK_SECRET", "") + self.packs = parse_packs(e.get("MDE_PACKS", "")) + self.rates = parse_rates(e.get("MDE_RATES", "")) + self.runners_open = e.get("MDE_RUNNERS_OPEN", "") == "1" + self.admin_token = e.get("MDE_ADMIN_TOKEN", "") + self.brand = "ForceField Silicon / MDEngine" + # RunPod launcher (mde_launcher.py). No RUNPOD_API_KEY -> no launcher -> `start` writes launch.env. + self.runpod_api_key = e.get("RUNPOD_API_KEY", "") + self.runner_image = e.get("MDE_RUNNER_IMAGE", "") # default in mde_launcher.DEFAULT_IMAGE + self.pod_disk_gb = int(e.get("MDE_POD_DISK_GB", "20") or 20) + self.min_cuda = e.get("MDE_MIN_CUDA", "12.4") + self.gpu_ladder = parse_ladder(e.get("MDE_GPU_LADDER", "")) + self.launch_timeout_s = int(e.get("MDE_LAUNCH_TIMEOUT_S", "600") or 600) + self.reaper_interval_s = int(e.get("MDE_REAPER_INTERVAL_S", "300") or 300) + +# ----------------------------------------------------------------------------- storage + +SCHEMA = """ +create table if not exists keys(key_id text primary key, key_hash text unique not null, email text, + created text not null, label text); +create table if not exists credits(session_id text primary key, key_id text not null, usd real not null, + gpu_s integer not null, price_id text, created text not null); +create table if not exists jobs(id text primary key, key_id text not null, token_hash text, spec text not null, + state text not null, created text not null, started text, finished text, gpu text, rate real, + billed_s integer default 0, thermo text default '[]', exitcode integer, error text, attempt integer default 1, + last_hb real, pod_id text, launched_at text); +create table if not exists pending_keys(session_id text primary key, key_id text not null, full_key text not null, + created text not null); +create table if not exists meta(k text primary key, v text not null); +create index if not exists jobs_key on jobs(key_id, created); +create index if not exists credits_key on credits(key_id); +""" +MIGRATIONS = ["alter table jobs add column pod_id text", "alter table jobs add column launched_at text"] + +class DB: + """sqlite wrapper shared by the service and the admin CLI. All access under one RLock.""" + def __init__(self, path): + os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) + self.lock = threading.RLock() + self.c = sqlite3.connect(path, check_same_thread=False, timeout=10) + self.c.row_factory = sqlite3.Row + with self.lock: + self.c.execute("pragma journal_mode=wal"); self.c.executescript(SCHEMA) + for m in MIGRATIONS: # idempotent: pre-launcher databases gain the new columns + try: self.c.execute(m); self.c.commit() + except sqlite3.OperationalError as e: + if "duplicate column" not in str(e): raise + if not self.meta("blob_secret"): self.set_meta("blob_secret", secrets.token_hex(32)) + + def q(self, sql, *args): + with self.lock: return self.c.execute(sql, args).fetchall() + def one(self, sql, *args): + with self.lock: return self.c.execute(sql, args).fetchone() + def x(self, sql, *args): + with self.lock: self.c.execute(sql, args); self.c.commit() + def meta(self, k): + r = self.one("select v from meta where k=?", k); return r["v"] if r else None + def set_meta(self, k, v): self.x("insert or replace into meta(k,v) values(?,?)", k, v) + + # keys + def key_by_hash(self, h): return self.one("select * from keys where key_hash=?", h) + def key_by_id(self, kid): return self.one("select * from keys where key_id=?", kid) + def create_key(self, email=None, label=None): + full, kid, h = new_api_key() + self.x("insert into keys(key_id,key_hash,email,created,label) values(?,?,?,?,?)", kid, h, email, now(), label) + return full, kid + + # ledger + def add_credit(self, session_id, key_id, usd, gpu_s, price_id): + """Idempotent on session_id. Returns True if inserted, False if it already existed.""" + with self.lock: + if self.one("select 1 from credits where session_id=?", session_id): return False + self.x("insert into credits(session_id,key_id,usd,gpu_s,price_id,created) values(?,?,?,?,?,?)", + session_id, key_id, float(usd), int(gpu_s), price_id, now()) + return True + + def billed_usd(self, key_id): + total = 0.0 + for j in self.q("select * from jobs where key_id=? and state not in ('created','uploaded','queued','launching')", key_id): + total += job_cost(j) + return total + + def credited_usd(self, key_id): + r = self.one("select coalesce(sum(usd),0) s from credits where key_id=?", key_id); return float(r["s"]) + + def balance(self, key_id): return round(self.credited_usd(key_id) - self.billed_usd(key_id), 6) + + # jobs + def job(self, jid): return self.one("select * from jobs where id=?", jid) + def job_by_pod(self, pod_id): return self.one("select * from jobs where pod_id=?", pod_id) if pod_id else None + def set_job(self, jid, **kw): + cols = ", ".join(f"{k}=?" for k in kw); self.x(f"update jobs set {cols} where id=?", *kw.values(), jid) + +def billed_seconds(j): + """Live seconds for a running job; stored value otherwise.""" + if j["state"] == "running" and j["started"]: + return max(int(time.time() - parse_ts(j["started"])), j["billed_s"] or 0) + return j["billed_s"] or 0 + +def job_cost(j): + if j["error"] in NOT_BILLED_ERRORS: return 0.0 + return round(billed_seconds(j) * (j["rate"] or 0) / 3600, 6) + +def status(j): + spec = json.loads(j["spec"]); billed = billed_seconds(j) + return {"id": j["id"], "state": j["state"], "states": "|".join(STATES), "created": j["created"], + "started": j["started"], "finished": j["finished"], "gpu": j["gpu"], "rate_usd_per_h": j["rate"], + "billed_s": billed, "cost_usd": round(job_cost(j), 4), "thermo_tail": json.loads(j["thermo"] or "[]"), + "exitcode": j["exitcode"], "error": j["error"], "attempt": j["attempt"], "label": spec.get("label"), + "pod_id": j["pod_id"]} + +# ----------------------------------------------------------------------------- Stripe + +def stripe_fetch_session(session_id, secret_key): + """GET the Checkout Session with line_items expanded. Module-level so tests can monkeypatch it.""" + url = "https://api.stripe.com/v1/checkout/sessions/%s?expand[]=line_items" % urllib.parse.quote(session_id, safe="") + req = urllib.request.Request(url, headers={"Authorization": "Bearer " + secret_key, "User-Agent": "mde-endpoint/" + VERSION}) + with urllib.request.urlopen(req, timeout=20) as r: + return json.loads(r.read().decode()) + +def verify_stripe_signature(payload: bytes, header: str, secret: str, tolerance=STRIPE_TOLERANCE_S, now_ts=None): + """Stripe-Signature: t=,v1=[,v1=...]; v1 = HMAC-SHA256(secret, f"{t}.{payload}").""" + if not header or not secret: return False + parts = dict(p.split("=", 1) for p in header.split(",") if "=" in p) + t = parts.get("t"); sigs = [p.split("=", 1)[1] for p in header.split(",") if p.startswith("v1=")] + if not t or not sigs or not t.isdigit(): return False + if abs((now_ts or time.time()) - int(t)) > tolerance: return False + expected = hmac.new(secret.encode(), f"{t}.".encode() + payload, hashlib.sha256).hexdigest() + return any(hmac.compare_digest(expected, s) for s in sigs) + +def stripe_sign(payload: bytes, secret: str, ts=None): + """Produce a Stripe-Signature header (used by tests and the README's rotation check).""" + ts = int(ts or time.time()); sig = hmac.new(secret.encode(), f"{ts}.".encode() + payload, hashlib.sha256).hexdigest() + return f"t={ts},v1={sig}" + +# ----------------------------------------------------------------------------- application + +class App: + def __init__(self, cfg: Config): + self.cfg = cfg; self.db = DB(cfg.db) + os.makedirs(cfg.blobs, exist_ok=True) + self.grant_lock = threading.Lock() + self.public_url = cfg.public_url or "http://" + cfg.bind + self.launch_env = os.path.join(os.path.dirname(os.path.abspath(cfg.blobs)), "launch.env") + self.launcher = RunPodLauncher(cfg, public_url=self.public_url) if cfg.runpod_api_key else None # tests inject FakeLauncher + self.reaper_fail = {} # pod_id -> consecutive reaper delete failures (in memory; 3 = reaper.stuck) + self.last_reap = 0.0 + self._bg = []; self._bg_lock = threading.Lock() + + # -- background work (pod create/delete never blocks an HTTP response); join_bg() is for tests + def spawn(self, target, *args): + t = threading.Thread(target=target, args=args, daemon=True) + with self._bg_lock: + self._bg = [x for x in self._bg if x.is_alive()]; self._bg.append(t) + t.start(); return t + def join_bg(self, timeout=10): + deadline = time.time() + timeout + while True: + with self._bg_lock: live = [x for x in self._bg if x.is_alive()] + if not live or time.time() > deadline: return not live + live[0].join(max(0.01, deadline - time.time())) + + # -- blob URLs stand in for presigned object-storage URLs: HMAC over name|exp + def blob_path(self, name): return os.path.join(self.cfg.blobs, name) + def blob_sig(self, name, exp): return hmac.new(self.db.meta("blob_secret").encode(), f"{name}|{exp}".encode(), hashlib.sha256).hexdigest()[:32] + def blob_url(self, name, ttl=BLOB_TTL_S): + exp = int(time.time()) + ttl; return f"{self.public_url}/blob/{name}?exp={exp}&sig={self.blob_sig(name, exp)}", exp + def blob_ok(self, name, query): + q = parse_qs(query); exp = (q.get("exp") or [""])[0]; sig = (q.get("sig") or [""])[0] + if not exp.isdigit() or int(exp) < time.time(): return False + return hmac.compare_digest(self.blob_sig(name, int(exp)), sig) + + # -- purchase: one idempotent handler for /welcome and the webhook + def grant(self, sess, source): + """Credit a paid Checkout Session once. Returns dict(kind, key_id, full_key|None, usd, gpu_h, balance).""" + sid = sess.get("id"); + if not sid: raise ValueError("session has no id") + if sess.get("payment_status") != "paid": raise ValueError("payment_status=%s" % sess.get("payment_status")) + with self.grant_lock: + existing = self.db.one("select * from credits where session_id=?", sid) + if existing: + pend = self.db.one("select * from pending_keys where session_id=?", sid) + full = None + if pend and source == "welcome": # first visit after a webhook-first grant: show once + full = pend["full_key"]; self.db.x("delete from pending_keys where session_id=?", sid) + log("key.revealed", key_id=existing["key_id"], session=sid) + return {"kind": "already", "key_id": existing["key_id"], "full_key": full, "usd": existing["usd"], + "gpu_h": existing["gpu_s"] / 3600, "balance": self.db.balance(existing["key_id"])} + price_id, usd, gpu_h = self.pack_for(sess) + ref = (sess.get("client_reference_id") or "").strip() + details = sess.get("customer_details") or {} + email = details.get("email") or sess.get("customer_email") + k = self.db.key_by_id(ref) if ref else None + if k: + self.db.add_credit(sid, k["key_id"], usd, gpu_h * 3600, price_id) + log("credit.topup", key_id=k["key_id"], usd=usd, gpu_h=gpu_h, price_id=price_id, session=sid, source=source) + return {"kind": "topup", "key_id": k["key_id"], "full_key": None, "usd": usd, "gpu_h": gpu_h, "balance": self.db.balance(k["key_id"])} + full, kid = self.db.create_key(email=email, label="stripe:" + sid[-8:]) + self.db.add_credit(sid, kid, usd, gpu_h * 3600, price_id) + if source != "welcome": # buyer has not seen the key yet: hold it for the welcome page + self.db.x("insert or replace into pending_keys(session_id,key_id,full_key,created) values(?,?,?,?)", sid, kid, full, now()) + log("credit.new_key", key_id=kid, usd=usd, gpu_h=gpu_h, price_id=price_id, session=sid, source=source, has_email=bool(email)) + return {"kind": "new", "key_id": kid, "full_key": full, "usd": usd, "gpu_h": gpu_h, "balance": self.db.balance(kid)} + + def pack_for(self, sess): + """price id decides the hours (coupons do not reduce them); usd credited = hours * rate.""" + rate = self.cfg.rates.get("any", 2.0) + price_id = None + try: price_id = sess["line_items"]["data"][0]["price"]["id"] + except (KeyError, IndexError, TypeError): pass + if price_id in self.cfg.packs: + _, gpu_h = self.cfg.packs[price_id]; return price_id, round(gpu_h * rate, 2), gpu_h + # Unknown price id: honor the amount actually paid at the base rate, and say so in the log. + usd = (sess.get("amount_total") or 0) / 100.0 + log("pack.unknown_price", price_id=price_id, amount_total=sess.get("amount_total")) + return price_id, round(usd, 2), round(usd / rate, 4) + + def purge_pending(self): + cutoff = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - PENDING_KEY_TTL_S)) + self.db.x("delete from pending_keys where created < ?", cutoff) + + # -- pods: one per job, created after `start`, deleted at every terminal state + def launch_job(self, jid, token, wall, gpu): + """Background: walk the launcher's ladder; queued -> launching(pod_id) or failed:no_capacity (unbilled).""" + try: pod_id = self.launcher.create(jid, token, wall, gpu) + except NoCapacity as e: + log("job.no_capacity", job=jid, error=str(e)); self.fail_unlaunched(jid); return + except Exception as e: + log("launch.error", job=jid, error=repr(e)); self.fail_unlaunched(jid); return + j = self.db.job(jid) + if not j: self.release_pod(jid, pod_id, "job_vanished"); return + if j["state"] == "queued": self.db.set_job(jid, state="launching", pod_id=pod_id, launched_at=now()) + elif j["state"] in ("launching", "running"): self.db.set_job(jid, pod_id=pod_id, launched_at=j["launched_at"] or now()) + else: # cancelled while the create call was in flight: no orphan + self.db.set_job(jid, pod_id=pod_id); self.release_pod(jid, pod_id, "terminal_during_launch"); return + log("job.launching", job=jid, key_id=j["key_id"], pod_id=pod_id) + + def fail_unlaunched(self, jid): + j = self.db.job(jid) + if j and j["state"] in ("queued", "launching"): + self.db.set_job(jid, state="failed", finished=now(), error="no_capacity", token_hash=None) + + def release_pod(self, jid, pod_id, reason): + """Delete the job's pod in the background (never from the request thread).""" + if not self.launcher or not pod_id: return + def run(): + try: self.launcher.delete(pod_id); log("pod.deleted", job=jid, pod_id=pod_id, reason=reason) + except Exception as e: log("pod.delete_failed", job=jid, pod_id=pod_id, reason=reason, error=repr(e)) + self.spawn(run) + + def finish_job(self, j, log_ev, **fields): + """Terminal write + pod release + log, in that order (state first so a crash mid-way leaves the reaper a terminal job).""" + self.db.set_job(j["id"], **fields) + cur = self.db.job(j["id"]) + log(log_ev, job=j["id"], key_id=j["key_id"], **{k: v for k, v in fields.items() if k in ("state", "error", "billed_s", "exitcode")}, + cost_usd=round(job_cost(cur), 4), pod_id=cur["pod_id"]) + self.release_pod(j["id"], cur["pod_id"], fields.get("error") or fields.get("state")) + return cur + + # -- watchdog: lost heartbeats -> pod_lost; launching too long -> no_capacity (neither is billed) + def watchdog_once(self): + cutoff = time.time() - HEARTBEAT_LOST_S + for j in self.db.q("select * from jobs where state='running' and last_hb is not null and last_hb < ?", cutoff): + self.finish_job(j, "job.pod_lost", state="failed", finished=now(), error="pod_lost", token_hash=None) + lcut = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(time.time() - self.cfg.launch_timeout_s)) + for j in self.db.q("select * from jobs where state='launching' and last_hb is null and launched_at is not null and launched_at < ?", lcut): + self.finish_job(j, "job.launch_timeout", state="failed", finished=now(), error="no_capacity", token_hash=None) + + # -- reaper (CONTRACT "Pod lifecycle", GJOB-099): a pod exists only while a job is launching/running + def reaper_once(self): + if not self.launcher: return 0 + try: pods = self.launcher.list_pods() + except Exception as e: log("reaper.list_failed", error=repr(e)); return 0 + self.last_reap = time.time(); n = 0 + for p in pods: + pid = p.get("id"); name = p.get("name") or "" + if not pid or not name.startswith(POD_PREFIX): continue + j = self.db.job_by_pod(pid) or self.db.job(name[len(POD_PREFIX):]) + reason = None + if not j: reason = "no_job" + elif j["state"] in TERMINAL: reason = "job_terminal" + else: + age = pod_age_s(p); wall = int((json.loads(j["spec"]).get("wall_limit_s")) or 86400) + if age is not None and age > wall + REAPER_GRACE_S: reason = "overage" + if not reason: self.reaper_fail.pop(pid, None); continue + try: + self.launcher.delete(pid); self.reaper_fail.pop(pid, None); n += 1 + log("reaper.deleted", pod_id=pid, name=name, job=j["id"] if j else None, reason=reason, status=p.get("status")) + except Exception as e: + k = self.reaper_fail[pid] = self.reaper_fail.get(pid, 0) + 1 + log("reaper.delete_failed", pod_id=pid, name=name, reason=reason, failures=k, error=repr(e)) + if k >= 3: log("reaper.stuck", pod_id=pid, name=name, reason=reason, failures=k) + return n + + def watchdog_loop(self, stop): + while not stop.wait(30): + try: self.watchdog_once(); self.purge_pending() + except Exception as e: log("watchdog.error", error=repr(e)) + if self.launcher and time.time() - self.last_reap >= self.cfg.reaper_interval_s: + try: self.reaper_once() + except Exception as e: log("reaper.error", error=repr(e)) + +# ----------------------------------------------------------------------------- HTTP + +PAGE = """ +{brand}
{brand}
{body}
""" + +def page(brand, body): return PAGE.format(brand=brand, body=body).encode() +def esc(s): return str(s).replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) + +class Handler(BaseHTTPRequestHandler): + app: App + server_version = "mde-endpoint/" + VERSION; sys_version = "" + def log_message(self, *a): pass + def send(self, code, obj=None, raw=None, ctype="application/json"): + body = raw if raw is not None else (json.dumps(obj).encode() if obj is not None else b"") + self.send_response(code); self.send_header("Content-Type", ctype); self.send_header("Content-Length", str(len(body))) + self.send_header("Cache-Control", "no-store"); self.end_headers(); self.wfile.write(body) + def html(self, code, body): return self.send(code, raw=page(self.app.cfg.brand, body), ctype="text/html; charset=utf-8") + def body(self): + n = int(self.headers.get("Content-Length") or 0); return self.rfile.read(n) if n else b"" + def bearer(self): + a = self.headers.get("Authorization", ""); return a[7:].strip() if a.startswith("Bearer ") else None + def api_key(self): + b = self.bearer() + if not b or not b.startswith("mde_"): return None + return self.app.db.key_by_hash(sha256(b)) + def pod_job(self, jid): + j = self.app.db.job(jid); b = self.bearer() + if not j or not b or not j["token_hash"] or not hmac.compare_digest(j["token_hash"], sha256(b)): return None + return j + def json_body(self): + try: return json.loads(self.body() or b"{}") + except json.JSONDecodeError: return None + + # ------------------------------------------------------------------ GET + def do_GET(self): + u = urlparse(self.path); p = u.path.rstrip("/") or "/"; parts = p.split("/") + db = self.app.db + if p == "/v1/health": + return self.send(200, {"ok": True, "version": VERSION, "runners": "open" if self.app.cfg.runners_open else "closed", + "launcher": self.app.launcher.kind if self.app.launcher else "none"}) + if p == "/welcome": return self.welcome(parse_qs(u.query)) + if parts[1] == "blob" and len(parts) == 3: + name = parts[2] + if not self.app.blob_ok(name, u.query): return self.send(403, {"error": "bad or expired blob url"}) + f = self.app.blob_path(name) + if not os.path.exists(f): return self.send(404, {"error": "no blob"}) + self.send_response(200); self.send_header("Content-Type", "application/gzip"); self.send_header("Content-Length", str(os.path.getsize(f))); self.end_headers() + with open(f, "rb") as fh: + while chunk := fh.read(1 << 20): self.wfile.write(chunk) + return + if parts[1] == "internal": # pod side + j = self.pod_job(parts[3]) if len(parts) == 4 and parts[2] == "jobs" else None + if not j: return self.send(401, {"error": "bad token"}) + spec = json.loads(j["spec"]) + in_url, _ = self.app.blob_url(f"{j['id']}.in.tar.gz"); out_url, _ = self.app.blob_url(f"{j['id']}.out.tar.gz") + return self.send(200, {"input_url": in_url, "input": spec["input"], "launch": spec.get("launch", "default"), + "wall_limit_s": spec.get("wall_limit_s", 86400), "results_put_url": out_url}) + if parts[1] == "v1" and len(parts) > 2 and parts[2] == "admin": return self.admin(parts[3:]) + k = self.api_key() + if not k: return self.send(401, {"error": "bad api key"}) + if p == "/v1/me": + return self.send(200, {"balance_usd": db.balance(k["key_id"]), "rate_table": self.app.cfg.rates, "keys_created": k["created"], "key_id": k["key_id"]}) + if p == "/v1/jobs": + q = parse_qs(u.query); limit = max(1, min(int((q.get("limit") or ["50"])[0] or 50), 200)); cursor = (q.get("cursor") or [None])[0] + if cursor and (cj := db.job(cursor)) and cj["key_id"] == k["key_id"]: + rows = db.q("select * from jobs where key_id=? and (created= 4 and parts[1] == "v1" and parts[2] == "jobs": + j = db.job(parts[3]) + if not j or j["key_id"] != k["key_id"]: return self.send(404, {"error": "no job"}) + if len(parts) == 4: return self.send(200, status(j)) + if parts[4] == "results": + if j["state"] not in ("done", "failed"): return self.send(409, {"error": "not finished"}) + f = self.app.blob_path(f"{j['id']}.out.tar.gz"); sz = os.path.getsize(f) if os.path.exists(f) else 0 + url, exp = self.app.blob_url(f"{j['id']}.out.tar.gz") + return self.send(200, {"download_url": url, "expires": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(exp)), "bytes": sz}) + self.send(404, {"error": "no route"}) + + def admin(self, parts): + tok = self.app.cfg.admin_token + if not tok or not self.bearer() or not hmac.compare_digest(tok, self.bearer()): return self.send(403, {"error": "forbidden"}) + db = self.app.db + if parts == ["stats"]: + by_state = {r["state"]: r["n"] for r in db.q("select state, count(*) n from jobs group by state")} + return self.send(200, {"keys": db.one("select count(*) n from keys")["n"], "credits_usd": db.one("select coalesce(sum(usd),0) s from credits")["s"], + "credit_rows": db.one("select count(*) n from credits")["n"], "jobs": by_state, "pending_keys": db.one("select count(*) n from pending_keys")["n"]}) + return self.send(404, {"error": "no route"}) + + def welcome(self, q): + sid = (q.get("session_id") or [""])[0].strip() + if not sid or len(sid) > 200: return self.html(400, "

Missing session

This page is reached from the payment confirmation link.

") + if not self.app.cfg.stripe_secret: log("welcome.no_stripe_key", session=sid); return self.html(503, "

Not configured

Payment lookup is unavailable right now. Your payment is safe; try again shortly.

") + try: + sess = stripe_fetch_session(sid, self.app.cfg.stripe_secret) + except urllib.error.HTTPError as e: + log("welcome.stripe_http", session=sid, code=e.code); return self.html(404 if e.code == 404 else 502, "

Session not found

We could not find that payment session. If you were charged, reply to your Stripe receipt.

") + except Exception as e: + log("welcome.stripe_error", session=sid, error=repr(e)); return self.html(502, "

Temporary problem

Could not reach the payment provider. Reload in a minute; your payment is safe.

") + if sess.get("payment_status") != "paid": + return self.html(402, "

Payment not completed

Stripe reports this session as %s. Once it is paid, reload this page.

" % esc(sess.get("payment_status"))) + try: g = self.app.grant(sess, "welcome") + except Exception as e: + log("welcome.grant_error", session=sid, error=repr(e)); return self.html(500, "

Something went wrong

Your payment is recorded; reply to your Stripe receipt and we will fix it.

") + bal = "$%.2f" % g["balance"]; hours = ("%.4g" % g["gpu_h"]) + if g["full_key"]: + body = f"""

Your MDEngine key

Credited {hours} GPU-hours (${g['usd']:.2f}). Balance {bal}

+

Shown once. Copy it now; it is stored hashed and cannot be displayed again.

{esc(g['full_key'])}
+

Key id {esc(g['key_id'])} (this short id is what appears in support and top-ups).

+

Command line

mdengine login {esc(g['full_key'])}
+mdengine run --gpu in.lmp
+

App

  1. MDEngine ▸ Settings ▸ Accelerated (API key)
  2. Paste the key
  3. File ▸ Run Accelerated… (⇧⌘R)
+

MCP: submit_lammps host=cloud once the key is saved by mdengine login.

+

Credits never expire. Check the balance any time with mdengine account.

""" + return self.html(200, body) + if g["kind"] == "topup": + return self.html(200, f"""

Credits added

Credited {hours} GPU-hours (${g['usd']:.2f}) to your existing key {esc(g['key_id'])}.

+

New balance {bal}

mdengine account shows the same figure.

""") + return self.html(200, f"""

Already issued

This purchase was already credited to key {esc(g['key_id'])}; balance {bal}.

+

The full key is shown only once, right after payment. If you did not save it, mdengine account works if you already logged in; otherwise reply to your Stripe receipt and quote the key id above and we will issue a replacement.

""") + + # ------------------------------------------------------------------ PUT (blob upload) + def do_PUT(self): + u = urlparse(self.path); parts = u.path.split("/") + if len(parts) != 3 or parts[1] != "blob": return self.send(404, {"error": "no route"}) + name = parts[2] + if not self.app.blob_ok(name, u.query): return self.send(403, {"error": "bad or expired blob url"}) + n = int(self.headers.get("Content-Length") or 0) + if n > MAX_BLOB: return self.send(413, {"error": "tarball exceeds 2 GB"}) + jid = name.split(".")[0]; j = self.app.db.job(jid) + if not j: return self.send(404, {"error": "no job"}) + if name.endswith(".in.tar.gz") and j["state"] != "created": return self.send(409, {"error": f"state is {j['state']}"}) + tmp = self.app.blob_path(name + ".part"); left = n + with open(tmp, "wb") as fh: + while left > 0: + chunk = self.rfile.read(min(left, 1 << 20)) + if not chunk: break + fh.write(chunk); left -= len(chunk) + if left: os.remove(tmp); return self.send(400, {"error": "short body"}) + os.replace(tmp, self.app.blob_path(name)) + if name.endswith(".in.tar.gz"): self.app.db.set_job(jid, state="uploaded") + log("blob.put", job=jid, name=name, bytes=n) + self.send(200, {"ok": True}) + + # ------------------------------------------------------------------ POST + def do_POST(self): + p = urlparse(self.path).path.rstrip("/"); parts = p.split("/"); db = self.app.db + if p == "/v1/stripe/webhook": return self.webhook() + if parts[1] == "internal": # pod side: heartbeat / done + j = self.pod_job(parts[3]) if len(parts) == 5 and parts[2] == "jobs" else None + if not j: return self.send(401, {"error": "bad token"}) + b = self.json_body() + if b is None: return self.send(400, {"error": "bad json"}) + if parts[4] == "heartbeat": + kw = {"thermo": json.dumps([str(x) for x in (b.get("thermo_tail") or [])][-20:]), "last_hb": time.time()} + if j["state"] in ("queued", "launching"): kw.update(state="running", started=now()); log("job.running", job=j["id"]) + db.set_job(j["id"], **kw); return self.send(200, {"ok": True}) + if parts[4] == "done": + if j["state"] in TERMINAL: return self.send(200, {"ok": True, "state": j["state"]}) + rc = int(b.get("exitcode", 1)); err = b.get("error") + started = j["started"] or now(); billed = int(b.get("elapsed_s", 0)) + if j["started"]: billed = min(billed, int(time.time() - parse_ts(j["started"])) + 60) # pod cannot bill more than wall time + have = os.path.exists(self.app.blob_path(f"{j['id']}.out.tar.gz")) + state = "done" if rc == 0 and have else "failed" + if not have and not err: err = "no_results" + self.app.finish_job(j, "job.finished", state=state, finished=now(), started=started, exitcode=rc, error=err, billed_s=billed, token_hash=None) + return self.send(200, {"ok": True, "state": state}) + return self.send(404, {"error": "no route"}) + k = self.api_key() + if not k: return self.send(401, {"error": "bad api key"}) + if p == "/v1/jobs": + spec = self.json_body() + if spec is None or "input" not in spec: return self.send(400, {"error": "input required"}) + gpu = spec.get("gpu", "any"); rate = self.app.cfg.rates.get(gpu) + if rate is None: return self.send(400, {"error": "unknown gpu"}) + try: est = int(spec.get("estimate_s", 0)); wall = int(spec.get("wall_limit_s", 14400)) + except (TypeError, ValueError): return self.send(400, {"error": "bad numbers"}) + if wall > 86400 or wall <= 0: return self.send(400, {"error": "wall_limit_s must be 1..86400"}) + if len(str(spec.get("label") or "")) > 120: return self.send(400, {"error": "label too long"}) + if db.balance(k["key_id"]) < rate * max(est, 900) / 3600: return self.send(402, {"error": "insufficient balance"}) + if not self.app.cfg.runners_open: + log("job.refused_closed", key_id=k["key_id"]) + return self.send(503, {"error": "gpu_runners_open_soon", "message": "GPU runners open this week; your credits are safe and never expire."}) + jid = job_id() + spec["wall_limit_s"] = wall + db.x("insert into jobs(id,key_id,spec,state,created,gpu,rate) values(?,?,?,?,?,?,?)", + jid, k["key_id"], json.dumps(spec), "created", now(), gpu, rate) + url, exp = self.app.blob_url(f"{jid}.in.tar.gz", ttl=3600) + log("job.created", job=jid, key_id=k["key_id"], gpu=gpu) + return self.send(201, {"id": jid, "upload_url": url, "upload_expires": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(exp))}) + if len(parts) == 5 and parts[1] == "v1" and parts[2] == "jobs" and parts[4] == "start": + j = db.job(parts[3]) + if not j or j["key_id"] != k["key_id"]: return self.send(404, {"error": "no job"}) + if j["state"] != "uploaded": return self.send(409, {"error": f"state is {j['state']}"}) + tok = "jt_" + secrets.token_hex(16) # minted here, stored hashed, handed to the pod only + db.set_job(j["id"], state="queued", token_hash=sha256(tok)) + spec = json.loads(j["spec"]); wall = int(spec.get("wall_limit_s") or 86400) + if self.app.launcher: # background: queued -> launching | failed:no_capacity + self.app.spawn(self.app.launch_job, j["id"], tok, wall, j["gpu"]) + else: # dev path: launch env on disk for a hand-run pod + with open(self.app.launch_env, "w", opener=lambda f, fl: os.open(f, fl, 0o600)) as fh: + fh.write(f"MDE_ENDPOINT={self.app.public_url}\nMDE_JOB_ID={j['id']}\nMDE_JOB_TOKEN={tok}\n") + log("job.queued", job=j["id"], key_id=k["key_id"], launcher=bool(self.app.launcher)) + return self.send(202, {"id": j["id"], "state": "queued"}) + self.send(404, {"error": "no route"}) + + def webhook(self): + payload = self.body(); cfg = self.app.cfg + if not verify_stripe_signature(payload, self.headers.get("Stripe-Signature", ""), cfg.webhook_secret): + log("webhook.bad_signature"); return self.send(400, {"error": "bad signature"}) + try: event = json.loads(payload) + except json.JSONDecodeError: return self.send(400, {"error": "bad json"}) + etype = event.get("type"); obj = (event.get("data") or {}).get("object") or {} + if etype != "checkout.session.completed": return self.send(200, {"ok": True, "ignored": etype}) + sid = obj.get("id") + if not sid: return self.send(400, {"error": "no session id"}) + sess = obj + if cfg.stripe_secret: # webhook payloads carry no line_items; fetch to learn the price id + try: sess = stripe_fetch_session(sid, cfg.stripe_secret) + except Exception as e: log("webhook.stripe_fetch_failed", session=sid, error=repr(e)) + if sess.get("payment_status") != "paid": return self.send(200, {"ok": True, "ignored": "unpaid"}) + try: g = self.app.grant(sess, "webhook") + except Exception as e: + log("webhook.grant_error", session=sid, error=repr(e)); return self.send(500, {"error": "grant failed"}) + return self.send(200, {"ok": True, "result": g["kind"], "key_id": g["key_id"]}) + + # ------------------------------------------------------------------ DELETE (cancel) + def do_DELETE(self): + parts = urlparse(self.path).path.rstrip("/").split("/"); db = self.app.db + k = self.api_key() + if not k: return self.send(401, {"error": "bad api key"}) + if len(parts) != 4 or parts[1] != "v1" or parts[2] != "jobs": return self.send(404, {"error": "no route"}) + j = db.job(parts[3]) + if not j or j["key_id"] != k["key_id"]: return self.send(404, {"error": "no job"}) + if j["state"] in TERMINAL: return self.send(409, {"error": "terminal"}) + billed = billed_seconds(j) if j["state"] == "running" else 0 + cur = self.app.finish_job(j, "job.cancelled", state="cancelled", finished=now(), error="cancelled", billed_s=billed, token_hash=None) + self.send(202, status(cur)) + +# ----------------------------------------------------------------------------- server + +def make_server(cfg: Config): + """Bind and return (server, app). Port 0 in MDE_BIND picks a free port (tests).""" + app = App(cfg) + handler = type("BoundHandler", (Handler,), {"app": app}) + host, _, port = cfg.bind.rpartition(":") + srv = ThreadingHTTPServer((host or "127.0.0.1", int(port or 8080)), handler) + srv.daemon_threads = True + if not cfg.public_url: app.public_url = "http://%s:%d" % srv.server_address[:2] + if app.launcher: app.launcher.public_url = app.public_url + return srv, app + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--env-file", default=os.environ.get("MDE_ENV_FILE"), help="KEY=VALUE file; existing env wins") + a = ap.parse_args(argv); load_env_file(a.env_file) + cfg = Config(); srv, app = make_server(cfg) + log("start", version=VERSION, bind=cfg.bind, public_url=app.public_url, runners="open" if cfg.runners_open else "closed", + packs=len(cfg.packs), stripe=bool(cfg.stripe_secret), webhook=bool(cfg.webhook_secret), db=cfg.db, + launcher=app.launcher.kind if app.launcher else "none", ladder=cfg.gpu_ladder if app.launcher else None) + if app.launcher: # boot-time reap: an endpoint outage must not leave orphans behind it + try: log("reaper.boot", deleted=app.reaper_once()) + except Exception as e: log("reaper.error", error=repr(e)) + stop = threading.Event(); threading.Thread(target=app.watchdog_loop, args=(stop,), daemon=True).start() + try: srv.serve_forever() + except KeyboardInterrupt: pass + finally: stop.set(); srv.server_close(); log("stop") + +if __name__ == "__main__": + main() diff --git a/hosted/endpoint/mde_launcher.py b/hosted/endpoint/mde_launcher.py new file mode 100644 index 0000000..7634b1b --- /dev/null +++ b/hosted/endpoint/mde_launcher.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""RunPod pod launcher for the MDEngine hosted endpoint (hosted/CONTRACT.md "Pod lifecycle"). Stdlib only. + +One pod per job. The pod is the runner image with MDE_ENDPOINT / MDE_JOB_ID / MDE_JOB_TOKEN in its env; +docker/runner-gpu/start.sh execs runner.sh when MDE_JOB_ID is set. The endpoint deletes the pod at every +terminal job state and the reaper (mde_endpoint.App.reaper_once) clears anything that slipped through. + +RunPod REST v2 (https://api.runpod.io/v2/openapi.json): + POST /v2/pods 201 Pod create; any other status = this ladder rung failed + GET /v2/pods/{id} 200 Pod status in PROVISIONING STARTING RUNNING EXITED ERROR TERMINATED + DELETE /v2/pods/{id} 204 (404 = gone) idempotent + GET /v2/pods 200 {"pods":[Pod]} (live) | {"items":[Pod]} (docs) | bare list + +Logging discipline: the API key is only ever a header; request bodies (they carry the job token) are never +logged; Pod objects returned by RunPod carry the pod env (job token) and are never logged either -- only +error bodies, truncated to LOG_BODY_MAX. +""" +import json, threading, time, urllib.error, urllib.request +from datetime import datetime, timezone + +RUNPOD_BASE = "https://api.runpod.io" +DEFAULT_IMAGE = "ghcr.io/forcefieldsilicon/mdengine-runner-gpu:ADA89" +DEFAULT_LADDER = [("COMMUNITY", "NVIDIA GeForce RTX 4090"), ("SECURE", "NVIDIA GeForce RTX 4090")] +CLOUDS = ("COMMUNITY", "SECURE") +HTTP_TIMEOUT_S = 30 +LOG_BODY_MAX = 300 +POD_NAME_PREFIX = "mde-" + +class LauncherError(Exception): + """Transport or protocol failure talking to RunPod (never carries the API key or a token).""" + +class NoCapacity(LauncherError): + """Every rung of the fallback ladder refused to create a pod.""" + +def log(ev, **kw): + rec = {"ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "ev": ev}; rec.update(kw) + print(json.dumps(rec, separators=(",", ":"), default=str), flush=True) + +def trunc(s, n=LOG_BODY_MAX): + s = s if isinstance(s, str) else str(s) + return s if len(s) <= n else s[:n] + "...(%d more)" % (len(s) - n) + +def parse_ladder(s): + """MDE_GPU_LADDER="COMMUNITY:NVIDIA GeForce RTX 4090,SECURE:NVIDIA GeForce RTX 4090" -> [(cloud, gpu_id)].""" + if not (s or "").strip(): return list(DEFAULT_LADDER) + out = [] + for item in s.split(","): + item = item.strip() + if not item: continue + if ":" not in item: raise ValueError("MDE_GPU_LADDER entry needs CLOUD:gpu id, got %r" % item) + cloud, gid = item.split(":", 1); cloud = cloud.strip().upper(); gid = gid.strip() + if cloud not in CLOUDS or not gid: raise ValueError("bad MDE_GPU_LADDER entry %r" % item) + out.append((cloud, gid)) + if not out: raise ValueError("MDE_GPU_LADDER is empty") + return out + +def parse_iso(s): + """RunPod createdAt ('2026-09-06T12:34:56.789Z' or with an offset) -> unix seconds, or None.""" + if not s: return None + try: + if s.endswith("Z"): s = s[:-1] + "+00:00" + d = datetime.fromisoformat(s) + if d.tzinfo is None: d = d.replace(tzinfo=timezone.utc) + return d.timestamp() + except ValueError: return None + +def pod_age_s(pod, now_ts=None): + t = parse_iso((pod or {}).get("createdAt")); return None if t is None else max(0.0, (now_ts or time.time()) - t) + +def iso_now(ts=None): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(ts)) + +# ----------------------------------------------------------------------------- RunPod + +class RunPodLauncher: + kind = "runpod" + def __init__(self, cfg, public_url=None): + self.api_key = cfg.runpod_api_key + if not self.api_key: raise ValueError("RUNPOD_API_KEY is empty") + self.image = getattr(cfg, "runner_image", "") or DEFAULT_IMAGE + self.disk_gb = int(getattr(cfg, "pod_disk_gb", 20) or 20) + self.min_cuda = getattr(cfg, "min_cuda", "") or "12.4" + self.ladder = list(getattr(cfg, "gpu_ladder", None) or DEFAULT_LADDER) + self.public_url = (public_url or cfg.public_url or "").rstrip("/") + self.base = getattr(cfg, "runpod_base", "") or RUNPOD_BASE + self.timeout = HTTP_TIMEOUT_S + self.backoff_s = 1.0 # delete() retry base; tests set 0 + self._lock = threading.Lock() # serialises pod creation (one launch thread at a time is plenty) + + # -- transport (split so tests can monkeypatch either layer) ------------------------------------ + def _open(self, req): + """urlopen -> (status, text). HTTP errors are returned, not raised; transport errors raise LauncherError.""" + try: + with urllib.request.urlopen(req, timeout=self.timeout) as r: return r.status, r.read().decode(errors="replace") + except urllib.error.HTTPError as e: + with e: return e.code, e.read().decode(errors="replace") + except (urllib.error.URLError, OSError, TimeoutError) as e: + raise LauncherError("runpod transport: %s" % trunc(repr(e), 120)) from None + + def _request(self, method, path, body=None): + """(status, text) for METHOD {base}{path}. Authorization is set here and nowhere else.""" + data = json.dumps(body).encode() if body is not None else None + req = urllib.request.Request(self.base + path, data=data, method=method, + headers={"Authorization": "Bearer " + self.api_key, "Accept": "application/json", + "User-Agent": "mde-endpoint-launcher"}) + if data is not None: req.add_header("Content-Type", "application/json") + return self._open(req) + + # -- interface --------------------------------------------------------------------------------- + def pod_body(self, job_id, token, cloud, gpu_id, wall_limit_s=86400): + return {"name": POD_NAME_PREFIX + job_id, "image": self.image, "cloud": cloud, + "gpu": {"id": gpu_id, "count": 1, "minCudaVersion": self.min_cuda}, "disk": self.disk_gb, + "env": {"MDE_ENDPOINT": self.public_url, "MDE_JOB_ID": job_id, "MDE_JOB_TOKEN": token, + "MDE_WALL_LIMIT_S": str(int(wall_limit_s))}} # pod-side TTL (start.sh): wall + 600 s + + def create(self, job_id, token, wall_limit_s, gpu): + """Walk the ladder until one POST /v2/pods returns 201; return the pod id. Raises NoCapacity.""" + with self._lock: + for rung, (cloud, gpu_id) in enumerate(self.ladder, 1): + try: code, text = self._request("POST", "/v2/pods", self.pod_body(job_id, token, cloud, gpu_id, wall_limit_s)) + except LauncherError as e: code, text = 0, str(e) + pod_id = None + if code == 201: # a Pod body: carries env, so never logged + try: pod_id = json.loads(text).get("id") + except (ValueError, AttributeError): pod_id = None + text = "<201 without pod id>" + log("launch.attempt", job=job_id, rung=rung, cloud=cloud, gpu_id=gpu_id, gpu=gpu, wall_limit_s=wall_limit_s, + code=code, ok=bool(pod_id), **({} if pod_id else {"body": trunc(text)})) + if pod_id: return pod_id + raise NoCapacity("no rung of %d accepted job %s" % (len(self.ladder), job_id)) + + def delete(self, pod_id, retries=3): + """DELETE /v2/pods/{id}; 204 and 404 are success. 429/5xx/transport errors retry with backoff.""" + last = None + for i in range(retries): + try: code, text = self._request("DELETE", "/v2/pods/" + pod_id) + except LauncherError as e: code, text = 0, str(e) + if code in (200, 202, 204, 404): return True + last = "%d %s" % (code, trunc(text)) + if code and code != 429 and code < 500: break # other 4xx: not retryable + if i + 1 < retries: time.sleep(self.backoff_s * (2 ** i)) + raise LauncherError("delete %s failed: %s" % (pod_id, last)) + + def get(self, pod_id): + code, text = self._request("GET", "/v2/pods/" + pod_id) + if code == 404: return None + if code != 200: raise LauncherError("get %s: %d %s" % (pod_id, code, trunc(text))) + return json.loads(text) + + def list_pods(self): + code, text = self._request("GET", "/v2/pods") + if code != 200: raise LauncherError("list pods: %d %s" % (code, trunc(text))) + data = json.loads(text) + # live API (2026-09) wraps as {"pods":[...]}; OpenAPI example says {"items":[...]}; older: bare list + items = (data.get("pods") or data.get("items")) if isinstance(data, dict) else data + return list(items or []) + +# ----------------------------------------------------------------------------- fake (tests / dry runs) + +class FakeLauncher: + """Same interface, in memory. `fail_create=True` raises NoCapacity; pod ids in `fail_delete` refuse deletion.""" + kind = "fake" + def __init__(self, public_url=""): + self.public_url = public_url; self.pods = {}; self.deleted = []; self.calls = [] + self.fail_create = False; self.fail_delete = set(); self._n = 0; self._lock = threading.Lock() + + def add_pod(self, name, created_at=None, pod_id=None, env=None): + with self._lock: + self._n += 1; pid = pod_id or "fakepod%d" % self._n + self.pods[pid] = {"id": pid, "name": name, "status": "RUNNING", "cloud": "COMMUNITY", + "createdAt": created_at or iso_now(), "env": dict(env or {})} + return pid + + def create(self, job_id, token, wall_limit_s, gpu): + self.calls.append(("create", job_id)) + log("launch.attempt", job=job_id, rung=1, cloud="FAKE", gpu_id="fake", gpu=gpu, wall_limit_s=wall_limit_s, ok=not self.fail_create) + if self.fail_create: raise NoCapacity("fake: no capacity") + return self.add_pod(POD_NAME_PREFIX + job_id, env={"MDE_ENDPOINT": self.public_url, "MDE_JOB_ID": job_id, "MDE_JOB_TOKEN": token}) + + def delete(self, pod_id): + self.calls.append(("delete", pod_id)) + if pod_id in self.fail_delete: raise LauncherError("fake: delete refused for %s" % pod_id) + with self._lock: self.pods.pop(pod_id, None); self.deleted.append(pod_id) + return True + + def get(self, pod_id): return self.pods.get(pod_id) + def list_pods(self): return [dict(p) for p in self.pods.values()] diff --git a/hosted/endpoint/test_endpoint.py b/hosted/endpoint/test_endpoint.py new file mode 100644 index 0000000..e14ebe3 --- /dev/null +++ b/hosted/endpoint/test_endpoint.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Tests for mde_endpoint.py / mde_admin.py. Stdlib unittest; starts the server on a random port with a +temp db and a fake Stripe (stripe_fetch_session monkeypatched). No network beyond 127.0.0.1. + + python3 hosted/endpoint/test_endpoint.py -v +""" +import io, json, os, re, shutil, sys, tempfile, threading, time, unittest, urllib.error, urllib.request +from contextlib import redirect_stdout + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import mde_endpoint as E +import mde_admin as A +import mde_launcher as L + +PRICE_STARTER, PRICE_LAB = "price_test_starter", "price_test_lab" +WEBHOOK_SECRET = "whsec_test_" + "a" * 20 + +FAKE_SESSIONS = {} + +def fake_session(sid, price=PRICE_STARTER, paid=True, amount=2500, ref=None, email="buyer@example.test"): + s = {"id": sid, "object": "checkout.session", "payment_status": "paid" if paid else "unpaid", "amount_total": amount, + "client_reference_id": ref, "customer_details": {"email": email}, + "line_items": {"data": [{"price": {"id": price}}]}} + FAKE_SESSIONS[sid] = s; return s + +def fake_fetch(sid, secret): + if sid not in FAKE_SESSIONS: raise urllib.error.HTTPError("https://api.stripe.com", 404, "no such session", {}, None) + return FAKE_SESSIONS[sid] + +class Server: + def __init__(self, runners_open=False, launcher=None): + self.dir = tempfile.mkdtemp(prefix="mde-test-") + env = {"MDE_DB": os.path.join(self.dir, "t.sqlite"), "MDE_BLOBS": os.path.join(self.dir, "blobs"), "MDE_BIND": "127.0.0.1:0", + "STRIPE_SECRET_KEY": "sk_test_fake", "STRIPE_WEBHOOK_SECRET": WEBHOOK_SECRET, + "MDE_PACKS": f"{PRICE_STARTER}:25:12.5,{PRICE_LAB}:100:50", "MDE_RUNNERS_OPEN": "1" if runners_open else "", + "MDE_ADMIN_TOKEN": "admintok"} + self.cfg = E.Config(env); self.srv, self.app = E.make_server(self.cfg) + self.base = self.app.public_url + if launcher is not None: launcher.public_url = self.base; self.app.launcher = launcher + self.t = threading.Thread(target=self.srv.serve_forever, daemon=True); self.t.start() + def close(self): + self.srv.shutdown(); self.srv.server_close(); self.app.db.c.close(); shutil.rmtree(self.dir, ignore_errors=True) + + def req(self, method, path, body=None, key=None, headers=None, raw=None): + data = raw if raw is not None else (json.dumps(body).encode() if body is not None else None) + url = path if path.startswith("http") else self.base + path + r = urllib.request.Request(url, data=data, method=method) + if key: r.add_header("Authorization", "Bearer " + key) + if body is not None: r.add_header("Content-Type", "application/json") + for k, v in (headers or {}).items(): r.add_header(k, v) + try: + with urllib.request.urlopen(r, timeout=10) as resp: return resp.status, resp.read(), resp.headers + except urllib.error.HTTPError as e: + with e: return e.code, e.read(), e.headers + + def js(self, *a, **kw): + code, body, _ = self.req(*a, **kw); return code, json.loads(body) + + def admin_key(self, credit, email="op@example.test"): + """Create a key via the admin CLI against the same sqlite file; returns (full_key, key_id).""" + out = io.StringIO() + with redirect_stdout(out): A.main(["--db", self.cfg.db, "--env-file", "/nonexistent", "key", "new", "--email", email, "--credit", str(credit)]) + full = re.search(r"\b(mde_[0-9a-f]{32})\b", out.getvalue()).group(1) + kid = re.search(r"key_id\s+(\w+)", out.getvalue()).group(1) + return full, kid + +class Base(unittest.TestCase): + runners_open = False + fake_launcher = False + def setUp(self): + self._orig = E.stripe_fetch_session; E.stripe_fetch_session = fake_fetch; FAKE_SESSIONS.clear() + self.fake = L.FakeLauncher() if self.fake_launcher else None + self.s = Server(self.runners_open, launcher=self.fake); self._log = io.StringIO(); self._logpatch = redirect_stdout(self._log); self._logpatch.__enter__() + def tearDown(self): + self._logpatch.__exit__(None, None, None); self.s.close(); E.stripe_fetch_session = self._orig + +class TestAuthAndBalance(Base): + def test_health(self): + code, j = self.s.js("GET", "/v1/health") + self.assertEqual(code, 200); self.assertTrue(j["ok"]); self.assertEqual(j["runners"], "closed"); self.assertEqual(j["version"], E.VERSION) + + def test_me_and_key_hashing(self): + full, kid = self.s.admin_key(20) + code, j = self.s.js("GET", "/v1/me", key=full) + self.assertEqual(code, 200); self.assertEqual(j["balance_usd"], 20.0); self.assertEqual(j["key_id"], kid) + self.assertIn("rate_table", j); self.assertIn("keys_created", j) + # Only the hash is at rest. + row = self.s.app.db.one("select * from keys where key_id=?", kid) + self.assertEqual(row["key_hash"], E.sha256(full)); self.assertNotIn(full, json.dumps(dict(row))) + self.assertEqual(kid, E.sha256(full)[:8]) + # Bad / missing keys. + self.assertEqual(self.s.js("GET", "/v1/me", key="mde_" + "0" * 32)[0], 401) + self.assertEqual(self.s.js("GET", "/v1/me")[0], 401) + # Log output never contains the full key. + self.assertNotIn(full, self._log.getvalue()) + + def test_submit_closed_503_after_auth_and_balance(self): + full, _ = self.s.admin_key(20) + spec = {"input": "in.lmp", "gpu": "any", "estimate_s": 3600, "wall_limit_s": 7200} + code, j = self.s.js("POST", "/v1/jobs", body=spec, key=full) + self.assertEqual(code, 503); self.assertEqual(j["error"], "gpu_runners_open_soon"); self.assertIn("never expire", j["message"]) + # Auth failure wins over the flag, and so does an insufficient balance. + self.assertEqual(self.s.js("POST", "/v1/jobs", body=spec, key="mde_" + "f" * 32)[0], 401) + poor, _ = self.s.admin_key(0.10) + self.assertEqual(self.s.js("POST", "/v1/jobs", body=spec, key=poor)[0], 402) + self.assertEqual(self.s.js("GET", "/v1/jobs", key=full)[1], {"jobs": [], "next_cursor": None}) + +class TestPurchase(Base): + def welcome(self, sid): + code, body, _ = self.s.req("GET", "/welcome?session_id=" + sid); return code, body.decode() + + def test_first_purchase_creates_key(self): + fake_session("cs_test_1") + code, html = self.welcome("cs_test_1") + self.assertEqual(code, 200); self.assertIn("ForceField Silicon / MDEngine", html); self.assertIn("Shown once", html) + full = re.search(r"mdengine login (mde_[0-9a-f]{32})", html).group(1) + code, j = self.s.js("GET", "/v1/me", key=full) + self.assertEqual(code, 200); self.assertEqual(j["balance_usd"], 25.0) # 12.5 h * $2 + c = self.s.app.db.one("select * from credits where session_id='cs_test_1'") + self.assertEqual(c["gpu_s"], 45000); self.assertEqual(c["price_id"], PRICE_STARTER); self.assertEqual(c["key_id"], j["key_id"]) + self.assertEqual(self.s.app.db.key_by_id(j["key_id"])["email"], "buyer@example.test") + self.assertNotIn(full, self._log.getvalue()) + # Revisit: idempotent, no key shown. + code, html2 = self.welcome("cs_test_1") + self.assertEqual(code, 200); self.assertIn("Already issued", html2); self.assertNotIn(full, html2); self.assertIn(j["key_id"], html2) + self.assertEqual(self.s.app.db.one("select count(*) n from credits")["n"], 1) + self.assertEqual(self.s.app.db.one("select count(*) n from keys")["n"], 1) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 25.0) + + def test_coupon_discount_still_grants_full_hours(self): + fake_session("cs_coupon", price=PRICE_LAB, amount=1) # paid $0.01 with a coupon + code, html = self.welcome("cs_coupon") + full = re.search(r"mdengine login (mde_[0-9a-f]{32})", html).group(1) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 100.0) # 50 h * $2 + + def test_topup_existing_key_by_client_reference_id(self): + full, kid = self.s.admin_key(5) + fake_session("cs_topup", ref=kid) + code, html = self.welcome("cs_topup") + self.assertEqual(code, 200); self.assertIn("existing key", html); self.assertIn(kid, html); self.assertIn("$30.00", html) + self.assertNotRegex(html, r"mde_[0-9a-f]{32}") + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 30.0) + self.assertEqual(self.s.app.db.one("select count(*) n from keys")["n"], 1) + + def test_unknown_reference_id_creates_new_key(self): + fake_session("cs_badref", ref="deadbeef") + code, html = self.welcome("cs_badref") + self.assertEqual(code, 200); self.assertRegex(html, r"mde_[0-9a-f]{32}") + + def test_unpaid_and_missing_sessions(self): + fake_session("cs_unpaid", paid=False) + self.assertEqual(self.welcome("cs_unpaid")[0], 402) + self.assertEqual(self.welcome("cs_nope")[0], 404) + self.assertEqual(self.s.req("GET", "/welcome")[0], 400) + self.assertEqual(self.s.app.db.one("select count(*) n from credits")["n"], 0) + +class TestWebhook(Base): + def event(self, sid, etype="checkout.session.completed"): + return json.dumps({"id": "evt_1", "type": etype, "data": {"object": {"id": sid, "object": "checkout.session", "payment_status": "paid"}}}).encode() + + def post(self, payload, sig): + return self.s.js("POST", "/v1/stripe/webhook", raw=payload, headers={"Stripe-Signature": sig, "Content-Type": "application/json"}) + + def test_signature_reject(self): + fake_session("cs_wh"); p = self.event("cs_wh") + self.assertEqual(self.post(p, "")[0], 400) + self.assertEqual(self.post(p, E.stripe_sign(p, "whsec_wrong"))[0], 400) + self.assertEqual(self.post(p, E.stripe_sign(p, WEBHOOK_SECRET, ts=time.time() - 600))[0], 400) # stale + self.assertEqual(self.post(p + b" ", E.stripe_sign(p, WEBHOOK_SECRET))[0], 400) # tampered body + self.assertEqual(self.s.app.db.one("select count(*) n from credits")["n"], 0) + + def test_accept_then_welcome_reveals_once(self): + fake_session("cs_wh"); p = self.event("cs_wh") + code, j = self.post(p, E.stripe_sign(p, WEBHOOK_SECRET)) + self.assertEqual(code, 200); self.assertEqual(j["result"], "new") + # Duplicate delivery: 200, nothing double-credited. + code, j2 = self.post(p, E.stripe_sign(p, WEBHOOK_SECRET)) + self.assertEqual(code, 200); self.assertEqual(j2["result"], "already"); self.assertEqual(j2["key_id"], j["key_id"]) + self.assertEqual(self.s.app.db.one("select count(*) n from credits")["n"], 1) + # The buyer lands on /welcome after the webhook already ran: key shown exactly once. + code, body, _ = self.s.req("GET", "/welcome?session_id=cs_wh"); html = body.decode() + self.assertEqual(code, 200); full = re.search(r"mdengine login (mde_[0-9a-f]{32})", html).group(1) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 25.0) + self.assertEqual(self.s.app.db.one("select count(*) n from pending_keys")["n"], 0) + code, body, _ = self.s.req("GET", "/welcome?session_id=cs_wh") + self.assertIn("Already issued", body.decode()); self.assertNotIn(full, body.decode()) + + def test_welcome_then_webhook_is_idempotent(self): + fake_session("cs_both") + code, body, _ = self.s.req("GET", "/welcome?session_id=cs_both") + full = re.search(r"mdengine login (mde_[0-9a-f]{32})", body.decode()).group(1) + p = self.event("cs_both"); code, j = self.post(p, E.stripe_sign(p, WEBHOOK_SECRET)) + self.assertEqual(code, 200); self.assertEqual(j["result"], "already") + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 25.0) + self.assertEqual(self.s.app.db.one("select count(*) n from pending_keys")["n"], 0) + + def test_other_events_ignored(self): + p = self.event("cs_x", etype="payment_intent.succeeded") + code, j = self.post(p, E.stripe_sign(p, WEBHOOK_SECRET)); self.assertEqual(code, 200); self.assertEqual(j["ignored"], "payment_intent.succeeded") + +class TestAdmin(Base): + def run_admin(self, *args): + out = io.StringIO() + with redirect_stdout(out): A.main(["--db", self.s.cfg.db, "--env-file", "/nonexistent", *args]) + return out.getvalue() + + def test_key_new_and_credit_add(self): + full, kid = self.s.admin_key(10, email="a@example.test") + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 10.0) + out = self.run_admin("credit", "add", "--key", kid, "--usd", "15") + self.assertIn("balance $25.00", out) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 25.0) + lst = self.run_admin("key", "list"); self.assertIn(kid, lst); self.assertIn("a@example.test", lst); self.assertNotIn(full, lst) + led = self.run_admin("ledger"); self.assertEqual(led.count("admin-"), 2) + st = self.run_admin("stats"); self.assertIn("keys 1", st); self.assertIn("credited usd 25.00", st) + with self.assertRaises(SystemExit): self.run_admin("credit", "add", "--key", "nokey000", "--usd", "1") + + def test_admin_http_stats(self): + self.assertEqual(self.s.js("GET", "/v1/admin/stats")[0], 403) + code, j = self.s.js("GET", "/v1/admin/stats", key="admintok"); self.assertEqual(code, 200); self.assertEqual(j["keys"], 0) + +class TestJobFlowWhenOpen(Base): + """The mock's job logic, behind the flag, end to end with a stand-in pod.""" + runners_open = True + def launch_token(self): + with open(self.s.app.launch_env) as fh: return dict(l.split("=", 1) for l in fh.read().splitlines())["MDE_JOB_TOKEN"] + + def test_submit_upload_start_pod_done_billed(self): + full, kid = self.s.admin_key(20) + self.assertEqual(self.s.js("GET", "/v1/health")[1]["runners"], "open") + code, j = self.s.js("POST", "/v1/jobs", body={"input": "in.lmp", "gpu": "rtx4090", "estimate_s": 600, "label": "t"}, key=full) + self.assertEqual(code, 201); jid = j["id"]; self.assertRegex(jid, r"^MDJOB-\d{8}-[A-Z0-9]{6}$") + up = j["upload_url"]; self.assertIn("sig=", up) + # Unsigned PUT refused; signed PUT accepted; state -> uploaded. + self.assertEqual(self.s.req("PUT", up.split("?")[0], raw=b"x")[0], 403) + self.assertEqual(self.s.req("PUT", up, raw=b"deck-tarball")[0], 200) + self.assertEqual(self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1]["state"], "uploaded") + code, j = self.s.js("POST", f"/v1/jobs/{jid}/start", key=full); self.assertEqual(code, 202); self.assertEqual(j["state"], "queued") + tok = self.launch_token(); self.assertTrue(tok.startswith("jt_")) + self.assertEqual(self.s.app.db.job(jid)["token_hash"], E.sha256(tok)) + # Pod side. + self.assertEqual(self.s.js("GET", f"/internal/jobs/{jid}", key="jt_" + "0" * 32)[0], 401) + code, spec = self.s.js("GET", f"/internal/jobs/{jid}", key=tok) + self.assertEqual(code, 200); self.assertEqual(spec["input"], "in.lmp"); self.assertEqual(spec["launch"], "default") + self.assertEqual(self.s.req("GET", spec["input_url"])[1], b"deck-tarball") + self.assertEqual(self.s.js("POST", f"/internal/jobs/{jid}/heartbeat", body={"thermo_tail": ["Step Temp", "1 300"], "elapsed_s": 1}, key=tok)[0], 200) + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1]; self.assertEqual(st["state"], "running"); self.assertEqual(st["thermo_tail"], ["Step Temp", "1 300"]) + self.assertEqual(self.s.req("PUT", spec["results_put_url"], raw=b"results-tarball")[0], 200) + code, j = self.s.js("POST", f"/internal/jobs/{jid}/done", body={"exitcode": 0, "elapsed_s": 36, "results_bytes": 15}, key=tok) + self.assertEqual(code, 200); self.assertEqual(j["state"], "done") + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual(st["state"], "done"); self.assertEqual(st["billed_s"], 36); self.assertEqual(st["cost_usd"], 0.02); self.assertEqual(st["exitcode"], 0) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 19.98) + # Token invalidated after terminal state; results fetchable. + self.assertEqual(self.s.js("GET", f"/internal/jobs/{jid}", key=tok)[0], 401) + code, r = self.s.js("GET", f"/v1/jobs/{jid}/results", key=full); self.assertEqual(code, 200); self.assertEqual(r["bytes"], 15) + self.assertEqual(self.s.req("GET", r["download_url"])[1], b"results-tarball") + lst = self.s.js("GET", "/v1/jobs", key=full)[1]; self.assertEqual([x["id"] for x in lst["jobs"]], [jid]) + self.assertEqual(self.s.js("DELETE", f"/v1/jobs/{jid}", key=full)[0], 409) + # Ledger from the CLI agrees. + out = io.StringIO() + with redirect_stdout(out): A.main(["--db", self.s.cfg.db, "--env-file", "/nonexistent", "ledger", "--key", kid]) + self.assertIn(jid, out.getvalue()); self.assertIn("-$ 0.0200", out.getvalue()) + + def test_pod_lost_not_billed_and_cancel(self): + full, kid = self.s.admin_key(20) + jid = self.s.js("POST", "/v1/jobs", body={"input": "in.lmp"}, key=full)[1]["id"] + up = self.s.app.blob_url(f"{jid}.in.tar.gz")[0]; self.s.req("PUT", up, raw=b"x"); self.s.js("POST", f"/v1/jobs/{jid}/start", key=full) + tok = self.launch_token() + self.s.js("POST", f"/internal/jobs/{jid}/heartbeat", body={"thermo_tail": [], "elapsed_s": 0}, key=tok) + self.s.app.db.set_job(jid, last_hb=time.time() - 1000, started="2020-01-01T00:00:00Z") + self.s.app.watchdog_once() + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual((st["state"], st["error"], st["cost_usd"]), ("failed", "pod_lost", 0.0)) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 20.0) + # Cancel a fresh job. + j2 = self.s.js("POST", "/v1/jobs", body={"input": "in.lmp"}, key=full)[1]["id"] + code, st = self.s.js("DELETE", f"/v1/jobs/{j2}", key=full); self.assertEqual(code, 202); self.assertEqual(st["state"], "cancelled") + +class TestLauncherFlow(Base): + """Job lifecycle with a FakeLauncher injected: pods are created at start and gone at every terminal state.""" + runners_open = True; fake_launcher = True + + def started_job(self, credit=20): + full, kid = self.s.admin_key(credit) + jid = self.s.js("POST", "/v1/jobs", body={"input": "in.lmp", "gpu": "rtx4090", "wall_limit_s": 3600}, key=full)[1]["id"] + self.s.req("PUT", self.s.app.blob_url(f"{jid}.in.tar.gz")[0], raw=b"deck") + code, j = self.s.js("POST", f"/v1/jobs/{jid}/start", key=full); self.assertEqual((code, j["state"]), (202, "queued")) + self.assertTrue(self.s.app.join_bg()); return full, jid + + def pod_token(self, jid): + pod = self.s.app.db.job(jid)["pod_id"]; return self.fake.pods[pod]["env"]["MDE_JOB_TOKEN"] + + def test_start_launches_pod_then_done_deletes_it(self): + self.assertEqual(self.s.js("GET", "/v1/health")[1]["launcher"], "fake") + full, jid = self.started_job() + self.assertFalse(os.path.exists(self.s.app.launch_env)) # no dev launch.env with a launcher + row = self.s.app.db.job(jid); st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual(st["state"], "launching"); self.assertEqual(st["pod_id"], row["pod_id"]); self.assertTrue(row["launched_at"]) + pod = self.fake.pods[row["pod_id"]] + self.assertEqual(pod["name"], "mde-" + jid); self.assertEqual(pod["env"]["MDE_JOB_ID"], jid); self.assertEqual(pod["env"]["MDE_ENDPOINT"], self.s.base) + tok = pod["env"]["MDE_JOB_TOKEN"]; self.assertTrue(tok.startswith("jt_")); self.assertEqual(row["token_hash"], E.sha256(tok)) + self.assertNotIn(tok, self._log.getvalue()) + # Pod boots, heartbeats -> running; results; done -> pod deleted in the background. + code, spec = self.s.js("GET", f"/internal/jobs/{jid}", key=tok); self.assertEqual(code, 200); self.assertEqual(spec["wall_limit_s"], 3600) + self.s.js("POST", f"/internal/jobs/{jid}/heartbeat", body={"thermo_tail": ["x"], "elapsed_s": 1}, key=tok) + self.assertEqual(self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1]["state"], "running") + self.s.req("PUT", spec["results_put_url"], raw=b"out") + code, j = self.s.js("POST", f"/internal/jobs/{jid}/done", body={"exitcode": 0, "elapsed_s": 10, "results_bytes": 3}, key=tok) + self.assertEqual((code, j["state"]), (200, "done")); self.assertTrue(self.s.app.join_bg()) + self.assertEqual(self.fake.deleted, [row["pod_id"]]); self.assertEqual(self.fake.pods, {}) + self.assertIn('"ev":"pod.deleted"', self._log.getvalue()) + self.assertEqual(self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1]["billed_s"], 10) + + def test_cancel_deletes_pod(self): + full, jid = self.started_job(); pod = self.s.app.db.job(jid)["pod_id"] + code, st = self.s.js("DELETE", f"/v1/jobs/{jid}", key=full); self.assertEqual((code, st["state"]), (202, "cancelled")) + self.assertTrue(self.s.app.join_bg()); self.assertEqual(self.fake.deleted, [pod]); self.assertEqual(self.fake.pods, {}) + + def test_pod_lost_deletes_pod_unbilled(self): + full, jid = self.started_job(); pod = self.s.app.db.job(jid)["pod_id"]; tok = self.pod_token(jid) + self.s.js("POST", f"/internal/jobs/{jid}/heartbeat", body={"thermo_tail": [], "elapsed_s": 0}, key=tok) + self.s.app.db.set_job(jid, last_hb=time.time() - 1000, started="2020-01-01T00:00:00Z") + self.s.app.watchdog_once(); self.assertTrue(self.s.app.join_bg()) + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual((st["state"], st["error"], st["cost_usd"]), ("failed", "pod_lost", 0.0)) + self.assertEqual(self.fake.deleted, [pod]); self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 20.0) + + def test_launch_timeout_is_no_capacity_and_deletes_pod(self): + full, jid = self.started_job(); pod = self.s.app.db.job(jid)["pod_id"] + self.s.app.watchdog_once(); self.assertEqual(self.s.app.db.job(jid)["state"], "launching") # fresh: untouched + self.s.app.db.set_job(jid, launched_at="2020-01-01T00:00:00Z") + self.s.app.watchdog_once(); self.assertTrue(self.s.app.join_bg()) + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual((st["state"], st["error"], st["cost_usd"]), ("failed", "no_capacity", 0.0)) + self.assertEqual(self.fake.deleted, [pod]); self.assertIsNone(self.s.app.db.job(jid)["token_hash"]) + self.assertIn('"ev":"job.launch_timeout"', self._log.getvalue()) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 20.0) + + def test_no_capacity_from_launcher(self): + self.fake.fail_create = True + full, jid = self.started_job() + st = self.s.js("GET", f"/v1/jobs/{jid}", key=full)[1] + self.assertEqual((st["state"], st["error"], st["cost_usd"], st["pod_id"]), ("failed", "no_capacity", 0.0, None)) + self.assertEqual(self.fake.pods, {}); self.assertIsNone(self.s.app.db.job(jid)["token_hash"]) + self.assertIn('"ev":"job.no_capacity"', self._log.getvalue()) + self.assertEqual(self.s.js("GET", "/v1/me", key=full)[1]["balance_usd"], 20.0) + + def test_reaper(self): + full, jid = self.started_job(); live = self.s.app.db.job(jid)["pod_id"] + orphan = self.fake.add_pod("mde-MDJOB-20260101-NOJOB1") + other = self.fake.add_pod("someone-elses-pod") # not ours: never touched + # A terminal job whose pod delete failed earlier (simulate by re-adding the pod after cancel). + full2, jid2 = self.started_job(); self.s.js("DELETE", f"/v1/jobs/{jid2}", key=full2); self.s.app.join_bg() + stale = self.fake.add_pod("mde-" + jid2, pod_id=self.s.app.db.job(jid2)["pod_id"]); self.fake.deleted.clear() + # A running job whose pod is older than wall_limit_s + 20 min. + full3, jid3 = self.started_job(); old = self.s.app.db.job(jid3)["pod_id"] + self.fake.pods[old]["createdAt"] = "2020-01-01T00:00:00.000Z" + n = self.s.app.reaper_once() + self.assertEqual(n, 3); self.assertEqual(sorted(self.fake.deleted), sorted([orphan, stale, old])) + self.assertIn(live, self.fake.pods); self.assertIn(other, self.fake.pods) + logs = self._log.getvalue() + for reason in ("no_job", "job_terminal", "overage"): self.assertIn('"reason":"%s"' % reason, logs) + # Stuck pod: three failed passes -> reaper.stuck. + stuck = self.fake.add_pod("mde-MDJOB-20260101-STUCK1"); self.fake.fail_delete.add(stuck) + for _ in range(3): self.s.app.reaper_once() + self.assertEqual(self.s.app.reaper_fail[stuck], 3); self.assertIn('"ev":"reaper.stuck"', self._log.getvalue()) + self.fake.fail_delete.clear(); self.s.app.reaper_once(); self.assertNotIn(stuck, self.s.app.reaper_fail) + +class TestRunPodLauncher(unittest.TestCase): + """Request shaping against a monkeypatched transport; no network.""" + KEY = "rpa_TESTKEY_" + "z" * 24 + def launcher(self, **env): + base = {"RUNPOD_API_KEY": self.KEY, "MDE_PUBLIC_URL": "https://api.example.test/", "MDE_DB": "/nonexistent/x.sqlite"} + base.update(env); return L.RunPodLauncher(E.Config(base)) + + def test_create_body_and_ladder(self): + lc = self.launcher(); calls = [] + def fake_request(method, path, body=None): + calls.append((method, path, body)); return (500, '{"error":"no capacity"}') if len(calls) == 1 else (201, '{"id":"pod123","name":"x"}') + lc._request = fake_request + log = io.StringIO() + with redirect_stdout(log): pid = lc.create("MDJOB-20260906-ABC123", "jt_" + "0" * 32, 3600, "rtx4090") + self.assertEqual(pid, "pod123"); self.assertEqual(len(calls), 2) + for (m, p, b), (cloud, gid) in zip(calls, L.DEFAULT_LADDER): + self.assertEqual((m, p), ("POST", "/v2/pods")); self.assertEqual(b["cloud"], cloud) + self.assertEqual(b["gpu"], {"id": gid, "count": 1, "minCudaVersion": "12.4"}); self.assertEqual(b["disk"], 20) + self.assertEqual(b["name"], "mde-MDJOB-20260906-ABC123"); self.assertEqual(b["image"], L.DEFAULT_IMAGE) + self.assertEqual(b["env"], {"MDE_ENDPOINT": "https://api.example.test", "MDE_JOB_ID": "MDJOB-20260906-ABC123", "MDE_JOB_TOKEN": "jt_" + "0" * 32, "MDE_WALL_LIMIT_S": "3600"}) + self.assertNotIn("dataCenterIds", b) + out = log.getvalue(); self.assertEqual(out.count('"ev":"launch.attempt"'), 2); self.assertIn('"code":500', out) + self.assertNotIn(self.KEY, out); self.assertNotIn("jt_" + "0" * 32, out) + + def test_config_overrides(self): + lc = self.launcher(MDE_RUNNER_IMAGE="ghcr.io/x/y:z", MDE_POD_DISK_GB="40", MDE_MIN_CUDA="12.8", + MDE_GPU_LADDER="SECURE:NVIDIA A100 80GB PCIe, community:NVIDIA GeForce RTX 4090") + self.assertEqual(lc.ladder, [("SECURE", "NVIDIA A100 80GB PCIe"), ("COMMUNITY", "NVIDIA GeForce RTX 4090")]) + b = lc.pod_body("J", "t", "SECURE", "NVIDIA A100 80GB PCIe") + self.assertEqual((b["image"], b["disk"], b["gpu"]["minCudaVersion"]), ("ghcr.io/x/y:z", 40, "12.8")) + with self.assertRaises(ValueError): L.parse_ladder("PRIVATE:foo") + with self.assertRaises(ValueError): L.parse_ladder("nocolon") + + def test_all_rungs_fail_is_no_capacity(self): + lc = self.launcher(); lc._request = lambda m, p, body=None: (422, '{"error":"bad"}') + with redirect_stdout(io.StringIO()), self.assertRaises(L.NoCapacity): lc.create("J", "t", 60, "any") + # 201 without an id is also a failed rung, and its body is not echoed. + lc._request = lambda m, p, body=None: (201, '{"env":{"MDE_JOB_TOKEN":"jt_secret"}}'); out = io.StringIO() + with redirect_stdout(out), self.assertRaises(L.NoCapacity): lc.create("J", "t", 60, "any") + self.assertNotIn("jt_secret", out.getvalue()) + + def test_delete_semantics(self): + lc = self.launcher(); lc.backoff_s = 0; seq = [] + def scripted(codes): + it = iter(codes) + def f(m, p, body=None): + seq.append((m, p)); return next(it), "{}" + return f + lc._request = scripted([204]); self.assertTrue(lc.delete("p1")) + lc._request = scripted([404]); self.assertTrue(lc.delete("p1")) + lc._request = scripted([429, 500, 204]); self.assertTrue(lc.delete("p1")) + lc._request = scripted([500, 500, 500]) + with self.assertRaises(L.LauncherError): lc.delete("p1") + seq.clear(); lc._request = scripted([403]) + with self.assertRaises(L.LauncherError): lc.delete("p1") + self.assertEqual(seq, [("DELETE", "/v2/pods/p1")]) # other 4xx: no retry + + def test_get_and_list(self): + lc = self.launcher() + lc._request = lambda m, p, body=None: (200, '{"items":[{"id":"a","name":"mde-x"}]}') + self.assertEqual([p["id"] for p in lc.list_pods()], ["a"]) + lc._request = lambda m, p, body=None: (200, '[{"id":"b"}]'); self.assertEqual(lc.list_pods()[0]["id"], "b") + lc._request = lambda m, p, body=None: (200, '{"pods":[{"id":"c","name":"mde-y"}]}') # LIVE shape (2026-09-06) + self.assertEqual([p["id"] for p in lc.list_pods()], ["c"]) + lc._request = lambda m, p, body=None: (404, ""); self.assertIsNone(lc.get("zz")) + lc._request = lambda m, p, body=None: (200, '{"id":"zz","status":"RUNNING"}'); self.assertEqual(lc.get("zz")["status"], "RUNNING") + lc._request = lambda m, p, body=None: (500, "boom") + with self.assertRaises(L.LauncherError): lc.list_pods() + + def test_authorization_header_and_url(self): + lc = self.launcher(); seen = [] + def fake_open(req): + seen.append(req); return 204, "" + lc._open = fake_open; lc.delete("p9") + req = seen[0]; self.assertEqual(req.full_url, "https://api.runpod.io/v2/pods/p9"); self.assertEqual(req.get_method(), "DELETE") + auth = req.get_header("Authorization"); self.assertTrue(auth and auth.startswith("Bearer ") and len(auth) == len("Bearer ") + len(self.KEY)) + self.assertEqual(req.timeout if hasattr(req, "timeout") else lc.timeout, lc.timeout) + with self.assertRaises(ValueError): L.RunPodLauncher(E.Config({"RUNPOD_API_KEY": "", "MDE_DB": "/nonexistent/x"})) + + def test_pod_age(self): + self.assertAlmostEqual(L.pod_age_s({"createdAt": "2020-01-01T00:00:00.000Z"}, now_ts=1577836800 + 90), 90, places=3) + self.assertAlmostEqual(L.pod_age_s({"createdAt": "2020-01-01T01:00:00+01:00"}, now_ts=1577836800 + 5), 5, places=3) + self.assertIsNone(L.pod_age_s({"createdAt": "garbage"})); self.assertIsNone(L.pod_age_s({})) + +class TestUnits(unittest.TestCase): + def test_schema_migration_adds_pod_columns(self): + d = tempfile.mkdtemp(prefix="mde-mig-"); path = os.path.join(d, "old.sqlite") + import sqlite3 + c = sqlite3.connect(path) + c.executescript(E.SCHEMA.replace(", pod_id text, launched_at text", "")); c.close() + db = E.DB(path); cols = {r["name"] for r in db.q("pragma table_info(jobs)")} + self.assertIn("pod_id", cols); self.assertIn("launched_at", cols) + db2 = E.DB(path); db.c.close(); db2.c.close(); shutil.rmtree(d, ignore_errors=True) # second open: duplicate column ignored + + def test_signature_roundtrip(self): + p = b'{"a":1}'; h = E.stripe_sign(p, "s") + self.assertTrue(E.verify_stripe_signature(p, h, "s")); self.assertFalse(E.verify_stripe_signature(p, h, "t")) + self.assertFalse(E.verify_stripe_signature(p, "t=abc,v1=00", "s")); self.assertFalse(E.verify_stripe_signature(p, h, "")) + def test_parse_packs_rates(self): + self.assertEqual(E.parse_packs("price_a:25:12.5, price_b:100:50"), {"price_a": (25.0, 12.5), "price_b": (100.0, 50.0)}) + self.assertEqual(E.parse_rates(""), {"any": 2.0, "rtx4090": 2.0}); self.assertEqual(E.parse_rates("any:2,a100:4"), {"any": 2.0, "a100": 4.0}) + def test_env_file(self): + with tempfile.NamedTemporaryFile("w", suffix=".env", delete=False) as f: f.write("# c\nX_MDE_T=\"v 1\"\nY_MDE_T=2\n") + os.environ.pop("X_MDE_T", None); os.environ["Y_MDE_T"] = "keep"; E.load_env_file(f.name); os.unlink(f.name) + self.assertEqual(os.environ["X_MDE_T"], "v 1"); self.assertEqual(os.environ["Y_MDE_T"], "keep") + +if __name__ == "__main__": + unittest.main() diff --git a/hosted/mock/mock_endpoint.py b/hosted/mock/mock_endpoint.py new file mode 100644 index 0000000..73488aa --- /dev/null +++ b/hosted/mock/mock_endpoint.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Local mock of the MDEngine hosted endpoint (hosted/CONTRACT.md v1). Stdlib only. + +Purpose: let the pull-runner (docker/runner-gpu/runner.sh) and the clients (CLI/MCP/app) be +developed and tested offline. Same routes and JSON as the real endpoint; storage = sqlite + a local +blob dir served at /blob/ (GET/PUT) standing in for presigned object-storage URLs. NOT the +production server: no TLS, no key hashing, no pod launcher (a "launch" here just marks the job +queued and prints the env a pod would get). + + python3 mock_endpoint.py --port 8787 --data /tmp/mde-mock + curl -H 'Authorization: Bearer mde_test' localhost:8787/v1/me +""" +import argparse, json, os, secrets, sqlite3, time, uuid +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlparse + +RATES = {"any": 2.0, "rtx4090": 2.0, "a100": 4.0} # $/GPU-h placeholder until GJOB-092 +STATES = "created uploaded queued launching running uploading done failed cancelled".split() + +def now(): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) +def job_id(): return "MDJOB-%s-%s" % (time.strftime("%Y%m%d", time.gmtime()), secrets.token_hex(3).upper()) + +class DB: + def __init__(self, path): + self.c = sqlite3.connect(path, check_same_thread=False); self.c.row_factory = sqlite3.Row + self.c.executescript(""" + create table if not exists keys(key text primary key, balance_usd real, created text); + create table if not exists jobs(id text primary key, key text, token text, spec text, state text, + created text, started text, finished text, gpu text, rate real, billed_s int default 0, + thermo text default '[]', exitcode int, error text, attempt int default 1, last_hb real); + """) + if not self.c.execute("select 1 from keys where key='mde_test'").fetchone(): + self.c.execute("insert into keys values('mde_test', 20.0, ?)", (now(),)); self.c.commit() + def key(self, k): return self.c.execute("select * from keys where key=?", (k,)).fetchone() + def job(self, jid): return self.c.execute("select * from jobs where id=?", (jid,)).fetchone() + def set(self, jid, **kw): + cols = ", ".join(f"{k}=?" for k in kw); self.c.execute(f"update jobs set {cols} where id=?", (*kw.values(), jid)); self.c.commit() + +def status(j): + spec = json.loads(j["spec"]) + billed = j["billed_s"] + if j["state"] == "running" and j["started"]: + billed = int(time.time() - time.mktime(time.strptime(j["started"], "%Y-%m-%dT%H:%M:%SZ")) + time.timezone) + return {"id": j["id"], "state": j["state"], "states": "|".join(STATES), "created": j["created"], + "started": j["started"], "finished": j["finished"], "gpu": j["gpu"], "rate_usd_per_h": j["rate"], + "billed_s": billed, "cost_usd": round(billed * j["rate"] / 3600, 4), "thermo_tail": json.loads(j["thermo"]), + "exitcode": j["exitcode"], "error": j["error"], "attempt": j["attempt"], "label": spec.get("label")} + +class H(BaseHTTPRequestHandler): + db: DB; data: str; base: str + def log_message(self, *a): pass + def send(self, code, obj=None, raw=None, ctype="application/json"): + body = raw if raw is not None else (json.dumps(obj).encode() if obj is not None else b"") + self.send_response(code); self.send_header("Content-Type", ctype); self.send_header("Content-Length", str(len(body))) + self.end_headers(); self.wfile.write(body) + def body(self): + n = int(self.headers.get("Content-Length") or 0); return self.rfile.read(n) if n else b"" + def bearer(self): + a = self.headers.get("Authorization", ""); return a[7:] if a.startswith("Bearer ") else None + def blob(self, name): return os.path.join(self.data, "blobs", name) + + def do_GET(self): + p = urlparse(self.path).path.rstrip("/"); parts = p.split("/") + if parts[1] == "blob": + f = self.blob(parts[2]) + if not os.path.exists(f): return self.send(404, {"error": "no blob"}) + with open(f, "rb") as fh: return self.send(200, raw=fh.read(), ctype="application/gzip") + if parts[1] == "internal": # pod side + j = self.db.job(parts[3]) if len(parts) > 3 else None + if not j or self.bearer() != j["token"]: return self.send(401, {"error": "bad token"}) + spec = json.loads(j["spec"]) + return self.send(200, {"input_url": f"{self.base}/blob/{j['id']}.in.tar.gz", "input": spec["input"], + "launch": spec.get("launch", "default"), "wall_limit_s": spec.get("wall_limit_s", 86400), + "results_put_url": f"{self.base}/blob/{j['id']}.out.tar.gz"}) + k = self.db.key(self.bearer() or "") + if not k: return self.send(401, {"error": "bad api key"}) + if p == "/v1/me": return self.send(200, {"balance_usd": k["balance_usd"], "rate_table": RATES, "keys_created": k["created"]}) + if p == "/v1/jobs": + rows = self.db.c.execute("select * from jobs where key=? order by created desc limit 50", (k["key"],)).fetchall() + return self.send(200, {"jobs": [status(j) for j in rows]}) + if len(parts) >= 4 and parts[2] == "jobs": + j = self.db.job(parts[3]) + if not j or j["key"] != k["key"]: return self.send(404, {"error": "no job"}) + if len(parts) == 4: return self.send(200, status(j)) + if parts[4] == "results": + if j["state"] not in ("done", "failed"): return self.send(409, {"error": "not finished"}) + f = self.blob(f"{j['id']}.out.tar.gz"); sz = os.path.getsize(f) if os.path.exists(f) else 0 + return self.send(200, {"download_url": f"{self.base}/blob/{j['id']}.out.tar.gz", "expires": None, "bytes": sz}) + self.send(404, {"error": "no route"}) + + def do_PUT(self): + parts = urlparse(self.path).path.split("/") + if parts[1] != "blob": return self.send(404, {"error": "no route"}) + os.makedirs(os.path.dirname(self.blob("x")), exist_ok=True) + with open(self.blob(parts[2]), "wb") as fh: fh.write(self.body()) + jid = parts[2].split(".")[0] + if parts[2].endswith(".in.tar.gz") and self.db.job(jid) and self.db.job(jid)["state"] == "created": self.db.set(jid, state="uploaded") + self.send(200, {"ok": True}) + + def do_POST(self): + p = urlparse(self.path).path.rstrip("/"); parts = p.split("/") + if parts[1] == "internal": # pod side: heartbeat / done + j = self.db.job(parts[3]) + if not j or self.bearer() != j["token"]: return self.send(401, {"error": "bad token"}) + b = json.loads(self.body() or b"{}") + if parts[4] == "heartbeat": + kw = {"thermo": json.dumps(b.get("thermo_tail", [])[-20:]), "last_hb": time.time()} + if j["state"] in ("queued", "launching"): kw.update(state="running", started=now()) + self.db.set(j["id"], **kw); return self.send(200, {"ok": True}) + if parts[4] == "done": + rc = int(b.get("exitcode", 1)); err = b.get("error") + started = j["started"] or now(); billed = int(b.get("elapsed_s", 0)) + have = os.path.exists(self.blob(f"{j['id']}.out.tar.gz")) + state = "done" if rc == 0 and have else "failed" + if not have and not err: err = "no_results" + self.db.set(j["id"], state=state, finished=now(), started=started, exitcode=rc, error=err, billed_s=billed, token=secrets.token_hex(4)) # token invalidated + cost = billed * j["rate"] / 3600 + if err not in ("pod_lost",): self.db.c.execute("update keys set balance_usd=balance_usd-? where key=?", (cost, j["key"])); self.db.c.commit() + print(f"[mock] {j['id']} -> {state} rc={rc} err={err} billed={billed}s cost=${cost:.4f}", flush=True) + return self.send(200, {"ok": True, "state": state}) + k = self.db.key(self.bearer() or "") + if not k: return self.send(401, {"error": "bad api key"}) + if p == "/v1/jobs": + spec = json.loads(self.body() or b"{}") + if "input" not in spec: return self.send(400, {"error": "input required"}) + gpu = spec.get("gpu", "any"); rate = RATES.get(gpu) + if rate is None: return self.send(400, {"error": "unknown gpu"}) + if k["balance_usd"] < rate * max(int(spec.get("estimate_s", 0)), 900) / 3600: return self.send(402, {"error": "insufficient balance"}) + jid = job_id(); tok = "jt_" + secrets.token_hex(16) + self.db.c.execute("insert into jobs(id,key,token,spec,state,created,gpu,rate) values(?,?,?,?,?,?,?,?)", + (jid, k["key"], tok, json.dumps(spec), "created", now(), gpu, rate)); self.db.c.commit() + return self.send(201, {"id": jid, "upload_url": f"{self.base}/blob/{jid}.in.tar.gz", "upload_expires": None}) + if len(parts) == 5 and parts[2] == "jobs" and parts[4] == "start": + j = self.db.job(parts[3]) + if not j or j["key"] != k["key"]: return self.send(404, {"error": "no job"}) + if j["state"] != "uploaded": return self.send(409, {"error": f"state is {j['state']}"}) + self.db.set(j["id"], state="queued") + # The real endpoint launches a pod here. The mock prints what the launcher would inject. + print(f"[mock] LAUNCH {j['id']}: MDE_ENDPOINT={self.base} MDE_JOB_ID={j['id']} MDE_JOB_TOKEN={j['token']}", flush=True) + with open(os.path.join(self.data, "launch.env"), "w") as fh: + fh.write(f"MDE_ENDPOINT={self.base}\nMDE_JOB_ID={j['id']}\nMDE_JOB_TOKEN={j['token']}\n") + return self.send(202, {"id": j["id"], "state": "queued"}) + self.send(404, {"error": "no route"}) + + def do_DELETE(self): + parts = urlparse(self.path).path.rstrip("/").split("/") + k = self.db.key(self.bearer() or "") + if not k or len(parts) != 4: return self.send(401, {"error": "bad api key"}) + j = self.db.job(parts[3]) + if not j or j["key"] != k["key"]: return self.send(404, {"error": "no job"}) + if j["state"] in ("done", "failed", "cancelled"): return self.send(409, {"error": "terminal"}) + self.db.set(j["id"], state="cancelled", finished=now(), error="cancelled"); self.send(202, {"id": j["id"], "state": "cancelled"}) + +if __name__ == "__main__": + ap = argparse.ArgumentParser(); ap.add_argument("--port", type=int, default=8787); ap.add_argument("--data", default="/tmp/mde-mock") + a = ap.parse_args(); os.makedirs(os.path.join(a.data, "blobs"), exist_ok=True) + H.db = DB(os.path.join(a.data, "mock.sqlite")); H.data = a.data; H.base = f"http://127.0.0.1:{a.port}" + print(f"[mock] listening on {H.base} data={a.data} test key: mde_test ($20)", flush=True) + ThreadingHTTPServer(("127.0.0.1", a.port), H).serve_forever()