Files
mlx-knife/scripts/build.sh
T
Nikesh Parajuli 6ea9bd33f5 patches
2026-02-24 22:19:59 -05:00

552 lines
17 KiB
Bash
Executable File

#!/usr/bin/env bash
set -euo pipefail
# Build and bundle a macOS arm64 (Apple Silicon) CLI tarball using PyInstaller.
# Invoked as scripts/build.sh
#
# Outputs: artifacts/<name>.tgz containing dist/<name>/...
# Optional env:
# ONEFILE=1 -> build single-file binary (defaults to onedir for reliability)
# VERSION=1.1 -> append version to tarball name
# BIN_NAME=my-cli -> override binary/folder name (default: msty-mlx-studio)
# USE_HF_XET=0 -> skip installing Hugging Face Xet plugin
#
# Requirements (run on an Apple Silicon Mac):
# - Xcode Command Line Tools
# - Python 3.9+ available as `python3`
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"
# -----------------------------------------------------------------------------
# Temporary version override (VERSION env)
# -----------------------------------------------------------------------------
ORIG_VERSION="$(python3 - <<'PY'
import pathlib, re
text = pathlib.Path("mlxk2/__init__.py").read_text()
m = re.search(r'__version__\s*=\s*"([^"]+)"', text)
print(m.group(1) if m else "0.0.0")
PY
)"
TARGET_VERSION="${VERSION:-$ORIG_VERSION}"
TMP_INIT=""
TMP_README=""
RESTORE_VERSION=0
APPLIED_PATCHES=()
cleanup_version() {
if [[ "$RESTORE_VERSION" == "1" ]]; then
if [[ -n "$TMP_INIT" && -f "$TMP_INIT" ]]; then
cp "$TMP_INIT" mlxk2/__init__.py
rm -f "$TMP_INIT"
fi
if [[ -n "$TMP_README" && -f "$TMP_README" ]]; then
cp "$TMP_README" README.md
rm -f "$TMP_README"
fi
fi
}
cleanup_patches() {
if [[ "${#APPLIED_PATCHES[@]}" -gt 0 ]]; then
for pf in "${APPLIED_PATCHES[@]}"; do
patch -l -p1 -R < "$pf" >/dev/null 2>&1 || true
done
fi
}
cleanup_all() {
cleanup_version
cleanup_patches
}
trap cleanup_all EXIT
if [[ "$TARGET_VERSION" != "$ORIG_VERSION" ]]; then
echo "[version] Updating project version: $ORIG_VERSION -> $TARGET_VERSION"
TMP_INIT="$(mktemp)"
cp mlxk2/__init__.py "$TMP_INIT"
if [[ -f README.md ]]; then
TMP_README="$(mktemp)"
cp README.md "$TMP_README"
fi
RESTORE_VERSION=1
python3 - "$ORIG_VERSION" "$TARGET_VERSION" <<'PY'
import sys
import pathlib
import re
old, new = sys.argv[1:3]
def tag(ver: str) -> str:
if 'b' in ver:
base, beta = ver.split('b', 1)
if beta and beta.isdigit():
return f"v{base}-beta.{beta}"
return f"v{ver}"
init_path = pathlib.Path("mlxk2/__init__.py")
text = init_path.read_text()
text = re.sub(r'__version__\s*=\s*".*?"', f'__version__ = "{new}"', text)
init_path.write_text(text)
readme_path = pathlib.Path("README.md")
if readme_path.exists():
readme = readme_path.read_text()
readme = readme.replace(old, new)
readme = readme.replace(tag(old), tag(new))
readme_path.write_text(readme)
PY
fi
# -----------------------------------------------------------------------------
# Apply local CLI patch (idempotent)
# -----------------------------------------------------------------------------
apply_patch_with_cleanup() {
local patch_file="$1"
[[ -f "$patch_file" ]] || return
if patch --dry-run -l -p1 < "$patch_file" >/dev/null 2>&1; then
echo "[patch] Applying patch from $patch_file"
patch -l -p1 < "$patch_file"
APPLIED_PATCHES+=("$patch_file")
elif patch --dry-run -R -l -p1 < "$patch_file" >/dev/null 2>&1; then
echo "[patch] Patch already applied; skipping ($patch_file)"
else
echo "ERROR: Failed to apply patch: $patch_file" >&2
exit 1
fi
}
apply_patch_with_cleanup "$ROOT_DIR/scripts/patches/cli.patch"
apply_patch_with_cleanup "$ROOT_DIR/scripts/patches/server_models.patch"
apply_patch_with_cleanup "$ROOT_DIR/scripts/patches/server_streaming_usage.patch"
apply_patch_with_cleanup "$ROOT_DIR/scripts/patches/runner_decode.patch"
# -----------------------------------------------------------------------------
# CLI flags (custom wheels)
# -----------------------------------------------------------------------------
while [[ $# -gt 0 ]]; do
case "$1" in
--mlx-wheel)
shift
if [[ -n "${1:-}" ]]; then
export MLX_WHEEL="${1:-}"
shift || true
else
echo "ERROR: --mlx-wheel requires a path" >&2
exit 1
fi
;;
--mlx-lm-wheel)
shift
if [[ -n "${1:-}" ]]; then
export MLX_LM_WHEEL="${1:-}"
shift || true
else
echo "ERROR: --mlx-lm-wheel requires a path" >&2
exit 1
fi
;;
*)
echo "Unknown flag: $1" >&2
exit 1
;;
esac
done
# -----------------------------------------------------------------------------
# Platform checks
# -----------------------------------------------------------------------------
if [[ "$(uname -s)" != "Darwin" ]]; then
echo "This script must run on macOS." >&2
exit 1
fi
if [[ "$(uname -m)" != "arm64" ]]; then
echo "This script targets Apple Silicon (arm64) only." >&2
exit 1
fi
APP_NAME="${BIN_NAME:-${APP_NAME:-msty-mlx-studio}}"
BUILD_VENV=".venv-build"
DIST_DIR="$ROOT_DIR/dist"
ART_DIR="$ROOT_DIR/artifacts"
SPEC_DIR="$ROOT_DIR/.pyi-spec"
mkdir -p "$ART_DIR" "$SPEC_DIR"
# -----------------------------------------------------------------------------
# Create build virtualenv
# -----------------------------------------------------------------------------
echo "[1/3] Creating build venv..."
CUSTOM_PY_DIR="$ROOT_DIR/python"
if [[ -n "${PYTHON_EXE:-}" ]]; then
echo "Using override Python (PYTHON_EXE): $PYTHON_EXE"
elif [[ -x "$CUSTOM_PY_DIR/bin/python3" ]]; then
PYTHON_EXE="$CUSTOM_PY_DIR/bin/python3"
echo "Using project Python: $PYTHON_EXE"
else
PYTHON_EXE="$(command -v python3)"
echo "Using system Python: $PYTHON_EXE"
fi
"$PYTHON_EXE" -m venv "$BUILD_VENV"
source "$BUILD_VENV/bin/activate"
if ! python -m pip --version >/dev/null 2>&1; then
echo "Bootstrapping pip via ensurepip..."
python -m ensurepip --upgrade || true
fi
python -m pip install --upgrade pip setuptools wheel
# -----------------------------------------------------------------------------
# Install dependencies and project
# -----------------------------------------------------------------------------
echo "[2/3] Installing project, deps and build tools..."
CUSTOM_MLX_WHEEL="${MLX_WHEEL:-}"
if [[ -z "$CUSTOM_MLX_WHEEL" && -f "$ROOT_DIR/scripts/mlx-custom.whl" ]]; then
CUSTOM_MLX_WHEEL="$ROOT_DIR/scripts/mlx-custom.whl"
fi
if [[ -n "$CUSTOM_MLX_WHEEL" ]]; then
echo "[opt] Using custom MLX wheel: $CUSTOM_MLX_WHEEL"
[[ -f "$CUSTOM_MLX_WHEEL" ]] || { echo "ERROR: MLX wheel not found: $CUSTOM_MLX_WHEEL" >&2; exit 1; }
fi
CUSTOM_MLX_LM_WHEEL="${MLX_LM_WHEEL:-}"
if [[ -z "$CUSTOM_MLX_LM_WHEEL" && -f "$ROOT_DIR/scripts/mlx-lm-custom.whl" ]]; then
CUSTOM_MLX_LM_WHEEL="$ROOT_DIR/scripts/mlx-lm-custom.whl"
fi
if [[ -n "$CUSTOM_MLX_LM_WHEEL" ]]; then
echo "[opt] Using custom MLX-LM wheel: $CUSTOM_MLX_LM_WHEEL"
[[ -f "$CUSTOM_MLX_LM_WHEEL" ]] || { echo "ERROR: MLX-LM wheel not found: $CUSTOM_MLX_LM_WHEEL" >&2; exit 1; }
fi
if [[ -f requirements.txt ]]; then
if [[ -n "$CUSTOM_MLX_WHEEL" || -n "$CUSTOM_MLX_LM_WHEEL" ]]; then
REQ_TMP="$(mktemp)"
awk 'BEGIN{IGNORECASE=0} {
line=$0; trimmed=line; gsub(/^[ \t]+/,"",trimmed);
if (trimmed ~ /^#/ || trimmed ~ /^$/) { print line; next }
if (ENVIRON["FILTER_MLX"] && trimmed ~ /^mlx([ \t]|[><=]|$)/) { next }
if (ENVIRON["FILTER_MLX_LM"] && trimmed ~ /^mlx-lm([ \t]|[><=]|$)/) { next }
print line
}' FILTER_MLX=$([[ -n "$CUSTOM_MLX_WHEEL" ]] && echo 1) FILTER_MLX_LM=$([[ -n "$CUSTOM_MLX_LM_WHEEL" ]] && echo 1) requirements.txt > "$REQ_TMP"
python -m pip install ${PIP_ARGS:-} -r "$REQ_TMP"
rm -f "$REQ_TMP"
else
python -m pip install ${PIP_ARGS:-} -r requirements.txt
fi
fi
if [[ -n "${MLX_REPO_REF:-}" ]]; then
echo "[opt] Overriding MLX from: ${MLX_REPO_REF}"
python -m pip install ${PIP_ARGS:-} --no-deps "${MLX_REPO_REF}"
fi
if [[ -n "${MLX_LM_REPO_REF:-}" ]]; then
echo "[opt] Overriding MLX-LM from: ${MLX_LM_REPO_REF}"
python -m pip install ${PIP_ARGS:-} --no-deps "${MLX_LM_REPO_REF}"
fi
python -m pip install ${PIP_ARGS:-} .
if [[ "${USE_HF_XET:-1}" == "1" ]]; then
echo "[opt] Installing Hugging Face Xet plugin..."
python -m pip install "huggingface_hub[hf_xet]" || true
fi
if [[ -n "$CUSTOM_MLX_WHEEL" ]]; then
echo "[opt] Installing custom MLX wheel: $CUSTOM_MLX_WHEEL"
python -m pip install ${PIP_ARGS:-} --force-reinstall "$CUSTOM_MLX_WHEEL"
fi
if [[ -n "$CUSTOM_MLX_LM_WHEEL" ]]; then
echo "[opt] Installing custom MLX-LM wheel: $CUSTOM_MLX_LM_WHEEL"
python -m pip install ${PIP_ARGS:-} --force-reinstall "$CUSTOM_MLX_LM_WHEEL"
fi
python -m pip install "pyinstaller>=6.6"
# -----------------------------------------------------------------------------
# Entry script for PyInstaller
# -----------------------------------------------------------------------------
ENTRY_SCRIPT="$ROOT_DIR/.entry_mlxk_build.py"
cat > "$ENTRY_SCRIPT" <<'PY'
"""
Custom entry shim for PyInstaller bundle.
Implements supervised shutdown by default in a frozen binary:
- Parent process patches mlxk2.operations.serve.start_server to spawn a child
process (the same frozen binary) that runs the server in-process.
- Parent supervises: SIGTERM on first Ctrl-C, SIGKILL on timeout/second Ctrl-C.
Also supports a child mode keyed by MLXK2_CHILD_SERVER=1 to run the server
directly without going through CLI argument parsing.
"""
import os
import json
import signal
import subprocess
import sys
import time
def _to_bool(val: str) -> bool:
return str(val).lower() in ("1", "true", "yes", "on")
# Fast path: strip CLI name from version output for compatibility
if "--version" in sys.argv and "--json" not in sys.argv and "--python-info" not in sys.argv:
try:
from mlxk2 import __version__ as _ver
except Exception:
_ver = "0.0.0"
print(_ver)
sys.exit(0)
# Diagnostic: print embedded Python details
if "--python-info" in sys.argv or os.environ.get("MLXK2_PY_INFO") == "1":
info = {
"python_version": sys.version.split(" (", 1)[0],
"executable": sys.executable,
"prefix": sys.prefix,
"base_prefix": getattr(sys, "base_prefix", None),
"frozen": bool(getattr(sys, "frozen", False)),
"meipass": getattr(sys, "_MEIPASS", None),
"platform": sys.platform,
}
print(json.dumps(info, indent=2))
sys.exit(0)
# Diagnostic: package info helpers
if "--mlx-info" in sys.argv or "--pkg-info" in sys.argv:
try:
if "--mlx-info" in sys.argv:
pkg = "mlx"
else:
idx = sys.argv.index("--pkg-info")
pkg = sys.argv[idx + 1]
import importlib
mod = importlib.import_module(pkg)
version = None
try:
import importlib.metadata as md
try:
version = md.version(pkg)
except Exception:
version = None
except Exception:
pass
if version is None:
version = getattr(mod, "__version__", None)
dist_info_dir = None
try:
import importlib.metadata as md
dist = md.distribution(pkg)
files = list(dist.files or [])
for f in files:
if str(f).endswith("METADATA"):
dist_info_dir = os.fspath(dist.locate_file(f))
dist_info_dir = os.path.dirname(dist_info_dir)
break
except Exception:
dist_info_dir = None
out = {
"package": pkg,
"version": version,
"module_file": getattr(mod, "__file__", None),
"dist_info": dist_info_dir,
}
print(json.dumps(out, indent=2))
sys.exit(0)
except Exception as e:
print(json.dumps({"error": str(e)}))
sys.exit(2)
# Diagnostic: verify MLX stack imports cleanly
if "--check-mlx-stack" in sys.argv:
out = {"ok": False, "mlx": None, "mlx_lm": None, "error": None}
try:
import importlib
mlx = importlib.import_module("mlx")
out["mlx"] = {
"version": getattr(mlx, "__version__", None),
"file": getattr(mlx, "__file__", None),
}
try:
mlx_lm = importlib.import_module("mlx_lm")
out["mlx_lm"] = {
"version": getattr(mlx_lm, "__version__", None),
"file": getattr(mlx_lm, "__file__", None),
}
except Exception as e:
out["error"] = f"import mlx_lm failed: {e}"
print(json.dumps(out, indent=2))
sys.exit(2)
out["ok"] = True
print(json.dumps(out, indent=2))
sys.exit(0)
except Exception as e:
out["error"] = str(e)
print(json.dumps(out, indent=2))
sys.exit(2)
# Child mode: run the server in-process (bypasses CLI parsing)
if os.environ.get("MLXK2_CHILD_SERVER") == "1":
from mlxk2.core.server_base import run_server as _run_server
host = os.environ.get("MLXK2_HOST", "127.0.0.1")
port = int(os.environ.get("MLXK2_PORT", "8000"))
log_level = os.environ.get("MLXK2_LOG_LEVEL", "info")
reload = _to_bool(os.environ.get("MLXK2_RELOAD", "0"))
_mt = os.environ.get("MLXK2_MAX_TOKENS", "")
max_tokens = int(_mt) if _mt.strip() else None
_run_server(host=host, port=port, max_tokens=max_tokens, reload=reload, log_level=log_level)
sys.exit(0)
# Parent mode: patch start_server to supervise a child process
from mlxk2.core.server_base import run_server as _run_server
import mlxk2.operations.serve as _serve
def _run_supervised_server(*, host: str, port: int, log_level: str, reload: bool, max_tokens):
env = os.environ.copy()
env["MLXK2_CHILD_SERVER"] = "1"
env["MLXK2_HOST"] = host
env["MLXK2_PORT"] = str(port)
env["MLXK2_LOG_LEVEL"] = log_level
env["MLXK2_RELOAD"] = "1" if reload else "0"
if max_tokens is not None:
env["MLXK2_MAX_TOKENS"] = str(max_tokens)
proc = subprocess.Popen([sys.executable], start_new_session=True, env=env)
shutting_down = {"done": False}
def _graceful_shutdown(_sig=None, _frame=None):
if shutting_down["done"]:
return
shutting_down["done"] = True
try:
os.killpg(proc.pid, signal.SIGTERM)
except Exception:
pass
deadline = time.time() + 5.0
while time.time() < deadline:
ret = proc.poll()
if ret is not None:
return
try:
time.sleep(0.1)
except Exception:
pass
try:
os.killpg(proc.pid, signal.SIGKILL)
except Exception:
pass
try:
signal.signal(signal.SIGINT, _graceful_shutdown)
except Exception:
pass
try:
signal.signal(signal.SIGTERM, _graceful_shutdown)
except Exception:
pass
try:
return proc.wait()
except KeyboardInterrupt:
previous = signal.signal(signal.SIGINT, signal.SIG_IGN)
try:
try:
os.killpg(proc.pid, signal.SIGTERM)
except Exception:
pass
deadline = time.time() + 5.0
while time.time() < deadline:
ret = proc.poll()
if ret is not None:
return ret
try:
time.sleep(0.1)
except KeyboardInterrupt:
break
try:
os.killpg(proc.pid, signal.SIGKILL)
except Exception:
pass
while True:
ret = proc.poll()
if ret is not None:
return ret
time.sleep(0.05)
finally:
try:
signal.signal(signal.SIGINT, previous)
except Exception:
pass
def _patched_start_server(*, model=None, port=8000, host="127.0.0.1", max_tokens=None,
reload=False, log_level="info", verbose=False, supervise=True):
return _run_server(host=host, port=port, max_tokens=max_tokens, reload=reload, log_level=log_level)
_serve.start_server = _patched_start_server
if __name__ == "__main__":
from mlxk2.cli import main
main()
PY
# -----------------------------------------------------------------------------
# Build with PyInstaller
# -----------------------------------------------------------------------------
echo "[3/3] Building and packaging..."
PYI_OPTS=(
--clean
--noconfirm
--name "$APP_NAME"
--log-level=WARN
--specpath "$SPEC_DIR"
--collect-all mlx
--collect-all mlx_lm
--collect-all huggingface_hub
--collect-all requests
--collect-all fastapi
--collect-all pydantic
--collect-all uvicorn
)
if [[ "${ONEFILE:-0}" == "1" ]]; then
PYI_OPTS+=(--onefile)
else
PYI_OPTS+=(--onedir)
fi
pyinstaller "${PYI_OPTS[@]}" "$ENTRY_SCRIPT"
# -----------------------------------------------------------------------------
# Package archive
# -----------------------------------------------------------------------------
TAR_BASE="$APP_NAME"
if [[ -n "${VERSION:-}" ]]; then
TAR_BASE="$TAR_BASE-${VERSION}"
fi
TAR_NAME="${ARCHIVE_NAME:-${TAR_BASE}.tgz}"
PACK_DIR="$ROOT_DIR/.pack-tmp"
rm -rf "$PACK_DIR"
mkdir -p "$PACK_DIR/$APP_NAME"
if [[ "${ONEFILE:-0}" == "1" ]]; then
cp "$DIST_DIR/$APP_NAME" "$PACK_DIR/$APP_NAME/$APP_NAME"
else
rsync -a --delete "$DIST_DIR/$APP_NAME/" "$PACK_DIR/$APP_NAME/"
fi
cp -f LICENSE "$PACK_DIR/$APP_NAME/" 2>/dev/null || true
tar -C "$PACK_DIR" -czf "$ART_DIR/$TAR_NAME" "$APP_NAME"
echo "Created: $ART_DIR/$TAR_NAME"
echo "Build complete. Binaries are in dist/, archive in artifacts/."
rm -f "$ENTRY_SCRIPT"