How we got Linux VM container workloads running natively on the Mac so they can use the Apple Silicon GPU (mps/MLX): how we did it and the sandbox tradeoff
(Disclosure: I work at Union AI and maintain Flyte.)
Flyte is an open source AI runtime with a durable execution engine where you can train, fine-tune, and serve AI models and agents.
Problem: Flyte's local devbox runs every task and app in a container. On a Mac that container is a Linux VM, and a Linux VM has no Metal. So torch.backends.mps.is_available() is False, MLX doesn't run, and any ML code you iterate on locally falls back to the CPU. You can't pass the GPU through: there's nothing in the VM to pass it to.
What we did: a task or app that requests flyte.Resources(gpu="AppleSilicon") runs its code natively on the Mac instead. The container starts as normal, then its entrypoint hands its own invocation to a bridge running on the host. The bridge builds a cached uv environment from the image's pip packages, rewrites the in-cluster endpoints to the devbox's host ports, and runs it. The container streams the logs, mirrors the exit code, forwards aborts, and proxies the app port, so the console and the run look normal.
On an M5 Max: a small CNN trained on MNIST in 7 s on mps; a Qwen2.5-0.5B app served from mps at roughly 84–97 tokens/s; a 4096×4096 fp32 matmul ran at 13.4 TFLOPS on mps vs 1.6 on the CPU.
Things worth knowing:
- Running workload code on your laptop means running it outside any sandbox. So it's opt-in (flyte start devbox --apple-silicon), and the flag binds every devbox port to 127.0.0.1.
- Only pip packages carry over to the Mac, not apt packages or the rest of the image. Mounted k8s secrets don't either.
- Asking for the GPU on a devbox without the flag, an Intel Mac, Linux, or a real cluster fails before the task starts with a named error. We chose that over a quiet CPU fallback.
- If you only need completions and want to keep container isolation, there's a container mode that gives the task a host mlx_lm.server behind OPENAI_BASE_URL.
PRs for implementation details:
- flyte#8181 (backend)
- flyte-sdk#1685 (SDK, examples, proof runbook)
- flyte2-ui#42 (console).
Check out the docs to give it a try: union.ai
Curious whether others have gone the "run it on the host" route for local GPU access, or found a way to keep the sandbox.
5 points · 2 comments
Mentions
How we got Linux VM container workloads running natively on the Mac so they can use the Apple Silicon GPU (mps/MLX): how we did it and the sandbox tradeoff (Disclosure: I work at Union AI and maintain Flyte.) Flyte is an open source AI runtime with a durable execution engine where you can train, fine-tune, and serve AI models and agents. Problem: Flyte's local devbox runs every task and app in a container. On a Mac that container is a Linux VM, and a Linux VM has no Metal. So torch.backends.mps.is_available() is False, MLX doesn't run, and any ML code you iterate on locally falls back to the CPU. You can't pass the GPU through: there's nothing in the VM to pass it to. What we did: a task or app that requests flyte.Resources(gpu="AppleSilicon") runs its code natively on the Mac instead. The container starts as normal, then its entrypoint hands its own invocation to a bridge running on the host. The bridge builds a cached uv environment from the image's pip packages, rewrites the in-cluster endpoints to the devbox's host ports, and runs it. The container streams the logs, mirrors the exit code, forwards aborts, and proxies the app port, so the console and the run look normal. On an M5 Max: a small CNN trained on MNIST in 7 s on mps; a Qwen2.5-0.5B app served from mps at roughly 84–97 tokens/s; a 4096×4096 fp32 matmul ran at 13.4 TFLOPS on mps vs 1.6 on the CPU. Things worth knowing: - Running workload code on your laptop means running it outside any sandbox. So it's opt-in (flyte start devbox --apple-silicon), and the flag binds every devbox port to 127.0.0.1. - Only pip packages carry over to the Mac, not apt packages or the rest of the image. Mounted k8s secrets don't either. - Asking for the GPU on a devbox without the flag, an Intel Mac, Linux, or a real cluster fails before the task starts with a named error. We chose that over a quiet CPU fallback. - If you only need completions and want to keep container isolation, there's a container mode that gives the task a host mlx_lm.server behind OPENAI_BASE_URL. PRs for implementation details: - flyte#8181 (backend) - flyte-sdk#1685 (SDK, examples, proof runbook) - flyte2-ui#42 (console). Check out the docs to give it a try: union.ai Curious whether others have gone the "run it on the host" route for local GPU access, or found a way to keep the sandbox.