| """Install and launch the pinned Swift MLX runtime using an existing local model.""" |
|
|
| import argparse |
| import hashlib |
| import inspect |
| import json |
| import os |
| import platform |
| import shlex |
| import shutil |
| import subprocess |
| import sys |
| import urllib.request |
| import venv |
| from pathlib import Path |
|
|
| BASE = "c69d1288440a0dc4e6401fc417098b07598dccd5" |
| REVISIONS = { |
| 4: "730aab9b0395b26f7d9cf4b0dfa6a4c788aff6fd", |
| 5: "8aff72b145212e62c15146e41dcbef35ea5fa9ff", |
| } |
| PATCHES = { |
| "swift15-mlx-lm.patch": "f6f1d0bdafa45863bfbf93dac0398c481c993ea04fdf38b9bae98c643f89eaec", |
| "enable-5bit.patch": "b985961eac3035e05ca4c9f3a8b283c26e6dd69bab4d113997fc04fe2a4f99cc", |
| "swift15-server-cache.patch": "a7fbc0f0524ee7864d9f41a98a2e35acf0e63f53e369ee5b9a87ca12926a9eab", |
| } |
| MODEL_HASHES = { |
| 4: "f559e4559946bfad4ce20b93906028449b76e007f42f36bb83d2c2c5f6cb95a9", |
| 5: "ecda666d6d5f1059388a2ce4267d9ff8625398bb925d6c0840d0059ce17b676f", |
| } |
| SERVER_HASH = "791f8c1eb3431c4d24b4f9084f67b665c38e3d547ee4a307e52c5b2b6f12aba0" |
|
|
|
|
| def digest(path): |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
|
|
| def local_config(model): |
| config = json.loads((model / "config.json").read_text()) |
| quant = config.get("quantization", {}) |
| bits = quant.get("bits") |
| if bits not in REVISIONS or quant != { |
| "bits": bits, |
| "group_size": 64, |
| "mode": "affine", |
| }: |
| raise ValueError( |
| "Expected the Swift 1.5 affine 4-bit or 5-bit / group-size-64 checkpoint." |
| ) |
| if config.get("language_model_only") is not False: |
| raise ValueError("Expected the complete Swift architecture configuration.") |
| for name in ("tokenizer.json", "tokenizer_config.json", "chat_template.jinja"): |
| if not (model / name).is_file(): |
| raise ValueError(f"Missing local model asset: {name}") |
| index = json.loads((model / "model.safetensors.index.json").read_text()) |
| shards = set(index["weight_map"].values()) |
| if not shards: |
| raise ValueError("The checkpoint index is empty.") |
| for name in shards: |
| if Path(name).name != name or not name.endswith(".safetensors"): |
| raise ValueError("The checkpoint index contains an invalid shard path.") |
| if not (model / name).is_file(): |
| raise ValueError( |
| f"Missing local weight shard: {name}; complete the existing download first." |
| ) |
| return config, bits |
|
|
|
|
| def runtime_check(config, bits): |
| from importlib.metadata import version |
|
|
| from mlx_lm import server |
| from mlx_lm.utils import _get_classes |
|
|
| cls, args = _get_classes(config) |
| args.from_dict(config) |
| if cls.__module__ != "mlx_lm.models.qwen3_5_full": |
| raise ValueError( |
| "Swift architecture patch is missing. Use the serve launcher created by setup." |
| ) |
| if digest(inspect.getsourcefile(cls)) != MODEL_HASHES[bits]: |
| raise ValueError( |
| "The installed architecture differs from the pinned Swift implementation." |
| ) |
| if digest(server.__file__) != SERVER_HASH: |
| raise ValueError("The tested server cache patch is missing or has changed.") |
| versions = { |
| name: version(name) |
| for name in ("mlx", "mlx-metal", "mlx-lm", "transformers", "huggingface_hub") |
| } |
| expected = { |
| "mlx": "0.32.2", |
| "mlx-metal": "0.32.2", |
| "mlx-lm": "0.32.0", |
| "transformers": "5.14.1", |
| "huggingface_hub": "1.31.0", |
| } |
| if versions != expected: |
| raise ValueError( |
| f"Installed package versions differ from the tested runtime: {versions}" |
| ) |
| print( |
| json.dumps( |
| { |
| "status": "PASS_RUNTIME_IDENTITY", |
| "python": sys.executable, |
| "model_class": cls.__module__ + "." + cls.__name__, |
| "server": server.__file__, |
| "versions": versions, |
| "weights_loaded": False, |
| } |
| ), |
| flush=True, |
| ) |
|
|
|
|
| def setup(model, bits, destination): |
| if shutil.which("git") is None: |
| raise ValueError( |
| "Git is required. Install the macOS command-line tools, then retry." |
| ) |
| if destination.exists(): |
| raise ValueError( |
| f"Runtime directory already exists: {destination}. Use its serve launcher, or choose a new --runtime-dir." |
| ) |
| destination.mkdir(parents=True) |
| source = destination / "mlx-lm" |
| envdir = destination / ".venv" |
| python = envdir / "bin/python" |
| patchdir = destination / "patches" |
| patchdir.mkdir() |
| patch_names = ["swift15-mlx-lm.patch"] |
| if bits == 5: |
| patch_names.append("enable-5bit.patch") |
| patch_names.append("swift15-server-cache.patch") |
| for name in patch_names: |
| url = f"https://huggingface.co/ukisai/Swift-1.5-{bits}bit-MLX/resolve/{REVISIONS[bits]}/compatibility/{name}" |
| with urllib.request.urlopen(url, timeout=60) as response: |
| data = response.read() |
| if hashlib.sha256(data).hexdigest() != PATCHES[name]: |
| raise ValueError(f"Patch checksum mismatch: {name}") |
| (patchdir / name).write_bytes(data) |
| subprocess.run(["git", "init", "--quiet", str(source)], check=True) |
| git = ["git", "-C", str(source)] |
| subprocess.run( |
| git + ["remote", "add", "origin", "https://github.com/ml-explore/mlx-lm.git"], |
| check=True, |
| ) |
| subprocess.run(git + ["fetch", "--depth", "1", "origin", BASE], check=True) |
| subprocess.run(git + ["checkout", "--detach", BASE], check=True) |
| for name in patch_names: |
| subprocess.run(git + ["apply", "--check", str(patchdir / name)], check=True) |
| subprocess.run(git + ["apply", str(patchdir / name)], check=True) |
| venv.EnvBuilder(with_pip=True).create(envdir) |
| subprocess.run( |
| [ |
| str(python), |
| "-I", |
| "-m", |
| "pip", |
| "install", |
| "mlx==0.32.2", |
| "mlx-metal==0.32.2", |
| "transformers==5.14.1", |
| "huggingface_hub==1.31.0", |
| "pillow==12.3.0", |
| "safetensors==0.8.0", |
| "-e", |
| str(source), |
| ], |
| check=True, |
| ) |
| subprocess.run([str(python), "-I", "-m", "pip", "check"], check=True) |
| installed = destination / "swift_runtime.py" |
| shutil.copyfile(__file__, installed) |
| command = [str(python), "-I", str(installed), "start", "--model", str(model)] |
| subprocess.run(command + ["--check-only"], check=True) |
| launcher = destination / "serve" |
| launcher.write_text("#!/bin/sh\nexec " + shlex.join(command) + ' "$@"\n') |
| launcher.chmod(0o755) |
| print( |
| f"\nSetup complete. Start or restart with this exact command:\n{shlex.quote(str(launcher))}" |
| ) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| sub = parser.add_subparsers(dest="action", required=True) |
| install = sub.add_parser( |
| "setup", |
| help="Install into a new isolated directory; reuse existing model files.", |
| ) |
| install.add_argument("--model", required=True, type=Path) |
| install.add_argument("--runtime-dir", type=Path) |
| start = sub.add_parser( |
| "start", help="Verify the selected Python runtime and serve a local checkpoint." |
| ) |
| start.add_argument("--model", required=True, type=Path) |
| start.add_argument("--check-only", action="store_true") |
| args, extra = parser.parse_known_args() |
| if args.action == "setup" and extra: |
| parser.error("Unknown setup arguments: " + " ".join(extra)) |
| if sys.version_info[:2] != (3, 12): |
| parser.error( |
| "Use python3.12 for this tested runtime. Homebrew's Python 3.14 environment is separate." |
| ) |
| if platform.system() != "Darwin" or platform.machine() != "arm64": |
| parser.error("This installer and launcher target native Apple Silicon macOS.") |
| model = args.model.expanduser().resolve() |
| try: |
| config, bits = local_config(model) |
| if args.action == "setup": |
| target = ( |
| args.runtime_dir |
| or Path.home() / ".local/share/swift15-mlx" / f"{bits}bit" |
| ) |
| setup(model, bits, target.expanduser().resolve()) |
| return |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" |
| runtime_check(config, bits) |
| if args.check_only: |
| return |
| os.execv( |
| sys.executable, |
| [ |
| sys.executable, |
| "-I", |
| "-m", |
| "mlx_lm.server", |
| "--model", |
| str(model), |
| "--host", |
| "127.0.0.1", |
| "--port", |
| "8080", |
| *extra, |
| ], |
| ) |
| except ( |
| OSError, |
| ValueError, |
| KeyError, |
| ImportError, |
| subprocess.CalledProcessError, |
| ) as exc: |
| raise SystemExit(f"Swift runtime setup/start failed: {exc}") from None |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
