From 53662ff8048b9230b26c403c6fd84200c6c159f7 Mon Sep 17 00:00:00 2001 From: Travis James Date: Wed, 26 Aug 2026 10:30:39 -0500 Subject: [PATCH] feat(embeddings): add native MLX executor --- .../surreal-memory/src/embeddings/candle.rs | 37 ++- docs/lessons.md | 5 + executors/mlx/.gitignore | 1 + executors/mlx/Package.resolved | 114 +++++++ executors/mlx/Package.swift | 51 ++++ .../MLXEmbeddingService.swift | 282 ++++++++++++++++++ .../SurrealMemoryMLXExecutor/main.swift | 123 ++++++++ .../Planner.swift | 93 ++++++ .../ProgressHeartbeat.swift | 45 +++ .../Protocol.swift | 160 ++++++++++ .../ProtocolTests.swift | 102 +++++++ .../mlx/Tests/fixtures/parity-corpus.json | 26 ++ executors/mlx/scripts/compare_embeddings.py | 170 +++++++++++ src/config.rs | 47 ++- src/executor.rs | 245 +++++++++++++-- src/main.rs | 108 ++++++- tests/executor_recovery.rs | 36 +++ 17 files changed, 1608 insertions(+), 37 deletions(-) create mode 100644 executors/mlx/.gitignore create mode 100644 executors/mlx/Package.resolved create mode 100644 executors/mlx/Package.swift create mode 100644 executors/mlx/Sources/SurrealMemoryMLXExecutor/MLXEmbeddingService.swift create mode 100644 executors/mlx/Sources/SurrealMemoryMLXExecutor/main.swift create mode 100644 executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Planner.swift create mode 100644 executors/mlx/Sources/SurrealMemoryMLXExecutorCore/ProgressHeartbeat.swift create mode 100644 executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Protocol.swift create mode 100644 executors/mlx/Tests/SurrealMemoryMLXExecutorCoreTests/ProtocolTests.swift create mode 100644 executors/mlx/Tests/fixtures/parity-corpus.json create mode 100755 executors/mlx/scripts/compare_embeddings.py diff --git a/crates/surreal-memory/src/embeddings/candle.rs b/crates/surreal-memory/src/embeddings/candle.rs index 7a53217..185eb60 100644 --- a/crates/surreal-memory/src/embeddings/candle.rs +++ b/crates/surreal-memory/src/embeddings/candle.rs @@ -31,6 +31,7 @@ struct CandleEmbeddingsInner { pub struct CandleEmbeddings { inner: OnceCell>>, model_id: String, + model_revision: String, cache_dir: String, expected_dimensions: usize, } @@ -52,6 +53,8 @@ impl CandleEmbeddings { Ok(Self { inner: OnceCell::new(), model_id: model_id.to_string(), + model_revision: std::env::var("LOCAL_EMBEDDING_MODEL_REVISION") + .unwrap_or_else(|_| "main".to_string()), cache_dir: cache_dir.to_string(), expected_dimensions, }) @@ -98,7 +101,7 @@ impl CandleEmbeddings { tracing::info!("Using device: {:?}", device); let (config_path, tokenizer_path, weights_path) = - Self::download_model(&self.model_id, &self.cache_dir) + Self::download_model(&self.model_id, &self.model_revision, &self.cache_dir) .await .context("Failed to download model files")?; @@ -195,6 +198,12 @@ impl CandleEmbeddings { } fn get_device() -> Result { + let device_preference = std::env::var("LOCAL_EMBEDDING_DEVICE").ok(); + if force_cpu(device_preference.as_deref())? { + tracing::warn!("LOCAL_EMBEDDING_DEVICE=cpu: using the explicit degraded CPU backend"); + return Ok(Device::Cpu); + } + #[cfg(feature = "cuda")] { if candle_core::utils::cuda_is_available() { @@ -237,6 +246,7 @@ impl CandleEmbeddings { async fn download_model( model_id: &str, + model_revision: &str, cache_dir: &str, ) -> Result<(PathBuf, PathBuf, PathBuf)> { // hf-hub keeps repositories under `/hub`: both @@ -251,8 +261,9 @@ impl CandleEmbeddings { let hub_dir = Self::hub_cache_dir(cache_dir); tracing::info!( - "Resolving model from Hugging Face: {} (cache: {})", + "Resolving model from Hugging Face: {}@{} (cache: {})", model_id, + model_revision, hub_dir.display() ); @@ -263,7 +274,11 @@ impl CandleEmbeddings { .with_cache_dir(hub_dir) .build() .context("build Hugging Face API client")?; - let repo = api.repo(Repo::new(model_id.to_string(), RepoType::Model)); + let repo = api.repo(Repo::with_revision( + model_id.to_string(), + RepoType::Model, + model_revision.to_string(), + )); // Download config, tokenizer, and weights concurrently. The weights // future tries safetensors first and falls back to the PyTorch file. @@ -444,6 +459,14 @@ impl CandleEmbeddings { } } +fn force_cpu(value: Option<&str>) -> Result { + match value.unwrap_or("auto").trim().to_ascii_lowercase().as_str() { + "auto" => Ok(false), + "cpu" => Ok(true), + other => anyhow::bail!("LOCAL_EMBEDDING_DEVICE must be 'auto' or 'cpu', got '{other}'"), + } +} + #[async_trait] impl EmbeddingService for CandleEmbeddings { async fn embed(&self, text: &str) -> Result { @@ -643,6 +666,14 @@ mod tests { assert_eq!(estimate("sentence-transformers/all-MiniLM-L6-v2"), 384); } + #[test] + fn device_preference_is_explicit_and_fail_closed() { + assert!(!force_cpu(None).unwrap()); + assert!(!force_cpu(Some("auto")).unwrap()); + assert!(force_cpu(Some("CPU")).unwrap()); + assert!(force_cpu(Some("metal")).is_err()); + } + #[test] fn model_boundary_guard_accepts_below_and_at_capacity_only() { assert!(validate_model_input_len(510, 512).is_ok()); diff --git a/docs/lessons.md b/docs/lessons.md index a39867e..c4ef556 100644 --- a/docs/lessons.md +++ b/docs/lessons.md @@ -37,6 +37,11 @@ Format: `YYYY-MM-DD — Rule — *(context: what went wrong)*` - 2026-08-26 — Check `uptime` before diagnosing a timing bug. Metal init measured 1.2s, 14s, 36s, then never — I attributed that to code before noticing load average 269 and a static `/health` handler taking 21s. Timing symptoms on a saturated machine are not evidence about code. - 2026-08-26 — A watchdog must not cover a phase in which the watched process is structurally incapable of reporting progress. The executor child could not heartbeat until after its first request read, yet the 30s per-request watchdog was armed at spawn. Startup needs its own budget and an explicit readiness signal. - 2026-08-26 — Never stack a PR onto another PR's branch when the base will merge first. PR #9 targeted PR #8's branch; #8 merged to main, then #9 merged into an orphaned branch and GitHub reported "MERGED" while main had none of it. Target `main` and rebase, or the fix silently never ships. +- 2026-08-26 — When adding a typed field to a struct that derives Serde, derive the same Serde direction on the field type before the first compile. *(Context: `LocalEmbeddingBackend` was added to deserializable `Config` without `Deserialize`.)* +- 2026-08-26 — Do not run the aggregate `prometheus-rust-auditor audit` command in this repository: its pipeline generates a GitHub Actions workflow, which violates the local-only validation policy. Run the read-only enforcement, format, dependency, inventory, and partition checks individually. +- 2026-08-26 — A copied SwiftPM executable is not self-contained when a C target ships resources: install `mlx-swift_Cmlx.bundle` beside the executable and smoke-test the copied path, because the build-tree binary can hide a missing `default.metallib` deployment. +- 2026-08-26 — Persisted executor generations span server processes, but a newly pre-warmed child starts with process-local numbering. Adopt the healthy child into the durable generation sequence; do not kill it merely because its initial number is lower, or queue recovery turns into an unnecessary cold model launch. +- 2026-08-26 — A liveness heartbeat must run independently of the work it supervises. An async Swift `Task` heartbeat can be starved while MLX/Metal synchronously occupies a cooperative executor; use a dedicated Dispatch timer and synchronously drain it before writing the terminal protocol message. ## Schema / migrations *(carried forward from CLAUDE.md operational rules)* diff --git a/executors/mlx/.gitignore b/executors/mlx/.gitignore new file mode 100644 index 0000000..30bcfa4 --- /dev/null +++ b/executors/mlx/.gitignore @@ -0,0 +1 @@ +.build/ diff --git a/executors/mlx/Package.resolved b/executors/mlx/Package.resolved new file mode 100644 index 0000000..04b832e --- /dev/null +++ b/executors/mlx/Package.resolved @@ -0,0 +1,114 @@ +{ + "originHash" : "06b5628b56d2ca3741635554cce26e6f45b53bcfcefda4c3212a028436dbf9ff", + "pins" : [ + { + "identity" : "eventsource", + "kind" : "remoteSourceControl", + "location" : "https://github.com/mattt/EventSource.git", + "state" : { + "revision" : "86b5096ac59ab46e66bd1f6377c604bc1dab0bc2", + "version" : "1.5.1" + } + }, + { + "identity" : "mlx-swift", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift", + "state" : { + "revision" : "dc43e62d7055353c7f99fa071a4e71d29dfddc44", + "version" : "0.31.4" + } + }, + { + "identity" : "mlx-swift-lm", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ml-explore/mlx-swift-lm", + "state" : { + "revision" : "bd4b7434e6bdb588c7ef55706ff8904cb7fd4c57", + "version" : "3.31.4" + } + }, + { + "identity" : "swift-asn1", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-asn1.git", + "state" : { + "revision" : "a9a5efd40eaf558a2bcd48d64b1d1646be686008", + "version" : "1.7.1" + } + }, + { + "identity" : "swift-collections", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-collections.git", + "state" : { + "revision" : "a0cb0954ecb21e4e31b0070e6ed5674e8556685a", + "version" : "1.6.0" + } + }, + { + "identity" : "swift-crypto", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-crypto.git", + "state" : { + "revision" : "47d3869a7291f085c1fb9fb1e6d3b97a793f45c6", + "version" : "4.5.1" + } + }, + { + "identity" : "swift-huggingface", + "kind" : "remoteSourceControl", + "location" : "https://github.com/huggingface/swift-huggingface", + "state" : { + "revision" : "b721959445b617d0bf03910b2b4aced345fd93bf", + "version" : "0.9.0" + } + }, + { + "identity" : "swift-jinja", + "kind" : "remoteSourceControl", + "location" : "https://github.com/huggingface/swift-jinja.git", + "state" : { + "revision" : "7d0b8880ef8e567dd4e0089f8b99fb354129017c", + "version" : "2.4.2" + } + }, + { + "identity" : "swift-numerics", + "kind" : "remoteSourceControl", + "location" : "https://github.com/apple/swift-numerics", + "state" : { + "revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2", + "version" : "1.1.1" + } + }, + { + "identity" : "swift-syntax", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swiftlang/swift-syntax.git", + "state" : { + "revision" : "79e4b74a295b6eb74a8b585e3a39d29e70c1dbd1", + "version" : "603.0.2" + } + }, + { + "identity" : "swift-transformers", + "kind" : "remoteSourceControl", + "location" : "https://github.com/huggingface/swift-transformers", + "state" : { + "revision" : "b38443e44d93eca770f2eb68e2a4d0fa100f9aa2", + "version" : "1.3.0" + } + }, + { + "identity" : "yyjson", + "kind" : "remoteSourceControl", + "location" : "https://github.com/ibireme/yyjson.git", + "state" : { + "revision" : "8b4a38dc994a110abaec8a400615567bd996105f", + "version" : "0.12.0" + } + } + ], + "version" : 3 +} diff --git a/executors/mlx/Package.swift b/executors/mlx/Package.swift new file mode 100644 index 0000000..d740628 --- /dev/null +++ b/executors/mlx/Package.swift @@ -0,0 +1,51 @@ +// swift-tools-version: 6.1 + +import PackageDescription + +let package = Package( + name: "surreal-memory-mlx-executor", + platforms: [.macOS(.v14)], + products: [ + .executable( + name: "surreal-memory-mlx-executor", + targets: ["SurrealMemoryMLXExecutor"] + ) + ], + dependencies: [ + .package( + url: "https://github.com/ml-explore/mlx-swift-lm", + exact: "3.31.4" + ), + .package( + url: "https://github.com/ml-explore/mlx-swift", + exact: "0.31.4" + ), + .package( + url: "https://github.com/huggingface/swift-huggingface", + exact: "0.9.0" + ), + .package( + url: "https://github.com/huggingface/swift-transformers", + exact: "1.3.0" + ) + ], + targets: [ + .target(name: "SurrealMemoryMLXExecutorCore"), + .executableTarget( + name: "SurrealMemoryMLXExecutor", + dependencies: [ + "SurrealMemoryMLXExecutorCore", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXEmbedders", package: "mlx-swift-lm"), + .product(name: "MLXLMCommon", package: "mlx-swift-lm"), + .product(name: "MLXHuggingFace", package: "mlx-swift-lm"), + .product(name: "HuggingFace", package: "swift-huggingface"), + .product(name: "Tokenizers", package: "swift-transformers") + ] + ), + .testTarget( + name: "SurrealMemoryMLXExecutorCoreTests", + dependencies: ["SurrealMemoryMLXExecutorCore"] + ) + ] +) diff --git a/executors/mlx/Sources/SurrealMemoryMLXExecutor/MLXEmbeddingService.swift b/executors/mlx/Sources/SurrealMemoryMLXExecutor/MLXEmbeddingService.swift new file mode 100644 index 0000000..9b9efbb --- /dev/null +++ b/executors/mlx/Sources/SurrealMemoryMLXExecutor/MLXEmbeddingService.swift @@ -0,0 +1,282 @@ +import Foundation +import HuggingFace +import MLX +import MLXEmbedders +import MLXHuggingFace +import MLXLMCommon +import SurrealMemoryMLXExecutorCore +import Tokenizers + +struct ExecutorSettings: Sendable { + static let defaultModelID = "BAAI/bge-small-en-v1.5" + static let defaultRevision = "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a" + static let expectedDimensions = 384 + + let modelID: String + let modelRevision: String + let dimensions: Int + let hubCache: URL + + static func fromEnvironment() throws -> ExecutorSettings { + let environment = ProcessInfo.processInfo.environment + let modelID = environment["LOCAL_EMBEDDING_MODEL"] ?? defaultModelID + let revision = environment["LOCAL_EMBEDDING_MODEL_REVISION"] ?? defaultRevision + let dimensions = Int(environment["LOCAL_EMBEDDING_DIMENSIONS"] ?? "") + ?? expectedDimensions + + guard modelID == defaultModelID else { + throw ExecutorError.unsupportedModel(modelID) + } + guard revision == defaultRevision else { + throw ExecutorError.unsupportedRevision(revision) + } + guard dimensions == expectedDimensions else { + throw ExecutorError.dimensionMismatch( + expected: expectedDimensions, + actual: dimensions + ) + } + + let hubCache: URL + if let configured = environment["HF_HUB_CACHE"] { + hubCache = URL(filePath: configured, directoryHint: .isDirectory) + } else if let home = environment["MODEL_CACHE_DIR"] ?? environment["HF_HOME"] { + hubCache = URL(filePath: home, directoryHint: .isDirectory) + .appending(path: "hub", directoryHint: .isDirectory) + } else { + hubCache = FileManager.default.homeDirectoryForCurrentUser + .appending(path: ".cache/huggingface/hub", directoryHint: .isDirectory) + } + + return ExecutorSettings( + modelID: modelID, + modelRevision: revision, + dimensions: dimensions, + hubCache: hubCache + ) + } + + var snapshotDirectory: URL { + let repository = "models--" + modelID.replacingOccurrences(of: "/", with: "--") + return hubCache + .appending(path: repository, directoryHint: .isDirectory) + .appending(path: "snapshots", directoryHint: .isDirectory) + .appending(path: modelRevision, directoryHint: .isDirectory) + } +} + +enum ExecutorError: Error, LocalizedError { + case unsupportedModel(String) + case unsupportedRevision(String) + case dimensionMismatch(expected: Int, actual: Int) + case missingSnapshot(URL) + case snapshotRevisionMismatch(expected: String, actual: String) + case missingModelFile(String) + case invalidModelConfiguration + case inputTooLong(actual: Int, maximum: Int) + + var errorDescription: String? { + switch self { + case .unsupportedModel(let model): + "MLX executor is certified only for \(ExecutorSettings.defaultModelID), got \(model)" + case .unsupportedRevision(let revision): + "MLX executor is certified only for revision \(ExecutorSettings.defaultRevision), got \(revision)" + case .dimensionMismatch(let expected, let actual): + "embedding dimension mismatch: expected \(expected), got \(actual)" + case .missingSnapshot(let path): + "pinned model snapshot is not cached at \(path.path); run --prefetch first" + case .snapshotRevisionMismatch(let expected, let actual): + "download resolved revision \(actual), expected \(expected)" + case .missingModelFile(let file): + "pinned model snapshot is missing \(file); run --prefetch first" + case .invalidModelConfiguration: + "model config.json does not contain a positive max_position_embeddings" + case .inputTooLong(let actual, let maximum): + "input_too_long: tokenizer produced \(actual) tokens for model capacity \(maximum)" + } + } +} + +final class MLXEmbeddingService: Sendable { + let settings: ExecutorSettings + let container: EmbedderModelContainer + let maxInputTokens: Int + + private init( + settings: ExecutorSettings, + container: EmbedderModelContainer, + maxInputTokens: Int + ) { + self.settings = settings + self.container = container + self.maxInputTokens = maxInputTokens + } + + static func loadCached(settings: ExecutorSettings) async throws -> MLXEmbeddingService { + let directory = settings.snapshotDirectory + try validateSnapshot(directory, expectedRevision: settings.modelRevision) + let container = try await EmbedderModelFactory.shared.loadContainer( + from: directory, + using: #huggingFaceTokenizerLoader() + ) + let maximum = try readMaximumInputTokens(from: directory) + return MLXEmbeddingService( + settings: settings, + container: container, + maxInputTokens: maximum + ) + } + + static func prefetch(settings: ExecutorSettings) async throws -> MLXEmbeddingService { + try FileManager.default.createDirectory( + at: settings.hubCache, + withIntermediateDirectories: true + ) + let hub = HubClient(cache: HubCache(cacheDirectory: settings.hubCache)) + let configuration = ModelConfiguration( + id: settings.modelID, + revision: settings.modelRevision + ) + let container = try await EmbedderModelFactory.shared.loadContainer( + from: #hubDownloader(hub), + using: #huggingFaceTokenizerLoader(), + configuration: configuration, + useLatest: false + ) + let resolvedDirectory = try await container.modelDirectory.resolvingSymlinksInPath() + guard resolvedDirectory.lastPathComponent == settings.modelRevision else { + throw ExecutorError.snapshotRevisionMismatch( + expected: settings.modelRevision, + actual: resolvedDirectory.lastPathComponent + ) + } + try validateSnapshot(resolvedDirectory, expectedRevision: settings.modelRevision) + return MLXEmbeddingService( + settings: settings, + container: container, + maxInputTokens: try readMaximumInputTokens(from: resolvedDirectory) + ) + } + + func warmup() async throws -> [Float] { + let embedding = try await embedBatch(["warmup"])[0] + guard embedding.count == settings.dimensions else { + throw ExecutorError.dimensionMismatch( + expected: settings.dimensions, + actual: embedding.count + ) + } + return embedding + } + + func plan(text: String) async throws -> [EmbeddingPlanPart] { + try await container.perform { context in + try EmbeddingPlanner.plan( + text: text, + maxInputTokens: maxInputTokens, + encode: { value, addSpecialTokens in + context.tokenizer.encode( + text: value, + addSpecialTokens: addSpecialTokens + ) + }, + decode: { tokenIDs, skipSpecialTokens in + context.tokenizer.decode( + tokenIds: tokenIDs, + skipSpecialTokens: skipSpecialTokens + ) + } + ) + } + } + + func embedBatch(_ texts: [String]) async throws -> [[Float]] { + guard !texts.isEmpty else { return [] } + let embeddings = try await container.perform { context in + let encoded = texts.map { + context.tokenizer.encode(text: $0, addSpecialTokens: true) + } + if let tooLong = encoded.first(where: { $0.count > maxInputTokens }) { + throw ExecutorError.inputTooLong( + actual: tooLong.count, + maximum: maxInputTokens + ) + } + + let maximum = encoded.map(\.count).max() ?? 0 + let paddingToken = context.tokenizer.convertTokenToId("[PAD]") ?? 0 + let paddedTokens = encoded.map { + $0 + Array(repeating: paddingToken, count: maximum - $0.count) + } + let masks = encoded.map { + Array(repeating: Float(1), count: $0.count) + + Array(repeating: Float(0), count: maximum - $0.count) + } + let inputs = stacked(paddedTokens.map { MLXArray($0) }) + let attentionMask = stacked(masks.map { MLXArray($0) }) + let tokenTypes = MLXArray.zeros(like: inputs) + let output = context.model( + inputs, + positionIds: nil, + tokenTypeIds: tokenTypes, + attentionMask: attentionMask + ) + let pooled = Pooling( + strategy: .mean, + dimension: settings.dimensions + )( + output, + mask: attentionMask, + normalize: true, + applyLayerNorm: false + ) + pooled.eval() + return pooled.map { $0.asArray(Float.self) } + } + + for embedding in embeddings where embedding.count != settings.dimensions { + throw ExecutorError.dimensionMismatch( + expected: settings.dimensions, + actual: embedding.count + ) + } + return embeddings + } + + private static func validateSnapshot( + _ directory: URL, + expectedRevision: String + ) throws { + var isDirectory: ObjCBool = false + guard FileManager.default.fileExists( + atPath: directory.path, + isDirectory: &isDirectory + ), isDirectory.boolValue else { + throw ExecutorError.missingSnapshot(directory) + } + guard directory.lastPathComponent == expectedRevision else { + throw ExecutorError.snapshotRevisionMismatch( + expected: expectedRevision, + actual: directory.lastPathComponent + ) + } + for file in ["config.json", "tokenizer.json", "model.safetensors"] { + guard FileManager.default.fileExists( + atPath: directory.appending(path: file).path + ) else { + throw ExecutorError.missingModelFile(file) + } + } + } + + private static func readMaximumInputTokens(from directory: URL) throws -> Int { + let data = try Data(contentsOf: directory.appending(path: "config.json")) + let object = try JSONSerialization.jsonObject(with: data) + guard let values = object as? [String: Any], + let maximum = values["max_position_embeddings"] as? Int, + maximum > 0 else { + throw ExecutorError.invalidModelConfiguration + } + return maximum + } +} diff --git a/executors/mlx/Sources/SurrealMemoryMLXExecutor/main.swift b/executors/mlx/Sources/SurrealMemoryMLXExecutor/main.swift new file mode 100644 index 0000000..3c5c679 --- /dev/null +++ b/executors/mlx/Sources/SurrealMemoryMLXExecutor/main.swift @@ -0,0 +1,123 @@ +import Foundation +import SurrealMemoryMLXExecutorCore + +final class OutputWriter: @unchecked Sendable { + private let lock = NSLock() + + func write(_ message: ExecutorMessage) throws { + lock.lock() + defer { lock.unlock() } + var data = try JSONEncoder().encode(message) + data.append(0x0A) + FileHandle.standardOutput.write(data) + } +} +@main +enum SurrealMemoryMLXExecutorMain { + static func main() async { + do { + try await run() + } catch { + writeStandardError("surreal-memory-mlx-executor: \(error.localizedDescription)\n") + Foundation.exit(1) + } + } + + private static func run() async throws { + let arguments = Array(CommandLine.arguments.dropFirst()) + let settings = try ExecutorSettings.fromEnvironment() + switch arguments.first { + case "--version", "-V": + print("surreal-memory-mlx-executor 1.0.0 mlx-swift-lm 3.31.4") + case "--prefetch": + let service = try await MLXEmbeddingService.prefetch(settings: settings) + try await printSmoke(service: service, mode: "prefetch") + case "--smoke": + let service = try await MLXEmbeddingService.loadCached(settings: settings) + try await printSmoke(service: service, mode: "smoke") + case "embedding-executor", nil: + let service = try await MLXEmbeddingService.loadCached(settings: settings) + _ = try await service.warmup() + try await runProtocol(service: service) + default: + throw CocoaError(.executableLoad) + } + } + + private static func printSmoke( + service: MLXEmbeddingService, + mode: String + ) async throws { + let embedding = try await service.warmup() + let norm = sqrt(embedding.map { $0 * $0 }.reduce(0, +)) + let output: [String: Any] = [ + "backend": "mlx", + "dimensions": embedding.count, + "mode": mode, + "model_id": service.settings.modelID, + "model_revision": service.settings.modelRevision, + "norm": norm, + "status": "ok" + ] + let data = try JSONSerialization.data(withJSONObject: output, options: [.sortedKeys]) + print(String(decoding: data, as: UTF8.self)) + } + + private static func runProtocol(service: MLXEmbeddingService) async throws { + let writer = OutputWriter() + try writer.write( + .ready( + backend: "mlx", + modelID: service.settings.modelID, + modelRevision: service.settings.modelRevision, + dimensions: service.settings.dimensions + ) + ) + + while let line = readLine(strippingNewline: true) { + let request: ExecutorRequest + do { + request = try JSONDecoder().decode( + ExecutorRequest.self, + from: Data(line.utf8) + ) + } catch { + try writer.write(.failed(requestID: 0, error: "decode request: \(error)")) + continue + } + + try writer.write(.progress(requestID: request.requestID, phase: "accepted")) + // MLX evaluation can synchronously occupy a Swift cooperative + // executor thread while Metal finishes the graph. Drive watchdog + // progress from a dedicated Dispatch queue so a busy inference + // cannot starve its own heartbeat and be mistaken for a hang. + let heartbeat = ProgressHeartbeat(interval: 0.25) { + try? writer.write( + .progress(requestID: request.requestID, phase: "working") + ) + } + + let message: ExecutorMessage + do { + let result: ExecutorResult + switch request.command { + case .plan(let text): + result = .plan(parts: try await service.plan(text: text)) + case .embed(_, let text): + result = .embedding(try await service.embedBatch([text])[0]) + case .embedBatch(let texts): + result = .batch(try await service.embedBatch(texts)) + } + message = .completed(requestID: request.requestID, result: result) + } catch { + message = .failed(requestID: request.requestID, error: String(describing: error)) + } + heartbeat.stop() + try writer.write(message) + } + } + + private static func writeStandardError(_ text: String) { + FileHandle.standardError.write(Data(text.utf8)) + } +} diff --git a/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Planner.swift b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Planner.swift new file mode 100644 index 0000000..4d43c51 --- /dev/null +++ b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Planner.swift @@ -0,0 +1,93 @@ +import CryptoKit +import Foundation + +public enum EmbeddingPlannerError: Error, LocalizedError, Equatable { + case noContentCapacity(maximum: Int, specialTokens: Int) + case unableToConstructWindow + + public var errorDescription: String? { + switch self { + case .noContentCapacity(let maximum, let specialTokens): + "model capacity \(maximum) does not leave room after \(specialTokens) special tokens" + case .unableToConstructWindow: + "unable to construct a model-safe token window" + } + } +} +public enum EmbeddingPlanner { + public static func plan( + text: String, + maxInputTokens: Int, + encode: (String, Bool) throws -> [Int], + decode: ([Int], Bool) throws -> String + ) throws -> [EmbeddingPlanPart] { + let specialTokens = try encode("", true).count + let usable = maxInputTokens - specialTokens + guard usable > 0 else { + throw EmbeddingPlannerError.noContentCapacity( + maximum: maxInputTokens, + specialTokens: specialTokens + ) + } + + let source = try encode(text, false) + if source.count + specialTokens <= maxInputTokens { + return [part(index: 0, start: 0, end: source.count, ids: source, content: text)] + } + + let overlap = min(32, max(0, usable - 1)) + let step = usable - overlap + var parts: [EmbeddingPlanPart] = [] + var start = 0 + while start < source.count { + var end = min(start + usable, source.count) + var content = try decode(Array(source[start.. maxInputTokens { + end -= 1 + guard end > start else { + throw EmbeddingPlannerError.unableToConstructWindow + } + content = try decode(Array(source[start.. EmbeddingPlanPart { + var bytes = Data(capacity: ids.count * MemoryLayout.size) + for id in ids { + var littleEndian = UInt32(id).littleEndian + withUnsafeBytes(of: &littleEndian) { bytes.append(contentsOf: $0) } + } + let digest = SHA256.hash(data: bytes).map { String(format: "%02x", $0) }.joined() + return EmbeddingPlanPart( + partIndex: index, + tokenStart: start, + tokenEnd: end, + tokenCount: end - start, + tokenHash: digest, + content: content + ) + } +} diff --git a/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/ProgressHeartbeat.swift b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/ProgressHeartbeat.swift new file mode 100644 index 0000000..b039519 --- /dev/null +++ b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/ProgressHeartbeat.swift @@ -0,0 +1,45 @@ +import Foundation + +/// A progress timer backed by its own Dispatch queue. +/// +/// MLX/Metal evaluation may synchronously occupy a Swift cooperative executor +/// thread. A heartbeat implemented as another `Task` can therefore be starved +/// alongside the work it is meant to supervise. Dispatch keeps the liveness +/// signal independent, and `stop()` synchronously drains any in-flight callback +/// before the terminal protocol message is written. +public final class ProgressHeartbeat: @unchecked Sendable { + private let queue = DispatchQueue(label: "ai.prometheus.surreal-memory.mlx-heartbeat") + private let timer: DispatchSourceTimer + private let stateLock = NSLock() + private var stopped = false + + public init( + interval: TimeInterval, + action: @escaping @Sendable () -> Void + ) { + precondition(interval > 0) + timer = DispatchSource.makeTimerSource(queue: queue) + timer.schedule(deadline: .now() + interval, repeating: interval) + timer.setEventHandler(handler: action) + timer.resume() + } + + public func stop() { + stateLock.lock() + guard !stopped else { + stateLock.unlock() + return + } + stopped = true + stateLock.unlock() + + queue.sync { + timer.setEventHandler {} + timer.cancel() + } + } + + deinit { + stop() + } +} diff --git a/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Protocol.swift b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Protocol.swift new file mode 100644 index 0000000..6c36ea0 --- /dev/null +++ b/executors/mlx/Sources/SurrealMemoryMLXExecutorCore/Protocol.swift @@ -0,0 +1,160 @@ +import Foundation + +public let executorProtocolVersion = 1 + +public struct ExecutorRequest: Decodable, Sendable { + public let requestID: UInt64 + public let operationID: String? + public let command: ExecutorCommand + + enum CodingKeys: String, CodingKey { + case requestID = "request_id" + case operationID = "operation_id" + case command + } +} +public enum ExecutorCommand: Decodable, Sendable { + case plan(text: String) + case embed(partIndex: Int, text: String) + case embedBatch(texts: [String]) + + enum CodingKeys: String, CodingKey { + case command + case text + case partIndex = "part_index" + case texts + } + + public init(from decoder: Decoder) throws { + let values = try decoder.container(keyedBy: CodingKeys.self) + switch try values.decode(String.self, forKey: .command) { + case "plan": + self = .plan(text: try values.decode(String.self, forKey: .text)) + case "embed": + self = .embed( + partIndex: try values.decode(Int.self, forKey: .partIndex), + text: try values.decode(String.self, forKey: .text) + ) + case "embed_batch": + self = .embedBatch(texts: try values.decode([String].self, forKey: .texts)) + case let command: + throw DecodingError.dataCorruptedError( + forKey: .command, + in: values, + debugDescription: "unsupported executor command '\(command)'" + ) + } + } +} + +public struct EmbeddingPlanPart: Codable, Equatable, Sendable { + public let partIndex: Int + public let tokenStart: Int + public let tokenEnd: Int + public let tokenCount: Int + public let tokenHash: String + public let content: String + + public init( + partIndex: Int, + tokenStart: Int, + tokenEnd: Int, + tokenCount: Int, + tokenHash: String, + content: String + ) { + self.partIndex = partIndex + self.tokenStart = tokenStart + self.tokenEnd = tokenEnd + self.tokenCount = tokenCount + self.tokenHash = tokenHash + self.content = content + } + + enum CodingKeys: String, CodingKey { + case partIndex = "part_index" + case tokenStart = "token_start" + case tokenEnd = "token_end" + case tokenCount = "token_count" + case tokenHash = "token_hash" + case content + } +} + +public enum ExecutorResult: Encodable, Sendable { + case plan(parts: [EmbeddingPlanPart]) + case embedding([Float]) + case batch([[Float]]) + + enum CodingKeys: String, CodingKey { + case result + case parts + case embedding + case embeddings + } + + public func encode(to encoder: Encoder) throws { + var values = encoder.container(keyedBy: CodingKeys.self) + switch self { + case .plan(let parts): + try values.encode("plan", forKey: .result) + try values.encode(parts, forKey: .parts) + case .embedding(let embedding): + try values.encode("embedding", forKey: .result) + try values.encode(embedding, forKey: .embedding) + case .batch(let embeddings): + try values.encode("batch", forKey: .result) + try values.encode(embeddings, forKey: .embeddings) + } + } +} + +public enum ExecutorMessage: Encodable, Sendable { + case ready( + backend: String, + modelID: String, + modelRevision: String, + dimensions: Int + ) + case progress(requestID: UInt64, phase: String) + case completed(requestID: UInt64, result: ExecutorResult) + case failed(requestID: UInt64, error: String) + + enum CodingKeys: String, CodingKey { + case message + case protocolVersion = "protocol_version" + case backend + case modelID = "model_id" + case modelRevision = "model_revision" + case dimensions + case requestID = "request_id" + case phase + case result + case error + } + + public func encode(to encoder: Encoder) throws { + var values = encoder.container(keyedBy: CodingKeys.self) + switch self { + case .ready(let backend, let modelID, let modelRevision, let dimensions): + try values.encode("ready", forKey: .message) + try values.encode(executorProtocolVersion, forKey: .protocolVersion) + try values.encode(backend, forKey: .backend) + try values.encode(modelID, forKey: .modelID) + try values.encode(modelRevision, forKey: .modelRevision) + try values.encode(dimensions, forKey: .dimensions) + case .progress(let requestID, let phase): + try values.encode("progress", forKey: .message) + try values.encode(requestID, forKey: .requestID) + try values.encode(phase, forKey: .phase) + case .completed(let requestID, let result): + try values.encode("completed", forKey: .message) + try values.encode(requestID, forKey: .requestID) + try values.encode(result, forKey: .result) + case .failed(let requestID, let error): + try values.encode("failed", forKey: .message) + try values.encode(requestID, forKey: .requestID) + try values.encode(error, forKey: .error) + } + } +} diff --git a/executors/mlx/Tests/SurrealMemoryMLXExecutorCoreTests/ProtocolTests.swift b/executors/mlx/Tests/SurrealMemoryMLXExecutorCoreTests/ProtocolTests.swift new file mode 100644 index 0000000..b2883bc --- /dev/null +++ b/executors/mlx/Tests/SurrealMemoryMLXExecutorCoreTests/ProtocolTests.swift @@ -0,0 +1,102 @@ +import Foundation +import Testing +@testable import SurrealMemoryMLXExecutorCore + +@Test func decodesEmbedRequest() throws { + let json = #"{"request_id":7,"operation_id":"op-1","command":{"command":"embed","part_index":2,"text":"hello"}}"# + let request = try JSONDecoder().decode(ExecutorRequest.self, from: Data(json.utf8)) + #expect(request.requestID == 7) + #expect(request.operationID == "op-1") + guard case .embed(let partIndex, let text) = request.command else { + Issue.record("expected embed command") + return + } + #expect(partIndex == 2) + #expect(text == "hello") +} + +@Test func readyMessageCarriesCertifiedIdentity() throws { + let data = try JSONEncoder().encode( + ExecutorMessage.ready( + backend: "mlx", + modelID: "BAAI/bge-small-en-v1.5", + modelRevision: "revision", + dimensions: 384 + ) + ) + let object = try #require(JSONSerialization.jsonObject(with: data) as? [String: Any]) + #expect(object["message"] as? String == "ready") + #expect(object["protocol_version"] as? Int == executorProtocolVersion) + #expect(object["backend"] as? String == "mlx") + #expect(object["dimensions"] as? Int == 384) +} + +@Test func plannerUsesDeterministicOverlappingTokenWindows() throws { + func encode(_ text: String, _ special: Bool) -> [Int] { + let values = text.split(separator: " ").compactMap { Int($0) } + return special ? [101] + values + [102] : values + } + func decode(_ ids: [Int], _: Bool) -> String { + ids.map(String.init).joined(separator: " ") + } + + let input = (1...12).map(String.init).joined(separator: " ") + let first = try EmbeddingPlanner.plan( + text: input, + maxInputTokens: 8, + encode: encode, + decode: decode + ) + let second = try EmbeddingPlanner.plan( + text: input, + maxInputTokens: 8, + encode: encode, + decode: decode + ) + + #expect(first == second) + #expect(first.count == 7) + #expect(first[0].tokenStart == 0) + #expect(first[0].tokenEnd == 6) + #expect(first[1].tokenStart == 1) + #expect(first.allSatisfy { $0.tokenHash.count == 64 }) +} + +@Test func malformedCommandIsRejected() { + let json = #"{"request_id":7,"command":{"command":"unknown"}}"# + #expect(throws: DecodingError.self) { + try JSONDecoder().decode(ExecutorRequest.self, from: Data(json.utf8)) + } +} + +private final class HeartbeatCounter: @unchecked Sendable { + private let lock = NSLock() + private var value = 0 + + func increment() { + lock.lock() + value += 1 + lock.unlock() + } + + func read() -> Int { + lock.lock() + defer { lock.unlock() } + return value + } +} + +@Test func progressHeartbeatRunsIndependentlyAndStopsSynchronously() async throws { + let counter = HeartbeatCounter() + let heartbeat = ProgressHeartbeat(interval: 0.01) { + counter.increment() + } + + try await Task.sleep(for: .milliseconds(80)) + heartbeat.stop() + let stoppedAt = counter.read() + #expect(stoppedAt > 0) + + try await Task.sleep(for: .milliseconds(40)) + #expect(counter.read() == stoppedAt) +} diff --git a/executors/mlx/Tests/fixtures/parity-corpus.json b/executors/mlx/Tests/fixtures/parity-corpus.json new file mode 100644 index 0000000..a1dec29 --- /dev/null +++ b/executors/mlx/Tests/fixtures/parity-corpus.json @@ -0,0 +1,26 @@ +{ + "queries": [ + "How does the learning worker resume accepted memory operations?", + "What protects the embedding subprocess from protocol desynchronization?", + "Which SurrealDB version is installed by the service pack?", + "How are skills distributed to supported coding tools?", + "Why did the Metal embedding backend stall during startup?" + ], + "documents": [ + "The learning worker reconciles durable receipts and resumes unfinished operation parts after restart.", + "Accepted memory operations remain in the durable ledger until every embedding part commits.", + "The executor supervisor correlates JSONL messages by request ID and respawns after malformed or mismatched frames.", + "A fresh subprocess generation is used after an executor exits or violates the wire protocol.", + "The native database service is certified against SurrealDB 3.2.4 and surrealdb-types 3.2.4.", + "SurrealDB stores memory records and vector indexes in the shared memory namespace and mcp database.", + "An immutable signed skill generation is activated across fourteen supported tool targets.", + "Receipts verify every installed skill placement and preserve generation provenance.", + "Candle completed the model download but stalled while compiling Metal kernels for the first BERT forward pass.", + "MLX uses Apple Silicon unified memory and a different Metal execution stack for local inference.", + "Prompt hooks retrieve bounded project and global knowledge context before the agent turn.", + "Stop hooks enqueue one idempotent Karpathy learning job per completed session.", + "The service installer ad-hoc signs copied Apple Silicon binaries before launchd starts them.", + "The BGE small English model produces normalized vectors with 384 dimensions.", + "Changing embedding models would require a controlled full re-embedding migration." + ] +} diff --git a/executors/mlx/scripts/compare_embeddings.py b/executors/mlx/scripts/compare_embeddings.py new file mode 100755 index 0000000..32babd7 --- /dev/null +++ b/executors/mlx/scripts/compare_embeddings.py @@ -0,0 +1,170 @@ +#!/usr/bin/env python3 +"""Certify MLX embeddings against the Candle/CPU BGE implementation.""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import subprocess +from pathlib import Path + +MODEL_ID = "BAAI/bge-small-en-v1.5" +MODEL_REVISION = "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a" +DIMENSIONS = 384 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--candle", required=True, type=Path) + parser.add_argument("--mlx", required=True, type=Path) + parser.add_argument( + "--corpus", + type=Path, + default=Path(__file__).parents[1] / "Tests/fixtures/parity-corpus.json", + ) + return parser.parse_args() + + +def read_message(process: subprocess.Popen[str]) -> dict[str, object]: + assert process.stdout is not None + line = process.stdout.readline() + if not line: + stderr = process.stderr.read() if process.stderr else "" + raise RuntimeError(f"executor closed its output: {stderr}") + return json.loads(line) + + +def embed(binary: Path, backend: str, texts: list[str]) -> list[list[float]]: + environment = os.environ.copy() + environment.update( + { + "EMBEDDING_PROVIDER": "local", + "LOCAL_EMBEDDING_BACKEND": backend, + "LOCAL_EMBEDDING_MODEL": MODEL_ID, + "LOCAL_EMBEDDING_MODEL_REVISION": MODEL_REVISION, + "LOCAL_EMBEDDING_DIMENSIONS": str(DIMENSIONS), + "MODEL_CACHE_DIR": str(Path.home() / ".cache/huggingface"), + } + ) + process = subprocess.Popen( + [str(binary), "embedding-executor"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=environment, + ) + try: + ready = read_message(process) + if ready.get("message") != "ready" or ready.get("backend") != backend: + raise RuntimeError(f"unexpected {backend} ready message: {ready}") + request = { + "request_id": 1, + "operation_id": None, + "command": {"command": "embed_batch", "texts": texts}, + } + assert process.stdin is not None + process.stdin.write(json.dumps(request) + "\n") + process.stdin.flush() + while True: + message = read_message(process) + if message.get("message") == "progress": + continue + if message.get("message") == "failed": + raise RuntimeError(f"{backend} inference failed: {message.get('error')}") + if message.get("message") == "completed": + result = message["result"] + assert isinstance(result, dict) + embeddings = result["embeddings"] + assert isinstance(embeddings, list) + return embeddings + raise RuntimeError(f"unexpected {backend} message: {message}") + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) + + +def dot(left: list[float], right: list[float]) -> float: + return sum(a * b for a, b in zip(left, right, strict=True)) + + +def norm(vector: list[float]) -> float: + return math.sqrt(dot(vector, vector)) + + +def cosine(left: list[float], right: list[float]) -> float: + return dot(left, right) / (norm(left) * norm(right)) + + +def ranking(query: list[float], documents: list[list[float]]) -> list[int]: + return sorted( + range(len(documents)), + key=lambda index: cosine(query, documents[index]), + reverse=True, + ) + + +def main() -> None: + arguments = parse_args() + corpus = json.loads(arguments.corpus.read_text()) + queries = corpus["queries"] + documents = corpus["documents"] + texts = queries + documents + candle = embed(arguments.candle, "candle", texts) + mlx = embed(arguments.mlx, "mlx", texts) + + if len(candle) != len(texts) or len(mlx) != len(texts): + raise SystemExit("embedding count mismatch") + for backend, vectors in (("candle", candle), ("mlx", mlx)): + for index, vector in enumerate(vectors): + if len(vector) != DIMENSIONS: + raise SystemExit( + f"{backend} vector {index} has {len(vector)} dimensions, expected {DIMENSIONS}" + ) + magnitude = norm(vector) + if not 0.999 <= magnitude <= 1.001: + raise SystemExit(f"{backend} vector {index} has norm {magnitude}") + + paired = [cosine(left, right) for left, right in zip(candle, mlx, strict=True)] + if min(paired) < 0.999: + raise SystemExit(f"paired cosine minimum {min(paired):.9f} is below 0.999") + + query_count = len(queries) + candle_documents = candle[query_count:] + mlx_documents = mlx[query_count:] + top_one_matches = 0 + top_five_overlap = 0 + for index in range(query_count): + candle_rank = ranking(candle[index], candle_documents) + mlx_rank = ranking(mlx[index], mlx_documents) + top_one_matches += candle_rank[0] == mlx_rank[0] + top_five_overlap += len(set(candle_rank[:5]) & set(mlx_rank[:5])) + + if top_one_matches != query_count: + raise SystemExit(f"top-1 parity failed for {query_count - top_one_matches} queries") + overlap_ratio = top_five_overlap / (query_count * 5) + if overlap_ratio < 0.95: + raise SystemExit(f"aggregate top-5 overlap {overlap_ratio:.1%} is below 95%") + + print( + json.dumps( + { + "dimensions": DIMENSIONS, + "paired_cosine_min": min(paired), + "status": "ok", + "top_1_matches": f"{top_one_matches}/{query_count}", + "top_5_overlap": overlap_ratio, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/src/config.rs b/src/config.rs index a0e74f1..8ccb599 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,25 @@ use crate::embeddings::EmbeddingProvider; use serde::Deserialize; -use std::env; +use std::{env, path::PathBuf}; + +pub const DEFAULT_LOCAL_EMBEDDING_MODEL: &str = "BAAI/bge-small-en-v1.5"; +pub const DEFAULT_LOCAL_EMBEDDING_REVISION: &str = "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum LocalEmbeddingBackend { + Candle, + Mlx, +} + +impl LocalEmbeddingBackend { + pub fn as_str(self) -> &'static str { + match self { + Self::Candle => "candle", + Self::Mlx => "mlx", + } + } +} #[derive(Debug, Clone, Deserialize)] pub struct Config { @@ -12,6 +31,10 @@ pub struct Config { pub surreal_password: Option, pub embedded_path: Option, pub embedding_provider: EmbeddingProvider, + pub local_embedding_backend: LocalEmbeddingBackend, + pub local_embedding_executor: Option, + pub local_embedding_model_revision: String, + pub local_embedding_dimensions: usize, /// When true, the embedding model is loaded eagerly at startup so the first /// user-facing write does not pay the cold-load latency. Defaults to true; /// set `EMBEDDING_WARMUP=false` to restore purely lazy loading. @@ -57,7 +80,7 @@ impl Config { _ => { // Default to local embeddings let model_id = env::var("LOCAL_EMBEDDING_MODEL") - .unwrap_or_else(|_| "BAAI/bge-small-en-v1.5".to_string()); + .unwrap_or_else(|_| DEFAULT_LOCAL_EMBEDDING_MODEL.to_string()); let model_path = env::var("MODEL_CACHE_DIR").ok(); EmbeddingProvider::Local { model_id, @@ -78,6 +101,26 @@ impl Config { .ok() .or_else(|| Some("./data/memory.db".to_string())), embedding_provider, + local_embedding_backend: match env::var("LOCAL_EMBEDDING_BACKEND") + .unwrap_or_else(|_| "candle".to_string()) + .to_ascii_lowercase() + .as_str() + { + "candle" => LocalEmbeddingBackend::Candle, + "mlx" => LocalEmbeddingBackend::Mlx, + value => anyhow::bail!( + "LOCAL_EMBEDDING_BACKEND must be 'candle' or 'mlx', got '{value}'" + ), + }, + local_embedding_executor: env::var_os("LOCAL_EMBEDDING_EXECUTOR").map(PathBuf::from), + local_embedding_model_revision: env::var("LOCAL_EMBEDDING_MODEL_REVISION") + .unwrap_or_else(|_| DEFAULT_LOCAL_EMBEDDING_REVISION.to_string()), + local_embedding_dimensions: env::var("LOCAL_EMBEDDING_DIMENSIONS") + .ok() + .map(|value| value.parse::()) + .transpose() + .map_err(|error| anyhow::anyhow!("invalid LOCAL_EMBEDDING_DIMENSIONS: {error}"))? + .unwrap_or(384), // Default on: with warmup off, the first user-facing request paid // the entire cold model load, which is exactly the window in which // the supervisor watchdog used to kill the executor. Warmup moves diff --git a/src/executor.rs b/src/executor.rs index 086d6bf..00fe7f5 100644 --- a/src/executor.rs +++ b/src/executor.rs @@ -45,7 +45,18 @@ enum ExecutorMessage { /// embedding service is constructed. Until this arrives the child cannot /// emit anything at all, so the parent must not hold it to the per-request /// progress watchdog. - Ready, + Ready { + #[serde(default, skip_serializing_if = "Option::is_none")] + protocol_version: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + backend: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + model_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + model_revision: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + dimensions: Option, + }, Progress { request_id: u64, phase: String, @@ -110,6 +121,70 @@ struct OperationBaseline { error: Option, } +const EXECUTOR_PROTOCOL_VERSION: u32 = 1; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecutorIdentity { + pub backend: String, + pub model_id: String, + pub model_revision: String, + pub dimensions: usize, +} + +impl ExecutorIdentity { + fn validate_ready( + &self, + protocol_version: Option, + backend: Option<&str>, + model_id: Option<&str>, + model_revision: Option<&str>, + dimensions: Option, + ) -> Result<()> { + let protocol_version = + protocol_version.context("executor ready missing protocol_version")?; + if protocol_version != EXECUTOR_PROTOCOL_VERSION { + anyhow::bail!( + "executor protocol mismatch: expected {}, got {}", + EXECUTOR_PROTOCOL_VERSION, + protocol_version + ); + } + let backend = backend.context("executor ready missing backend")?; + if backend != self.backend { + anyhow::bail!( + "executor backend mismatch: expected '{}', got '{}'", + self.backend, + backend + ); + } + let model_id = model_id.context("executor ready missing model_id")?; + if model_id != self.model_id { + anyhow::bail!( + "executor model mismatch: expected '{}', got '{}'", + self.model_id, + model_id + ); + } + let model_revision = model_revision.context("executor ready missing model_revision")?; + if model_revision != self.model_revision { + anyhow::bail!( + "executor model revision mismatch: expected '{}', got '{}'", + self.model_revision, + model_revision + ); + } + let dimensions = dimensions.context("executor ready missing dimensions")?; + if dimensions != self.dimensions { + anyhow::bail!( + "executor dimension mismatch: expected {}, got {}", + self.dimensions, + dimensions + ); + } + Ok(()) + } +} + /// Long-lived embedding subprocess supervised by progress messages. The /// watchdog only decides whether the child is responsive; operation success is /// still determined exclusively by durable ledger and part states. @@ -140,11 +215,12 @@ pub struct SupervisedEmbeddingService { ready: AtomicBool, events_tx: broadcast::Sender, child_env: Vec<(String, String)>, + expected_identity: Option, } impl SupervisedEmbeddingService { pub fn new(executable: PathBuf, dimensions: usize, watchdog: Duration) -> Self { - Self::with_child_env(executable, dimensions, watchdog, Vec::new()) + Self::with_child_env_and_identity(executable, dimensions, watchdog, Vec::new(), None) } pub fn with_child_env( @@ -152,6 +228,16 @@ impl SupervisedEmbeddingService { dimensions: usize, watchdog: Duration, child_env: Vec<(String, String)>, + ) -> Self { + Self::with_child_env_and_identity(executable, dimensions, watchdog, child_env, None) + } + + pub fn with_child_env_and_identity( + executable: PathBuf, + dimensions: usize, + watchdog: Duration, + child_env: Vec<(String, String)>, + expected_identity: Option, ) -> Self { let (events_tx, _) = broadcast::channel(1024); let startup_budget = Duration::from_millis( @@ -178,6 +264,7 @@ impl SupervisedEmbeddingService { ready: AtomicBool::new(false), events_tx, child_env, + expected_identity, } } @@ -277,10 +364,27 @@ impl SupervisedEmbeddingService { line.trim() ) })?; - if !matches!(message, ExecutorMessage::Ready) { - anyhow::bail!( + match message { + ExecutorMessage::Ready { + protocol_version, + backend, + model_id, + model_revision, + dimensions, + } => { + if let Some(expected) = &self.expected_identity { + expected.validate_ready( + protocol_version, + backend.as_deref(), + model_id.as_deref(), + model_revision.as_deref(), + dimensions, + )?; + } + } + message => anyhow::bail!( "embedding executor generation {generation} sent {message:?} before readiness" - ); + ), } } Ok(Err(error)) => { @@ -700,16 +804,18 @@ impl EmbeddingService for SupervisedEmbeddingService { .fetch_max(previous.generation.saturating_add(1), Ordering::SeqCst); let mut state = self.state.lock().await; if let Some(child) = state.as_mut() - && child.generation < previous.generation + && child.generation <= previous.generation { - self.stop_child( - child, - None, - ExecutorEventKind::Exited, - "executor generation preceded durable ledger generation".to_owned(), - ) - .await?; - *state = None; + // Generation numbers are supervisor bookkeeping, not part of the + // child protocol. A restarted server begins counting at one, so a + // healthy, pre-warmed child can legitimately have a lower number + // than an operation persisted by an earlier server process. + // Killing that child forced an unnecessary cold MLX launch during + // queue recovery. Adopt it into the durable sequence instead. + let adopted_generation = self.next_generation.fetch_add(1, Ordering::SeqCst); + child.generation = adopted_generation; + self.current_generation + .store(adopted_generation, Ordering::SeqCst); } drop(state); @@ -753,12 +859,56 @@ pub async fn run_embedding_executor() -> Result<()> { #[cfg(not(debug_assertions))] let fixture = false; - let service: Arc = if fixture { - Arc::new(FixtureEmbeddingService) + let (service, identity): (Arc, ExecutorIdentity) = if fixture { + ( + Arc::new(FixtureEmbeddingService), + ExecutorIdentity { + backend: "fixture".to_owned(), + model_id: "fixture".to_owned(), + model_revision: "fixture".to_owned(), + dimensions: 2, + }, + ) } else { let config = crate::config::Config::from_env()?; - Arc::from( + let (backend, model_id, model_revision) = match &config.embedding_provider { + surreal_memory::embeddings::EmbeddingProvider::Local { model_id, .. } => { + if config.local_embedding_backend != crate::config::LocalEmbeddingBackend::Candle { + anyhow::bail!( + "the Rust embedding-executor only supports LOCAL_EMBEDDING_BACKEND=candle" + ); + } + ( + "candle".to_owned(), + model_id.clone(), + config.local_embedding_model_revision.clone(), + ) + } + surreal_memory::embeddings::EmbeddingProvider::OpenAI { model, .. } => { + ("openai".to_owned(), model.clone(), "api".to_owned()) + } + surreal_memory::embeddings::EmbeddingProvider::Cohere { model, .. } => { + ("cohere".to_owned(), model.clone(), "api".to_owned()) + } + #[cfg(feature = "palace")] + surreal_memory::embeddings::EmbeddingProvider::Fast => ( + "fast".to_owned(), + "fastembed".to_owned(), + "bundled".to_owned(), + ), + }; + let service: Arc = Arc::from( surreal_memory::embeddings::create_embedding_service(config.embedding_provider).await?, + ); + let dimensions = service.dimensions(); + ( + service, + ExecutorIdentity { + backend, + model_id, + model_revision, + dimensions, + }, ) }; let stdin = tokio::io::stdin(); @@ -770,7 +920,17 @@ pub async fn run_embedding_executor() -> Result<()> { // process exec, dynamic linking of a large GPU-linked binary, config // parsing, service construction — happens while the child is structurally // incapable of producing output. The parent times that phase separately. - write_message(&mut writer, &ExecutorMessage::Ready).await?; + write_message( + &mut writer, + &ExecutorMessage::Ready { + protocol_version: Some(EXECUTOR_PROTOCOL_VERSION), + backend: Some(identity.backend), + model_id: Some(identity.model_id), + model_revision: Some(identity.model_revision), + dimensions: Some(identity.dimensions), + }, + ) + .await?; while let Some(line) = lines.next_line().await? { let request: ExecutorRequest = serde_json::from_str(&line)?; @@ -914,3 +1074,52 @@ fn exit_fixture_once(text: &str) -> Result<()> { Err(error) => Err(error.into()), } } + +#[cfg(test)] +mod identity_tests { + use super::*; + + fn identity() -> ExecutorIdentity { + ExecutorIdentity { + backend: "mlx".to_owned(), + model_id: "BAAI/bge-small-en-v1.5".to_owned(), + model_revision: "revision".to_owned(), + dimensions: 384, + } + } + + #[test] + fn accepts_exact_executor_identity() { + identity() + .validate_ready( + Some(EXECUTOR_PROTOCOL_VERSION), + Some("mlx"), + Some("BAAI/bge-small-en-v1.5"), + Some("revision"), + Some(384), + ) + .unwrap(); + } + + #[test] + fn rejects_model_revision_mismatch() { + let error = identity() + .validate_ready( + Some(EXECUTOR_PROTOCOL_VERSION), + Some("mlx"), + Some("BAAI/bge-small-en-v1.5"), + Some("other"), + Some(384), + ) + .unwrap_err(); + assert!(error.to_string().contains("model revision mismatch")); + } + + #[test] + fn rejects_missing_identity_fields() { + let error = identity() + .validate_ready(None, None, None, None, None) + .unwrap_err(); + assert!(error.to_string().contains("protocol_version")); + } +} diff --git a/src/main.rs b/src/main.rs index 556e669..ca22a41 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,9 +3,9 @@ use std::sync::Arc; use surreal_memory::storage::migrations::{inspect_legacy_enum_data, repair_legacy_enum_data}; use surreal_memory_server::{ api, - config::Config, + config::{Config, LocalEmbeddingBackend}, embeddings::{self, EmbeddingService, create_embedding_service}, - executor::{SupervisedEmbeddingService, run_embedding_executor}, + executor::{ExecutorIdentity, SupervisedEmbeddingService, run_embedding_executor}, mcp::MemoryMcpServer, storage::{MemoryStorage, create_storage}, workers, @@ -104,8 +104,19 @@ async fn main() -> Result<()> { let embedding_service = init_embedding_service(&config).await?; let retry_config = parse_retry_config_from_env(); + let mlx_backend = matches!( + &config.embedding_provider, + embeddings::EmbeddingProvider::Local { .. } + ) && config.local_embedding_backend == LocalEmbeddingBackend::Mlx; if config.embedding_warmup { - warmup_embedding(Arc::clone(&embedding_service)).await; + let warmup = warmup_embedding(Arc::clone(&embedding_service)).await; + if mlx_backend { + warmup.context("MLX embedding warmup failed; refusing to open the API")?; + } else if let Err(error) = warmup { + tracing::warn!(%error, "Embedding warmup failed; model will load lazily on first use"); + } + } else if mlx_backend { + anyhow::bail!("EMBEDDING_WARMUP must be enabled for the MLX backend"); } let api_port: u16 = std::env::var("API_PORT") @@ -443,7 +454,12 @@ async fn load_config() -> Result { tracing::info!(" Embedding: Cohere ({})", model); } embeddings::EmbeddingProvider::Local { model_id, .. } => { - tracing::info!(" Embedding: Local ({})", model_id); + tracing::info!( + " Embedding: Local ({}, backend={}, revision={})", + model_id, + config.local_embedding_backend.as_str(), + config.local_embedding_model_revision + ); } #[cfg(feature = "palace")] embeddings::EmbeddingProvider::Fast => { @@ -460,14 +476,76 @@ async fn init_embedding_service(config: &Config) -> Result { + if dimensions != config.local_embedding_dimensions { + anyhow::bail!( + "configured local embedding dimensions {} do not match model dimensions {}", + config.local_embedding_dimensions, + dimensions + ); + } + let executable = match config.local_embedding_backend { + LocalEmbeddingBackend::Candle => std::env::current_exe() + .context("resolve server executable for Candle embedding supervisor")?, + LocalEmbeddingBackend::Mlx => { + let path = config.local_embedding_executor.as_ref().context( + "LOCAL_EMBEDDING_EXECUTOR is required when LOCAL_EMBEDDING_BACKEND=mlx", + )?; + std::fs::canonicalize(path).with_context(|| { + format!("resolve MLX embedding executor '{}'", path.display()) + })? + } + }; + ( + executable, + ExecutorIdentity { + backend: config.local_embedding_backend.as_str().to_owned(), + model_id: model_id.clone(), + model_revision: config.local_embedding_model_revision.clone(), + dimensions, + }, + ) + } + embeddings::EmbeddingProvider::OpenAI { model, .. } => ( + std::env::current_exe().context("resolve server embedding executor")?, + ExecutorIdentity { + backend: "openai".to_owned(), + model_id: model.clone(), + model_revision: "api".to_owned(), + dimensions, + }, + ), + embeddings::EmbeddingProvider::Cohere { model, .. } => ( + std::env::current_exe().context("resolve server embedding executor")?, + ExecutorIdentity { + backend: "cohere".to_owned(), + model_id: model.clone(), + model_revision: "api".to_owned(), + dimensions, + }, + ), + #[cfg(feature = "palace")] + embeddings::EmbeddingProvider::Fast => ( + std::env::current_exe().context("resolve server embedding executor")?, + ExecutorIdentity { + backend: "fast".to_owned(), + model_id: "fastembed".to_owned(), + model_revision: "bundled".to_owned(), + dimensions, + }, + ), + }; let watchdog_ms = std::env::var("SURREAL_EXECUTOR_WATCHDOG_MS") .ok() .and_then(|value| value.parse().ok()) .unwrap_or(30_000); - let service = SupervisedEmbeddingService::new( - std::env::current_exe().context("resolve server executable for embedding supervisor")?, + let service = SupervisedEmbeddingService::with_child_env_and_identity( + executable, dimensions, std::time::Duration::from_millis(watchdog_ms), + Vec::new(), + Some(identity), ); tracing::info!( "🧠 Embedding service configured ({} dimensions); model loads on warmup or first use", @@ -500,16 +578,18 @@ async fn run_mcp_server(storage: Arc) -> Result<()> { /// model will be retried lazily on first use. Readiness is reported via /// `EmbeddingService::is_ready()`, so `/health` stays accurate regardless of /// whether warmup ran or succeeded. -async fn warmup_embedding(service: Arc) { +async fn warmup_embedding(service: Arc) -> Result<()> { tracing::info!("🔥 Warming up embedding model..."); - match service.embed("warmup").await { - Ok(_) => { - tracing::info!("✅ Embedding model warmed up and ready"); - } - Err(e) => { - tracing::warn!(error = %e, "Embedding warmup failed; model will load lazily on first use"); - } + let embedding = service.embed("warmup").await?; + if embedding.len() != service.dimensions() { + anyhow::bail!( + "embedding warmup dimension mismatch: expected {}, got {}", + service.dimensions(), + embedding.len() + ); } + tracing::info!("✅ Embedding model warmed up and ready"); + Ok(()) } async fn run_api_server( diff --git a/tests/executor_recovery.rs b/tests/executor_recovery.rs index 1e22b7f..6b802ff 100644 --- a/tests/executor_recovery.rs +++ b/tests/executor_recovery.rs @@ -3,6 +3,7 @@ use std::{path::PathBuf, sync::Arc, time::Duration}; use futures_util::StreamExt; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; +use surreal_memory::embeddings::ExecutorSnapshot; use surreal_memory::{ EmbeddingService, MemoryStorage, SurrealStorage, embeddings::ExecutorEventKind, }; @@ -303,3 +304,38 @@ async fn server_restart_advances_the_persisted_executor_generation() { second_executor.terminate_idle_executor().await.unwrap(); std::fs::remove_dir_all(marker_dir).unwrap(); } + +#[tokio::test] +async fn durable_generation_adopts_the_pre_warmed_child_without_restarting_it() { + let executable = PathBuf::from(env!("CARGO_BIN_EXE_surreal-memory-server")); + let executor = SupervisedEmbeddingService::with_child_env( + executable, + 2, + Duration::from_secs(10), + vec![("SURREAL_EXECUTOR_FIXTURE".to_owned(), "1".to_owned())], + ); + + executor.embed("warmup").await.unwrap(); + let warm = executor.executor_snapshot().unwrap(); + let persisted = ExecutorSnapshot { + generation: warm.generation + 10, + progress_seq: warm.progress_seq, + exit_count: warm.exit_count, + last_exit: None, + error: None, + }; + + executor + .prepare_operation("resume-with-warm-child", &persisted) + .await + .unwrap(); + let adopted = executor.executor_snapshot().unwrap(); + assert!(adopted.generation > persisted.generation); + assert_eq!(adopted.exit_count, warm.exit_count); + + executor.embed("still-warm").await.unwrap(); + let after = executor.executor_snapshot().unwrap(); + assert_eq!(after.generation, adopted.generation); + assert_eq!(after.exit_count, warm.exit_count); + executor.terminate_idle_executor().await.unwrap(); +}