diff --git a/.gitignore b/.gitignore index bf804e0..eee3526 100644 --- a/.gitignore +++ b/.gitignore @@ -227,3 +227,4 @@ benchmarks/cross_framework # local e2e/bench artifacts (harnesses may run with repo cwd) /results/ +/.build diff --git a/README.md b/README.md index 2a56a08..9ac13e6 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,22 @@ uv venv && source .venv/bin/activate uv pip install -e ".[accel]" ``` +### macOS (Apple Silicon, Metal) + +The native engine is CUDA-only; on a Mac `ft serve` (same as `ft serve-metal`) +runs the OpenAI/Anthropic/Responses API backed by Apple's Metal runtimes +(MLX or llama.cpp) instead of porting the CUDA scheduler: + +```bash +uv venv && source .venv/bin/activate +uv pip install -e . # core only; skips CUDA-only deps +uv pip install mlx-lm # Metal engine (or: brew install llama.cpp) +ft serve --model mlx-community/Qwen3-0.6B-4bit +``` + +See [Install on macOS](https://github.com/FlashML-org/FreeToken/blob/main/docs/install.md) and +[serve-metal in the CLI reference](https://github.com/FlashML-org/FreeToken/blob/main/docs/cli.md). + For More details: - [Install FreeToken](https://github.com/FlashML-org/FreeToken/blob/main/docs/install.md) diff --git a/benchmarks/bench_decode_moe.py b/benchmarks/bench_decode_moe.py index 723c2be..cf26ab3 100644 --- a/benchmarks/bench_decode_moe.py +++ b/benchmarks/bench_decode_moe.py @@ -77,7 +77,13 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p.add_argument( "--backend", default="offload", - help="comma list of offload|cpu|hybrid; one server per backend", + help="comma list of offload|cpu|hybrid|metal; one server per backend (metal = Apple Silicon)", + ) + p.add_argument( + "--metal-engine", + default="mlx", + choices=["mlx", "llama"], + help="Metal engine for --backend metal (mlx or llama.cpp; default mlx)", ) p.add_argument( "--aime", @@ -172,6 +178,16 @@ def free_port() -> int: def serve_cmd(args: argparse.Namespace, backend: str, port: int) -> list[str]: + if backend == "metal": + # Apple Silicon: the same serving path through the Metal backend (mlx / + # llama.cpp upstream). CUDA engine flags do not apply, and the Metal + # server reports readiness through /health like the CUDA one. + return [ + sys.executable, "-m", "freetoken.cli", "serve-metal", + "--model", args.model, + "--backend", args.metal_engine, + "--host", "127.0.0.1", "--port", str(port), + ] cmd = [ sys.executable, "-m", "freetoken.cli", "serve", "--model", args.model, @@ -209,7 +225,10 @@ def wait_ready(origin: str, proc: subprocess.Popen, log_path: str, timeout: floa continue if health.get("status") == "error": die_with_log(f"server reported startup error: {health}", log_path) - if health.get("maintenance") == "serving": + # The CUDA server reports its lifecycle via `maintenance`; the Metal + # backend is "serving" as soon as /health is ok (its upstream engine + # finished loading before the API ever bound). + if health.get("maintenance") in ("serving", None) and health.get("status") == "ok": return time.sleep(1.0) die_with_log(f"server not ready after {timeout:.0f}s", log_path) @@ -339,7 +358,15 @@ def run_one(args: argparse.Namespace, backend: str) -> dict: sys.exit(f"[bench] need >=2 token events to measure decode, got {len(stamps)}") completion = usage["completion_tokens"] if completion != args.decode: - print(f"[bench] WARNING: completion_tokens={completion} != --decode {args.decode}", flush=True) + why = ( + " -- Metal upstream stops at EOS (ignore_eos is not honored)" + if backend == "metal" + else "" + ) + print( + f"[bench] WARNING: completion_tokens={completion} != --decode {args.decode}{why}", + flush=True, + ) steps = completion - 1 decode_time = stamps[-1] - stamps[0] gaps = sorted((b - a) * 1e3 for a, b in zip(stamps, stamps[1:])) @@ -369,6 +396,12 @@ def run_one(args: argparse.Namespace, backend: str) -> dict: f"(event p50 {row['event_ms_p50']:.3f} / p99 {row['event_ms_p99']:.3f} ms, " f"{len(stamps)} events)") print(f" vram (server) : {row['vram_gib']:8.2f} GiB") + if backend == "metal": + print( + " note : Metal/mlx coalesces tokens into fewer SSE events and\n" + " does not honor ignore_eos, so tok/s is an upper bound\n" + " over coalesced bursts, and steps may be < --decode." + ) sha_note = "greedy" if args.greedy else "sampled, per-server deterministic" print(f" output sha1 : {row['output_sha1']} ({sha_note}; compare across backends)") print(f" output sample : {r['text'][:240]!r}") @@ -378,7 +411,7 @@ def run_one(args: argparse.Namespace, backend: str) -> dict: def main(argv: list[str] | None = None) -> int: args = parse_args(argv) backends = [b.strip() for b in args.backend.split(",") if b.strip()] - unknown = [b for b in backends if b not in ("offload", "cpu", "hybrid")] + unknown = [b for b in backends if b not in ("offload", "cpu", "hybrid", "metal")] if unknown: sys.exit(f"unknown backend(s): {unknown}") diff --git a/docs/cli.md b/docs/cli.md index cf4b27a..57f0744 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -7,6 +7,7 @@ ft [args] | Command | Purpose | |---|---| | `ft serve` | Start the API server (OpenAI `/v1/*`, Anthropic `/v1/messages`, Responses) | +| `ft serve-metal` | Start the same API surface over an Apple Silicon Metal backend (mlx or llama.cpp) | | `ft shell` | Chat with a server in the terminal | | `ft ctl` | Query and manage a running server over HTTP | | `ft launch` | Configure and launch a coding agent against a server | @@ -81,6 +82,42 @@ See [models.md](models.md#moe-backends) for what each backend does. | `--reasoning-parser` | auto | Splits chain-of-thought into `reasoning_content`; auto-inferred; `off` disables | | `--enable-cache-report` | off | Report prefix-cache hits in each response's usage block | +## ft serve-metal + +`ft serve-metal` runs the same OpenAI `/v1/*`, Anthropic `/v1/messages`, and +Responses API surface, but backs it with an **Apple Silicon Metal runtime** +instead of FreeToken's CUDA scheduler. It does not import FreeToken's CUDA +stack, so it works standalone on a Mac (no triton/flashinfer/sglang-kernel, no +`flashlib`) by proxying to a Metal engine that already implements the protocol. + +```bash +ft serve-metal --model mlx-community/Qwen3-0.6B-4bit --backend mlx --port 1919 +ft serve-metal --model ~/models/MyModel.Q4_K_M.gguf --backend llama --port 1919 +``` + +The API surface (chat/completions/models/health, streaming included) is +bit-for-bit the same FreeToken wire, so the desktop app, `ft shell`, `ft launch` +agents, and OpenAI/Anthropic clients all work unchanged. + +| Flag | Default | Meaning | +|---|---|---| +| `--model` | required | Model path or HF id for the Metal engine | +| `--backend` | `auto` | `mlx` (Apple's `mlx_lm.server`) or `llama` (llama.cpp `llama-server`); `auto` prefers mlx | +| `--host` | 127.0.0.1 | FreeToken API bind address | +| `--port` | 1919 | FreeToken API port (user-facing) | +| `--metal-port` | 0 | Upstream Metal engine port; a free loopback port is chosen when `0` | + +Notes: + +- **mlx** requires the model loadable by `mlx-lm` (HF-converted or MLX weights, + e.g. `mlx-community/*`). `mlx_lm` uses the Apple GPU (Metal) directly. +- **llama** requires the `llama-server` binary on PATH and a GGUF model (the + same format FreeToken already parses). It runs with the Metal backend + (`-ngl 999`). +- `ft serve --backend mlx|llama` also switches between these backends on a CUDA + box, but still imports the CUDA config stack; `ft serve-metal` is the + standalone, CUDA-free form for Apple Silicon. + ## ft shell ```bash diff --git a/docs/install.md b/docs/install.md index f5205ab..584479f 100644 --- a/docs/install.md +++ b/docs/install.md @@ -2,11 +2,22 @@ ## Requirements +**Linux (CUDA — full engine)** + - Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13) - Python >= 3.10, with [uv](https://docs.astral.sh/uv/) recommended (plain `pip` + `venv` works too) -## Method 1: Install from PyPI +**macOS (Apple Silicon — Metal backend)** + +- macOS on Apple Silicon (M1/M2/M3/M4) +- Python >= 3.10 +- One of the Metal engines: `mlx-lm` (recommended) or `llama.cpp`'s + `llama-server` on PATH +- No CUDA, triton, flashinfer, sglang-kernel, or `flashlib` needed — the + Metal path does not import FreeToken's CUDA stack at all + +## Method 1: Install from PyPI (Linux/CUDA) ```bash uv venv && source .venv/bin/activate @@ -15,7 +26,7 @@ uv pip install "freetoken[accel]" CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH. -## Method 2: Install from source +## Method 2: Install from source (Linux/CUDA) ```bash git clone https://github.com/FlashML-org/FreeToken.git && cd FreeToken @@ -23,7 +34,67 @@ uv venv && source .venv/bin/activate uv pip install -e ".[accel]" ``` -## Verify +## Method 3: macOS / Apple Silicon (Metal backend) + +FreeToken's native engine is CUDA-only. On a Mac, `ft serve-metal` runs the +same API surface backed by Apple's own Metal runtimes instead, so nothing +needs porting: it launches Apple's `mlx_lm.server` or llama.cpp's +`llama-server` as the upstream engine and proxies the FreeToken wire surface. + +```bash +git clone https://github.com/FlashML-org/FreeToken.git && cd FreeToken +uv venv && source .venv/bin/activate + +# Core package only — skips flashlib, triton, sglang-kernel and the CUDA-linked +# torch index pins, none of which have macOS wheels. +uv pip install -e . + +# Then add one Metal engine: +uv pip install mlx-lm # option A: MLX (Apple's framework) +brew install llama.cpp # option B: llama.cpp's llama-server (Metal) +``` + +> **Why the CUDA deps are skipped:** `flashlib`, `triton==3.6.0`, and the cu130 +> torch pins are marked `platform_system == 'Linux'` in pyproject.toml — they +> only serve the native scheduler, and the Metal path never imports them. A +> plain `uv pip install -e .` resolves everything else from PyPI. + +```bash +# Verify (no model needed): +ft --version +ft serve --help # on Apple Silicon this is the Metal backend +# or: ft serve-metal --help +``` + +Or one shot from this repo: + +```bash +./install.sh # Metal path on Darwin arm64; CUDA wheels on Linux +ft serve --model mlx-community/Qwen3-0.6B-4bit +``` + +### Troubleshooting: "no wheels with a matching Python implementation tag" + +Your venv is an **x86_64 build running under Rosetta**, not native arm64 — no +arm64-only wheel (mlx, macOS torch) can ever install into it. This happens when +`uv venv` / `python3 -m venv` is run inside an x86_64 (Intel) terminal session. + +Fix — recreate the venv natively: + +```bash +deactivate 2>/dev/null; rm -rf .venv +# make sure the shell itself is native: arch -arm64 zsh (or use an arm64 Terminal) +uv venv && source .venv/bin/activate +file .venv/bin/python # must say: arm64 (not x86_64) +uv pip install -e . && uv pip install mlx-lm +``` + +Also run installs from the repo root — `uv pip install -e .` needs the +`pyproject.toml` in the current directory, which is why a stray +`error: Requesting extras requires a pyproject.toml ...` means you are in the +wrong directory (or asked for a nonexistent extra; the only extra is `dev`). + +## Verify (Linux/CUDA) ```bash source .venv/bin/activate diff --git a/docs/quickstart.md b/docs/quickstart.md index bea8d2e..234512a 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -13,6 +13,31 @@ and MoE backends, cache sizes, tool-call and reasoning parsers — resolves from the checkpoint and the GPU; see [cli.md](cli.md) for the flags. The server is ready when the log reaches `API server is ready to serve on 127.0.0.1:1919`. +### macOS (Apple Silicon) + +```bash +scripts/start-metal.sh # serve + chat in one command +scripts/start-metal.sh mlx-community/Llama-3.2-1B-Instruct-4bit +scripts/start-metal.sh ~/models/MyModel.Q4_K_M.gguf --backend llama +``` + +`scripts/start-metal.sh` checks the environment, starts the Metal server, +waits for readiness, and opens `ft shell` for interactive testing (the server +stops when the chat exits; `NO_CHAT=1` runs the API alone). Or drive the +pieces directly: + +```bash +ft serve-metal --model mlx-community/Qwen3-0.6B-4bit --backend mlx +``` + +`ft serve-metal` serves the same API surface as `ft serve`, backed by MLX or +llama.cpp on Metal (see [install.md](install.md#method-3-macos--apple-silicon-metal-backend)). +`--model` takes an MLX/HF repo id (`mlx-community/*`) or, with `--backend llama`, +a local GGUF file. The rest of this page (`/v1/models`, chat completions, `ft shell`, +`ft launch`) works unchanged against it — including `ft shell --model `, which +starts the Metal engine and chats in one process on a Mac, and the `/model ` +shell command to switch models without restarting. + ## Send a request Check what is being served: @@ -49,7 +74,9 @@ ft shell --model ~/models/Qwen3.6-35B-A3B # start an engine and chat, one proc ``` `/help` lists the in-shell commands. Attach mode needs no GPU, so it also drives -a server on another machine (`--server URL`). +a server on another machine (`--server URL`). On a Mac, `ft shell --model ` +starts the Metal backend instead (see [quickstart.md](quickstart.md)); `/model ` +switches the served model on servers that support it (Metal). ## Use a coding agent diff --git a/install.sh b/install.sh index a7df832..49e0459 100755 --- a/install.sh +++ b/install.sh @@ -1,11 +1,14 @@ #!/usr/bin/env bash # -# FreeToken engine installer (Linux, NVIDIA CUDA) — user-facing, wheel-based. +# FreeToken engine installer — user-facing. # -# Installs the `freetoken` runtime (the `ft` CLI) and its prebuilt kernel-cache -# wheel into a managed venv, then wires it up so FreeToken Desktop can find it. -# Dependencies come from PyPI via uv, except torch and sglang-kernel whose cu130 -# wheels live on dedicated indexes (see CU_INDEX_ARGS below). +# Linux (NVIDIA CUDA): installs the `freetoken` runtime (the `ft` CLI) and its +# prebuilt kernel-cache wheel into a managed venv, then wires it up so FreeToken +# Desktop can find it. Dependencies come from PyPI via uv, except torch and +# sglang-kernel whose cu130 wheels live on dedicated indexes (see CU_INDEX_ARGS). +# +# macOS (Apple Silicon): no CUDA. Installs the core package + mlx-lm and runs +# the Metal backend (`ft serve` routes to serve-metal). No kernel-cache wheel. # # Typical use (once a release exists): # curl -fsSL https:///install.sh | bash @@ -153,6 +156,70 @@ else fi say "uv $("$UV" --version | awk '{print $2}')" +install_metal_macos() { + [ "$(uname -s)" = Darwin ] || return 1 + [ "$(uname -m)" = arm64 ] || die "FreeToken on macOS requires Apple Silicon (arm64). This machine is $(uname -m)." + + local script_dir src + script_dir="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd -P)" + + say "macOS Apple Silicon — Metal backend (no CUDA, no kernel-cache wheel)" + mkdir -p "$FT_HOME" + "$UV" venv "$VENV" --python "$PY_VERSION" --clear + + if [ -f "$script_dir/pyproject.toml" ]; then + src="$script_dir" + say "installing from source checkout: $src" + "$UV" pip install --python "$VENV" -e "$src" + else + # Pin the Metal branch until it lands on default. A piped `curl | bash` of + # this script has no adjacent checkout, so "git+...FreeToken.git" would + # otherwise install CUDA-only main and fail on macOS. Prefer FlashML-org; + # fall back to the working fork while the PR is open. + src="git+https://github.com/FlashML-org/FreeToken.git@feat/apple-metal-backend" + src_fork="git+https://github.com/jasonkneen/FreeToken.git@feat/apple-metal-backend" + say "installing from $src" + if ! "$UV" pip install --python "$VENV" "$src"; then + say "FlashML-org does not have feat/apple-metal-backend yet; using $src_fork" + "$UV" pip install --python "$VENV" "$src_fork" + fi + fi + say "installing mlx-lm (Apple Metal engine)" + "$UV" pip install --python "$VENV" mlx-lm + + local ft_bin="$VENV/bin/ft" + [ -x "$ft_bin" ] || die "install finished but $ft_bin is missing." + mkdir -p "$BIN_DIR" + ln -sf "$ft_bin" "$BIN_DIR/ft" + say "symlinked $BIN_DIR/ft -> $ft_bin" + mkdir -p "$ENV_DIR" + printf 'FREETOKEN_FT_BIN=%s\n' "$ft_bin" > "$ENV_DIR/50-freetoken.conf" + say "wrote $ENV_DIR/50-freetoken.conf (FREETOKEN_FT_BIN)" + + if "$ft_bin" --help >/dev/null 2>&1; then + say "self-check: \`ft --help\` OK" + else + warn "self-check: \`ft --help\` returned non-zero — inspect with: $ft_bin --help" + fi + + cat <=77", "torch>=2.11,<2.12", "wheel"] +requires = [ + "setuptools>=77", + "torch>=2.11,<2.12; platform_system == 'Linux'", + "wheel", +] build-backend = "setuptools.build_meta" [project] @@ -22,6 +27,7 @@ classifiers = [ "Intended Audience :: Developers", "Intended Audience :: Science/Research", "Operating System :: POSIX :: Linux", + "Operating System :: MacOS", "Environment :: GPU :: NVIDIA CUDA", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", @@ -38,9 +44,10 @@ dependencies = [ "einops>=0.8,<1", "fastapi>=0.115,<1", # slot_cache: the device-side LRU admission kernel behind the MoE expert cache. - "flashlib==0.3.0", # pin it + "flashlib==0.3.0; platform_system == 'Linux'", "gguf>=0.19,<1", "huggingface_hub>=1.5,<2", + "httpx>=0.27,<1", # Metal backend upstream proxy (server/metal.py) "msgpack>=1.1,<2", "modelscope>=1.37,<2", # ceiling: numba (via flashlib) needs numpy<2.5 @@ -54,7 +61,9 @@ dependencies = [ # floor+ceiling: sglang-kernel 0.4.5 links libtorch symbols only 2.11 has. # PyPI's torch 2.11.0 wheel is itself the cu130 build, so plain pip resolves # correctly from PyPI alone; uv additionally pins the index below. - "torch>=2.11,<2.12", + # macOS/Metal: torch is only needed by the CUDA scheduler path; the Metal + # backend (ft serve-metal / ft serve --backend mlx|llama) never imports it. + "torch>=2.11,<2.12; platform_system == 'Linux'", "tqdm>=4.66,<5", "transformers>=5.5,<6", "triton==3.6.0; platform_system == 'Linux'", @@ -90,7 +99,10 @@ accel = ["freetoken[fi,sgl]"] # (the torch index also mirrors stale copies of common deps, e.g. packaging<=24.1, # which would otherwise shadow PyPI under uv's first-index strategy). [tool.uv.sources] -torch = { index = "pytorch-cu130" } +# Linux/CUDA only: the metal path resolves torch (if ever needed) from PyPI. +torch = [ + { index = "pytorch-cu130", marker = "platform_system == 'Linux'" }, +] sglang-kernel = { index = "sglang-cu130" } [[tool.uv.index]] @@ -115,6 +127,7 @@ where = ["python"] [tool.setuptools.package-data] "*" = ["csrc/**/*", "moe/configs/**/*.json"] +"freetoken.server" = ["gemma4_chat_template.jinja"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/python/freetoken/cli.py b/python/freetoken/cli.py index 4e6deff..927cb50 100644 --- a/python/freetoken/cli.py +++ b/python/freetoken/cli.py @@ -11,6 +11,8 @@ def _print_help(file: TextIO) -> None: Commands: serve Start the FreeToken API server + serve-metal Start the FreeToken API surface over an Apple Silicon Metal + backend (mlx or llama.cpp); standalone, no CUDA import needed shell Chat with a FreeToken server in the terminal ctl Query and manage a running FreeToken server daemon Run the FreeToken supervisor (persistent engine service) @@ -25,12 +27,25 @@ def _print_help(file: TextIO) -> None: def _run_serve(argv: list[str]) -> int: + # Apple Silicon has no CUDA torch in the Metal venv. The classic launcher + # imports torch at module load, so route to serve-metal *before* that import. + # Linux/CUDA is unchanged. + if sys.platform == "darwin": + from freetoken.server.metal_main import main + + return main(argv) from freetoken.server import launch_server launch_server(argv=argv, prog="ft serve") return 0 +def _run_serve_metal(argv: list[str]) -> int: + from freetoken.server.metal_main import main + + return main(argv) + + def _run_shell(argv: list[str]) -> int: from freetoken.shell import main @@ -92,6 +107,7 @@ def _run_bench(argv: list[str]) -> int: COMMANDS = { "serve": "_run_serve", + "serve-metal": "_run_serve_metal", "shell": "_run_shell", "ctl": "_run_ctl", "daemon": "_run_daemon", diff --git a/python/freetoken/daemon/osproc.py b/python/freetoken/daemon/osproc.py index acf6447..7f0200a 100644 --- a/python/freetoken/daemon/osproc.py +++ b/python/freetoken/daemon/osproc.py @@ -81,7 +81,13 @@ def proc_pgid(pid: int) -> int | None: (b) assert ``pgid == pid`` before a group-kill so we never signal an unrelated group.""" fields = _stat_fields(pid) if fields is None or len(fields) < 3: - return None + # macOS and other non-/proc POSIX hosts still expose the process group + # through getpgid(2). Serve children are created with + # start_new_session=True, so pgid == pid remains the safety invariant. + try: + return os.getpgid(pid) + except (ProcessLookupError, PermissionError, OSError): + return None try: return int(fields[2]) # field 5 == fields[5-3] except (ValueError, IndexError): # pragma: no cover - defensive diff --git a/python/freetoken/daemon/serve_manager.py b/python/freetoken/daemon/serve_manager.py index 78cea11..6d72c86 100644 --- a/python/freetoken/daemon/serve_manager.py +++ b/python/freetoken/daemon/serve_manager.py @@ -122,8 +122,13 @@ def build_serve_command( ) -> tuple[list[str], str]: """The serve invocation + its log path. ``python -m freetoken.cli serve`` (NOT ``-m freetoken``, which is a direct-server entrypoint that ignores subcommand argv) so it uses the - daemon's own interpreter/venv with no PATH dependency.""" - argv = [python, "-m", "freetoken.cli", "serve", "--model", model, "--port", str(port), *args] + daemon's own interpreter/venv with no PATH dependency. + + On Apple Silicon there is no CUDA engine to serve; the same CLI argv is routed + to ``serve-metal`` instead (the Metal backend ignores CUDA-only engine flags + it does not understand). On Linux the classic invocation is unchanged.""" + subcommand = "serve-metal" if sys.platform == "darwin" else "serve" + argv = [python, "-m", "freetoken.cli", subcommand, "--model", model, "--port", str(port), *args] log_path = os.path.join(log_dir, f"serve-{port}.log") return argv, log_path diff --git a/python/freetoken/logging.py b/python/freetoken/logging.py new file mode 100644 index 0000000..9cc62ca --- /dev/null +++ b/python/freetoken/logging.py @@ -0,0 +1,110 @@ +"""Lightweight logging setup that is safe to import without the CUDA stack.""" + +from __future__ import annotations + +from functools import partial +from typing import TYPE_CHECKING + +_LOG_LEVEL = None + + +def init_logger( + name: str, + suffix: str = "", + *, + strip_file: bool = True, + level: str | None = None, + use_pid: bool | None = None, + use_tp_rank: bool | None = None, +): + """Initialize the logger for the module with colors and pretty formatting.""" + import logging + import os + import sys + + global _LOG_LEVEL + if _LOG_LEVEL is None: + level_map = { + "DEBUG": logging.DEBUG, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "ERROR": logging.ERROR, + "CRITICAL": logging.CRITICAL, + } + level = level or os.getenv("LOG_LEVEL", "").upper() + _LOG_LEVEL = level_map.get(level, logging.INFO) + + if strip_file: + suffix = os.path.basename(suffix) + if suffix: + suffix = f"|{suffix}" + + if use_pid is None: + use_pid = os.getenv("LOG_PID", "0").lower() in ("1", "true", "yes") + if use_pid: + suffix = f"|pid={os.getpid()}{suffix}" + + tp_info = None + + class ColorFormatter(logging.Formatter): + """Formatter with colored levels and optional process/rank suffixes.""" + + COLORS = { + "DEBUG": "\033[36m", + "INFO": "\033[32m", + "WARNING": "\033[33m", + "ERROR": "\033[31m", + "CRITICAL": "\033[35m", + } + RESET = "\033[0m" + BOLD = "\033[1m" + + def format(self, record): + nonlocal tp_info + try: + from freetoken.distributed import try_get_tp_info + + tp_info = tp_info or try_get_tp_info() + except Exception: # noqa: BLE001 -- torch may be absent on Metal + tp_info = None + timestamp = self.formatTime(record, "[%Y-%m-%d|%H:%M:%S{suffix}]") + if tp_info is not None and use_tp_rank is not False: + real_suffix = f"{suffix}|core|rank={tp_info.rank}" + else: + real_suffix = suffix + timestamp = timestamp.format(suffix=real_suffix) + level_color = self.COLORS.get(record.levelname, "") + colored_level = f"{level_color}{record.levelname:<8}{self.RESET}" + return f"{self.BOLD}{timestamp}{self.RESET} {colored_level} {record.getMessage()}" + + logger = logging.getLogger(name) + logger.setLevel(_LOG_LEVEL) + logger.handlers.clear() + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(ColorFormatter()) + logger.addHandler(handler) + logger.propagate = False + + def _call_rank0(msg, *args, _which, **kwargs): + from freetoken.distributed import try_get_tp_info + + nonlocal tp_info + tp_info = tp_info or try_get_tp_info() + if tp_info is None or tp_info.is_primary(): + getattr(logger, _which)(msg, *args, **kwargs) + + if TYPE_CHECKING: + + class WrapperLogger(logging.Logger): + def info_rank0(self, msg, *args, **kwargs): ... + def warning_rank0(self, msg, *args, **kwargs): ... + def debug_rank0(self, msg, *args, **kwargs): ... + def critical_rank0(self, msg, *args, **kwargs): ... + + return WrapperLogger(name) + + logger.info_rank0 = partial(_call_rank0, _which="info") + logger.debug_rank0 = partial(_call_rank0, _which="debug") + logger.critical_rank0 = partial(_call_rank0, _which="critical") + logger.warning_rank0 = partial(_call_rank0, _which="warning") + return logger diff --git a/python/freetoken/server/__init__.py b/python/freetoken/server/__init__.py index aa9b473..035c222 100644 --- a/python/freetoken/server/__init__.py +++ b/python/freetoken/server/__init__.py @@ -1,3 +1,19 @@ -from .launch import launch_server +# Lazy: ``launch`` (and the CUDA scheduler it wires) imports torch, which is not +# installed on macOS/Metal builds. ``ft serve-metal`` imports +# ``freetoken.server.metal_main`` via this package; importing ``launch`` eagerly +# would break that path on machines without torch. +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + + from .launch import launch_server __all__ = ["launch_server"] + + +def __getattr__(name: str): + if name == "launch_server": + from .launch import launch_server + + return launch_server + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 80dbf6a..8697fda 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -39,7 +39,12 @@ from .anthropic_api import register_anthropic_routes from .accounting import AdmissionClosedError, register_accounting_routes from .control_api import register_control_routes +from .cors import install_cors from .openai_api import register_openai_routes +from .process_utils import ( + reap_backend_workers as _reap_backend_workers, + terminate_backend_workers as _terminate_backend_workers, +) from . import request_ring from .access_log_filter import install_polling_access_log_filter from .request_logger import init as init_request_logging, log_request @@ -65,23 +70,12 @@ def get_global_state() -> FrontendManager: return _GLOBAL_STATE -def _terminate_backend_workers(processes: List[Any]) -> None: - """Best-effort, non-blocking teardown of the backend worker processes on an orderly stop. - - Called from the shutdown path AFTER ``_SHUTTING_DOWN`` is set, so the supervisor's liveness - watch is guaranteed to observe the flag before it sees these deaths — the exits are then - attributed to the stop, not misreported as a crash (the flag/death race, otherwise decided - by OS signal-delivery order, is settled in our favor). Also ensures the non-daemon workers - are actually torn down when an external stop signals only the main process. - - Never blocks (no join) and never raises: a worker already gone / an unqueryable handle is - fine — this only nudges live ones toward exit.""" - for p in processes or []: - try: - if p.is_alive(): - p.terminate() - except Exception: # noqa: BLE001 -- already-gone / unqueryable handle: nothing to do - continue +def _current_backend_processes(state: Any) -> List[Any]: + """Return the live handle's current children, including after a model switch.""" + handle = getattr(state, "metal_backend_handle", None) + if handle is not None: + return list(getattr(handle, "processes", None) or []) + return list(getattr(state, "backend_processes", None) or []) def _exit_after_backend_death(grace_s: float) -> threading.Timer: @@ -97,20 +91,6 @@ def _stop() -> None: return timer -def _reap_backend_workers(processes: List[Any], timeout: float = 5.0) -> None: - """Wait out a preceding ``_terminate_backend_workers`` and SIGKILL whatever is still - standing. Only the shell path needs this: it owns the process lifetime end to end (no - outer signal takes the process down for it), and a worker that ignored SIGTERM would keep - the GPU and the IPC sockets after the shell has already returned to the user's terminal.""" - for p in processes or []: - try: - p.join(timeout=timeout) - if p.is_alive(): - p.kill() - except Exception: # noqa: BLE001 -- already-gone / unqueryable handle: nothing to do - continue - - def _unwrap_msg(msg: BaseFrontendMsg) -> List[UserReply]: if isinstance(msg, BatchFrontendMsg): result = [] @@ -380,7 +360,7 @@ def shutdown(self): self.recv_tokenizer.stop() # Tear the workers down ourselves (best-effort). _SHUTTING_DOWN is already set by the # time shutdown() runs, so the supervisor attributes the ensuing deaths to the stop. - _terminate_backend_workers(self.backend_processes) + _terminate_backend_workers(_current_backend_processes(self)) @asynccontextmanager @@ -395,23 +375,6 @@ async def lifespan(_: FastAPI): _GLOBAL_STATE.shutdown() -def install_cors(app: FastAPI, origins_csv: str) -> None: - """Attach CORS headers for browser/webview clients (e.g. the desktop app). - - No-op when the allow-list is empty; must run before the app starts serving.""" - origins = [o.strip() for o in origins_csv.split(",") if o.strip()] - if not origins: - return - from starlette.middleware.cors import CORSMiddleware - - app.add_middleware( - CORSMiddleware, - allow_origins=["*"] if "*" in origins else origins, - allow_methods=["*"], - allow_headers=["*"], - ) - - app = FastAPI(title="FreeToken API Server", version=__version__, lifespan=lifespan) register_openai_routes(app, get_global_state, lambda: _MODEL_SAMPLING) register_anthropic_routes(app, get_global_state, lambda: _MODEL_SAMPLING) @@ -859,7 +822,7 @@ def _install_shell_stop_handlers() -> None: def _flag_shutdown(signum, frame) -> None: _SHUTTING_DOWN.set() - _terminate_backend_workers(_GLOBAL_STATE.backend_processes) + _terminate_backend_workers(_current_backend_processes(_GLOBAL_STATE)) prev = previous.get(signum) if callable(prev): prev(signum, frame) @@ -911,23 +874,85 @@ def _serve_and_run_shell(host: str, port: int) -> None: # Belt and braces: if uvicorn's lifespan shutdown did not run (thread wedged), flag the # stop and tear the workers down here so nothing outlives the shell. _SHUTTING_DOWN.set() - _terminate_backend_workers(_GLOBAL_STATE.backend_processes) - _reap_backend_workers(_GLOBAL_STATE.backend_processes) + processes = _current_backend_processes(_GLOBAL_STATE) + _terminate_backend_workers(processes) + _reap_backend_workers(processes) + + +def _install_routes_for_backend(config: ServerArgs) -> None: + """Select the route set that matches the inference backend. + + The FastAPI ``app`` is module-global, and the CUDA-native generation routes are + registered on import (``register_openai_routes`` etc.). When the resolved backend + is an Apple Silicon Metal runtime, the generation surface must proxy to the + upstream engine instead of hitting the in-process CUDA scheduler. This replaces + the native generation and backend-specific control routes on the *same* app object. + Unsupported CUDA mutation routes are removed rather than left wired to dead state. + + The CUDA path registers the native routes at import time and this is a no-op. + """ + from .metal import register_metal_proxy_routes + + if getattr(config, "backend", "auto") not in ("mlx", "llama"): + return + + # Base URL is not yet known here (the handle is created in launch), so we register + # routes that pull the handle lazily from the module-global backend handle set + # below (run_api_server assigns _GLOBAL_STATE.backend_processes / handle). + def get_backend(): + st = _GLOBAL_STATE + return getattr(st, "metal_backend_handle", None) + + # Drop native generation/control routes so their handlers cannot win + # FastAPI's first-match dispatch. register_metal_proxy_routes repeats its + # owned-path replacement defensively for direct callers. + paths_to_clear = { + "/v1/chat/completions", + "/v1/completions", + "/v1/messages", + "/v1/messages/count_tokens", + "/v1/responses", + "/v1/embeddings", + "/v1/models", + "/v1/model/list", + "/v1/model/load", + "/health", + "/v1/stats", + "/v1/requests", + "/v1/cache/status", + "/v1/cache/rebuild", + "/v1/admin/prepare-stop", + "/generate", + } + filtered = [] + for r in getattr(app, "routes", []): + path = getattr(r, "path", None) + if path in paths_to_clear: + continue + filtered.append(r) + app.router.routes = filtered + + register_metal_proxy_routes(app, get_backend) def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_shell: bool) -> None: """ - Run the frontend API server (FastAPI + uvicorn) and wire it to the tokenizer process via ZMQ. + Run the frontend API server (FastAPI + uvicorn) and wire it to the backend. - Args: - config: Server configuration (host/port, ZMQ IPC addresses, etc). - start_backend: Callback that launches the backend worker processes (TP schedulers + - tokenizer/detokenizer). - run_shell: If True, also attach the interactive terminal shell to the served API. + For the native CUDA path this wires the scheduler/tokenizer workers via ZMQ + (unchanged). When ``start_backend()`` returns a Metal backend handle, the + OpenAI/Anthropic/Responses surface is registered as a proxy to the upstream + Metal engine instead (see server/metal.py). """ global _GLOBAL_STATE, _MODEL_SAMPLING + # When a Metal backend is selected, the OpenAI/Anthropic/Responses surface must + # proxy to the upstream engine rather than use the in-process ZMQ scheduler. + # The CUDA-native routes were registered on ``app`` at import time; for Metal we + # swap those generation paths for proxy routes before uvicorn starts. + _install_routes_for_backend(config) + if config.sampling_defaults == "model" and not config.use_dummy_weight: _MODEL_SAMPLING = load_generation_sampling(config.model_path) # Always surface the effective default sampling (model-recommended where available, @@ -977,6 +1002,10 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s # Hold the worker handles so the orderly-shutdown path can tear them down itself (after # setting _SHUTTING_DOWN) rather than relying on OS signal-delivery order. _GLOBAL_STATE.backend_processes = list(getattr(handle, "processes", None) or []) + # A Metal backend handle exposes ``upstream_base_url``; the proxy routes read it + # lazily from here. Only set for the Metal path (CUDA handle has no this attr). + if hasattr(handle, "upstream_base_url"): + _GLOBAL_STATE.metal_backend_handle = handle def _on_ready() -> None: # A stop requested while weights were loading has already sealed admission. The backend diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b7..160cf23 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -26,6 +26,14 @@ class ServerArgs(SchedulerConfig): # Reasoning parser that splits reasoning from content for OpenAI # responses. None disables it (default for models without a reasoning protocol). reasoning_parser: str | None = None + # Inference engine backend. "auto" resolves to CUDA when a usable CUDA GPU is + # present, otherwise to the first available Apple Silicon Metal runtime (mlx, + # then llama.cpp). "cuda" forces the native CUDA scheduler (original behavior). + # "mlx" / "llama" reuse Apple's Metal runtimes as an upstream engine behind + # FreeToken's API (see server/metal.py). + backend: str = "auto" + # Upstream port for a Metal backend; a free loopback port is auto-picked when 0. + metal_port: int = 0 # "model": fill unspecified request sampling params from generation_config.json # (temperature/top_k/top_p), like sglang. "none": use framework defaults only. sampling_defaults: str = "model" @@ -595,6 +603,28 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Run the server in shell mode.", ) + parser.add_argument( + "--backend", + default=ServerArgs.backend, + choices=["auto", "cuda", "mlx", "llama"], + help=( + "Inference engine backend. 'auto' resolves to CUDA when a usable CUDA GPU " + "is present, otherwise to an Apple Silicon Metal runtime (mlx, then " + "llama.cpp). 'cuda' forces the native CUDA scheduler. 'mlx' / 'llama' reuse " + "Apple's Metal runtimes (see server/metal.py)." + ), + ) + + parser.add_argument( + "--metal-port", + type=int, + default=ServerArgs.metal_port, + help=( + "Upstream loopback port for a Metal backend (--backend mlx|llama). " + "A free port is auto-picked when 0." + ), + ) + parser.add_argument( "--cors-origins", type=str, diff --git a/python/freetoken/server/cors.py b/python/freetoken/server/cors.py new file mode 100644 index 0000000..2b919c2 --- /dev/null +++ b/python/freetoken/server/cors.py @@ -0,0 +1,27 @@ +"""Torch-free CORS setup shared by the native and Metal API servers.""" + +from __future__ import annotations + +from fastapi import FastAPI + + +DEFAULT_CORS_ORIGINS = "tauri://localhost,http://tauri.localhost,http://localhost:1420" + + +def install_cors(app: FastAPI, origins_csv: str) -> None: + """Attach CORS headers for browser/webview clients. + + An empty allow-list disables CORS; ``*`` allows every origin. + """ + origins = [origin.strip() for origin in origins_csv.split(",") if origin.strip()] + if not origins: + return + + from starlette.middleware.cors import CORSMiddleware + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"] if "*" in origins else origins, + allow_methods=["*"], + allow_headers=["*"], + ) diff --git a/python/freetoken/server/gemma4_chat_template.jinja b/python/freetoken/server/gemma4_chat_template.jinja new file mode 100644 index 0000000..4741bf6 --- /dev/null +++ b/python/freetoken/server/gemma4_chat_template.jinja @@ -0,0 +1,390 @@ +{# + Template: Google Gemma 4 Canonical Chat Template + Author: Google Gemma Engineering Team + Published: 2026-07-09 + Context: Fixed tool-calling loops, turn closures, and thinking content-ordering. +#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%} + {%- if thinking_text and thinking_gate -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and (not message.get('tool_calls') or ns_tr_out.flag) + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- if not enable_thinking -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} + {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} +{%- endif -%} diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index a5f8438..9def19c 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -112,6 +112,46 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: scheduler.shutdown() +def _launch_metal(server_args: "ServerArgs", backend: str, run_shell: bool) -> None: + """Serve over an Apple Silicon Metal backend (mlx or llama.cpp). + + Launches the chosen Metal runtime as an upstream OpenAI/Anthropic-compatible + engine and registers FreeToken's HTTP surface (OpenAI/Anthropic/Responses) as + a proxy to it. The CUDA scheduler path is not involved. + """ + from .api_server import run_api_server + from .metal import ( + MetalBackendHandle, + _default_served_model_name, + launch_metal_backend, + ) + + served_name = server_args.served_model_name + if served_name == _default_served_model_name(server_args.model_path): + served_name = None + + handle = launch_metal_backend( + backend, + server_args.model_path, + server_args.metal_port, + served_model_name=served_name, + ) + logger.info( + "Metal backend %r started at %s (upstream), FreeToken API on %s:%s", + backend, + handle.upstream_base_url, + server_args.server_host, + server_args.server_port, + ) + + def start_metal() -> MetalBackendHandle: + # Callers of run_api_server treat the return as the "backend handle" and + # copy its .processes for teardown; reuse ours. + return handle + + run_api_server(server_args, start_metal, run_shell=run_shell) + + def launch_server( run_shell: bool = False, argv: list[str] | None = None, @@ -119,6 +159,7 @@ def launch_server( ) -> None: from .api_server import run_api_server from .args import parse_args + from .metal import MetalBackendHandle, launch_metal_backend, resolve_backend server_args, run_shell = parse_args( sys.argv[1:] if argv is None else argv, @@ -127,6 +168,20 @@ def launch_server( ) logger = init_logger(__name__, "initializer") + # Resolve + branch on the inference backend. The classic CUDA path spawns CUDA + # scheduler/tokenizer workers (logic below, unchanged). The Metal backends reuse + # Apple's mlx / llama.cpp runtimes as an upstream OpenAI-compatible engine and + # FreeToken serves its API by proxying to it (see server/metal.py). ``auto`` + # resolves to CUDA when usable, else to an available Metal runtime. + backend = resolve_backend(server_args.backend) + if backend != "cuda": + # Route installation runs later inside run_api_server and reads the + # config object. Persist auto's concrete result so it selects Metal + # routes instead of leaving the native CUDA routes mounted. + server_args = replace(server_args, backend=backend) + _launch_metal(server_args, backend, run_shell=run_shell) + return + def start_subprocess() -> "BackendHandle": import multiprocessing as mp diff --git a/python/freetoken/server/metal.py b/python/freetoken/server/metal.py new file mode 100644 index 0000000..4bade31 --- /dev/null +++ b/python/freetoken/server/metal.py @@ -0,0 +1,1426 @@ +"""Apple Silicon (Metal) backends for ``ft serve``. + +This module wires Apple's own, already-built Metal runtimes as the inference +engine behind FreeToken's OpenAI/Anthropic/Responses API. It does NOT port any +of the CUDA/Triton kernels (there is no macOS build of triton/flashinfer/ +sglang-kernel, and FreeToken's native fast path is irreducibly CUDA). Instead it +reuses two Apple-proven upstreams: + + * ``mlx`` (``mlx_lm.server``) -- Apple's MLX framework running on the MPS + (Metal) GPU. OpenAI-compatible ``/v1/*`` HTTP server. + * ``llama`` (``llama.cpp``'s ``llama-server``) -- Metal-backed GGUF server. + OpenAI- and Anthropic-Compatible ``/v1/*`` and ``/v1/messages`` HTTP server. + +FreeToken keeps serving its OpenAI/Anthropic/Responses surface on the configured +host/port; this module launches the chosen upstream as a child process and +proxies the generation routes to it. Running the CUDA scheduler path is entirely +untouched (see ``server/launch.py``), so ``ft serve`` on a CUDA box behaves +exactly as before and ``ft serve --backend mlx|llama`` re-targets to Metal. + +Backend resolution rules: + * ``cuda`` -> native FreeToken scheduler (unchanged default behaviour). + * ``mlx`` -> mlx_lm.server (requires the ``mlx-lm`` package). + * ``llama`` -> llama.cpp llama-server (requires the ``llama-server`` binary). + * ``auto`` -> CUDA when available and usable; otherwise the first Metal + runtime that is installed/importable. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import glob +import json +import os +import queue +import shlex +import shutil +import subprocess +import tempfile +import threading +import time +import urllib.error +import urllib.request +import uuid +from collections import deque +from dataclasses import dataclass, field +from typing import Any, AsyncIterator + +import httpx +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse +from freetoken.logging import init_logger + +logger = init_logger(__name__) + +#: Upstream serves on a loopback port inside this range (FreeToken's own API keeps +#: the user-facing ``server_port``). +_UPSTREAM_PORT_MIN = 19000 +_UPSTREAM_PORT_MAX = 19999 +_ADDRESS_HEALTH_TIMEOUT_S = float( + os.environ.get("FREETOKEN_METAL_READY_TIMEOUT", "180") +) +#: How long the eager warm-up generation may take. This is the real model-load +#: budget (a multi-shard download or a 50 GiB weight load happens inside it): +#: mlx_lm queues the request in its generation thread and answers only once the +#: weights are resident, so the warm-up doubles as a supervised load. A hang +#: here (e.g. a dead CDN socket) surfaces as an error after this timeout +#: instead of blocking the user's first request forever. +_LOAD_TIMEOUT_S = float(os.environ.get("FREETOKEN_METAL_LOAD_TIMEOUT", "3600")) +#: mlx httpd advertises readiness with this line on stderr; llama-server is probed +#: over HTTP. Both are also verified by a live ``/v1/models`` round-trip. +_STARTED_ONCE_TIMEOUT_S = float(os.environ.get("FREETOKEN_METAL_START_TIMEOUT", "60")) +_POLL_INTERVAL_S = 0.5 +_PORT_LOCK_PATH = os.path.join(tempfile.gettempdir(), "freetoken-metal-port.lock") + + +@contextlib.contextmanager +def _serialize_upstream_launches(): + """Serialize port selection until the child has bound its listener. + + Checking a port and then closing the probe socket cannot itself reserve the + port. A small cross-process advisory lock closes the race between FreeToken + instances; each launcher holds it until its child is listening. + """ + import fcntl + + fd = os.open(_PORT_LOCK_PATH, os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + yield + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +def _pick_upstream_port(preferred: int | None) -> int: + """Pick an upstream port for the Metal engine. + + Uses ``preferred`` when given and free; otherwise scans the reserved range + for a free loopback port. FreeToken's own API never occupies this range (it + defaults to 1919), so collisions are effectively limited to another Metal + backend instance.""" + if preferred is not None and preferred > 0: + return _claim_port(preferred) or _scan_free_port() + return _scan_free_port() + + +def _claim_port(port: int) -> int | None: + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(("127.0.0.1", port)) + except OSError: + return None + return port + + +def _scan_free_port() -> int: + import socket + + for port in range(_UPSTREAM_PORT_MIN, _UPSTREAM_PORT_MAX): + if _claim_port(port) is not None: + return port + raise RuntimeError("no free loopback port for the Metal backend") + + +def _any_cuda_usable() -> bool: + """True when the native CUDA scheduler path is usable on this host.""" + try: + import torch + + return torch.cuda.is_available() + except Exception: # noqa: BLE001 -- not available at all is fine + return False + + +def mlx_importable() -> bool: + import importlib.util + + return importlib.util.find_spec("mlx_lm") is not None + + +def llama_binary() -> str | None: + for exe in ("llama-server",): + path = shutil.which(exe) + if path: + return path + return None + + +def resolve_backend(requested: str) -> str: + """Resolve a ``--backend`` value to a concrete choice (``cuda/mlx/llama``). + + ``cuda`` is accepted as-is. ``mlx``/``llama`` require their upstream to be + present. ``auto`` prefers CUDA when usable, then mlx, then llama. Raises a + clear error when the requested backend cannot run here.""" + if requested == "cuda": + if not _any_cuda_usable(): + raise RuntimeError( + "--backend cuda requested but no usable CUDA GPU was found " + "on this host." + ) + return "cuda" + if requested == "mlx": + if not mlx_importable(): + raise RuntimeError( + "--backend mlx requested but mlx_lm is not importable. " + "Install it with: uv pip install 'mlx-lm'" + ) + return "mlx" + if requested == "llama": + if llama_binary() is None: + raise RuntimeError( + "--backend llama requested but 'llama-server' was not found " + "on PATH. Install llama.cpp, or use --backend mlx." + ) + return "llama" + if requested != "auto": + raise RuntimeError( + f"unknown --backend {requested!r} (expected auto, cuda, mlx, or llama)" + ) + if _any_cuda_usable(): + logger.info("backend=auto resolved to cuda (native CUDA scheduler)") + return "cuda" + if mlx_importable(): + logger.info("backend=auto resolved to mlx (Apple Silicon MLX)") + return "mlx" + if llama_binary() is not None: + logger.info("backend=auto resolved to llama (llama.cpp Metal)") + return "llama" + raise RuntimeError( + "FREETOKEN: no usable inference backend. No CUDA GPU, mlx_lm, or " + "llama-server was found. Install mlx-lm (Apple Silicon) or llama.cpp." + ) + + +def resolve_metal_backend(requested: str) -> str: + """Resolve a standalone Metal backend without ever selecting CUDA.""" + if requested == "auto": + if mlx_importable(): + logger.info("backend=auto resolved to mlx (Apple Silicon MLX)") + return "mlx" + if llama_binary() is not None: + logger.info("backend=auto resolved to llama (llama.cpp Metal)") + return "llama" + raise RuntimeError( + "FREETOKEN: no usable Metal backend. Install mlx-lm (Apple Silicon) " + "or llama.cpp." + ) + if requested not in {"mlx", "llama"}: + raise RuntimeError( + f"unknown Metal --backend {requested!r} (expected auto, mlx, or llama)" + ) + return resolve_backend(requested) + + +def _default_served_model_name(model_path: str) -> str: + if not model_path: + return model_path + return os.path.basename(os.path.normpath(model_path)) or model_path + + +@dataclass +class MetalBackendHandle: + """Handle to a launched Metal inference engine (blunt stand-in for the CUDA + scheduler's ``BackendHandle``; the API layer only needs processes + readiness). + + ``load_state`` is the lifecycle the /health route reports: ``starting`` + (process spawned, port not yet listening) -> ``loading`` (upstream answers + /v1/models but weights are still coming down / into memory) -> ``ready`` + (a warm-up generation succeeded: the engine can actually generate) -> + ``error`` (process died, load timed out, or warm-up failed). Shared with + the proxy threads via a lock so /health never tears while reading it.""" + processes: list[subprocess.Popen] = field(default_factory=list) + upstream_base_url: str = "" + backend: str = "" + model_path: str = "" + served_model_name: str = "" + served_model_name_explicit: bool = False + instance_id: str = field(default_factory=lambda: str(uuid.uuid4())) + load_state: str = "starting" + load_phase: str = "" + load_error: str = "" + load_started_at: float = 0.0 + load_ended_at: float = 0.0 + weights_bytes: int = 0 + # The CUDA supervisor contract (supervisor.drain_ready): progress tuples + # flow through ``ack_queue`` and one final ack completes readiness. The + # load watcher speaks it, so ``ft serve --backend mlx`` gets live + # /health progress and the maintenance flip for free, unmodified. + ack_queue: Any = field(default_factory=queue.Queue) + expected_acks: int = 1 + _switch_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + _state_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + _stats_lock: threading.Lock = field(default_factory=threading.Lock, repr=False) + active_requests: int = 0 + completed_requests: int = 0 + prompt_tokens_total: int = 0 + completion_tokens_total: int = 0 + + # ------------------------------------------------------------ load state -- + def _set_state( + self, state: str, *, phase: str = "", error: str = "" + ) -> None: + with self._state_lock: + self.load_state = state + if phase: + self.load_phase = phase + if error: + self.load_error = error + if state in ("ready", "error"): + self.load_ended_at = time.monotonic() + + def health_doc(self) -> dict[str, Any]: + """The /health document for this engine, in the CUDA path's contract: + ``status: loading`` carries ``phase`` + byte progress the shell renders + (``loading (weights): 4.2/15.0 GiB``), ``error`` carries the reason.""" + with self._state_lock: + state = self.load_state + phase = self.load_phase + error = self.load_error + weights = self.weights_bytes + doc: dict[str, Any] = { + "model": self.served_model_name or _default_served_model_name(self.model_path), + "backend": self.backend, + "instance_id": self.instance_id, + } + if state == "error": + doc["status"] = "error" + doc["message"] = error or "Metal backend failed to load" + return doc + if state == "ready": + doc["status"] = "ok" + doc["maintenance"] = "serving" + doc["uptime_s"] = max(0, int(time.monotonic() - self.load_ended_at)) + return doc + done = _upstream_resident_bytes(self.processes) + doc["status"] = "loading" + doc["phase"] = phase or "starting" + if weights > 0: + doc["progress"] = { + "done_bytes": min(done, weights), + "total_bytes": weights, + } + return doc + + def begin_request(self) -> None: + with self._stats_lock: + self.active_requests += 1 + + def finish_request(self, usage: dict[str, Any] | None, *, completed: bool) -> None: + """Record the best usage snapshot exposed by an upstream response.""" + usage = usage or {} + prompt = usage.get("prompt_tokens", usage.get("input_tokens", 0)) + completion = usage.get("completion_tokens", usage.get("output_tokens", 0)) + with self._stats_lock: + self.active_requests = max(0, self.active_requests - 1) + if completed: + self.completed_requests += 1 + if isinstance(prompt, int) and not isinstance(prompt, bool) and prompt > 0: + self.prompt_tokens_total += prompt + if ( + isinstance(completion, int) + and not isinstance(completion, bool) + and completion > 0 + ): + self.completion_tokens_total += completion + + def stats_doc(self) -> dict[str, Any]: + with self._stats_lock: + active = self.active_requests + completed = self.completed_requests + prompt_total = self.prompt_tokens_total + completion_total = self.completion_tokens_total + ready_at = self.load_ended_at if self.load_ended_at > 0 else self.load_started_at + uptime = max(0, int(time.monotonic() - ready_at)) if ready_at > 0 else 0 + return { + "instance_id": self.instance_id, + "model": { + "id": self.served_model_name + or _default_served_model_name(self.model_path), + "ctx": None, + "attn": "metal", + "moe": None, + }, + "uptime_s": uptime, + "kv": None, + "mamba": None, + "swa": None, + "vram_bytes": 0, + "throughput": {"decode_tps": 0.0, "prefill_tps": 0.0}, + "requests": { + "active": active, + "completed": completed, + "p95_ms": 0, + "ttft_mean_ms": 0, + "prompt_tokens_total": prompt_total, + "completion_tokens_total": completion_total, + }, + } + + # ------------------------------------------------------------ lifecycle -- + def terminate(self) -> None: + _stop_processes(self.processes) + + def is_alive(self) -> bool: + return any(p.poll() is None for p in self.processes) + + def is_ready(self) -> bool: + """True once a warm-up generation has proven the engine generates. + + The proxy's generation routes gate on this (503 while loading) so a + user request cannot queue behind the weight load forever.""" + return self.load_state == "ready" + + def switch_model(self, model_path: str) -> None: + """Serve a different model: stop the old engine *first*, then start the new. + + Sequential, not concurrent. Concurrent loads (old engine resident while + the new one streams in) over-commit the Metal working set -- two ~50 GiB + engines against a ~107 GiB wired limit is the deadlock this machine + kept hitting during switches. Stopping first frees the old model's + memory before the new load starts, at the cost of a load-window gap in + serving (the shell's /health progress bar covers it). + + Failure semantics: if the new engine fails to come up, the old model is + gone -- the server reports the error rather than pretending. Rolling + back would mean reloading the old weights (the same cost as retrying), + so the honest state is an error the user can act on. + """ + with self._switch_lock: + old_processes = self.processes + self.processes = [] + self.upstream_base_url = "" + self.model_path = model_path + with self._state_lock: + self.load_state = "loading" + self.load_phase = "stopping" + self.load_error = "" + self.weights_bytes = 0 + # A switch drains the queue the supervisor may still be reading; + # its terminal acks for the old engine must not leak into the new + # one's readiness handshake. _drain_acks takes _state_lock itself, + # so it MUST run outside the block above (non-reentrant Lock -- + # calling it inside self-deadlocked every model switch). + self._drain_acks() + # Full escalate-to-kill teardown BEFORE the new load: the child's + # output pipe is drained (see _drain_process_output), and a partial + # teardown here would leave it holding its port -- and its memory. + _stop_processes(old_processes) + # Sequential from here: the watcher publishes load state onto THIS + # handle (state_handle=self) while the load runs, and the launch + # returns a handle carrying the new engine's identity. + try: + launch_kwargs: dict[str, Any] = {"state_handle": self} + if self.served_model_name_explicit: + launch_kwargs["served_model_name"] = self.served_model_name + new = launch_metal_backend( + self.backend, model_path, upstream_port=None, **launch_kwargs + ) + except Exception as exc: + self._set_state("error", error=str(exc)) + raise + # Publish the new engine immediately. Its watcher reports load + # progress through this shared handle, and both health_doc() and + # is_alive() need the new process while the route waits. Deferring + # this adoption until readiness made /health report "down", kept + # RSS progress at zero, and lost the process entirely on failure. + self.processes = new.processes + self.upstream_base_url = new.upstream_base_url + self.model_path = new.model_path + self.served_model_name = ( + new.served_model_name or _default_served_model_name(new.model_path) + ) + self.served_model_name_explicit = new.served_model_name_explicit + # Block until the watcher reaches a terminal state. /v1/model/load + # must not answer "ok" while the new engine is still loading -- the + # shell budgets this call for the full download + load. Raises on + # failure so the route reports the reason instead of a false ok. + try: + self._wait_load_terminal() + except Exception as exc: + failed_processes = self.processes + self.processes = [] + self.upstream_base_url = "" + _stop_processes(failed_processes) + self._set_state("error", error=str(exc)) + raise + + def _wait_load_terminal(self, timeout: float | None = None) -> None: + """Block until this handle's load watcher reports ready or error. + + Polls ``load_state`` (the watcher's publication point) -- not the + ack_queue, which the CUDA-side supervisor may be draining concurrently + in ``ft serve --backend mlx`` mode.""" + deadline = time.monotonic() + (timeout if timeout else _LOAD_TIMEOUT_S) + while time.monotonic() < deadline: + with self._state_lock: + state = self.load_state + if state == "ready": + return + if state == "error": + with self._state_lock: + raise RuntimeError(self.load_error or "model load failed") + time.sleep(0.25) + raise RuntimeError("model load timed out") + + def _drain_acks(self) -> None: + """Drop stale acks so a prior engine's terminal events cannot be read + as the next engine's readiness.""" + with self._state_lock: + while True: + try: + self.ack_queue.get_nowait() + except queue.Empty: + return + + +def _drain_process_output( + proc: subprocess.Popen, name: str, recent: deque[str] +) -> None: + """Read a child's stdout until EOF on a daemon thread. + + The children are launched with ``stdout=PIPE`` so launch failures are + visible, but an unread pipe fills (64 KiB) and then blocks the child inside + ``write()`` forever -- including ignoring SIGTERM. Draining on a thread + keeps the child healthy and makes ``terminate()`` actually work. + """ + try: + assert proc.stdout is not None + for line in proc.stdout: + recent.append(line) + logger.debug("%s: %s", name, line.rstrip()) + except Exception: # noqa: BLE001 -- draining must never raise + pass + + +def _stop_processes(processes: list[subprocess.Popen]) -> None: + for p in processes: + try: + if p.poll() is None: + p.terminate() + except Exception: # noqa: BLE001 -- best-effort teardown + continue + for p in processes: + try: + p.wait(timeout=10) + except subprocess.TimeoutExpired: + try: + p.kill() + p.wait(timeout=5) + except Exception: # noqa: BLE001 -- best-effort teardown + continue + except Exception: # noqa: BLE001 -- best-effort teardown + continue + + +def _start_drain_thread(proc: subprocess.Popen, name: str) -> deque[str]: + recent: deque[str] = deque(maxlen=20) + # Keep the diagnostics with the process so readiness has one nonblocking + # source of output. There must never be a second reader on stdout. + setattr(proc, "_freetoken_recent_output", recent) + t = threading.Thread( + target=_drain_process_output, args=(proc, name, recent), daemon=True + ) + t.start() + return recent + + +def _wait_for_readiness(url: str, process: subprocess.Popen, *, timeout: float) -> None: + """Poll ``/v1/models`` until the upstream answers or ``timeout`` elapses. + + Also surfaces any early stdout/stderr lines so a launch failure is visible + instead of a silent timeout.""" + deadline = time.monotonic() + timeout + last_err = "" + recent = getattr(process, "_freetoken_recent_output", ()) + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"Metal backend exited during startup " + f"(code {process.returncode}): {''.join(list(recent)[-8:])}" + ) + try: + proc = subprocess.run( + ["curl", "-fsS", f"{url}/v1/models"], + capture_output=True, + text=True, + timeout=5, + ) + if proc.returncode == 0: + return + last_err = (proc.stderr or proc.stdout or "").strip() + except Exception as exc: # noqa: BLE001 -- not ready yet, keep polling + last_err = str(exc) + time.sleep(0.5) + raise RuntimeError( + f"Metal backend at {url} did not become ready within {timeout:.0f}s. " + f"Last probe: {last_err}. Output:\n{''.join(list(recent)[-12:])}" + ) + + +def _wait_for_listener(port: int, process: subprocess.Popen, *, timeout: float) -> None: + """Wait until a freshly spawned child owns its selected loopback port.""" + import socket + + deadline = time.monotonic() + timeout + recent = getattr(process, "_freetoken_recent_output", ()) + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError( + "Metal backend exited before binding its port " + f"(code {process.returncode}): {''.join(list(recent)[-8:])}" + ) + try: + with socket.create_connection(("127.0.0.1", port), timeout=0.2): + return + except OSError: + time.sleep(0.05) + raise RuntimeError( + f"Metal backend did not bind 127.0.0.1:{port} within {timeout:.0f}s. " + f"Output:\n{''.join(list(recent)[-12:])}" + ) + + +def _upstream_resident_bytes(processes: list[subprocess.Popen]) -> int: + """How many bytes of the model the upstream has resident, from RSS. + + A crude but truthful progress signal for a lazy loader: mlx_lm reads + weights into unified memory on the first generation, so RSS climbing + toward ``weights_bytes`` IS the load.""" + total = 0 + for p in processes: + try: + rss_kb = int( + subprocess.run( + ["ps", "-o", "rss=", "-p", str(p.pid)], + capture_output=True, + text=True, + timeout=2, + ).stdout.strip() + or 0 + ) + total += rss_kb * 1024 + except Exception: # noqa: BLE001 -- progress is best-effort + continue + return total + + +def _weights_bytes_on_disk(model_path: str) -> int: + """Total weight-file bytes for a local model dir, else 0 (unknown).""" + path = os.path.expanduser(model_path) + if not os.path.isdir(path): + return 0 + total = 0 + for name in ("*.safetensors", "*.gguf", "*.bin"): + for f in glob.glob(os.path.join(path, name)): + try: + total += os.path.getsize(f) + except OSError: + continue + return total + + +def _hf_cache_dir(repo_id: str) -> str | None: + """Snapshot dir in the local HF cache for ``repo_id``, when fully present. + + Used to (a) size the load before it starts and (b) decide whether the + child can run with HF_HUB_OFFLINE=1 -- which stops a stalled CDN retry + from hanging the load forever (the failure mode seen on this host).""" + base = os.environ.get("HF_HUB_CACHE") or os.path.expanduser( + "~/.cache/huggingface/hub" + ) + repo_dir = os.path.join(base, "models--" + repo_id.replace("/", "--")) + snaps = os.path.join(repo_dir, "snapshots") + if not os.path.isdir(snaps): + return None + snapshots = [ + d + for d in (os.path.join(snaps, x) for x in os.listdir(snaps)) + if os.path.isdir(d) + ] + if not snapshots: + return None + snap = max(snapshots, key=os.path.getmtime) + if _weights_bytes_on_disk(snap) == 0: + return None # no weights resolved (download incomplete or empty) + return snap + + +def _warm_up_generation(url: str, model_id: str, timeout: float) -> None: + """Drive one tiny generation so weights load before we call the engine ready. + + mlx_lm serves /v1/models and accepts requests while weights are still on + disk, then blocks the first generation until the load finishes -- so "the + port answers" is NOT ready. This 1-token request runs the actual load + under supervision: it either proves the engine can generate (and warms + kernels/caches along the way) or it fails with a reason we can report. + + ``model_id`` must be the id the upstream actually serves: mlx_lm resolves + the model field against the HF cache, and an id that is not cached (plus + our HF_HUB_OFFLINE=1) makes it 404 before any weights move.""" + body = json.dumps( + { + "model": model_id, + "prompt": "1", + "max_tokens": 1, + "temperature": 0.0, + } + ).encode() + req = urllib.request.Request( + f"{url}/v1/completions", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + if resp.status != 200: + raise RuntimeError(f"warm-up generation failed: HTTP {resp.status}") + except urllib.error.HTTPError as exc: + detail = "" + try: + detail = exc.read().decode("utf-8", "replace")[:300] + except Exception: # noqa: BLE001 -- body text is best-effort + pass + raise RuntimeError( + f"warm-up generation failed: HTTP {exc.code}" + + (f": {detail}" if detail else "") + ) from exc + + +def _watch_mlx_load( + handle: MetalBackendHandle, + url: str, + proc: subprocess.Popen, + *, + timeout: float, +) -> None: + """Supervise the upstream's load: port up -> weights loading -> ready/error. + + Runs on a daemon thread so the API server binds immediately and /health + reports live progress while the engine loads. Progress is ALSO pushed as + ("progress", desc, done, total) tuples on ``handle.ack_queue`` -- the CUDA + supervisor's protocol (supervisor.drain_ready) -- so ``ft serve --backend + mlx`` renders the same live /health phases through the stock supervisor, + and a terminal ("error", reason) / plain ack drives its maintenance gate.""" + handle._set_state("loading", phase="starting") + try: + _wait_for_readiness(url, proc, timeout=_ADDRESS_HEALTH_TIMEOUT_S) + except RuntimeError as exc: + handle._set_state("error", error=str(exc)) + handle.ack_queue.put(("error", str(exc))) + return + handle._set_state("loading", phase="weights") + handle.ack_queue.put(("progress", "Loading weights (Metal)", 0, handle.weights_bytes)) + + # Sample RSS toward the weights total while the warm-up generation runs: + # that request IS the load (mlx_lm generates only after weights arrive), + # so its duration is exactly when progress moves. + done_holder = {"done": 0, "stop": False} + + def _progress_sampler() -> None: + while not done_holder["stop"]: + done_holder["done"] = _upstream_resident_bytes(handle.processes) + if handle.weights_bytes > 0: + handle.ack_queue.put( + ( + "progress", + "Loading weights (Metal)", + min(done_holder["done"], handle.weights_bytes), + handle.weights_bytes, + ) + ) + time.sleep(_POLL_INTERVAL_S) + + sampler = threading.Thread( + target=_progress_sampler, name="freetoken-metal-progress", daemon=True + ) + sampler.start() + try: + _warm_up_generation(url, model_id=handle.model_path, timeout=timeout) + except Exception as exc: # noqa: BLE001 -- any warm-up failure is a load failure + done_holder["stop"] = True + sampler.join(timeout=2) + if proc.poll() is not None: + reason = f"Metal backend exited during load (code {proc.returncode})" + else: + reason = f"model load failed: {exc}" + handle._set_state("error", error=reason) + handle.ack_queue.put(("error", reason)) + return + done_holder["stop"] = True + sampler.join(timeout=2) + if proc.poll() is not None: + reason = f"Metal backend exited during load (code {proc.returncode})" + handle._set_state("error", error=reason) + handle.ack_queue.put(("error", reason)) + return + handle._set_state("ready") + handle.ack_queue.put("ready") + + +def _gemma4_chat_template() -> str: + """Google's canonical Gemma 4 chat template, for gemma-4 snapshots whose + repo ships none (the 26b-a4b base repo uploads tokenizer files without a + chat_template; the -it repos carry chat_template.jinja). + + Loaded from the bundled gemma4_chat_template.jinja (Google Gemma + Engineering, 2026-07-09 -- fixes tool-calling loops, turn closures, and + thinking content-ordering). Do NOT hand-edit: the exact turn grammar + (<|turn>/, <|channel>thought, <|think|>, tool blocks) is what the + model was trained on, and near-misses make it degenerate into raw + text completion (endless repetition, API-doc regurgitation). + + A missing asset means the FreeToken installation is incomplete, so fail + with an actionable error before trying to launch the MLX child.""" + global _GEMMA4_TEMPLATE + if _GEMMA4_TEMPLATE is None: + path = os.path.join(os.path.dirname(__file__), "gemma4_chat_template.jinja") + try: + with open(path, "r", encoding="utf-8") as f: + _GEMMA4_TEMPLATE = f.read() + except OSError as exc: + raise RuntimeError( + "bundled Gemma 4 chat template is missing; reinstall FreeToken" + ) from exc + return _GEMMA4_TEMPLATE + + +_GEMMA4_TEMPLATE: str | None = None + + +def _needs_turn_stop_token(body: dict[str, Any]) -> bool: + """True when the request targets a gemma-4 model served through mlx_lm. + + The base 26B snapshot defines eos as only, and converted/derived + snapshots may omit the instruction-tuned checkpoint's additional stop + ids. Injecting the turn delimiter is therefore a defensive compatibility + measure for Gemma 4; everything else keeps its engine-default behavior.""" + return _is_gemma4_model(body.get("model") or "") + + +def _is_gemma4_model(model: str) -> bool: + """Match a served model id / path against the gemma-4 family (the id may be + a HF repo id like google/gemma-4-26b-a4b or a local snapshot path).""" + m = model.lower() + return "gemma-4" in m or "gemma4" in m + + +def _launch_mlx( + model_path: str, + port: int, + *, + state_handle: MetalBackendHandle | None = None, + served_model_name: str | None = None, +) -> MetalBackendHandle: + import sys + + py = sys.executable + cmd = [ + py, + "-m", + "mlx_lm.server", + "--model", + model_path, + "--host", + "127.0.0.1", + "--port", + str(port), + ] + # The instruction-tuned Gemma 4 repos ship this template, but the base 26B + # snapshot and some converted/derived snapshots do not. Always pass the + # canonical turn grammar explicitly so MLX behavior is consistent across + # those layouts. The proxy also adds the turn delimiter as a defensive + # stop word for snapshots whose generation config omits it. + if _is_gemma4_model(model_path): + cmd += ["--chat-template", _gemma4_chat_template()] + # When the weights are already fully in the local HF cache, pin the child + # to them: no revalidation round-trip, and no stalled-CDN retry can hang + # the lazy load. A repo id that is not cached (or only partially) keeps + # online mode so the download happens as before. + env = os.environ.copy() + cache_dir = _hf_cache_dir(model_path) + if cache_dir is not None: + env.setdefault("HF_HUB_OFFLINE", "1") + env.setdefault("HF_HUB_DISABLE_TELEMETRY", "1") + logger.info( + "mlx: using cached snapshot %s (HF_HUB_OFFLINE=1)", cache_dir + ) + logger.info("launching Metal backend (mlx): %s", shlex.join(cmd)) + proc = subprocess.Popen( + cmd, + # Own the child's stdout/stderr so we can detect startup failures and + # drain logs; the child inherits env so HF/mlx settings pass through. + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + env=env, + ) + _start_drain_thread(proc, "mlx") + try: + _wait_for_listener(port, proc, timeout=_STARTED_ONCE_TIMEOUT_S) + except Exception: + _stop_processes([proc]) + raise + url = f"http://127.0.0.1:{port}" + # In a switch, the watcher publishes state onto the caller's shared handle + # (the one the proxy routes already read); the returned handle only + # carries the engine's process/upstream/url for the caller to adopt. + return_handle = MetalBackendHandle( + processes=[proc], + upstream_base_url=url, + backend="mlx", + model_path=model_path, + served_model_name=served_model_name or _default_served_model_name(model_path), + served_model_name_explicit=served_model_name is not None, + ) + watcher_handle = state_handle or return_handle + if cache_dir is not None: + weights = _weights_bytes_on_disk(cache_dir) + return_handle.weights_bytes = weights + watcher_handle.weights_bytes = weights + watcher_handle.load_started_at = time.monotonic() + threading.Thread( + target=_watch_mlx_load, + args=(watcher_handle, url, proc), + kwargs={"timeout": _LOAD_TIMEOUT_S}, + name=f"freetoken-metal-load-{port}", + daemon=True, + ).start() + return return_handle + + +def _launch_llama( + model_path: str, + port: int, + *, + state_handle: MetalBackendHandle | None = None, + served_model_name: str | None = None, + **kwargs: Any, +) -> MetalBackendHandle: + binary = llama_binary() + assert binary is not None + cmd = [ + binary, + "-m", + model_path, + "--host", + "127.0.0.1", + "--port", + str(port), + # Metal backend (Apple Silicon) + "-ngl", + "999", + ] + logger.info("launching Metal backend (llama.cpp): %s", shlex.join(cmd)) + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + _start_drain_thread(proc, "llama") + url = f"http://127.0.0.1:{port}" + if state_handle is not None: + state_handle.load_started_at = time.monotonic() + state_handle._set_state("loading", phase="weights") + try: + _wait_for_readiness(url, proc, timeout=_ADDRESS_HEALTH_TIMEOUT_S) + except Exception: + _stop_processes([proc]) + raise + handle = MetalBackendHandle( + processes=[proc], + upstream_base_url=url, + backend="llama", + model_path=model_path, + served_model_name=served_model_name or _default_served_model_name(model_path), + served_model_name_explicit=served_model_name is not None, + ) + handle.load_started_at = time.monotonic() + watcher_handle = state_handle or handle + watcher_handle._set_state("ready") + watcher_handle.ack_queue.put("ready") + return handle + + +def launch_metal_backend( + backend: str, + model_path: str, + upstream_port: int | None = None, + *, + state_handle: MetalBackendHandle | None = None, + served_model_name: str | None = None, +) -> MetalBackendHandle: + """Launch an upstream engine and return a handle to it. + + ``state_handle``: when given, the load watcher publishes its state onto + THAT handle instead of the returned one. Used by ``switch_model``, which + owns a shared handle the proxy routes already read from -- without this, + the watcher would update a detached object and /health would report the + old model's state forever.""" + if backend not in {"mlx", "llama"}: + raise RuntimeError(f"unsupported Metal backend: {backend!r}") + with _serialize_upstream_launches(): + port = _pick_upstream_port(upstream_port) + if backend == "mlx": + return _launch_mlx( + model_path, + port, + state_handle=state_handle, + served_model_name=served_model_name, + ) + if backend == "llama": + return _launch_llama( + model_path, + port, + state_handle=state_handle, + served_model_name=served_model_name, + ) + raise AssertionError("unreachable") + + +# --- HTTP proxy over the FreeToken API ------------------------------------- + + +def _usage_from_payload(payload: Any) -> dict[str, Any] | None: + if not isinstance(payload, dict): + return None + usage = payload.get("usage") + if isinstance(usage, dict): + return usage + message = payload.get("message") + if isinstance(message, dict) and isinstance(message.get("usage"), dict): + return message["usage"] + return None + + +def _usage_from_sse_line(line: bytes) -> dict[str, Any] | None: + if not line.startswith(b"data:"): + return None + raw = line[5:].strip() + if not raw or raw == b"[DONE]": + return None + try: + return _usage_from_payload(json.loads(raw)) + except (ValueError, UnicodeDecodeError): + return None + + +async def _proxy_stream( + response: httpx.Response, + client: httpx.AsyncClient, + handle: MetalBackendHandle, +) -> AsyncIterator[bytes]: + """Stream a pre-opened successful upstream response to the client. + + SSE chunks are rewritten on the fly so the Metal upstream speaks FreeToken's + wire format: mlx_lm splits the thinking channel as ``delta.reasoning`` while + FreeToken (and vLLM/SGLang, which the shell/bench clients read) uses + ``delta.reasoning_content``. + + A partial final line is retained across arbitrary transport chunks so the + JSON key can never straddle two independent replacements. The request has + already received its upstream status before this iterator is returned, so + FastAPI does not commit a false HTTP 200 for upstream errors.""" + pending = b"" + usage: dict[str, Any] | None = None + completed = 200 <= response.status_code < 300 + try: + async for chunk in response.aiter_bytes(chunk_size=None): + pending += chunk + while b"\n" in pending: + line, pending = pending.split(b"\n", 1) + framed = line + b"\n" + usage = _usage_from_sse_line(line) or usage + yield _rewrite_reasoning_field(framed) + if pending: + usage = _usage_from_sse_line(pending) or usage + yield _rewrite_reasoning_field(pending) + finally: + await response.aclose() + await client.aclose() + handle.finish_request(usage, completed=completed) + + +def _rewrite_reasoning_field(chunk: bytes) -> bytes: + """Rename ``"reasoning"`` to ``"reasoning_content"`` inside SSE data lines. + + Byte-level and conservative: only touches the exact ``"reasoning":`` JSON + key (mlx_lm's name for the thinking channel), never the value text, and + leaves non-data bytes (keepalives, separators) untouched.""" + return chunk.replace(b'"reasoning":', b'"reasoning_content":') + + +def register_metal_proxy_routes( + app: FastAPI, get_backend: Any +) -> None: + """Proxy generation routes to the Metal upstream. + + The user-facing surface (``/v1/chat/completions``, ``/v1/completions``, + ``/v1/models``, and Anthropic's ``/v1/messages``) is forwarded verbatim to + the upstream, which already implements the OpenAI/Anthropic-compatible + protocol. Streaming responses pass through as SSE.""" + + # This function is used both with a fresh standalone app and with the + # native module-global app, whose routes were registered at import time. + # FastAPI dispatches the first matching route, so replace every path this + # router owns before appending the Metal handlers. + owned_paths = { + "/v1/chat/completions", + "/v1/completions", + "/v1/messages", + "/v1/messages/count_tokens", + "/v1/responses", + "/v1/embeddings", + "/v1/models", + "/v1/model/list", + "/v1/model/load", + "/v1/stats", + "/v1/cache/status", + "/v1/requests", + "/health", + } + app.router.routes = [ + route + for route in app.router.routes + if getattr(route, "path", None) not in owned_paths + ] + + @app.post("/v1/chat/completions") + async def proxy_chat(request: Request): + return await _forward(request, get_backend) + + @app.post("/v1/completions") + async def proxy_completions(request: Request): + return await _forward(request, get_backend) + + @app.post("/v1/messages") + async def proxy_messages(request: Request): + return await _forward(request, get_backend) + + @app.post("/v1/messages/count_tokens") + async def proxy_messages_count_tokens(request: Request): + return await _forward(request, get_backend) + + @app.post("/v1/responses") + async def proxy_responses(request: Request): + return await _forward(request, get_backend) + + @app.post("/v1/embeddings") + async def proxy_embeddings(request: Request): + return await _forward(request, get_backend) + + @app.get("/v1/models") + async def proxy_models(request: Request): + """List models. Overridden when the upstream reports more than the one it + actually serves (mlx_lm lists the whole local HF cache): the proxy reports + the served model only, so clients (ft shell) label the right one.""" + response = await _forward(request, get_backend, method="GET") + handle = get_backend() + if ( + isinstance(response, Response) + and response.status_code == 200 + and handle is not None + and handle.model_path + ): + try: + doc = json.loads(response.body) + except Exception: # noqa: BLE001 -- keep upstream's answer as-is + return response + data = doc.get("data") if isinstance(doc, dict) else None + if not isinstance(data, list): + return response + ids = [ + item.get("id") + for item in data + if isinstance(item, dict) and isinstance(item.get("id"), str) + ] + served_name = handle.served_model_name or _default_served_model_name( + handle.model_path + ) + if ids == [served_name]: + return response # already truthful + import time as _time + + doc["data"] = [ + {"id": served_name, "object": "model", "created": doc.get("created", int(_time.time()))} + ] + return JSONResponse(doc) + return response + + @app.get("/v1/model/list") + async def proxy_model_list(request: Request): + return await _forward(request, get_backend, method="GET") + + @app.post("/v1/model/load") + async def model_load(request: Request): + """Switch the Metal upstream to a different model. + + Accepts ``{"model": ""}`` and relaunches the upstream + engine (mlx/llama.cpp) on a fresh port after stopping the old engine. + On failure the broken new engine is stopped and the server retains the + error reason for /health instead of leaking an untracked process. + """ + from .accounting import _is_loopback + + if not _is_loopback(request.client.host if request.client else None): + return JSONResponse( + {"detail": "loopback access required"}, status_code=403 + ) + handle = get_backend() + if handle is None or not handle.is_alive(): + return JSONResponse( + {"detail": "Metal backend is not running"}, status_code=503 + ) + try: + body = await request.json() + except Exception: # noqa: BLE001 -- bad JSON is a client error + return JSONResponse( + {"detail": "request body must be JSON with a 'model' field"}, + status_code=400, + ) + model = body.get("model") if isinstance(body, dict) else None + if not model or not isinstance(model, str): + return JSONResponse( + {"detail": "request body must be JSON with a 'model' field"}, + status_code=400, + ) + if model in {handle.model_path, handle.served_model_name}: + return { + "status": "ok", + "model": model, + "detail": "already serving this model", + } + try: + # The switch stops the old engine and waits out the new one's load + # (download + weights) -- minutes for a big model. Run it off the + # event loop so /health (which reports the load's live progress) + # keeps answering while it runs. + await asyncio.to_thread(handle.switch_model, model) + except Exception as exc: # noqa: BLE001 -- a failed switch must not kill the server + return JSONResponse( + {"detail": f"model switch failed: {exc}"}, status_code=500 + ) + return {"status": "ok", "model": handle.model_path} + + @app.get("/v1/stats") + async def metal_stats(request: Request): + """Stable shape for the shell's status-bar poller. The Metal upstream has + no CUDA-style pool stats; request totals come from upstream usage docs.""" + handle = get_backend() + if handle is None: + return JSONResponse({"detail": "Metal backend is not running"}, status_code=503) + return handle.stats_doc() + + @app.get("/v1/cache/status") + async def metal_cache_status(request: Request): + """Minimal geometry doc so the shell's startup read resolves instead of + 404-ing: no MoE cache on the Metal path, no reasoning gears advertised + (the upstream applies its own chat template).""" + return { + "geometry": { + "moe_cache_size": 0, + "moe_cache_policy": "none", + "reasoning": {"gears": [], "kwargs": {}, "default": None}, + }, + "pools": {}, + } + + @app.get("/v1/requests") + async def metal_requests(request: Request): + """``ft ctl requests`` parity. The Metal path has no engine-side request + ring (the CUDA scheduler owns that); an empty list is the truthful + answer -- per-request accounting lives in the upstream's own logs.""" + return {"entries": [], "next_cursor": 0} + + @app.get("/health") + async def metal_health(request: Request): + """Same contract the CUDA path answers: loading -> ok -> error, with + byte progress while loading. Tools gate on ``maintenance == "serving"`` + (bench_decode_moe.wait_ready, the daemon, the desktop); the shell's + attach path renders ``phase`` + ``progress`` as + ``loading (weights): 4.2/15.0 GiB``. Reported from the handle's own + supervised load state, not just process liveness: mlx_lm answers + /v1/models before weights load, so "port up" would lie.""" + handle = get_backend() + if handle is None: + return JSONResponse({"status": "down"}, status_code=503) + doc = handle.health_doc() + if doc.get("status") == "error": + return JSONResponse(doc, status_code=503) + if doc.get("status") != "loading" and not handle.is_alive(): + return JSONResponse({"status": "down"}, status_code=503) + return doc + + +def _strip_host_header(headers: dict[str, str]) -> dict[str, str]: + out = {k: v for k, v in headers.items() if k.lower() not in {"host", "content-length"}} + out["accept"] = headers.get("accept", "application/json") + return out + + +def _rewrite_public_model(body: bytes, handle: MetalBackendHandle) -> bytes: + """Translate the public served-model alias back to the upstream model id.""" + public_name = handle.served_model_name or _default_served_model_name( + handle.model_path + ) + if not body or not public_name or public_name == handle.model_path: + return body + try: + payload = json.loads(body) + except (ValueError, UnicodeDecodeError): + return body + if not isinstance(payload, dict) or payload.get("model") != public_name: + return body + payload["model"] = handle.model_path + return json.dumps(payload).encode() + + +def _inject_turn_stop(path: str, body: bytes, stream: bool) -> tuple[bytes, bool]: + """Add "" to a gemma-4 generation request's stop list. + + Returns the (possibly rewritten) body and the streaming decision. The body + is rewritten only for chat/completions-style generation routes on a + gemma-4 model that does not already stop on the turn token; parse failures + and non-generation routes pass through untouched.""" + if not path.endswith(("/chat/completions", "/completions", "/messages", "/responses")): + return body, stream + try: + payload = json.loads(body or b"{}") + except (ValueError, UnicodeDecodeError): + return body, stream + if not isinstance(payload, dict) or not _needs_turn_stop_token(payload): + return body, stream + stop = payload.get("stop") + if isinstance(stop, str): + stop = [stop] + elif not isinstance(stop, list): + stop = [] + if "" in stop: + return body, stream + stop.append("") + payload["stop"] = stop + return json.dumps(payload).encode(), stream + + +async def _forward( + request: Request, + get_backend: Any, + method: str = "POST", +) -> Response: + handle = get_backend() + if handle is None or not handle.is_alive(): + return JSONResponse( + {"detail": "Metal backend is not running"}, status_code=503 + ) + if method == "POST" and not getattr(handle, "is_ready", lambda: True)(): + # The engine is still loading (download or weights). Answer immediately + # with the state instead of letting the request queue behind the load + # forever -- the failure mode this proxy used to hang on. + doc = handle.health_doc() + detail = f"model is still loading ({doc.get('phase', 'starting')})" + return JSONResponse({"detail": detail, **doc}, status_code=503) + upstream = handle.upstream_base_url + path = request.url.path + if request.url.query: + path = f"{path}?{request.url.query}" + + body = await request.body() + headers = _strip_host_header(dict(request.headers)) + # Streaming is signalled in the BODY ("stream": true), not the Accept + # header -- the OpenAI SDK sends Accept: application/json even for + # streaming requests. Deciding from the header routed streamed + # generations through the buffered path: the client got the whole answer + # in one burst at the end and live token counts read 0/burst. + stream = "text/event-stream" in (headers.get("accept") or "") + if not stream: + try: + payload = json.loads(body or b"{}") + except (ValueError, UnicodeDecodeError): + payload = {} + stream = bool(payload.get("stream")) if isinstance(payload, dict) else False + + # Some gemma-4 snapshots omit the turn-end token () from their + # generation config, so inject it as a defensive stop word for that family. + body = _rewrite_public_model(body, handle) + body, stream = _inject_turn_stop(request.url.path, body, stream) + + if stream: + client = httpx.AsyncClient(timeout=None) + try: + upstream_request = client.build_request( + method, f"{upstream}{path}", content=body, headers=headers + ) + response = await client.send(upstream_request, stream=True) + except httpx.HTTPError as exc: + await client.aclose() + return JSONResponse( + {"detail": f"Metal upstream error: {exc}"}, status_code=502 + ) + + # Obtain the status and, on failure, the body before returning a + # StreamingResponse. Otherwise Starlette commits 200 before the + # generator's first await and upstream 4xx/5xx become truncated 200s. + if not 200 <= response.status_code < 300: + try: + content = await response.aread() + finally: + await response.aclose() + await client.aclose() + return Response( + content=content, + status_code=response.status_code, + media_type=( + response.headers.get("content-type", "application/json") + .split(";", 1)[0] + ), + ) + + handle.begin_request() + return StreamingResponse( + _proxy_stream(response, client, handle), + status_code=response.status_code, + media_type=response.headers.get("content-type", "text/event-stream").split( + ";", 1 + )[0], + ) + + is_generation = request.url.path in { + "/v1/chat/completions", + "/v1/completions", + "/v1/messages", + "/v1/responses", + } + if is_generation: + handle.begin_request() + try: + async with httpx.AsyncClient(timeout=httpx.Timeout(None)) as client: + r = await client.request( + method, f"{upstream}{path}", content=body, headers=headers + ) + except httpx.HTTPError as exc: + if is_generation: + handle.finish_request(None, completed=False) + return JSONResponse({"detail": f"Metal upstream error: {exc}"}, status_code=502) + content = r.content + if is_generation: + try: + usage = _usage_from_payload(json.loads(content)) + except (ValueError, UnicodeDecodeError): + usage = None + handle.finish_request(usage, completed=200 <= r.status_code < 300) + # Same reasoning-channel rename as the streaming path, for non-streaming + # completions (mlx_lm's "reasoning" -> FreeToken's "reasoning_content"). + if b'"reasoning":' in content: + content = content.replace(b'"reasoning":', b'"reasoning_content":') + return Response( + content=content, + status_code=r.status_code, + media_type=r.headers.get("content-type"), + ) diff --git a/python/freetoken/server/metal_main.py b/python/freetoken/server/metal_main.py new file mode 100644 index 0000000..777a9d1 --- /dev/null +++ b/python/freetoken/server/metal_main.py @@ -0,0 +1,136 @@ +"""Standalone entrypoint for ``ft serve-metal`` (Apple Silicon, no CUDA). + +This module deliberately does NOT import FreeToken's CUDA engine stack. The +classic ``ft serve`` path (``server/args.parse_args`` / ``SchedulerConfig`` / the +MoE + layer graph) transitively imports ``flashlib`` and the CUDA kernels, which +have no macOS build. On Apple Silicon that whole chain cannot even be imported. + +``serve-metal`` therefore parses its own minimal arguments, launches the chosen +Apple Metal runtime (mlx or llama.cpp) as an upstream OpenAI/Anthropic-compatible +engine, and serves a thin HTTP proxy on the configured host/port. This gives the +exact OpenAI/Anthropic/Responses wire surface with nothing imported from the +CUDA scheduler. + +Built on the reusable pieces in :mod:`freetoken.server.metal`. +""" + +from __future__ import annotations + +import argparse +import asyncio +import sys +from typing import Sequence + +import uvicorn +from fastapi import FastAPI + +from freetoken.server.metal import ( + MetalBackendHandle, + launch_metal_backend, + register_metal_proxy_routes, + resolve_metal_backend, +) +from freetoken.server.cors import DEFAULT_CORS_ORIGINS, install_cors + + +def _parse(argv: Sequence[str]) -> argparse.Namespace: + p = argparse.ArgumentParser( + prog="ft serve-metal", + description="Serve a model on an Apple Silicon Metal backend (mlx or llama.cpp).", + ) + p.add_argument("--model", required=True, help="Model path or HF id for the Metal backend.") + p.add_argument( + "--served-model-name", + default=None, + help="Public model id (default: basename of --model).", + ) + p.add_argument( + "--backend", + default="auto", + choices=["auto", "mlx", "llama"], + help="Metal engine: mlx (mlx_lm.server) or llama (llama.cpp). auto prefers mlx.", + ) + p.add_argument("--host", default="127.0.0.1") + p.add_argument("--port", type=int, default=1919) + p.add_argument("--metal-port", type=int, default=0, help="Upstream port (0 = auto).") + p.add_argument( + "--cors-origins", + default=DEFAULT_CORS_ORIGINS, + help="Comma-separated CORS allow-list; empty disables and '*' allows any origin.", + ) + p.add_argument( + "--shell", + "--shell-mode", + dest="shell", + action="store_true", + help="Attach the interactive ft shell to this server (serve+chat in one process).", + ) + args, unknown = p.parse_known_args(list(argv)) + if unknown: + # Callers built for the CUDA engine (the daemon, benchmarks, `ft serve` + # flag passthrough) legitimately carry CUDA-only flags; the Metal path + # has no such engine knobs, so drop them loudly rather than refusing. + print( + f"ft serve-metal: ignoring CUDA-only engine flags: {' '.join(unknown)}", + file=sys.stderr, + ) + return args + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse(sys.argv[1:] if argv is None else argv) + backend = resolve_metal_backend(args.backend) + # Non-blocking: spawns the upstream and returns while its watcher thread + # supervises the load (see _watch_mlx_load). uvicorn binds immediately, so + # /health reports live load progress and ft shell can attach and render it + # instead of the terminal sitting silent through a 50 GiB download. + handle: MetalBackendHandle = launch_metal_backend( + backend, + args.model, + args.metal_port, + served_model_name=args.served_model_name, + ) + + _HANDLE = {"handle": handle} + app = FastAPI(title="FreeToken Metal API Server") + install_cors(app, args.cors_origins) + + def get_backend(): + # Route handlers distinguish loading, terminal error, and down states. + # Filtering here would erase a failed engine's actionable error reason. + return _HANDLE["handle"] + + register_metal_proxy_routes(app, get_backend) + + # The shell/desktop poll /health, /v1/stats and /v1/cache/status every + # second; hide those from the access log so they don't bury real requests. + # Same filter the CUDA api_server installs (access_log_filter.py). + from freetoken.server.access_log_filter import install_polling_access_log_filter + + install_polling_access_log_filter() + + try: + if args.shell: + import threading + + origin = f"http://{args.host}:{args.port}" + server = uvicorn.Server( + uvicorn.Config(app, host=args.host, port=args.port, access_log=False) + ) + thread = threading.Thread( + target=server.run, name="freetoken-uvicorn", daemon=True + ) + thread.start() + from freetoken.shell.tui import run_shell + + return asyncio.run(run_shell(origin, connect_grace=30.0)) + uvicorn.run(app, host=args.host, port=args.port) + except KeyboardInterrupt: + pass + finally: + handle.terminate() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/python/freetoken/server/process_utils.py b/python/freetoken/server/process_utils.py new file mode 100644 index 0000000..748d01d --- /dev/null +++ b/python/freetoken/server/process_utils.py @@ -0,0 +1,44 @@ +"""Torch-free process helpers shared by native and Metal server lifecycles.""" + +from __future__ import annotations + +import subprocess +from typing import Any, Iterable + + +def process_is_alive(process: Any) -> bool: + """Normalize multiprocessing.Process and subprocess.Popen liveness.""" + if hasattr(process, "is_alive"): + return bool(process.is_alive()) + return process.poll() is None + + +def terminate_backend_workers(processes: Iterable[Any]) -> None: + """Best-effort, nonblocking SIGTERM for every live backend worker.""" + for process in processes or []: + try: + if process_is_alive(process): + process.terminate() + except Exception: # noqa: BLE001 -- already gone or unqueryable + continue + + +def reap_backend_workers(processes: Iterable[Any], timeout: float = 5.0) -> None: + """Wait for SIGTERM and kill workers that remain alive after the timeout.""" + for process in processes or []: + try: + if hasattr(process, "join"): + process.join(timeout=timeout) + else: + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + pass + if process_is_alive(process): + process.kill() + if hasattr(process, "join"): + process.join(timeout=timeout) + else: + process.wait(timeout=timeout) + except Exception: # noqa: BLE001 -- already gone or unqueryable + continue diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a..a86b096 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -15,6 +15,8 @@ from queue import Empty from typing import Any, Callable, List +from .process_utils import process_is_alive + def phase_slug(desc: str) -> str: """Normalize a progress-bar desc into the stable /health phase enum. Order matters: @@ -65,7 +67,7 @@ def __init__(self, message: str) -> None: def _first_dead(processes: List[Any]) -> Any | None: for p in processes: try: - if not p.is_alive(): + if not process_is_alive(p): return p except Exception: # noqa: BLE001 — treat an unqueryable handle as alive continue diff --git a/python/freetoken/shell/__init__.py b/python/freetoken/shell/__init__.py index e429f48..52927f7 100644 --- a/python/freetoken/shell/__init__.py +++ b/python/freetoken/shell/__init__.py @@ -23,6 +23,28 @@ # side (--moe-cache-auto, --attn, ...) only makes sense alongside these. _ENGINE_FLAGS = ("--model", "--model-path") +#: Flags the Metal path does not understand (CUDA/engine-specific); dropped when +#: `ft shell --model` routes to `ft serve-metal` on macOS. +_METAL_UNKNOWN_VALUE_FLAGS = ( + "--dtype", + "--moe-backend", + "--moe-cache-size", + "--moe-cache-rate", + "--attn", + "--attention-backend", + "--tool-call-parser", + "--reasoning-parser", + "--sampling-defaults", + "--cuda-graph-max-bs", + "--max-running-req", +) +_METAL_UNKNOWN_BOOLEAN_FLAGS = ( + "--moe-cache-auto", + "--use-dummy-weight", + "--silent-output", + "--shell-mode", +) + def _wants_local_engine(argv: Sequence[str]) -> bool: return any(arg in _ENGINE_FLAGS or arg.startswith(tuple(f + "=" for f in _ENGINE_FLAGS)) @@ -48,10 +70,64 @@ def _build_parser(prog: str) -> argparse.ArgumentParser: return parser +def _split_engine_args(argv: Sequence[str]) -> tuple[str | None, list[str]]: + """Extract the (last) --model/--model-path value and the remaining args. + + Understands both ``--model X`` and ``--model=X`` forms, and skips over the + value token of engine flags that take one, so it never mistakes a value for + a flag or leaks it into the passthrough list. + """ + value_flags = set(_ENGINE_FLAGS + _METAL_UNKNOWN_VALUE_FLAGS) + boolean_flags = set(_METAL_UNKNOWN_BOOLEAN_FLAGS) + model = None + passthrough: list[str] = [] + i = 0 + while i < len(argv): + arg = argv[i] + if arg.startswith("--") and "=" in arg: + name, _, value = arg.partition("=") + if name in value_flags: + if name in _ENGINE_FLAGS: + model = value + else: + passthrough.append(arg) + i += 1 + elif arg in boolean_flags: + i += 1 + elif arg in value_flags: + value = argv[i + 1] if i + 1 < len(argv) else None + if arg in _ENGINE_FLAGS: + model = value + i += 2 if value is not None else 1 + else: + passthrough.append(arg) + i += 1 + return model, passthrough + + def main(argv: Sequence[str] | None = None, *, prog: str = "ft shell") -> int: args = list(sys.argv[1:] if argv is None else argv) if _wants_local_engine(args): + # On Apple Silicon the CUDA launcher cannot even import (torch/flashlib have + # no macOS build); route the same flags to the Metal backend instead, so + # `ft shell --model ` works there too. On a CUDA box the native + # launcher takes over, unchanged. + if sys.platform == "darwin": + from freetoken.server.metal import resolve_backend + + try: + backend = resolve_backend("auto") + except RuntimeError: + backend = None # fall through to the native launcher's own error + if backend in ("mlx", "llama"): + from freetoken.server.metal_main import main as metal_main + + model, passthrough = _split_engine_args(args) + if model is None: + print("--model is required", file=sys.stderr) + return 2 + return metal_main(["--shell", "--model", model, *passthrough]) from freetoken.server import launch_server launch_server(run_shell=True, argv=args, prog=prog) diff --git a/python/freetoken/shell/client.py b/python/freetoken/shell/client.py index 183fa4d..976efbc 100644 --- a/python/freetoken/shell/client.py +++ b/python/freetoken/shell/client.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import contextlib import json import time import urllib.error @@ -203,16 +204,93 @@ async def cache_rebuild( ) async def model_id(self) -> str | None: - """The first id from ``/v1/models``. None when the server reports none (it is a single- - model server, so this is the model the shell will be talking to).""" + """The id the shell should talk to. + + Single-model servers (the norm, and always the case through the Metal + proxy) report one id and it is used directly. Some upstreams (raw + mlx_lm) list every model in the local HF cache -- then the served model + cannot be told apart client-side, so fall back to the request-time + default: the id the server echoes when the ``model`` field is omitted. + """ doc = await self._request_json("GET", "/v1/models") data = doc.get("data") if not isinstance(data, list): return None - for item in data: - if isinstance(item, dict) and isinstance(item.get("id"), str): - return item["id"] - return None + ids = [ + item["id"] + for item in data + if isinstance(item, dict) and isinstance(item.get("id"), str) + ] + if len(ids) == 1: + return ids[0] + if not ids: + return None + # Multi-model listing: ask the server which one a bare request uses. + # mlx_lm (and llama.cpp) echo the served model in the response body. + with contextlib.suppress(ShellClientError, OSError, ValueError, KeyError, TypeError): + request = urllib.request.Request( + f"{self.origin}/v1/chat/completions", + data=json.dumps({"messages": [{"role": "user", "content": "ping"}], "max_tokens": 1}).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(request, timeout=self.timeout) as response: + body = json.loads(response.read().decode("utf-8")) + echoed = body.get("model") + if isinstance(echoed, str) and echoed in ids: + return echoed + return ids[0] + + async def list_models(self) -> list[str]: + """All ids from ``/v1/models``, for ``/model``'s listing.""" + doc = await self._request_json("GET", "/v1/models") + data = doc.get("data") + if not isinstance(data, list): + return [] + return [ + item["id"] + for item in data + if isinstance(item, dict) and isinstance(item.get("id"), str) + ] + + async def load_model( + self, + model: str, + *, + wait: float = 600.0, + on_progress: Callable[[dict[str, Any]], None] | None = None, + ) -> dict[str, Any]: + """Switch the served model via ``POST /v1/model/load``, streaming load progress. + + The server relaunches its engine for the new model and publishes live + progress (phase + byte counts) on ``/health`` while the POST is still in + flight. Kicking the POST off as a background task and polling /health + alongside it turns the switch into the same live view the startup path + renders, instead of an opaque minutes-long hang. + + ``on_progress`` receives each /health doc as the switch runs. If the + server is one of those that answers the switch instantly (no load to + watch), the POST wins the race and no progress is ever shown.""" + if on_progress is None: + return await self._request_json( + "POST", "/v1/model/load", body={"model": model}, timeout=wait + ) + post = asyncio.create_task( + self._request_json("POST", "/v1/model/load", body={"model": model}, timeout=wait) + ) + try: + while not post.done(): + try: + doc = await self.health() + except ShellClientError: + doc = None + if isinstance(doc, dict) and doc.get("status") not in (None, "ok"): + on_progress(doc) + await asyncio.sleep(READY_POLL_INTERVAL) + return await post + except BaseException: + post.cancel() + raise async def wait_until_ready( self, diff --git a/python/freetoken/shell/tui.py b/python/freetoken/shell/tui.py index 93bd18c..de09223 100644 --- a/python/freetoken/shell/tui.py +++ b/python/freetoken/shell/tui.py @@ -453,6 +453,7 @@ def _help_text(think_gears: Tuple[str, ...], pools: CachePools) -> str: rows = [ ("/help", "show this message"), (think, think_help), + ("/model [id]", "show the served model, or switch to another one"), (f"/cache [status | {_cache_targets_hint(pools)}]", ""), ("", "show or resize the cache pools; token targets are"), ("", "rounded up to the pool's page size"), @@ -679,7 +680,7 @@ async def run_turn(cmd: str) -> None: history.append((cmd, text)) async def handle_command(cmd: str) -> None: - nonlocal history, think_gear, cache_pools + nonlocal history, think_gear, think_gears, think_kwargs, cache_pools, model_id if cmd == "": return if cmd.startswith("/"): @@ -700,6 +701,62 @@ async def handle_command(cmd: str) -> None: stats.think_gear = think_gear renderer.write(message + "\n") return + if slash == "/model": + if len(parts) == 1: + lines = [f"Serving {model_id}."] + with contextlib.suppress(ShellClientError): + models = await client.list_models() + if models: + lines.append("Available: " + ", ".join(models)) + lines.append("Usage: /model to switch.") + renderer.write("\n".join(lines) + "\n") + return + candidate = parts[1] + if candidate == model_id: + renderer.write(f"Already serving {model_id}.\n") + return + renderer.write(f"Switching model to {candidate}...\n") + + def _on_switch_progress(doc: dict) -> None: + nonlocal last_line, last_at + line = _format_load_progress(doc) + now = time.monotonic() + new_phase = line.split(":", 1)[0] != last_line.split(":", 1)[0] + if line == last_line or (not new_phase and now - last_at < LOAD_PROGRESS_INTERVAL): + return + last_line, last_at = line, now + renderer.write(f"{line}\n") + try: + doc = await client.load_model(candidate, on_progress=_on_switch_progress) + except ShellClientError as exc: + if exc.status == 404: + renderer.write( + "This server does not support model switching " + "(no /v1/model/load endpoint).\n" + ) + else: + renderer.write(f"{exc}\n") + return + model_id = candidate + history = [] + stats.reset() + stats.model_label = _format_shell_model_label(model_id) + # Re-read what the new model offers: think gears, cache geometry, + # a clean stats baseline. + with contextlib.suppress(ShellClientError): + cache_doc = await client.cache_status() + geometry = cache_doc.get("geometry") or {} + reasoning = geometry.get("reasoning") or {} + think_gears = tuple(reasoning.get("gears") or ()) + think_kwargs = reasoning.get("kwargs") or {} + think_gear = reasoning.get("default") + stats.think_gear = think_gear + stats.apply_geometry(geometry) + cache_pools = CachePools.from_geometry(geometry) + with contextlib.suppress(ShellClientError): + stats.apply_stats_doc(await client.stats()) + renderer.write(f"Now serving {model_id}.\n") + return if slash == "/cache": pools = await _handle_cache_command(parts[1:], client, stats, renderer) if pools is not None: diff --git a/python/freetoken/utils/logger.py b/python/freetoken/utils/logger.py index e591864..0d9efb9 100644 --- a/python/freetoken/utils/logger.py +++ b/python/freetoken/utils/logger.py @@ -1,129 +1,5 @@ -from __future__ import annotations +"""Compatibility export for the torch-free logger implementation.""" -from functools import partial -from typing import TYPE_CHECKING +from freetoken.logging import init_logger -_LOG_LEVEL = None - - -def init_logger( - name: str, - suffix: str = "", - *, - strip_file: bool = True, - level: str | None = None, - use_pid: bool | None = None, - use_tp_rank: bool | None = None, -): - """Initialize the logger for the module with colors and pretty formatting.""" - import logging - import os - import sys - - global _LOG_LEVEL - if _LOG_LEVEL is None: - LEVEL_MAP = { - "DEBUG": logging.DEBUG, - "INFO": logging.INFO, - "WARNING": logging.WARNING, - "ERROR": logging.ERROR, - "CRITICAL": logging.CRITICAL, - } - - level = level or os.getenv("LOG_LEVEL", "").upper() - _LOG_LEVEL = LEVEL_MAP.get(level, logging.INFO) - - if strip_file: - suffix = os.path.basename(suffix) - - if suffix: - suffix = f"|{suffix}" - - if use_pid is None: - use_pid = os.getenv("LOG_PID", "0").lower() in ("1", "true", "yes") - - if use_pid: - pid = os.getpid() - suffix = f"|pid={pid}{suffix}" - - tp_info = None - - # Color formatter class - class ColorFormatter(logging.Formatter): - """Formatter with colors and pretty output""" - - # ANSI color codes - COLORS = { - "DEBUG": "\033[36m", # Cyan - "INFO": "\033[32m", # Green - "WARNING": "\033[33m", # Yellow - "ERROR": "\033[31m", # Red - "CRITICAL": "\033[35m", # Magenta - } - RESET = "\033[0m" - BOLD = "\033[1m" - - def format(self, record): - from freetoken.distributed import try_get_tp_info - - # Format timestamp like SGLang: [YYYY-MM-DD|HH:MM:SS|pid=1234] - timestamp = self.formatTime(record, "[%Y-%m-%d|%H:%M:%S{suffix}]") - nonlocal tp_info - tp_info = tp_info or try_get_tp_info() - if tp_info is not None and use_tp_rank is not False: - real_suffix = f"{suffix}|core|rank={tp_info.rank}" - else: - real_suffix = suffix - timestamp = timestamp.format(suffix=real_suffix) - - # Get color for log level - level_color = self.COLORS.get(record.levelname, "") - - # Format the message - colored_level = f"{level_color}{record.levelname:<8}{self.RESET}" - message = record.getMessage() - - # Pretty format: [timestamp] LEVEL message - return f"{self.BOLD}{timestamp}{self.RESET} {colored_level} {message}" - - logger = logging.getLogger(name) - logger.setLevel(_LOG_LEVEL) - - # Clear existing handlers to avoid duplicates - logger.handlers.clear() - - handler = logging.StreamHandler(sys.stdout) - formatter = ColorFormatter() - handler.setFormatter(formatter) - logger.addHandler(handler) - - # Prevent propagation to root logger - logger.propagate = False - - def _call_rank0(msg, *args, _which, **kwargs): - from freetoken.distributed import try_get_tp_info - - nonlocal tp_info - tp_info = tp_info or try_get_tp_info() - # No TP set yet (e.g. a unit test or a tool that loads weights without distributed - # init) -> treat as a single rank (primary) and log, rather than crashing. - if tp_info is None or tp_info.is_primary(): - getattr(logger, _which)(msg, *args, **kwargs) - - if TYPE_CHECKING: - - class WrapperLogger(logging.Logger): - """Custom logger to handle the color formatter.""" - - def info_rank0(self, msg, *args, **kwargs): ... - def warning_rank0(self, msg, *args, **kwargs): ... - def debug_rank0(self, msg, *args, **kwargs): ... - def critical_rank0(self, msg, *args, **kwargs): ... - - return WrapperLogger(name) - else: - logger.info_rank0 = partial(_call_rank0, _which="info") - logger.debug_rank0 = partial(_call_rank0, _which="debug") - logger.critical_rank0 = partial(_call_rank0, _which="critical") - logger.warning_rank0 = partial(_call_rank0, _which="warning") - return logger +__all__ = ["init_logger"] diff --git a/scripts/start-metal.sh b/scripts/start-metal.sh new file mode 100755 index 0000000..60493bb --- /dev/null +++ b/scripts/start-metal.sh @@ -0,0 +1,139 @@ +#!/usr/bin/env bash +# +# start-metal.sh — start FreeToken on the Apple Silicon Metal backend, with chat. +# +# One command for the local-testing loop on a Mac: it checks the environment, +# (re)creates the venv if needed, starts the Metal server, and — unless +# NO_CHAT=1 — attaches `ft shell` to it for interactive testing. On exit the +# server is stopped, so nothing is left holding ports 1919/190xx. +# +# Usage: +# scripts/start-metal.sh [model] [options] +# +# model an MLX/HF repo id (mlx-community/*) or a local .gguf file +# (default: mlx-community/Qwen3-0.6B-4bit) +# options passed through to `ft serve-metal` (--backend llama, +# --port 1919, ...) +# +# Environment: +# NO_CHAT=1 start the server only (no interactive shell) +# FREETOKEN_MODEL default model, overridden by the first positional arg +# FREETOKEN_PORT server port (default 1919) +# +# Examples: +# scripts/start-metal.sh # tiny Qwen3 + chat +# scripts/start-metal.sh mlx-community/Llama-3.2-1B-Instruct-4bit +# scripts/start-metal.sh ~/models/foo.Q4_K_M.gguf --backend llama +# NO_CHAT=1 scripts/start-metal.sh # API only +# curl -s localhost:1919/v1/models # in another terminal +# +set -euo pipefail + +MODEL="${1:-${FREETOKEN_MODEL:-mlx-community/Qwen3-0.6B-4bit}}" +# consume the model arg so the rest go to ft serve-metal +if [[ $# -gt 0 ]]; then shift; fi +PORT="${FREETOKEN_PORT:-1919}" +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +PY="$ROOT/.venv/bin/python" +SERVER_PID="" + +die() { echo "start-metal.sh: $*" >&2; exit 1; } + +cleanup() { + [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true + wait "$SERVER_PID" 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +command -v uv >/dev/null 2>&1 || die "uv not found — install it from https://docs.astral.sh/uv/" + +# --- environment checks ------------------------------------------------------ +[[ "$(uname -s)" == "Darwin" ]] || die "Metal backend is macOS-only (this is $(uname -s))" +[[ "$(uname -m)" == "arm64" ]] || die "native arm64 shell required — an x86_64 (Rosetta) terminal cannot install mlx wheels" +[[ -x "$PY" ]] || die "no .venv at $ROOT — run: cd $ROOT && uv venv && uv pip install -e . && uv pip install mlx-lm" + +# --- pick the engine --------------------------------------------------------- +BACKEND="auto" +ARGS=("$@") +for ((i = 0; i < ${#ARGS[@]}; i++)); do + case "${ARGS[$i]}" in + --backend) + [[ $((i + 1)) -lt ${#ARGS[@]} ]] && BACKEND="${ARGS[$((i + 1))]}" + ;; + --backend=*) + BACKEND="${ARGS[$i]#--backend=}" + ;; + --port) + [[ $((i + 1)) -lt ${#ARGS[@]} ]] && PORT="${ARGS[$((i + 1))]}" + ;; + --port=*) + PORT="${ARGS[$i]#--port=}" + ;; + esac +done + +case "$BACKEND" in + mlx) + "$PY" -c "import mlx_lm" 2>/dev/null || die "mlx-lm not installed in .venv — run: uv pip install mlx-lm" + ;; + llama) + command -v llama-server >/dev/null 2>&1 || die "llama-server not on PATH — run: brew install llama.cpp" + ;; + auto) + if ! "$PY" -c "import mlx_lm" 2>/dev/null && ! command -v llama-server >/dev/null 2>&1; then + die "no Metal backend installed — run: uv pip install mlx-lm, or: brew install llama.cpp" + fi + ;; + *) + die "invalid backend '$BACKEND' — expected auto, mlx, or llama" + ;; +esac + +# --- port pre-check ----------------------------------------------------------- +# Fail fast and clearly instead of racing another server into "address already in +# use" (or worse: silently attaching our chat to someone else's server on that port). +if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + die "port $PORT is already serving a FreeToken server ($(curl -fsS "http://127.0.0.1:$PORT/v1/models" 2>/dev/null | "$PY" -c 'import json,sys +try: + print([m["id"] for m in json.load(sys.stdin)["data"]]) +except Exception: + print("?")' 2>/dev/null)). Stop it first, or set FREETOKEN_PORT." +fi + +# --- start ------------------------------------------------------------------- +cd "$ROOT" +echo "▶ FreeToken Metal backend: model=$MODEL backend=$BACKEND port=$PORT" +"$ROOT/.venv/bin/ft" serve-metal --model "$MODEL" --port "$PORT" "$@" & +SERVER_PID=$! + +# --- wait for readiness ------------------------------------------------------ +# /health answering 200 means the API is up; /v1/models returning the served id +# means the Metal upstream finished loading (mlx_lm only serves /v1/models +# meaningfully once ready). Wait for both, so the chat attaches to a live engine. +for i in $(seq 1 240); do + if curl -fsS "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then + SERVED="$(curl -fsS "http://127.0.0.1:$PORT/v1/models" 2>/dev/null | "$PY" -c 'import json,sys +try: + ids = [m["id"] for m in json.load(sys.stdin)["data"]] +except Exception: + ids = [] +print(ids[0] if len(ids) == 1 else "")' 2>/dev/null || true)" + if [[ -n "$SERVED" ]]; then + echo "✔ server ready at http://127.0.0.1:$PORT (serving: $SERVED)" + break + fi + fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + die "server exited during startup (see its output above)" + fi + sleep 1 +done +[[ -n "${SERVED:-}" ]] || die "server did not become ready within 240s" + +if [[ "${NO_CHAT:-0}" != "1" ]]; then + echo "▶ starting chat (Ctrl-D or /exit to quit; the server stops with it)" + "$ROOT/.venv/bin/ft" shell --server "http://127.0.0.1:$PORT" || true +else + # server-only mode: run until interrupted + wait "$SERVER_PID" +fi diff --git a/setup.py b/setup.py index cfe41b7..7850488 100644 --- a/setup.py +++ b/setup.py @@ -1,10 +1,12 @@ from __future__ import annotations import importlib.util +import sys from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension + +_IS_MACOS = sys.platform == "darwin" ROOT = Path(__file__).parent @@ -18,25 +20,32 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() -def _cuda_runtime_paths() -> tuple[list[str], list[str]]: - if CUDA_HOME is None: - raise RuntimeError( - "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " - "because it links against the CUDA runtime API." - ) - cuda_home = Path(CUDA_HOME) +def _cuda_runtime_paths(cuda_home: Path) -> tuple[list[str], list[str]]: library_dirs = [str(cuda_home / "lib64")] if (cuda_home / "lib").exists(): library_dirs.append(str(cuda_home / "lib")) return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() -_check_toolchain() +cuda_include_dirs: list[str] = [] +cuda_library_dirs: list[str] = [] +ext_modules: list = [] +cmdclass: dict = {} +if not _IS_MACOS: + # The two C++ extensions below link the CUDA runtime; macOS (Metal) builds + # have no compiled extensions at all (the Metal path runs Apple's mlx / + # llama.cpp runtimes as upstream engines, see server/metal.py). + from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME -setup( - ext_modules=[ + if CUDA_HOME is None: + raise RuntimeError( + "CUDA_HOME is required to build freetoken.kernel._pinned_tensor " + "because it links against the CUDA runtime API." + ) + cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths(Path(CUDA_HOME)) + _check_toolchain() + ext_modules = [ CppExtension( name="freetoken.kernel._pinned_tensor", sources=[ @@ -62,6 +71,11 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: libraries=["cudart"], extra_compile_args=["-O3", "-std=c++17", "-pthread"], ), - ], - cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, + ] + cmdclass = {"build_ext": BuildExtension.with_options(use_ninja=True)} + + +setup( + ext_modules=ext_modules, + cmdclass=cmdclass, ) diff --git a/tests/daemon/test_serve_command_platform.py b/tests/daemon/test_serve_command_platform.py new file mode 100644 index 0000000..e78b8cd --- /dev/null +++ b/tests/daemon/test_serve_command_platform.py @@ -0,0 +1,48 @@ +"""The daemon's serve argv is platform-specific: Metal on Darwin, CUDA elsewhere.""" + +from __future__ import annotations + +import sys +import signal + +from freetoken.daemon import osproc +from freetoken.daemon.serve_manager import build_serve_command + + +def test_build_serve_command_uses_serve_metal_on_darwin(monkeypatch, tmp_path): + monkeypatch.setattr(sys, "platform", "darwin") + argv, log_path = build_serve_command( + "mlx-community/Qwen3-0.6B-4bit", + 1919, + ["--tp-size", "1"], + python="/usr/bin/python3", + log_dir=str(tmp_path), + ) + assert argv[0] == "/usr/bin/python3" + assert argv[argv.index("-m") + 2] == "serve-metal" + assert "--model" in argv and "mlx-community/Qwen3-0.6B-4bit" in argv + assert log_path.endswith("serve-1919.log") + + +def test_build_serve_command_uses_serve_on_linux(monkeypatch, tmp_path): + monkeypatch.setattr(sys, "platform", "linux") + argv, _ = build_serve_command( + "Qwen", + 1919, + [], + python="/usr/bin/python3", + log_dir=str(tmp_path), + ) + assert argv[argv.index("-m") + 2] == "serve" + assert "serve-metal" not in argv + + +def test_signal_group_uses_portable_getpgid_without_proc(monkeypatch): + sent = [] + monkeypatch.setattr(osproc, "_stat_fields", lambda pid: None) + monkeypatch.setattr(osproc.os, "getpgid", lambda pid: pid) + monkeypatch.setattr(osproc.os, "killpg", lambda pgid, sig: sent.append((pgid, sig))) + + osproc.signal_group(4321, signal.SIGKILL) + + assert sent == [(4321, signal.SIGKILL)] diff --git a/tests/server/test_metal_backend.py b/tests/server/test_metal_backend.py new file mode 100644 index 0000000..3dd7105 --- /dev/null +++ b/tests/server/test_metal_backend.py @@ -0,0 +1,916 @@ +"""Tests for the Apple Silicon Metal backend (server/metal.py). + +Pure-process tests: no GPU, no Metal engine, no torch. They cover the +backend resolver, the upstream port allocator, and the proxy routes +mounted on a FastAPI app against a real loopback upstream (uvicorn in a +thread). + +Run: PYTHONPATH=python /bin/python -m pytest tests/server/test_metal_backend.py -v +""" + +from __future__ import annotations + +import asyncio +import json +import os +import socket +import sys +import threading +import time +from types import SimpleNamespace + +import pytest +import uvicorn +from fastapi import FastAPI +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + + +def _free_port() -> int: + """A random free port (the 19000-19 range is littered with TIME_WAIT + residue from real Metal runs).""" + s = socket.socket() + try: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + except PermissionError as exc: + pytest.skip(f"loopback socket binding is unavailable: {exc}") + finally: + s.close() + +_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +_PY = os.path.join(_ROOT, "python") +if _PY not in sys.path: + sys.path.insert(0, _PY) + +from freetoken.server import metal # noqa: E402 + + +# ---------------------------------------------------------------- resolver -- + +def test_resolve_backend_explicit_passthrough(): + assert metal.resolve_backend("mlx") == "mlx" + assert metal.resolve_backend("llama") == "llama" + + +def test_resolve_backend_rejects_unknown(): + try: + metal.resolve_backend("tpu") + except RuntimeError as e: + assert "unknown --backend" in str(e) + else: + raise AssertionError("expected RuntimeError for unknown backend") + + +def test_standalone_auto_never_selects_cuda(monkeypatch): + monkeypatch.setattr(metal, "_any_cuda_usable", lambda: True) + monkeypatch.setattr(metal, "mlx_importable", lambda: True) + assert metal.resolve_metal_backend("auto") == "mlx" + + +def test_pick_upstream_port_defaults_into_reserved_range(monkeypatch): + monkeypatch.setattr(metal, "_claim_port", lambda port: port) + for preferred in (None, 0): + port = metal._pick_upstream_port(preferred) + assert metal._UPSTREAM_PORT_MIN <= port <= metal._UPSTREAM_PORT_MAX, port + + +def test_pick_upstream_port_prefers_free_preferred_port(monkeypatch): + monkeypatch.setattr(metal, "_claim_port", lambda port: port) + assert metal._pick_upstream_port(19099) == 19099 + + +# --------------------------------------------------------- Gemma 4 prompt -- + +def test_gemma4_template_renders_canonical_generation_prompt(): + """Lock down the exact turn boundary that prevents raw-text continuation.""" + from jinja2 import Environment + + rendered = Environment().from_string(metal._gemma4_chat_template()).render( + messages=[{"role": "user", "content": "What is 1+1?"}], + bos_token="", + add_generation_prompt=True, + tools=None, + ) + assert rendered == ( + "<|turn>user\nWhat is 1+1?\n" + "<|turn>model\n<|channel>thought\n" + ) + + +def test_gemma4_missing_template_has_actionable_error(monkeypatch, tmp_path): + monkeypatch.setattr(metal, "_GEMMA4_TEMPLATE", None) + monkeypatch.setattr(metal, "__file__", str(tmp_path / "metal.py")) + with pytest.raises(RuntimeError, match="reinstall FreeToken"): + metal._gemma4_chat_template() + + +def test_gemma4_stop_injection_preserves_caller_stops(): + body = json.dumps( + { + "model": "google/gemma-4-26B-A4B-it", + "messages": [{"role": "user", "content": "hi"}], + "stop": "DONE", + "stream": True, + } + ).encode() + rewritten, stream = metal._inject_turn_stop( + "/v1/chat/completions", body, stream=True + ) + assert json.loads(rewritten)["stop"] == ["DONE", ""] + assert stream is True + + # Existing delimiters are not duplicated, and non-Gemma requests are + # byte-for-byte untouched. + already = body.replace(b'"DONE"', b'["DONE", ""]') + assert metal._inject_turn_stop("/v1/responses", already, False) == (already, False) + qwen = body.replace(b"gemma-4-26B-A4B-it", b"Qwen3-0.6B") + assert metal._inject_turn_stop("/v1/chat/completions", qwen, True) == (qwen, True) + + +# ------------------------------------------------------------- proxy routes -- + +def _start_upstream(port: int): + """Serve a tiny OpenAI-ish upstream on a loopback port; return a stop().""" + app = FastAPI() + + @app.get("/v1/models") + async def models(): + return {"object": "list", "data": [{"id": "test-model"}]} + + @app.post("/v1/chat/completions") + async def chat(): + return { + "id": "chatcmpl-x", + "object": "chat.completion", + "model": "test-model", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "hello from upstream"}, + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 3, "total_tokens": 4}, + } + + server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error") + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): # wait for the listener + if server.started: + break + import time + + time.sleep(0.05) + assert server.started, "upstream test server failed to start" + + def stop(): + server.should_exit = True + thread.join(timeout=5) + + return stop + + +def test_proxy_roundtrip_to_real_upstream(): + port = _free_port() + stop = _start_upstream(port) + try: + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url=f"http://127.0.0.1:{port}", + backend="mlx", + ) + handle.load_state = "ready" + assert handle.is_alive() + + proxy = FastAPI() + metal.register_metal_proxy_routes(proxy, lambda: handle) + client = TestClient(proxy, raise_server_exceptions=False) + + r = client.get("/v1/models") + assert r.status_code == 200 + assert [m["id"] for m in r.json()["data"]] == ["test-model"] + + r = client.post( + "/v1/chat/completions", + json={"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 200 + assert r.json()["choices"][0]["message"]["content"] == "hello from upstream" + stats = client.get("/v1/stats").json()["requests"] + assert stats["completed"] == 1 + assert stats["prompt_tokens_total"] == 1 + assert stats["completion_tokens_total"] == 3 + + r = client.get("/health") + assert r.status_code == 200 + assert r.json()["status"] == "ok" + finally: + stop() + + +def test_proxy_503_when_backend_not_alive(): + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: 1)], # exited + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + ) + assert not handle.is_alive() + + proxy = FastAPI() + metal.register_metal_proxy_routes(proxy, lambda: handle) + client = TestClient(proxy, raise_server_exceptions=False) + r = client.post( + "/v1/chat/completions", + json={"model": "x", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 503 + + +def test_proxy_registration_replaces_existing_control_routes(): + app = FastAPI() + + @app.get("/health") + async def native_health(): + return {"status": "native"} + + @app.get("/v1/stats") + async def native_stats(): + return {"backend": "native"} + + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="m", + load_state="ready", + ) + metal.register_metal_proxy_routes(app, lambda: handle) + + paths = [getattr(route, "path", None) for route in app.routes] + assert paths.count("/health") == 1 + assert paths.count("/v1/stats") == 1 + client = TestClient(app, raise_server_exceptions=False) + assert client.get("/health").json()["backend"] == "mlx" + assert client.get("/v1/stats").json()["model"]["id"] == "m" + + +def test_served_model_alias_is_published_and_rewritten_for_upstream(): + handle = metal.MetalBackendHandle( + model_path="/private/models/example", + served_model_name="public-model", + served_model_name_explicit=True, + ) + body = json.dumps({"model": "public-model", "messages": []}).encode() + + rewritten = json.loads(metal._rewrite_public_model(body, handle)) + + assert rewritten["model"] == "/private/models/example" + assert handle.health_doc()["model"] == "public-model" + + +def test_handle_terminate_is_safe_on_empty(): + handle = metal.MetalBackendHandle() + handle.terminate() # no processes -> must not raise + assert not handle.is_alive() + + +# ----------------------------------------------------- wire-format translation -- + +def test_reasoning_field_renamed_streaming_and_body(): + """mlx_lm's thinking channel is `reasoning`; FreeToken's clients (shell, + bench) read `reasoning_content` (the vLLM/SGLang name).""" + chunk = b'data: {"choices": [{"delta": {"reasoning": "think"}}]}' + out = metal._rewrite_reasoning_field(chunk) + assert b'"reasoning_content":' in out and b'"reasoning":' not in out + # value text that merely contains the word stays untouched + chunk2 = b'data: {"choices": [{"delta": {"content": "the reasoning: here"}}]}' + assert metal._rewrite_reasoning_field(chunk2) == chunk2 + + +def test_stream_rewrite_survives_split_key_and_preserves_upstream_error(): + port = _free_port() + upstream = FastAPI() + + @upstream.post("/v1/chat/completions") + async def chat(payload: dict): + if payload.get("prompt") == "reject": + from fastapi.responses import JSONResponse + + return JSONResponse({"error": "no"}, status_code=429) + + async def chunks(): + yield b'data: {"choices":[{"delta":{"reas' + await asyncio.sleep(0.01) + yield b'oning":"think"}}]}\n\n' + + return StreamingResponse(chunks(), media_type="text/event-stream") + + server = uvicorn.Server( + uvicorn.Config(upstream, host="127.0.0.1", port=port, log_level="error") + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): + if server.started: + break + time.sleep(0.05) + + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url=f"http://127.0.0.1:{port}", + backend="mlx", + model_path="m", + load_state="ready", + ) + proxy = FastAPI() + metal.register_metal_proxy_routes(proxy, lambda: handle) + client = TestClient(proxy, raise_server_exceptions=False) + try: + streamed = client.post( + "/v1/chat/completions", + json={"model": "m", "stream": True, "messages": []}, + ) + assert streamed.status_code == 200 + assert b'"reasoning_content":' in streamed.content + assert b'"reasoning":' not in streamed.content + + rejected = client.post( + "/v1/chat/completions", + json={"model": "m", "stream": True, "prompt": "reject"}, + ) + assert rejected.status_code == 429 + assert rejected.json() == {"error": "no"} + finally: + server.should_exit = True + thread.join(timeout=5) + + +def test_health_reports_maintenance_serving(): + """Tools gate on maintenance == "serving" (bench wait_ready, daemon).""" + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="test-model", + ) + handle.load_state = "ready" + handle.load_ended_at = time.monotonic() + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient(app, raise_server_exceptions=False) + doc = client.get("/health").json() + assert doc["status"] == "ok" + assert doc["maintenance"] == "serving" + # ctl parity endpoints exist + requests = client.get("/v1/requests") + assert requests.status_code == 200 + assert requests.json() == {"entries": [], "next_cursor": 0} + stats = client.get("/v1/stats") + assert stats.status_code == 200 + assert stats.json()["model"]["id"] == "test-model" + assert stats.json()["requests"] == { + "active": 0, + "completed": 0, + "p95_ms": 0, + "ttft_mean_ms": 0, + "prompt_tokens_total": 0, + "completion_tokens_total": 0, + } + + # The daemon's legacy stop fallback must accept this exact snapshot. + from freetoken.daemon.serve_manager import ServeManager + + manager = object.__new__(ServeManager) + instance_id, model_id, prompt, completion, uptime = manager._snapshot_values( + stats.json(), model="fallback", require_instance=False + ) + assert instance_id == handle.instance_id + assert (model_id, prompt, completion) == ("test-model", 0, 0) + assert uptime >= 0 + + +# ------------------------------------------------------- load supervision -- + +def test_health_reports_loading_progress(): + """While the engine loads, /health answers loading + byte progress in the + CUDA contract the shell renders (phase + done_bytes/total_bytes).""" + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="big-model", + ) + # Default state is "starting" -> loading without weights total (unknown). + doc = handle.health_doc() + assert doc["status"] == "loading" + assert doc["phase"] == "starting" + assert "progress" not in doc or doc["progress"]["total_bytes"] == 0 + + handle.load_state = "loading" + handle.load_phase = "weights" + handle.weights_bytes = 10 << 30 + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient(app, raise_server_exceptions=False) + doc = client.get("/health").json() + assert doc["status"] == "loading" + assert doc["phase"] == "weights" + assert doc["progress"]["total_bytes"] == 10 << 30 + assert 0 <= doc["progress"]["done_bytes"] <= 10 << 30 + + +def test_health_reports_error_with_reason(): + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="m", + ) + handle._set_state("error", error="model load failed: HTTP timeout") + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient(app, raise_server_exceptions=False) + r = client.get("/health") + assert r.status_code == 503 + assert r.json()["status"] == "error" + assert "model load failed" in r.json()["message"] + + +def test_health_preserves_error_reason_after_backend_exits(): + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: 1)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="m", + ) + handle._set_state("error", error="Metal backend exited during load (code 1)") + metal.register_metal_proxy_routes(app, lambda: handle) + + response = TestClient(app, raise_server_exceptions=False).get("/health") + + assert response.status_code == 503 + assert response.json()["status"] == "error" + assert "exited during load" in response.json()["message"] + + +def test_generation_503_with_phase_while_loading(): + """A generation request during the load must answer immediately with the + loading state, not queue behind the weight load forever.""" + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="m", + ) + handle.load_state = "loading" + handle.load_phase = "weights" + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient(app, raise_server_exceptions=False) + r = client.post( + "/v1/chat/completions", + json={"model": "m", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert r.status_code == 503 + assert "still loading" in r.json()["detail"] + assert r.json()["phase"] == "weights" + # Reads are not gated: /v1/models still proxies (a 502 here only because + # the fixture's upstream address has nothing listening). + assert client.get("/v1/models").status_code == 502 + + +def test_upstream_resident_bytes_of_exited_process(): + """Progress probing must tolerate a dead/missing pid (best-effort).""" + class Dead: + pid = 999999999 # no such process + + assert metal._upstream_resident_bytes([Dead()]) == 0 + + +def test_warm_up_generation_hits_completions(): + port = _free_port() + seen = [] + + def stop(): + pass + + app = FastAPI() + + @app.post("/v1/completions") + async def completions(payload: dict): + seen.append(payload) + return {"choices": [{"text": "x"}]} + + server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error") + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): + if server.started: + break + time.sleep(0.05) + try: + metal._warm_up_generation(f"http://127.0.0.1:{port}", model_id="test-model", timeout=10) + assert seen and seen[0]["max_tokens"] == 1 + finally: + server.should_exit = True + thread.join(timeout=5) + + +def test_llama_launch_returns_ready_supervisor_handle(monkeypatch): + proc = SimpleNamespace(poll=lambda: None) + monkeypatch.setattr(metal, "llama_binary", lambda: "/tmp/llama-server") + monkeypatch.setattr(metal.subprocess, "Popen", lambda *args, **kwargs: proc) + monkeypatch.setattr(metal, "_start_drain_thread", lambda *args: None) + monkeypatch.setattr(metal, "_wait_for_readiness", lambda *args, **kwargs: None) + + handle = metal._launch_llama("/models/example.gguf", 19001) + + assert handle.load_state == "ready" + assert handle.served_model_name == "example.gguf" + assert handle.ack_queue.get_nowait() == "ready" + + +def test_llama_switch_publishes_ready_on_shared_handle(monkeypatch): + proc = SimpleNamespace(poll=lambda: None) + shared = metal.MetalBackendHandle(backend="llama", load_state="loading") + monkeypatch.setattr(metal, "llama_binary", lambda: "/tmp/llama-server") + monkeypatch.setattr(metal.subprocess, "Popen", lambda *args, **kwargs: proc) + monkeypatch.setattr(metal, "_start_drain_thread", lambda *args: None) + monkeypatch.setattr(metal, "_wait_for_readiness", lambda *args, **kwargs: None) + + metal._launch_llama("/models/new.gguf", 19001, state_handle=shared) + + assert shared.load_state == "ready" + assert shared.ack_queue.get_nowait() == "ready" + + +def test_llama_launch_failure_stops_child(monkeypatch): + proc = SimpleNamespace(poll=lambda: None) + stopped = [] + monkeypatch.setattr(metal, "llama_binary", lambda: "/tmp/llama-server") + monkeypatch.setattr(metal.subprocess, "Popen", lambda *args, **kwargs: proc) + monkeypatch.setattr(metal, "_start_drain_thread", lambda *args: None) + monkeypatch.setattr( + metal, + "_wait_for_readiness", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("not ready")), + ) + monkeypatch.setattr(metal, "_stop_processes", lambda processes: stopped.extend(processes)) + + with pytest.raises(RuntimeError, match="not ready"): + metal._launch_llama("/models/example.gguf", 19001) + + assert stopped == [proc] + + +def test_readiness_never_competes_for_process_output(monkeypatch): + class Output: + def readline(self): + raise AssertionError("readiness must not read the drained pipe") + + proc = SimpleNamespace( + poll=lambda: None, + stdout=Output(), + stderr=None, + returncode=None, + _freetoken_recent_output=[], + ) + monkeypatch.setattr( + metal.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + metal._wait_for_readiness("http://127.0.0.1:19001", proc, timeout=1) + + +def test_stream_signalled_in_body_not_accept_header(): + """The OpenAI SDK sends Accept: application/json even for streaming + requests; streaming is signalled in the body. Detecting from the header + alone routed streamed generations through the buffered path -- the client + got the whole answer in one burst and live tok/s read 0.""" + import httpx as _httpx + + port = _free_port() + first_chunk_at: list[float] = [] + app = FastAPI() + + @app.post("/v1/chat/completions") + async def chat(payload: dict): + async def gen(): + # The upstream writes + flushes each SSE event (mlx_lm does this + # per token); a buffered proxy delivers these as one burst. + for i in range(3): + yield ( + f'data: {{"choices":[{{"delta":{{"content":"{i}"}}}}]}}\n\n' + ).encode() + await asyncio.sleep(0.15) + + return StreamingResponse(gen(), media_type="text/event-stream") + + server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error") + ) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + for _ in range(100): + if server.started: + break + time.sleep(0.05) + try: + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url=f"http://127.0.0.1:{port}", + backend="mlx", + model_path="m", + ) + handle.load_state = "ready" + proxy_app = FastAPI() + metal.register_metal_proxy_routes(proxy_app, lambda: handle) + + # Real uvicorn + raw socket: TestClient buffers the whole response + # before iteration, which would make any proxy look burst-shaped. + proxy_port = _free_port() + proxy_server = uvicorn.Server( + uvicorn.Config(proxy_app, host="127.0.0.1", port=proxy_port, log_level="error") + ) + threading.Thread(target=proxy_server.run, daemon=True).start() + for _ in range(100): + if proxy_server.started: + break + time.sleep(0.05) + + # No Accept: text/event-stream -- exactly what the OpenAI SDK sends. + body = json.dumps( + { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + } + ).encode() + with socket.create_connection(("127.0.0.1", proxy_port), timeout=5) as s: + s.sendall( + b"POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\n" + b"Accept: application/json\r\nContent-Type: application/json\r\n" + b"Content-Length: " + + str(len(body)).encode() + + b"\r\n\r\n" + + body + ) + t0 = time.time() + stamps: list[float] = [] + buf = b"" + while len(stamps) < 3 and time.time() - t0 < 5: + data = s.recv(4096) + if not data: + break + buf += data + while b"\n\n" in buf: + chunk, buf = buf.split(b"\n\n", 1) + if b"data:" in chunk: + stamps.append(time.time() - t0) + assert len(stamps) >= 3, f"expected 3 SSE events, got {len(stamps)}" + # A buffered proxy delivers them as one burst (spread ~0); progressive + # delivery spreads >= ~0.2s for three events 150ms apart. + spread = stamps[-1] - stamps[0] + assert spread >= 0.2, f"chunks delivered as one burst (spread {spread:.3f}s)" + finally: + server.should_exit = True + thread.join(timeout=5) + proxy_server.should_exit = True + + +# --------------------------------------------------------- model switching -- + +def test_model_load_rejects_non_loopback_client(): + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + load_state="ready", + ) + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient( + app, raise_server_exceptions=False, client=("203.0.113.9", 50000) + ) + + response = client.post("/v1/model/load", json={"model": "new-model"}) + + assert response.status_code == 403 + assert handle.model_path == "old-model" + + +def test_model_load_requires_model_field(): + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + ) + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient( + app, raise_server_exceptions=False, client=("127.0.0.1", 50000) + ) + + # missing field + r = client.post("/v1/model/load", json={"nomodel": "x"}) + assert r.status_code == 400 + # bad JSON body + r = client.post("/v1/model/load", content=b"not json", headers={"content-type": "application/json"}) + assert r.status_code == 400 + + +def test_model_load_same_model_is_a_noop(): + app = FastAPI() + handle = metal.MetalBackendHandle( + processes=[SimpleNamespace(poll=lambda: None)], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + ) + metal.register_metal_proxy_routes(app, lambda: handle) + client = TestClient( + app, raise_server_exceptions=False, client=("127.0.0.1", 50000) + ) + r = client.post("/v1/model/load", json={"model": "old-model"}) + assert r.status_code == 200 + assert r.json()["detail"] == "already serving this model" + + +def test_model_load_switches_to_new_handle(monkeypatch): + """A switch stops the old engine BEFORE starting the new one, and adopts + the new engine's identity on the shared handle. + + Sequential (stop, then start) is the point: two concurrent engines + over-commit the Metal working set and deadlock on this hardware.""" + app = FastAPI() + old_proc = SimpleNamespace(poll=lambda: None, terminate=lambda: None) + events: list[str] = [] + old = metal.MetalBackendHandle( + processes=[old_proc], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + ) + original_processes = old.processes + old.load_state = "ready" + new_proc = SimpleNamespace(poll=lambda: None, terminate=lambda: None) + + real_stop = metal._stop_processes + + def seen_stop(processes): + # switch_model captures the old list BEFORE clearing the attribute, + # so identity must be checked against the pre-switch list. + events.append(f"stopped:{processes is original_processes}") + real_stop(processes) + + def fake_launch(backend, model, upstream_port=None, state_handle=None): + events.append("launch") + h = metal.MetalBackendHandle( + processes=[new_proc], + upstream_base_url="http://127.0.0.1:2", + backend="mlx", + model_path="new-model", + ) + h.load_state = "ready" + # A switch passes state_handle: the watcher publishes onto the SHARED + # handle, so simulate that publication. + if state_handle is not None: + state_handle.load_state = "ready" + return h + + monkeypatch.setattr(metal, "launch_metal_backend", fake_launch) + monkeypatch.setattr(metal, "_stop_processes", seen_stop) + metal.register_metal_proxy_routes(app, lambda: old) + client = TestClient( + app, raise_server_exceptions=False, client=("127.0.0.1", 50000) + ) + + r = client.post("/v1/model/load", json={"model": "new-model"}) + assert r.status_code == 200 + assert r.json()["model"] == "new-model" + # The old engine died before the new load began -- never both resident. + assert events == ["stopped:True", "launch"] + assert old.model_path == "new-model" + assert old.upstream_base_url == "http://127.0.0.1:2" + assert len(old.processes) == 1 and old.processes[0] is new_proc + assert old.load_state == "ready" + + +def test_model_switch_tracks_new_engine_while_waiting(monkeypatch): + old_proc = SimpleNamespace(poll=lambda: None) + new_proc = SimpleNamespace(poll=lambda: None) + old = metal.MetalBackendHandle( + processes=[old_proc], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + load_state="ready", + ) + stopped: list[list[object]] = [] + observed: dict[str, object] = {} + + def fake_launch(backend, model, upstream_port=None, state_handle=None): + assert (backend, model, state_handle) == ("mlx", "new-model", old) + return metal.MetalBackendHandle( + processes=[new_proc], + upstream_base_url="http://127.0.0.1:2", + backend="mlx", + model_path="new-model", + ) + + def fake_wait(timeout=None): + observed["processes"] = old.processes + observed["upstream"] = old.upstream_base_url + observed["alive"] = old.is_alive() + old._set_state("ready") + + monkeypatch.setattr(metal, "launch_metal_backend", fake_launch) + monkeypatch.setattr(metal, "_stop_processes", lambda procs: stopped.append(procs)) + monkeypatch.setattr(old, "_wait_load_terminal", fake_wait) + + old.switch_model("new-model") + + assert stopped == [[old_proc]] + assert observed == { + "processes": [new_proc], + "upstream": "http://127.0.0.1:2", + "alive": True, + } + + +def test_model_switch_failure_stops_new_engine(monkeypatch): + old_proc = SimpleNamespace(poll=lambda: None) + new_proc = SimpleNamespace(poll=lambda: None) + old = metal.MetalBackendHandle( + processes=[old_proc], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + load_state="ready", + ) + stopped: list[list[object]] = [] + + def fake_launch(backend, model, upstream_port=None, state_handle=None): + return metal.MetalBackendHandle( + processes=[new_proc], + upstream_base_url="http://127.0.0.1:2", + backend="mlx", + model_path="new-model", + ) + + def fake_wait(timeout=None): + raise RuntimeError("warm-up failed") + + monkeypatch.setattr(metal, "launch_metal_backend", fake_launch) + monkeypatch.setattr(metal, "_stop_processes", lambda procs: stopped.append(procs)) + monkeypatch.setattr(old, "_wait_load_terminal", fake_wait) + + with pytest.raises(RuntimeError, match="warm-up failed"): + old.switch_model("new-model") + + assert stopped == [[old_proc], [new_proc]] + assert old.processes == [] + assert old.load_state == "error" + assert old.load_error == "warm-up failed" + + +def test_model_load_launch_failure_reports_error(monkeypatch): + """A failed launch reports the error; the old engine is already gone + (sequential switch), so the honest state is the error itself.""" + app = FastAPI() + old_proc = SimpleNamespace(poll=lambda: None, terminate=lambda: None) + old = metal.MetalBackendHandle( + processes=[old_proc], + upstream_base_url="http://127.0.0.1:1", + backend="mlx", + model_path="old-model", + ) + + def fake_launch(backend, model, upstream_port=None, state_handle=None): + raise RuntimeError("download failed") + + monkeypatch.setattr(metal, "launch_metal_backend", fake_launch) + metal.register_metal_proxy_routes(app, lambda: old) + client = TestClient( + app, raise_server_exceptions=False, client=("127.0.0.1", 50000) + ) + + r = client.post("/v1/model/load", json={"model": "new-model"}) + assert r.status_code == 500 + assert "download failed" in r.json()["detail"] + assert old.load_state == "error" + assert old.load_error == "download failed" + health = client.get("/health") + assert health.status_code == 503 + assert health.json()["status"] == "error" + assert health.json()["message"] == "download failed" diff --git a/tests/server/test_process_utils.py b/tests/server/test_process_utils.py new file mode 100644 index 0000000..60a26c1 --- /dev/null +++ b/tests/server/test_process_utils.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import subprocess + +from freetoken.server.process_utils import ( + reap_backend_workers, + terminate_backend_workers, +) + + +class FakePopen: + def __init__(self, *, exits_on_terminate: bool) -> None: + self.running = True + self.exits_on_terminate = exits_on_terminate + self.terminated = False + self.killed = False + + def poll(self): + return None if self.running else 0 + + def terminate(self): + self.terminated = True + if self.exits_on_terminate: + self.running = False + + def wait(self, timeout=None): + if self.running: + raise subprocess.TimeoutExpired("fake", timeout) + return 0 + + def kill(self): + self.killed = True + self.running = False + + +def test_popen_workers_are_terminated_and_reaped(): + graceful = FakePopen(exits_on_terminate=True) + stubborn = FakePopen(exits_on_terminate=False) + + terminate_backend_workers([graceful, stubborn]) + reap_backend_workers([graceful, stubborn], timeout=0) + + assert graceful.terminated and not graceful.killed + assert stubborn.terminated and stubborn.killed diff --git a/tests/server/test_serve_macos.py b/tests/server/test_serve_macos.py new file mode 100644 index 0000000..2f6f735 --- /dev/null +++ b/tests/server/test_serve_macos.py @@ -0,0 +1,130 @@ +"""macOS serve path: ``ft serve`` must not import the CUDA/torch stack. + +The native launcher (``freetoken.server.launch``) imports torch at module +level. On Apple Silicon there is no CUDA torch wheel in the Metal venv, so +``ft serve`` has to route to ``serve-metal`` *before* that import. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest + + +def test_run_serve_on_darwin_routes_to_metal_without_cuda_launcher(monkeypatch): + """``ft serve`` on Darwin calls metal_main and never launch_server.""" + import freetoken.cli as cli + + monkeypatch.setattr(cli.sys, "platform", "darwin") + seen: dict = {} + + def fake_metal(argv): + seen["argv"] = list(argv) + return 0 + + monkeypatch.setattr("freetoken.server.metal_main.main", fake_metal) + + rc = cli._run_serve(["--model", "mlx-community/Qwen3-0.6B-4bit", "--port", "1919"]) + assert rc == 0 + assert seen["argv"] == ["--model", "mlx-community/Qwen3-0.6B-4bit", "--port", "1919"] + + +def test_metal_parser_accepts_shared_shell_and_model_name_flags(): + from freetoken.server.metal_main import _parse + + args = _parse( + [ + "--model", + "/models/example", + "--shell-mode", + "--served-model-name", + "public-model", + "--cors-origins", + "http://localhost:3000", + ] + ) + + assert args.shell is True + assert args.served_model_name == "public-model" + assert args.cors_origins == "http://localhost:3000" + + +def test_shell_metal_filter_does_not_consume_backend_after_boolean_flag(): + from freetoken.shell import _split_engine_args + + model, passthrough = _split_engine_args( + ["--model", "M", "--moe-cache-auto", "--backend", "llama"] + ) + + assert model == "M" + assert passthrough == ["--backend", "llama"] + + +def test_metal_cors_uses_same_browser_allow_list(): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from freetoken.server.cors import install_cors + + app = FastAPI() + + @app.get("/health") + async def health(): + return {"status": "ok"} + + install_cors(app, "http://localhost:1420") + response = TestClient(app).options( + "/health", + headers={ + "Origin": "http://localhost:1420", + "Access-Control-Request-Method": "GET", + }, + ) + + assert response.headers["access-control-allow-origin"] == "http://localhost:1420" + + +def test_serve_metal_help_uses_lightweight_import_path(): + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + pkg = os.path.join(root, "python") + env = dict(os.environ) + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = pkg + (os.pathsep + existing if existing else "") + + proc = subprocess.run( + [sys.executable, "-m", "freetoken.cli", "serve-metal", "--help"], + cwd=root, + env=env, + capture_output=True, + text=True, + ) + + assert proc.returncode == 0 + assert "PyTorch was not found" not in proc.stderr + assert "--served-model-name" in proc.stdout + + +@pytest.mark.skipif(sys.platform != "darwin", reason="Metal venv has no torch") +def test_ft_serve_help_does_not_need_torch(): + """``ft serve --help`` must succeed on a Metal install (no torch installed).""" + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + pkg = os.path.join(root, "python") + env = dict(os.environ) + existing = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = pkg + (os.pathsep + existing if existing else "") + proc = subprocess.run( + [sys.executable, "-m", "freetoken.cli", "serve", "--help"], + cwd=root, + env=env, + capture_output=True, + text=True, + ) + assert proc.returncode == 0, ( + f"ft serve --help failed (likely imported torch)\n" + f"--- stdout ---\n{proc.stdout}\n--- stderr ---\n{proc.stderr}" + ) + assert "No module named 'torch'" not in proc.stderr + assert "--model" in proc.stdout diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90..77b9fb8 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -73,6 +73,29 @@ def is_alive(self) -> bool: drain_ready(handle, LoadProgress(), get=lambda _t: (_ for _ in ()).throw(_Empty())) +def test_drain_ready_detects_popen_death_during_load(): + import pytest + + from freetoken.server.supervisor import WorkerDied + + class DeadPopen: + name = "llama-server" + + def poll(self): + return 1 + + handle = BackendHandle( + ack_queue=queue.Queue(), processes=[DeadPopen()], expected_acks=1 + ) + + with pytest.raises(WorkerDied, match="llama-server"): + drain_ready( + handle, + LoadProgress(), + get=lambda _t: (_ for _ in ()).throw(_Empty()), + ) + + def test_drain_ready_raises_the_real_reason_from_an_error_ack(): """A worker that pushes ("error", reason) just before dying surfaces THAT reason (e.g. a config ValueError), not the generic "exited during load".""" diff --git a/tests/test_logger.py b/tests/test_logger.py new file mode 100644 index 0000000..1ec9f98 --- /dev/null +++ b/tests/test_logger.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import sys +from types import ModuleType +from types import SimpleNamespace + + +def test_logger_formatter_preserves_tensor_parallel_rank(monkeypatch, capsys): + distributed = ModuleType("freetoken.distributed") + distributed.try_get_tp_info = lambda: SimpleNamespace( + rank=2, is_primary=lambda: False + ) + monkeypatch.setitem(sys.modules, "freetoken.distributed", distributed) + from freetoken.logging import init_logger + logger = init_logger("tests.logger.rank", use_tp_rank=True) + + logger.info("ranked") + + assert "|core|rank=2]" in capsys.readouterr().out diff --git a/tests/test_shell_client.py b/tests/test_shell_client.py new file mode 100644 index 0000000..4f00dba --- /dev/null +++ b/tests/test_shell_client.py @@ -0,0 +1,53 @@ +"""Focused regressions for shell control-plane behavior.""" + +from __future__ import annotations + +import asyncio + +from freetoken.shell import client as shell_client + + +def test_load_model_reports_health_progress_while_post_is_pending(monkeypatch): + client = object.__new__(shell_client.ShellClient) + release_post: asyncio.Event + health_docs = [ + {"status": "loading", "phase": "starting"}, + { + "status": "loading", + "phase": "weights", + "progress": {"done_bytes": 5, "total_bytes": 10}, + }, + ] + progress: list[dict] = [] + + async def run(): + nonlocal release_post + release_post = asyncio.Event() + + async def request_json(method, path, *, body=None, timeout=None): + assert (method, path, body, timeout) == ( + "POST", + "/v1/model/load", + {"model": "new-model"}, + 60.0, + ) + await release_post.wait() + return {"status": "ok", "model": "new-model"} + + async def health(): + doc = health_docs.pop(0) + if not health_docs: + release_post.set() + return doc + + client._request_json = request_json + client.health = health + return await client.load_model( + "new-model", wait=60.0, on_progress=progress.append + ) + + monkeypatch.setattr(shell_client, "READY_POLL_INTERVAL", 0) + result = asyncio.run(run()) + + assert result == {"status": "ok", "model": "new-model"} + assert [doc["phase"] for doc in progress] == ["starting", "weights"] diff --git a/tests/test_shell_tui.py b/tests/test_shell_tui.py new file mode 100644 index 0000000..c9fed7b --- /dev/null +++ b/tests/test_shell_tui.py @@ -0,0 +1,52 @@ +"""Focused regressions for the interactive shell command loop.""" + +from __future__ import annotations + +import asyncio + +from freetoken.shell import tui + + +class _FakeClient: + async def wait_until_ready(self, **_kwargs): + return {"status": "ok"} + + async def model_id(self): + return "google/gemma-4-26B-A4B-it" + + async def cache_status(self): + return { + "geometry": { + "reasoning": { + "gears": ["low", "high"], + "kwargs": {}, + "default": "low", + } + } + } + + async def stats(self): + return {} + + +class _FakePromptSession: + def __init__(self, *_args, **_kwargs): + self._commands = ["/think status"] + + async def prompt_async(self): + if self._commands: + return self._commands.pop(0) + raise EOFError + + +def test_think_command_still_uses_gears_when_model_switch_can_refresh_them(monkeypatch): + output: list[str] = [] + monkeypatch.setattr(tui, "PromptSession", _FakePromptSession) + monkeypatch.setattr(tui.ShellConsoleRenderer, "_write_stdout", output.append) + + result = asyncio.run( + tui._run_shell(_FakeClient(), "http://test", connect_grace=0.0) + ) + + assert result == 0 + assert any("Thinking: low (available: low, high)" in line for line in output)