#!/usr/bin/env bash
# Purpose: start the mlx-lm server on a local MLX model directory, so that the same measuring
#          scripts can be pointed at Apple's native engine and at llama-server without
#          changing anything but the port and the model name
# Platform: mac (mlx-lm runs on Apple silicon)
# Minimum memory: 12 GB
# Assumes: mlx-lm installed (Part 8's MLX lesson), an MLX model directory in $MODEL whose
#          config.json carries a "quantization" block, the port free and nothing else large
#          running; the server runs in the foreground so that Ctrl-C stops it
#
# Usage: MODEL=$HOME/models/mlx-community/Qwen3-8B-4bit bash serve-mlx.sh
#
# Environment:
#   MODEL              absolute path of the MLX model directory   (required)
#   HOST               address to bind                            (default: 127.0.0.1)
#   PORT               port to listen on                          (default: 8081)
#   PROMPT_CACHE_SIZE  distinct prompt KV caches the server keeps (default: 1)
#   LOGLEVEL           server log level                           (default: INFO)
#   PIDFILE            where to write the server's pid            (default: ./mlx-server.pid)
#
# Two things about mlx_lm.server 0.31.3, read from its source, shape this script:
#   * A request's "model" field is resolved as a model to load. Send the exact string passed
#     as --model (the absolute path this script prints), or the server tries to load
#     whatever name you sent.
#   * The server keeps up to --prompt-cache-size distinct KV caches from earlier requests
#     (default 10). This lab sends a different prompt every time, so at the default the
#     memory reading would include several stale caches. 1 keeps it comparable with a
#     single llama-server slot.
# The server's own documentation says it is not recommended for production because it only
# implements basic security checks, so it stays bound to the loopback address.

set -euo pipefail

MODEL="${MODEL:-}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8081}"
PROMPT_CACHE_SIZE="${PROMPT_CACHE_SIZE:-1}"
LOGLEVEL="${LOGLEVEL:-INFO}"
PIDFILE="${PIDFILE:-./mlx-server.pid}"

die() { echo "serve-mlx: $*" >&2; exit 1; }

[ "$(uname -s)" = "Darwin" ] || die "mlx-lm runs on Apple silicon; this machine reports $(uname -s)"
[ -n "$MODEL" ] || die "set MODEL to the absolute path of an MLX model directory"
case "$MODEL" in /*) ;; *) die "MODEL must be an absolute path (it is also the name requests send)" ;; esac
[ -f "$MODEL/config.json" ] || die "$MODEL/config.json not found"
grep -q '"quantization"' "$MODEL/config.json" || echo "serve-mlx: warning: no quantization block in config.json" >&2
command -v mlx_lm.server >/dev/null || die "mlx_lm.server not found; install mlx-lm as the MLX lesson describes"
if command -v curl >/dev/null && curl --silent --max-time 2 "http://$HOST:$PORT/health" >/dev/null 2>&1; then
  die "something is already answering on $HOST:$PORT; stop it first (one server at a time)"
fi

# The interpreter that runs mlx_lm.server is the one whose packages it uses.
PY="$(head -n 1 "$(command -v mlx_lm.server)" | sed 's/^#!//')"

echo "==> mlx_lm.server"
"$PY" -c 'import mlx_lm, mlx.core as mx; print("    mlx-lm", mlx_lm.__version__, "/ mlx", mx.__version__)' \
  || echo "    (could not read the mlx-lm version from $PY)"
echo "    model   $MODEL"
echo "    quant   $(grep -A3 '"quantization"' "$MODEL/config.json" | tr -d ' \n' | cut -c1-60)"
echo "    listen  http://$HOST:$PORT/v1"
echo "    prompt cache size $PROMPT_CACHE_SIZE"
echo "    pid written to $PIDFILE"
echo "    send \"model\": \"$MODEL\" in every request"

echo $$ > "$PIDFILE"

# exec keeps this shell's pid, so $PIDFILE names the server process itself.
exec mlx_lm.server \
    --model "$MODEL" \
    --host "$HOST" \
    --port "$PORT" \
    --prompt-cache-size "$PROMPT_CACHE_SIZE" \
    --log-level "$LOGLEVEL"