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/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 3f26703..92c8750 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -189,6 +189,30 @@ 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. +/// +/// 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]) { + 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 struct MLXServer: AsyncParsableCommand { static let configuration = CommandConfiguration( @@ -1026,24 +1050,38 @@ 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) 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...") + 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..e8300d8 100644 --- a/docs/AEGIS_INTEGRATION.md +++ b/docs/AEGIS_INTEGRATION.md @@ -47,6 +47,30 @@ 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 +when it receives SIGTERM/SIGINT: + +```json +{"event":"exiting","reason":"requested"} +``` + +`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. + --- ## 🧠 Running 122B+ MoE Models (Critical) 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"