Skip to content

fix(mlx-worker): ship mlx.metallib and keep MLX errors off the protocol channel - #47

Open
zxch3n wants to merge 1 commit into
mainfrom
fix/mlx-worker-metallib
Open

fix(mlx-worker): ship mlx.metallib and keep MLX errors off the protocol channel#47
zxch3n wants to merge 1 commit into
mainfrom
fix/mlx-worker-metallib

Conversation

@zxch3n

@zxch3n zxch3n commented Aug 18, 2026

Copy link
Copy Markdown
Member

Problem

The MLX VLM worker could never start from swift build output. Command-line SwiftPM cannot compile Metal shaders (xcodebuild-only, per the mlx-swift README), so no default.metallib existed anywhere in .build. The worker's first MLX call died in Device::Device with "Failed to load the default metallib", and mlx-c's default error handler printed MLX error: ... to stdout — the worker's line-delimited JSON protocol channel — before calling exit(-1). afterrayd surfaced this as model worker returned invalid output: MLX worker wrote non-protocol stdout, and the app fell back to "language model is not configured".

Changes

  • scripts/build-mlx-metallib.sh (new): compiles the mlx-swift checkout's generated Metal shaders into mlx.metallib with the same flags as mlx's CMake kernel build, then colocates it with the built worker binaries — the first place mlx's load_colocated_library looks. Rebuilds only when shaders or Package.resolved change. Preflights the Xcode 26+ Metal component with a clear remedy (xcodebuild -downloadComponent MetalToolchain).
  • scripts/run-v0.sh / scripts/build-release.sh: call the script after building afterray-mlx-vlm-worker and ship the library in the app bundle. It lives in Contents/Resources/mlx.metallib behind a Contents/Helpers/mlx.metallib symlink, because a data file directly inside Helpers fails codesign --verify --deep --strict as unsigned nested code (verified empirically).
  • swift/AfterRayMlxVlmWorker/Sources/WorkerCore.swift: installMlxErrorHandler() installs a process-wide mlx-c error sink that logs to stderr and exits, instead of mlx-c's default stdout printf. Uses the deprecated MLX.setErrorHandler deliberately — it is mlx-swift's only process-wide hook; task-scoped withError/withErrorHandler still take precedence and are unaffected. The executable installs it at startup.
  • Tests (MlxWorkerProcessSafetyTests): (1) an unscoped MLX broadcast error reaches the installed sink — previously it would have hit stdout and killed the process; (2) the colocated mlx.metallib actually loads the Metal default library and runs a GPU op (skipped with instructions when the script has not run).
  • AGENTS.md updates for scripts/, swift/AfterRayMlxVlmWorker/, apps/AfterRayMlxVlmWorker/.

Verification

  • swift test --filter MlxWorkerProcessSafetyTests — both new tests pass.
  • Protocol smoke test against the signed, bundled worker (generate without load, malformed line, bad model_dir): stdout stays pure JSON, logs on stderr.
  • scripts/run-v0.sh --build-only — full dev build, bundle assembly, and codesign --verify --deep --strict all pass with the metallib included.
  • make test (cargo workspace + swift test + capture shim) passes.
  • Not verified: real-model generation end to end — this machine has no verified Qwen3.5 snapshot yet (the 4B download is still partial). The Metal-side mechanism is covered by the new in-process test; the model path is unchanged.
  • Pre-existing, unrelated: cargo clippy --workspace --all-targets -- -D warnings is already red on main (cast_sign_loss in crates/afterray-codec/src/jpeg.rs). This PR touches no Rust code.

Confidence

Medium-high. Both root causes are verified at the mechanism level with tests and a full signed dev build. The release pipeline change mirrors the dev flow but build-release.sh itself was not run (needs notarization credentials); notarization of the symlinked resource layout follows standard bundle rules but has not been exercised.

…ol channel

The MLX worker never worked from `swift build` output: command-line SwiftPM
cannot compile Metal shaders (xcodebuild-only, per the mlx-swift README), so
no default.metallib existed anywhere and the first MLX call died in
Device::Device with "Failed to load the default metallib". mlx-c's default
error handler then printed `MLX error: ...` to stdout — the worker's
line-delimited JSON channel — and exit(-1)'d, which afterrayd reported as
"MLX worker wrote non-protocol stdout".

- scripts/build-mlx-metallib.sh compiles the checkout's generated shaders
  (same flags as mlx's CMake kernel build) into mlx.metallib and colocates it
  with the built worker binaries, where mlx's load order looks first.
  run-v0.sh / build-release.sh call it after building the worker. In the app
  bundle the library lives in Contents/Resources behind a Helpers symlink —
  a data file inside Helpers fails codesign --verify --deep --strict.
- installMlxErrorHandler() installs a process-wide mlx-c error sink that
  logs to stderr before exiting; scoped `withError` handling is unaffected.

Verified: new worker-core tests (metallib loads the Metal default library;
unscoped broadcast error reaches the sink), protocol smoke test on the
signed bundled worker, make test, and run-v0.sh --build-only including
codesign verification. Not verified: real-model generation — no verified
Qwen3.5 snapshot on this machine yet (4B download still partial).

Model: kimi-code/kimi-for-coding
Harness: lody
@zxch3n zxch3n added the risk:middle Medium correctness risk label Aug 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

risk:middle Medium correctness risk

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant