omlx-router: a stdlib-only proxy that swaps the resident large model in oMLX, so one Apple silicon box can serve several big models from any client

README.md

omlx-router

A ~200-line, stdlib-only proxy that lets you keep several large models configured in oMLX on one Apple silicon machine and switch between them from any client — a laptop, a phone, a cron job — without touching the machine itself.

The problem

Unified memory is not all yours. Metal reports a ceiling (max_recommended_working_set_size) well below installed RAM — on a 48GB M4 Pro it is 37.44GB — and oMLX refuses any prompt whose projected KV cache would push past ~90% of that.

Two 16–19GB models therefore load fine and then cannot serve. Worse, oMLX does not evict to make room: it loads the second model alongside the first and fails the prefill afterwards, so you find out at request time. Swapping has to be explicit, which normally means an SSH session or the admin UI — exactly the thing you were trying to avoid by putting the models on a server.

What this does

It sits between your clients and oMLX. When a request names a large model that is not the resident one, the router unloads the other large models, pins the requested one, and only then forwards the request. Everything else is proxied byte-for-byte, streaming included.

Clients need no changes beyond a port number. Model selection is the switch.

Install

Save omlx-router.py anywhere and make it executable:

chmod +x omlx-router.py
OMLX_BIG_MODELS="Model-A,Model-B,Model-C" ./omlx-router.py

OMLX_BIG_MODELS is the list of models that cannot be resident at the same time. Names must match oMLX's model ids exactly (GET /v1/models). Leave out anything small — embeddings, a 4B, a vision model — because those coexist happily and should never trigger a swap.

Then point your clients at :8010 instead of :8000. Over a tailnet:

tailscale serve --bg --https=9443 http://127.0.0.1:8010

Run it at login (macOS)

Save omlx-router.plist to ~/Library/LaunchAgents/, edit the two paths and the model list, then:

launchctl bootstrap gui/$(id -u) ~/Library/LaunchAgents/omlx-router.plist

It deliberately runs under /usr/bin/python3. The script is stdlib-only, and Apple's system python has a stable path that is already allow-listed in the macOS application firewall — a Homebrew python bumping 3.14.6 → 3.14.7 silently revokes inbound access, because the firewall's allow-list is keyed on the binary path.

Configuration

Variable	Default
OMLX_UPSTREAM	http://127.0.0.1:8000	oMLX base URL
OMLX_ROUTER_HOST	0.0.0.0	bind address
OMLX_ROUTER_PORT	8010	listen port
OMLX_BIG_MODELS	(empty)	mutually exclusive model ids; empty means proxy only, no swapping
OMLX_API_KEY	from ~/.omlx/settings.json	admin key
OMLX_ROUTER_DRAIN_TIMEOUT	900	seconds to wait for in-flight requests before unloading anyway

Three things that are easy to get wrong

Unloading a model that is mid-generation. oMLX aborts those requests outright (Aborted N requests: model is being unloaded). Someone switching model on their phone would kill a long generation running from a laptop. The router counts in-flight requests per model and waits for the outgoing model to drain before unloading it.

Buffering the token stream. read(8192) blocks until it has 8192 bytes or hits EOF, which batches SSE tokens into lumps and makes streaming feel broken. Event streams are read a line at a time so tokens forward as they land.

Assuming the body always completes. If upstream cuts a body short the naive proxy raises mid-stream and the client sees a chunk-decode error. The router catches it and terminates the chunked encoding cleanly, so the client sees a truncated response instead.

Caveats

• Anything that talks to oMLX directly on :8000 bypasses the router and can still wedge the machine. Binding oMLX to 127.0.0.1 and exposing only the router makes that structurally impossible.
• A swap costs a model load — a few seconds — paid by whoever asks first.
• oMLX admin endpoints take a session cookie, not the bearer token the OpenAI-compatible routes use. That's why there's a login step.

MIT.

omlx-router.plist

<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>Label</key>
    <string>omlx-router</string>
    <key>ProgramArguments</key>
    <array>
        <!-- Apple's system python, deliberately: it is already allow-listed in
         the macOS application firewall, and its path does not move when
         Homebrew bumps python (3.14.6 -> 3.14.7 silently broke inbound access
         on a tailnet). The router is stdlib-only, so 3.9 is fine. -->
        <string>/usr/bin/python3</string>
        <string>/PATH/TO/omlx-router.py</string>
    </array>
    <key>RunAtLoad</key>
    <true/>
    <key>KeepAlive</key>
    <true/>
    <key>ThrottleInterval</key>
    <integer>10</integer>
    <key>StandardOutPath</key>
    <string>/PATH/TO/router.log</string>
    <key>StandardErrorPath</key>
    <string>/PATH/TO/router.log</string>
    <key>EnvironmentVariables</key>
    <dict>
        <key>OMLX_UPSTREAM</key>
        <string>http://127.0.0.1:8000</string>
        <key>OMLX_ROUTER_PORT</key>
        <string>8010</string>
        <key>OMLX_BIG_MODELS</key>
        <string>Model-A,Model-B,Model-C</string>
    </dict>
</dict>
</plist>

omlx-router.py

#!/usr/bin/env python3
"""Transparent proxy in front of oMLX that swaps the resident large model.

On a machine with unified memory, only one large model fits at a time. Metal
caps the working set well below installed RAM (`max_recommended_working_set_size`
-- 37.44GB on a 48GB Apple silicon box) and oMLX rejects any prompt that would
push past ~90% of that. oMLX will not evict to make room: it loads the second
model *alongside* the first and then fails the prefill. So switching has to be
explicit, and doing it by hand means being sat at the machine -- which defeats
the point of reaching it from a laptop, a phone, or a cron job.

This sits in front of oMLX and does the switch in the request path. A request
naming a large model that is not the resident one triggers unload/pin/load
before being forwarded. Every other request is proxied untouched, including
streaming responses. Point clients at this port instead of oMLX's; nothing else
changes.

Requires nothing outside the standard library.

Setup
-----
1. List the models that must not coexist:

       export OMLX_BIG_MODELS="Model-A,Model-B,Model-C"

   Names must match oMLX's model ids exactly. Small models (embeddings, a 4B,
   a vision model) are left out -- they coexist happily and are proxied without
   any swap. With the variable unset the router still works, purely as a
   transparent proxy.

2. Run it:

       OMLX_BIG_MODELS="..." ./omlx-router.py

   It listens on :8010 and forwards to oMLX on :8000 by default. Both are
   configurable (see the environment variables below).

3. Point every client at the router's port instead of oMLX's. To serve it over
   TLS on a tailnet:

       tailscale serve --bg --https=9443 http://127.0.0.1:8010

Environment
-----------
    OMLX_UPSTREAM               oMLX base URL         (http://127.0.0.1:8000)
    OMLX_ROUTER_HOST            bind address          (0.0.0.0)
    OMLX_ROUTER_PORT            listen port           (8010)
    OMLX_BIG_MODELS             comma-separated ids   (empty: no swapping)
    OMLX_API_KEY                admin key             (else ~/.omlx/settings.json)
    OMLX_ROUTER_DRAIN_TIMEOUT   seconds to wait for
                                in-flight requests    (900)
"""
from __future__ import annotations

import collections
import http.client
import json
import logging
import os
import threading
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

UPSTREAM = os.environ.get("OMLX_UPSTREAM", "http://127.0.0.1:8000").rstrip("/")
LISTEN_HOST = os.environ.get("OMLX_ROUTER_HOST", "0.0.0.0")
LISTEN_PORT = int(os.environ.get("OMLX_ROUTER_PORT", "8010"))

# The models that cannot be resident at the same time. Anything not named here
# -- embeddings, small instruct models, a vision model -- coexists fine and is
# proxied without ever triggering a swap.
BIG_MODELS = {m.strip() for m in os.environ.get("OMLX_BIG_MODELS", "").split(",")
              if m.strip()}

SWAP_PATHS = ("/v1/chat/completions", "/v1/completions", "/v1/responses")
HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade", "content-length",
}

log = logging.getLogger("omlx-router")
_swap_lock = threading.Lock()

# Requests currently being served upstream, per model. A swap must not unload a
# model that is mid-generation: oMLX aborts those requests outright
# ("Aborted N requests: model is being unloaded"), so a long prefill started
# from one machine would be killed by someone else switching model on another.
_inflight: dict[str, int] = collections.defaultdict(int)
_inflight_lock = threading.Condition()
DRAIN_TIMEOUT = float(os.environ.get("OMLX_ROUTER_DRAIN_TIMEOUT", "900"))


def _enter(model: str | None) -> None:
    if not model:
        return
    with _inflight_lock:
        _inflight[model] += 1


def _leave(model: str | None) -> None:
    if not model:
        return
    with _inflight_lock:
        _inflight[model] -= 1
        if _inflight[model] <= 0:
            _inflight.pop(model, None)
        _inflight_lock.notify_all()


def _wait_drain(model: str, timeout: float = DRAIN_TIMEOUT) -> bool:
    """Block until `model` has no in-flight requests. False if it timed out."""
    deadline = time.monotonic() + timeout
    with _inflight_lock:
        while _inflight.get(model, 0) > 0:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return False
            log.info("waiting for %d in-flight request(s) on %s to finish",
                     _inflight[model], model)
            _inflight_lock.wait(min(remaining, 5.0))
    return True


def api_key() -> str:
    env = os.environ.get("OMLX_API_KEY")
    if env:
        return env
    try:
        cfg = json.loads((Path.home() / ".omlx" / "settings.json").read_text())
        return cfg.get("auth", {}).get("api_key", "")
    except OSError:
        return ""


def _admin_session() -> str | None:
    """Log in to the admin API and return the session cookie.

    oMLX's admin endpoints take a session cookie, not the bearer token the
    OpenAI-compatible routes use.
    """
    body = json.dumps({"api_key": api_key()}).encode()
    req = urllib.request.Request(
        f"{UPSTREAM}/admin/api/login", data=body,
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            cookie = resp.headers.get("Set-Cookie", "")
            return cookie.split(";", 1)[0] if cookie else None
    except (urllib.error.URLError, TimeoutError):
        return None


def _admin(path: str, cookie: str, method: str = "GET", payload: dict | None = None):
    req = urllib.request.Request(
        f"{UPSTREAM}{path}", method=method,
        data=json.dumps(payload).encode() if payload is not None else None,
        headers={"Content-Type": "application/json", "Cookie": cookie},
    )
    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            return json.loads(resp.read().decode() or "{}")
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as exc:
        log.warning("admin %s %s failed: %s", method, path, exc)
        return None


def active_big() -> str | None:
    """Which large model oMLX currently treats as default.

    `/health` is unauthenticated and cheap. The router sets `is_default`
    whenever it pins, so this is an accurate proxy for "the resident large
    model" without an admin login on every request. If the model has since
    idled out, oMLX reloads it on demand -- still no swap needed.
    """
    try:
        with urllib.request.urlopen(f"{UPSTREAM}/health", timeout=10) as resp:
            return json.loads(resp.read().decode()).get("default_model")
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError):
        return None


def ensure_resident(target: str) -> None:
    """Make `target` the one resident large model, if it is not already."""
    with _swap_lock:
        if active_big() == target:
            return
        cookie = _admin_session()
        if not cookie:
            log.error("cannot reach admin API; forwarding without swap")
            return
        for other in BIG_MODELS - {target}:
            # Let anything already generating on the outgoing model finish
            # before pulling it out from under them.
            if not _wait_drain(other):
                log.warning("%s still busy after drain timeout; unloading anyway", other)
            # Unpin first or the unload is refused.
            _admin(f"/admin/api/models/{other}/settings", cookie, "PUT",
                   {"is_pinned": False, "is_default": False})
            _admin(f"/admin/api/models/{other}/unload", cookie, "POST")
            log.info("unloaded %s", other)
        _admin(f"/admin/api/models/{target}/settings", cookie, "PUT",
               {"is_pinned": True, "is_default": True})
        log.info("pinned %s", target)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "omlx-router"

    def log_message(self, fmt, *args):
        log.info("%s - %s", self.address_string(), fmt % args)

    def _relay(self, method: str) -> None:
        tracked: str | None = None
        length = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(length) if length else None

        # The only thing this proxy actually inspects.
        if method == "POST" and self.path in SWAP_PATHS and body:
            try:
                model = json.loads(body).get("model")
            except (json.JSONDecodeError, UnicodeDecodeError):
                model = None
            if model in BIG_MODELS:
                ensure_resident(model)
                tracked = model

        headers = {k: v for k, v in self.headers.items()
                   if k.lower() not in HOP_BY_HOP}
        _enter(tracked)
        req = urllib.request.Request(
            UPSTREAM + self.path, data=body, headers=headers, method=method)
        try:
            resp = urllib.request.urlopen(req, timeout=1800)
        except urllib.error.HTTPError as e:
            resp = e
        except (urllib.error.URLError, TimeoutError) as e:
            _leave(tracked)
            self.send_error(502, f"upstream unreachable: {e}")
            return

        try:
            with resp:
                self.send_response(resp.status)
                for k, v in resp.headers.items():
                    if k.lower() not in HOP_BY_HOP:
                        self.send_header(k, v)
                self.send_header("Transfer-Encoding", "chunked")
                self.end_headers()
                # `read(n)` blocks until n bytes or EOF, which batches SSE
                # tokens into lumps and makes the client feel laggy. For event
                # streams read a line at a time so tokens forward as they land.
                ctype = resp.headers.get("Content-Type") or ""
                read = resp.readline if "text/event-stream" in ctype else (
                    lambda: resp.read(8192))
                try:
                    while chunk := read():
                        self.wfile.write(b"%X\r\n%s\r\n" % (len(chunk), chunk))
                        self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError):
                    return
                except (http.client.IncompleteRead, http.client.HTTPException,
                        ValueError, OSError) as exc:
                    # Upstream cut the body short -- usually because the model
                    # was unloaded mid-generation. Terminate the chunked stream
                    # cleanly so the client sees a truncated response rather
                    # than a chunk-decode error.
                    log.warning("upstream body truncated (%s): %s",
                                type(exc).__name__, exc)
                try:
                    self.wfile.write(b"0\r\n\r\n")
                    self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError, ValueError):
                    pass
        finally:
            _leave(tracked)

    def do_GET(self):    self._relay("GET")
    def do_POST(self):   self._relay("POST")
    def do_PUT(self):    self._relay("PUT")
    def do_DELETE(self): self._relay("DELETE")
    def do_PATCH(self):  self._relay("PATCH")
    def do_OPTIONS(self):self._relay("OPTIONS")


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    srv = ThreadingHTTPServer((LISTEN_HOST, LISTEN_PORT), Handler)
    srv.daemon_threads = True
    log.info("omlx-router on %s:%s -> %s", LISTEN_HOST, LISTEN_PORT, UPSTREAM)
    if BIG_MODELS:
        log.info("swapping between: %s", ", ".join(sorted(BIG_MODELS)))
    else:
        log.warning("OMLX_BIG_MODELS is empty: proxying without swapping")
    srv.serve_forever()


if __name__ == "__main__":
    main()