From bb1161854451b40694373919fde15b43f0b6966b Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sat, 29 Aug 2026 06:40:33 -0700 Subject: [PATCH 1/3] fix: wire up Gemma4MTPBench target + repair regressed mlx-swift-lm/mlx-swift submodule pins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mlx-swift-lm was uncommitted-pinned backward past the commit that introduced DualModelMTP/MTPTokenIterator/Gemma4AssistantModel, breaking Gemma4MTPBench and forcing MTP speculative decoding to be commented out in InferenceEngine.swift as a workaround. Repointing both submodules to their current origin/main tips (mlx-swift-lm past the merged DSA stage-2 PR #61, mlx-swift past the matching MLXFast.fromFp8 addition) restores those types and lets the benchmark build again — no source workaround needed. Also adds scripts/bootstrap_local_tests.sh, which mirrors CI's "Install MLX Metal library" step (pip install mlx, copy its bundled metallib into every built .xctest bundle) so `swift test` is runnable locally without the manual cmake+make dance. Addresses the Tier 3 item in #128. Co-Authored-By: Claude Sonnet 5 --- Package.swift | 15 ++++++++++++++- mlx-swift | 2 +- mlx-swift-lm | 2 +- scripts/bootstrap_local_tests.sh | 33 ++++++++++++++++++++++++++++++++ 4 files changed, 49 insertions(+), 3 deletions(-) create mode 100755 scripts/bootstrap_local_tests.sh diff --git a/Package.swift b/Package.swift index 9286564..21562c1 100644 --- a/Package.swift +++ b/Package.swift @@ -9,7 +9,8 @@ let package = Package( .library(name: "DFlash", targets: ["DFlash"]), .executable(name: "SwiftLM", targets: ["SwiftLM"]), .executable(name: "SwiftBuddy", targets: ["SwiftBuddy"]), - .executable(name: "DFlashKernelBench", targets: ["DFlashKernelBench"]) + .executable(name: "DFlashKernelBench", targets: ["DFlashKernelBench"]), + .executable(name: "Gemma4MTPBench", targets: ["Gemma4MTPBench"]) ], dependencies: [ // Local Apple MLX Swift fork for C++ extensions @@ -53,6 +54,18 @@ let package = Package( ], path: "Sources/DFlashKernelBench" ), + // ── Gemma4 MTP Speculative Decoding Benchmark ─────────────── + .executableTarget( + name: "Gemma4MTPBench", + dependencies: [ + "MLXInferenceCore", + .product(name: "MLX", package: "mlx-swift"), + .product(name: "MLXLLM", package: "mlx-swift-lm"), + .product(name: "MLXLMCommon", package: "mlx-swift-lm"), + .product(name: "ArgumentParser", package: "swift-argument-parser"), + ], + path: "Sources/Gemma4MTPBench" + ), // ── STFT Audio Profiling Testing Script (macOS only) ─────────── .executableTarget( name: "SwiftLMTestSTFT", diff --git a/mlx-swift b/mlx-swift index 133864c..5639a6d 160000 --- a/mlx-swift +++ b/mlx-swift @@ -1 +1 @@ -Subproject commit 133864c733c8d4178547f8fe92897da6a788368f +Subproject commit 5639a6d9e6a7ab785e102d88d741879f529fce56 diff --git a/mlx-swift-lm b/mlx-swift-lm index b9bf50b..f8da831 160000 --- a/mlx-swift-lm +++ b/mlx-swift-lm @@ -1 +1 @@ -Subproject commit b9bf50bdafef02fffd5b83598a61bbf7d47434f9 +Subproject commit f8da83166361e49a63fc3df690c92df0146d08ee diff --git a/scripts/bootstrap_local_tests.sh b/scripts/bootstrap_local_tests.sh new file mode 100755 index 0000000..f3770c6 --- /dev/null +++ b/scripts/bootstrap_local_tests.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Makes `swift test` runnable locally without CI's help. +# +# A bare `swift test` aborts with "Failed to load the default metallib" +# because Package.swift links MLX but nothing on a local machine ever builds +# or installs mlx.metallib. CI works around this in .github/workflows/ci.yml +# ("Install MLX Metal library" step) by pip-installing the `mlx` wheel and +# copying its bundled metallib into every built .xctest bundle. This script +# does the same thing locally. +set -eo pipefail + +VENV_DIR="${MLX_METALLIB_VENV:-/tmp/swiftlm_mlx_venv}" + +echo "=> Building test harness (swift build --build-tests)..." +swift build --build-tests + +echo "=> Installing MLX Metal library..." +if [ ! -d "$VENV_DIR" ]; then + python3 -m venv "$VENV_DIR" +fi +"$VENV_DIR/bin/pip" install --quiet --upgrade mlx + +METALLIB=$(find "$VENV_DIR" -name "mlx.metallib" | head -1) +if [ -z "$METALLIB" ]; then + echo "error: mlx.metallib not found after pip install mlx" >&2 + exit 1 +fi + +cp "$METALLIB" .build/debug/ 2>/dev/null || true +cp "$METALLIB" .build/release/ 2>/dev/null || true +find .build -type d -name "MacOS" -exec cp "$METALLIB" {}/ \; + +echo "=> Done. Run tests with: swift test --skip-build" From f89176ef322b562cd1705fd2c9473ababfcff339 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 30 Aug 2026 23:23:04 -0700 Subject: [PATCH 2/3] feat: emit exiting{reason:"requested"} on clean shutdown (Aegis Engine Protocol v1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aegis-AI's daemon (aegis-daemon, W7 Rust rewrite) needs to tell a clean, planned shutdown apart from a real crash. Today SwiftLM's SIGTERM/SIGINT handlers print plain text and exit(0) with no machine-readable signal at all. Adds a small shared emitEvent(_:) helper (used by both the existing ready event and this new one, replacing ready's inline JSONSerialization call) and emits {"event":"exiting","reason":"requested"} on both signal handlers before exiting. Consumers should ignore unrecognized future reason values rather than error on them. This is Phase 1 of a larger plan (see Aegis-AI's daemon/spec/engine-protocol.md §3/§4) — self-reporting a SPECIFIC failure reason on a real crash (OOM, model load failure, etc.) is separate, larger follow-up work, since none of those paths have any error handling to build on today. Verified: swift build --target SwiftLM (debug) and swift build -c release --product SwiftLM both green. Smoke-tested against a real Aegis-AI daemon build in an isolated sandbox: reached ready normally, and a real SIGTERM via the daemon's stop endpoint exited the process cleanly with the new event observed, without the daemon misclassifying it as a crash. --- Sources/SwiftLM/Server.swift | 30 +++++++++++++++++++++++++----- docs/AEGIS_INTEGRATION.md | 12 ++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 3f26703..7a4a98a 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -189,6 +189,21 @@ final class ProgressTracker { } } +/// Emit one machine-readable JSON-lines event on stdout for the Aegis-AI +/// daemon to consume (`docs/AEGIS_INTEGRATION.md`, and the daemon's own +/// `daemon/spec/engine-protocol.md`). Shared by the `ready` event and the +/// `exiting` event so the JSON-encode-and-flush boilerplate isn't +/// duplicated at each call site — this is deliberately the SAME manual +/// `JSONSerialization` shape the pre-existing `ready` event already used, +/// not a new encoding convention. +func emitEvent(_ payload: [String: Any]) { + if let data = try? JSONSerialization.data(withJSONObject: payload), + let json = String(data: data, encoding: .utf8) { + print(json) + fflush(stdout) + } +} + @main struct MLXServer: AsyncParsableCommand { static let configuration = CommandConfiguration( @@ -1026,13 +1041,16 @@ struct MLXServer: AsyncParsableCommand { } readyEvent["partition"] = info } - if let data = try? JSONSerialization.data(withJSONObject: readyEvent), - let json = String(data: data, encoding: .utf8) { - print(json) - fflush(stdout) - } + emitEvent(readyEvent) // ── Graceful shutdown on SIGTERM/SIGINT ── + // Engine Protocol v1 (`daemon/spec/engine-protocol.md` §3, TD-7): + // emit `exiting{reason:"requested"}` before exiting so the daemon + // can tell a planned, no-error stop apart from a real failure — + // `requested` has no `ExitClassification` equivalent on the daemon + // side (a daemon-requested stop never reaches that classifier at + // all), it exists purely for the daemon to positively confirm this + // was a clean shutdown, not infer it from absence of other signals. let shutdownSource = DispatchSource.makeSignalSource(signal: SIGTERM, queue: .main) let interruptSource = DispatchSource.makeSignalSource(signal: SIGINT, queue: .main) signal(SIGTERM, SIG_IGN) @@ -1040,10 +1058,12 @@ struct MLXServer: AsyncParsableCommand { shutdownSource.setEventHandler { print("\n[SwiftLM] Received SIGTERM, shutting down gracefully...") + emitEvent(["event": "exiting", "reason": "requested"]) Darwin.exit(0) } interruptSource.setEventHandler { print("\n[SwiftLM] Received SIGINT, shutting down gracefully...") + emitEvent(["event": "exiting", "reason": "requested"]) Darwin.exit(0) } shutdownSource.resume() diff --git a/docs/AEGIS_INTEGRATION.md b/docs/AEGIS_INTEGRATION.md index 1bd0927..9143d51 100644 --- a/docs/AEGIS_INTEGRATION.md +++ b/docs/AEGIS_INTEGRATION.md @@ -47,6 +47,18 @@ The server will emit a machine-readable JSON ready event on stdout when it is re Aegis-AI should **wait for this event** before routing any requests to the server. +The server also emits a machine-readable JSON `exiting` event on stdout +before a clean, planned shutdown (SIGTERM/SIGINT): + +```json +{"event":"exiting","reason":"requested"} +``` + +`reason: "requested"` means this shutdown was expected (e.g. Aegis-AI +called `stop`) — not a crash. This is the first of what may grow into a +small set of self-reported exit reasons; consumers should ignore any +`reason` value they don't recognize rather than treat it as an error. + --- ## 🧠 Running 122B+ MoE Models (Critical) From 9f8c2c71d3bc3fc5f4ad2c10531349fe7e4cdbff Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Mon, 31 Aug 2026 07:19:53 -0700 Subject: [PATCH 3/3] fix(code-review): 4 confirmed findings from PR #164's review pass - emitEvent() used to silently swallow any JSON encoding failure with no diagnostic trail, for a function that's now the sole path for a protocol-critical signal. Now logs a message to stderr on failure. - Ignore SIGPIPE before the SIGTERM/SIGINT handlers are wired up: if the daemon's read end of our stdout pipe is already gone by the time a shutdown signal arrives, the exiting-event print/fflush could raise SIGPIPE, whose default disposition kills the process via signal instead of reaching Darwin.exit(0) -- producing exactly the ambiguous "was this a crash?" signature this feature exists to eliminate. - docs/AEGIS_INTEGRATION.md: reworded reason:"requested" to not imply it proves the daemon itself initiated the shutdown (any SIGTERM/SIGINT sender produces the identical event), and added that the event is best-effort (dispatched on the main queue, so a busy queue delays emission) rather than a synchronous guarantee at signal-delivery time. Co-Authored-By: Claude Sonnet 5 --- Sources/SwiftLM/Server.swift | 28 +++++++++++++++++++++++----- docs/AEGIS_INTEGRATION.md | 22 +++++++++++++++++----- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 7a4a98a..92c8750 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -196,12 +196,21 @@ final class ProgressTracker { /// duplicated at each call site — this is deliberately the SAME manual /// `JSONSerialization` shape the pre-existing `ready` event already used, /// not a new encoding convention. +/// +/// Code-review finding: encoding failure used to be silently swallowed — +/// for a protocol-critical signal the daemon is meant to positively rely +/// on (not just infer), a dropped emit with zero diagnostic trail would be +/// hard to ever notice. Logs to stderr on failure instead. func emitEvent(_ payload: [String: Any]) { - if let data = try? JSONSerialization.data(withJSONObject: payload), - let json = String(data: data, encoding: .utf8) { - print(json) - fflush(stdout) - } + guard let data = try? JSONSerialization.data(withJSONObject: payload), + let json = String(data: data, encoding: .utf8) + else { + FileHandle.standardError.write( + Data("[SwiftLM] failed to encode event for stdout: \(payload)\n".utf8)) + return + } + print(json) + fflush(stdout) } @main @@ -1055,6 +1064,15 @@ struct MLXServer: AsyncParsableCommand { let interruptSource = DispatchSource.makeSignalSource(signal: SIGINT, queue: .main) signal(SIGTERM, SIG_IGN) signal(SIGINT, SIG_IGN) + // Code-review finding: if the daemon's read end of our stdout pipe + // is already gone by the time a shutdown signal arrives (e.g. the + // daemon itself already crashed), the exiting-event print()/fflush + // below can raise SIGPIPE — whose default disposition kills this + // process via signal instead of reaching Darwin.exit(0), producing + // exactly the ambiguous "was this a crash?" signature this feature + // exists to eliminate. Ignore SIGPIPE so a closed pipe surfaces as + // an ordinary EPIPE write error instead. + signal(SIGPIPE, SIG_IGN) shutdownSource.setEventHandler { print("\n[SwiftLM] Received SIGTERM, shutting down gracefully...") diff --git a/docs/AEGIS_INTEGRATION.md b/docs/AEGIS_INTEGRATION.md index 9143d51..e8300d8 100644 --- a/docs/AEGIS_INTEGRATION.md +++ b/docs/AEGIS_INTEGRATION.md @@ -48,16 +48,28 @@ The server will emit a machine-readable JSON ready event on stdout when it is re Aegis-AI should **wait for this event** before routing any requests to the server. The server also emits a machine-readable JSON `exiting` event on stdout -before a clean, planned shutdown (SIGTERM/SIGINT): +when it receives SIGTERM/SIGINT: ```json {"event":"exiting","reason":"requested"} ``` -`reason: "requested"` means this shutdown was expected (e.g. Aegis-AI -called `stop`) — not a crash. This is the first of what may grow into a -small set of self-reported exit reasons; consumers should ignore any -`reason` value they don't recognize rather than treat it as an error. +`reason: "requested"` means *some* SIGTERM/SIGINT was delivered — it does +not distinguish who sent it. The common case is Aegis-AI calling `stop`, +but a manual `kill`/`pkill` or an external process manager produces the +identical event; don't treat `requested` as proof the daemon itself +initiated the shutdown. + +This event is **best-effort, not guaranteed**: it's emitted from a signal +handler dispatched on the main queue, so if the main queue is busy (e.g. +mid-request) when the signal arrives, emission is delayed until the queue +frees up — it is not synchronous at signal-delivery time. A consumer that +times out waiting for it and force-kills the process should not treat the +absence of this event as proof the shutdown wasn't requested. + +This is the first of what may grow into a small set of self-reported exit +reasons; consumers should ignore any `reason` value they don't recognize +rather than treat it as an error. ---