feat: select and apply MLX backend per shard (#958)

[oliver0006]

Motivation

Addresses #958 ("Support all MLX backends").

MLX supports several computational backends (Metal, CUDA, CPU), and the placement layer already knows which backends every node supports (state.node_backends, INSTANCE_META_BACKENDS) — but that decision was never recorded on the instance, and runners never called mx.set_default_device(). Backend selection only worked implicitly, by whatever default the local MLX build picked. Recording the backend per shard and applying it early in the runner lifecycle unlocks explicit support for heterogeneous clusters (e.g. Macs + Linux CUDA/CPU nodes in one ring) without introducing a new runner, as suggested in the issue.

Changes

• ShardMetadata.backend: Backend | None (default None) — records the compute backend a shard should run on; None preserves today's behavior end to end.
• assign_shard_backends() in placement_utils.py — pure function that picks, for every runner, the first backend of the instance-meta preference order (Metal > Cuda > Cpu, intersected with the model card's backends) that its node supports.
• place_instance() calls it right after get_shard_assignments(), so all instance types get per-runner backends.
• apply_shard_backend() in the runner bootstrap.py — maps the assigned backend to an MLX device and calls mx.set_default_device() before the engine is built (MlxMetal/MlxCuda → mx.gpu, MlxCpu → mx.cpu, anything else → no-op).

Why It Works

The master is the only place that knows both the node capabilities and the model requirements, so the backend decision is made there and travels to the runner inside the shard metadata it already receives. Setting the default device at the very top of the runner entrypoint (before MlxBuilder/MfluxBuilder construction) means every subsequent MLX allocation lands on the assigned device. The field defaults to None and unassigned shards change nothing, so old serialized states and any path that doesn't set a backend behave exactly as before.

Test Plan

Manual Testing

Hardware: Fedora Linux, AMD Ryzen (16 threads) + Radeon 780M iGPU — deliberately a non-Apple box, since that's what this issue is about. Installed the mlx-cpu wheel (matching the mlx-cpu extra) and verified apply_shard_backend(Backend.MlxCpu) flips mx.default_device() to CPU on a Linux host.

Automated Testing

• test_assign_shard_backends_picks_first_preferred_backend_per_node — pure-function behavior incl. preference order and nodes with no matching backend.
• test_placement_assigns_a_backend_to_every_shard — end-to-end place_instance() on a heterogeneous 3-node ring (Metal/CUDA/CPU nodes) asserts each runner gets its node's backend.
• test_apply_shard_backend_* — MLX default device is set for MlxCpu and untouched for None/Vllm (device restored after each test).
• test_master expectation updated: the created instance's shard now carries backend=MlxMetal, which asserts the full master→instance propagation.

uv run basedpyright (0 errors), uv run ruff check, ruff format --check all pass; uv run pytest src → 344 passed, 5 skipped (the 13 collection errors are pre-existing on Linux: mflux and mlx_lm.models.deepseek_v4 aren't importable in this environment and are untouched by this change). nix fmt wasn't runnable on this machine (no nix without root) — ruff format was used instead.

---

Disclosure: co-written by a human (@oliver0006) directing the work and Claude AI (Fable 5) writing code, following AGENTS.md and .clauderules.

🤖 Generated with Claude Code