diff --git a/README.md b/README.md index f6578a5e..785091b8 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Modern dotfiles repository managed with [Chezmoi](https://chezmoi.io/), featurin - **Automated Validation**: Pre-commit hooks and validation scripts ensure configuration quality - **Windows Enterprise Detection**: Automatic detection of Entra ID (Azure AD) and Intune enrollment status - **Task Automation**: Integrated [Task](https://taskfile.dev/) runner for common operations (validation, testing, installation) +- **Global Tasks**: Shell-independent `task -g llm:chat` with automatic MLX setup on Apple Silicon; see [global tasks](docs/development-tools.md#global-tasks). - **Tool Version Management**: [mise](https://mise.jdx.dev/) for managing development tool versions ## 🔧 Chezmoi Variables diff --git a/docs/development-tools.md b/docs/development-tools.md index d053e0ad..8f38a529 100644 --- a/docs/development-tools.md +++ b/docs/development-tools.md @@ -38,6 +38,125 @@ task dev:setup # Complete dev environment setup task ci:validate # Run CI validation pipeline ``` +## Global Tasks + +Chezmoi deploys `home/Taskfile.yml` to `~/Taskfile.yml`, separate from the +repository's development `Taskfile.yml`. With Task on your PATH, `task -g` +selects the home Taskfile from **any directory and shell**, including Fish, +Bash, Zsh, and PowerShell. Ordinary repository `task` commands stay unchanged. + +```sh +task -g --list +task -g llm:chat +``` + +### Local MLX chat (Apple Silicon) + +`llm:chat` uses Apple's `mlx-lm` Python API directly, not an inference frontend. +It requires native Apple Silicon macOS; other platforms get an explicit error +only if they invoke this task. A Task dependency checks prerequisites before +launch, installing only what is missing. + +The default setup uses: + +- `~/mlx-qwen/.venv/bin/python` with `mlx-lm` 0.32.2 or newer (developed against 0.32.2). +- The `mlx-community/Qwen3.5-9B-4bit` model in + `/Volumes/External Disk Data/AI/huggingface/hub`. +- `/Volumes/External Disk Data/AI/huggingface/xet` on the same external SSD. + +**Setup order:** native Apple Silicon and the mounted SSD are checked first. +Cache directories are created on that SSD and checked for writability and +symlink redirection. A `status` check then verifies the venv and imports the +required MLX APIs without loading weights. A working installation skips +provisioning entirely: no Python replacement, pip upgrade, or package upgrade. + +If the venv is missing, Task creates it with native Python 3.10+ and `venv`. +If suitable Python is missing, it uses the same Homebrew `python@3` formula +as the repository's macOS package setup. Missing pip is bootstrapped with +`ensurepip`; missing requirements are installed with isolated +`python -m pip install 'mlx-lm>=0.32.2'`, without `--upgrade`. Pip and Homebrew +download caches, plus setup temporary files, stay on the SSD; the venv itself +remains at the configured `MLX_VENV` path. No uv or inference frontend is involved. +An existing broken/non-native venv is **not deleted**: repair it or move it +aside, then rerun. Installation failures and a failed post-install import +check stop the task before model launch; corrupt installed packages may +require manual repair rather than a forced reinstall. + +The launcher at `~/.config/mlx/chat.py` sets `HF_HOME`, `HF_HUB_CACHE`, and +`HF_XET_CACHE` to the external cache, overriding inherited Hugging Face +settings. It first resolves the existing cached model without network access. +If absent, it downloads the model there (approximately 6 GB); existing cached +weights are reused, not copied or updated. There is no internal-disk fallback. +Connect the SSD first and run from a terminal permitted to access it (for +example Terminal.app); this does not bypass macOS volume permissions. + +**Bootstrap boundary:** Task cannot install itself before the `task` command +exists. Full macOS dotfiles setup already includes `go-task` and Python. If +Task alone is missing, use the existing Homebrew setup and `brew install go-task`, +or your configured mise installation with `mise use --global task@latest`. +If neither suitable Python nor Homebrew is available, the task reports that +Homebrew must be installed through the existing dotfiles setup; it does not run +a new privileged bootstrap script. + +Thinking is disabled via `enable_thinking=False` in the model's chat template, +which supplies an empty closed `` block. The launcher uses temperature +0.7, top-p 0.8, top-k 20, at most 512 reply tokens, and 512-token prefill chunks. +It does not use the unsupported `mlx_lm.chat` flags or a `/nothink` soft switch. + +Type `q` to quit, `r` to reset, or `h` for help. Ctrl+C (including during +generation) and EOF exit cleanly. Each reply reports token speed and peak MLX +memory; conversation state stays in memory, with no saved transcripts. +The prompt cache is reused across turns, including the correct end-of-turn +boundary after a truncated reply. + +**Memory bound:** Qwen3.5's hybrid cache ignores `max_kv_size`. Instead, every +turn must fit within **8,192 total tokens**, counting cached history, the new +prompt and its delimiters, and 512 reserved reply tokens. Oversized turns are +rejected before generation: reset or shorten the input. History is never +silently discarded. This bounds context growth, not total process RAM; +the approximately 6 GB of weights and other applications still need memory. + +### Selective deployment + +Once these changes are present in your primary dotfiles checkout, run these +commands **from its repository root**. Only the global Taskfile and launcher +are selected; Chezmoi scripts are explicitly excluded, and neither the venv +nor model caches are managed by Git or Chezmoi. + +Create the launcher's parent directory first. When selecting individual files, +Chezmoi 2.73.0 requires that parent to exist; a fresh home otherwise fails with +`stat .../.config/mlx: no such file or directory`. Use `--verbose` to display +the proposed changes during the dry-run: + +```sh +mkdir -p "$HOME/.config/mlx" && + chezmoi --source=. apply --exclude=scripts --dry-run --verbose "$HOME/Taskfile.yml" "$HOME/.config/mlx/chat.py" +``` + +Review that preview first, especially if you already maintain a home Taskfile. +Then apply and launch only if deployment succeeds: + +```sh +chezmoi --source=. apply --exclude=scripts "$HOME/Taskfile.yml" "$HOME/.config/mlx/chat.py" && + task -g llm:chat +``` + +There is no need to activate the venv, set shell-specific environment variables, +or remove an earlier standalone launcher. For an intentionally relocated +existing setup, Task variables can be overridden explicitly: + +```sh +task -g llm:chat MLX_VENV="/path/to/.venv" MLX_VOLUME="/Volumes/Another SSD" +``` + +Focused validation uses fake MLX/Hugging Face modules, a mock venv, and an +isolated Chezmoi destination. It covers installed/missing/broken environments +and setup ordering without real provisioning, downloads, or inference: + +```sh +timeout 60 ./tests/bash/run-tests.sh --test global-mlx-chat.bats +``` + ## Mise (Tool Version Manager) Mise manages tool versions defined in diff --git a/home/.chezmoiignore b/home/.chezmoiignore index c267a3a9..d0566f91 100644 --- a/home/.chezmoiignore +++ b/home/.chezmoiignore @@ -5,7 +5,6 @@ LICENSE tests/ install.sh install.ps1 -Taskfile.yml .mise.toml CONTRIBUTING.md STRUCTURE.md diff --git a/home/Taskfile.yml b/home/Taskfile.yml new file mode 100644 index 00000000..6ed39800 --- /dev/null +++ b/home/Taskfile.yml @@ -0,0 +1,100 @@ +version: '3' + +vars: + MLX_VENV: '{{.ROOT_DIR}}/mlx-qwen/.venv' + MLX_VOLUME: /Volumes/External Disk Data + MLX_CACHE_DIR: '{{.MLX_VOLUME}}/AI/huggingface' + MLX_PYTHON_CHECK: 'import platform, sys, venv; assert sys.version_info >= (3, 10) and platform.machine() == "arm64"' + MLX_VENV_CHECK: '{{.MLX_PYTHON_CHECK}}; assert sys.prefix != sys.base_prefix' + MLX_CHECK_ENV: >- + "{{.MLX_VENV}}/bin/python" "{{.ROOT_DIR}}/.config/mlx/chat.py" + --cache-dir "{{.MLX_CACHE_DIR}}" --volume "{{.MLX_VOLUME}}" --check-env + +tasks: + llm:chat: + desc: Set up and chat with local Qwen3.5 (Apple Silicon, thinking disabled) + interactive: true + silent: true + deps: [llm:setup] + cmds: + - >- + "{{.MLX_VENV}}/bin/python" "{{.ROOT_DIR}}/.config/mlx/chat.py" + --cache-dir "{{.MLX_CACHE_DIR}}" --volume "{{.MLX_VOLUME}}" + + llm:storage: + internal: true + silent: true + preconditions: + - sh: test "$(uname -s)" = Darwin && test "$(uname -m)" = arm64 + msg: 'llm:chat requires native Apple Silicon macOS (not Rosetta).' + - sh: mount | grep -F -- " on {{.MLX_VOLUME}} (" >/dev/null + msg: 'Connect the SSD at {{.MLX_VOLUME}} before setup. No internal cache fallback is used.' + - sh: | + case "{{.MLX_CACHE_DIR}}" in + *"/../"*|*"/.."|*"/./"*|*"/.") exit 1 ;; + "{{.MLX_VOLUME}}"/*) exit 0 ;; + *) exit 1 ;; + esac + msg: 'MLX_CACHE_DIR must be a directory beneath MLX_VOLUME, without dot segments.' + - sh: test -f "{{.ROOT_DIR}}/.config/mlx/chat.py" + msg: 'Missing ~/.config/mlx/chat.py. Apply the global Taskfile and launcher together.' + cmds: + - mkdir -p "{{.MLX_CACHE_DIR}}/hub" "{{.MLX_CACHE_DIR}}/xet" "{{.MLX_CACHE_DIR}}/pip" "{{.MLX_CACHE_DIR}}/tmp" + + llm:setup: + internal: true + deps: [llm:storage] + preconditions: + - sh: | + for path in "{{.MLX_CACHE_DIR}}" "{{.MLX_CACHE_DIR}}/hub" "{{.MLX_CACHE_DIR}}/xet" "{{.MLX_CACHE_DIR}}/pip" "{{.MLX_CACHE_DIR}}/tmp"; do + test "$(cd "$path" && pwd -P)" = "$path" && test -w "$path" || exit 1 + done + msg: 'External cache directories must be writable and not redirected through symlinks. Check SSD permissions.' + status: + - '{{.MLX_CHECK_ENV}}' + cmds: + - task: llm:venv + - | + export TMPDIR="{{.MLX_CACHE_DIR}}/tmp" + "{{.MLX_VENV}}/bin/python" -m pip --version || + "{{.MLX_VENV}}/bin/python" -m ensurepip + - >- + env TMPDIR="{{.MLX_CACHE_DIR}}/tmp" + "{{.MLX_VENV}}/bin/python" -m pip --isolated --require-virtualenv + --disable-pip-version-check --no-input --cache-dir "{{.MLX_CACHE_DIR}}/pip" + install 'mlx-lm>=0.32.2' + - '{{.MLX_CHECK_ENV}}' + + llm:venv: + internal: true + preconditions: + - sh: >- + (test ! -e "{{.MLX_VENV}}" && test ! -L "{{.MLX_VENV}}") || + "{{.MLX_VENV}}/bin/python" -c '{{.MLX_VENV_CHECK}}' + msg: 'Existing {{.MLX_VENV}} is not a usable Python 3.10+ arm64 venv. Repair it or move it aside; it will not be deleted automatically.' + status: + - '"{{.MLX_VENV}}/bin/python" -c ''{{.MLX_VENV_CHECK}}''' + cmds: + - task: llm:python + - | + export TMPDIR="{{.MLX_CACHE_DIR}}/tmp" + if python3 -c '{{.MLX_PYTHON_CHECK}}' 2>/dev/null; then + python3 -m venv "{{.MLX_VENV}}" + else + "$(brew --prefix python@3)/libexec/bin/python3" -m venv "{{.MLX_VENV}}" + fi + + llm:python: + internal: true + preconditions: + - sh: python3 -c '{{.MLX_PYTHON_CHECK}}' 2>/dev/null || command -v brew + msg: 'A native Python 3.10+ or Homebrew is required. Install Homebrew using the existing dotfiles setup; this task cannot safely bootstrap Homebrew.' + status: + - | + python3 -c '{{.MLX_PYTHON_CHECK}}' 2>/dev/null || + (command -v brew >/dev/null && "$(brew --prefix python@3)/libexec/bin/python3" -c '{{.MLX_PYTHON_CHECK}}') + cmds: + - >- + brew list --versions python@3 || + env HOMEBREW_CACHE="{{.MLX_CACHE_DIR}}/homebrew" TMPDIR="{{.MLX_CACHE_DIR}}/tmp" + brew install python@3 diff --git a/home/dot_config/mlx/chat.py b/home/dot_config/mlx/chat.py new file mode 100644 index 00000000..e9b3add5 --- /dev/null +++ b/home/dot_config/mlx/chat.py @@ -0,0 +1,208 @@ +"""Bounded-memory Qwen3.5 chat using Apple's mlx-lm Python API.""" + +import argparse +from contextlib import closing +from importlib.metadata import version, PackageNotFoundError +import os +from pathlib import Path +import platform +import sys + + +MODEL = "mlx-community/Qwen3.5-9B-4bit" +CONTEXT_TOKENS = 8192 +REPLY_TOKENS = 512 +HELP = "q: quit | r: reset conversation | h: help" + + +def configure_cache(cache_dir, volume): + mount = Path(volume).expanduser() + if not mount.is_absolute() or not mount.is_mount(): + raise ValueError(f"Connect the external SSD at {mount}.") + mount = mount.resolve(strict=True) + root = Path(cache_dir).expanduser() + if not root.is_absolute(): + raise ValueError("The cache directory must be an absolute path.") + root = root.resolve(strict=True) + if not root.is_relative_to(mount) or root == mount: + raise ValueError("The cache directory must be beneath the external SSD.") + for name in ("hub", "xet"): + child = root / name + if not child.is_dir(): + raise ValueError(f"Missing {root / name}. Connect the external SSD.") + if not child.resolve(strict=True).is_relative_to(mount): + raise ValueError(f"{child} points outside the external SSD.") + + # Set these before importing Hugging Face/MLX: they cache environment values. + # Deliberately override inherited settings; never fall back to internal storage. + os.environ.update( + HF_HOME=str(root), + HF_HUB_CACHE=str(root / "hub"), + HF_XET_CACHE=str(root / "xet"), + HF_HUB_OFFLINE="0", + TRANSFORMERS_OFFLINE="0", + HF_HUB_DISABLE_TELEMETRY="1", + ) + return root + + +def check_environment(): + if sys.version_info < (3, 10) or sys.prefix == sys.base_prefix: + raise ValueError("A Python 3.10+ virtual environment is required.") + import mlx.core + from mlx_lm import load, stream_generate + from mlx_lm.models.cache import make_prompt_cache + from mlx_lm.sample_utils import make_sampler + from packaging.version import Version + + if Version(version("mlx-lm")) < Version("0.32.2"): + raise ValueError("mlx-lm 0.32.2 or newer is required for this launcher.") + + +def chat(model, tokenizer): + import mlx.core as mx + from mlx_lm import stream_generate + from mlx_lm.models.cache import make_prompt_cache + from mlx_lm.sample_utils import make_sampler + + end = tokenizer.encode("<|im_end|>", add_special_tokens=False) + if len(end) != 1 or end[0] not in tokenizer.eos_token_ids: + raise ValueError("Expected Qwen3.5's <|im_end|> end-of-turn token.") + newline = tokenizer.encode("\n", add_special_tokens=False) + sampler = make_sampler(temp=0.7, top_p=0.8, top_k=20) + cache = make_prompt_cache(model) + used_tokens = 0 + separator = [] + print(f"{HELP}\nContext limit: {CONTEXT_TOKENS} tokens; replies: {REPLY_TOKENS}.") + print("Thinking disabled. Conversation stays in memory; nothing is saved.") + try: + while True: + try: + query = input("\nYou: ").strip() + except EOFError: + print() + break + if query == "q": + break + if query == "h": + print(HELP) + continue + if query == "r": + cache.clear() + mx.clear_cache() + cache = make_prompt_cache(model) + used_tokens = 0 + separator = [] + print("Conversation reset.") + continue + if not query: + continue + + prompt = separator + tokenizer.apply_chat_template( + [{"role": "user", "content": query}], + tokenize=True, + add_generation_prompt=True, + enable_thinking=False, + ) + required = used_tokens + len(prompt) + REPLY_TOKENS + if required > CONTEXT_TOKENS: + print( + f"Context limit: this turn needs {required}/{CONTEXT_TOKENS} " + "tokens including the reserved reply. Enter r to reset " + "or shorten your input; history has not been dropped." + ) + continue + + tokenizer.detokenizer.reset() + print("\nAssistant: ", end="", flush=True) + last = None + with closing( + stream_generate( + model, + tokenizer, + prompt, + prompt_cache=cache, + sampler=sampler, + max_tokens=REPLY_TOKENS, + prefill_step_size=512, + ) + ) as responses: + for response in responses: + print(response.text, end="", flush=True) + last = response + if last is None or last.finish_reason not in ("stop", "length"): + raise RuntimeError("MLX generation ended without a completed response.") + + # mlx-lm 0.32.2 caches each yielded token, including EOS, before yielding. + # Append only the missing boundary, not a re-rendered copy of history. + used_tokens += len(prompt) + last.generation_tokens + separator = ([] if last.token == end[0] else end) + newline + print( + f"\n[{last.generation_tokens} tokens; " + f"{last.generation_tps:.1f} tok/s; " + f"peak {last.peak_memory:.2f} GB; " + f"context {used_tokens}/{CONTEXT_TOKENS}]" + ) + if last.finish_reason == "length": + print("Reply reached the 512-token limit; the next turn will close it.") + except KeyboardInterrupt: + print() + finally: + cache.clear() + mx.clear_cache() + print("Bye.") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache-dir", required=True, help="Existing Hugging Face cache root") + parser.add_argument("--volume", required=True, help="Mounted external SSD") + parser.add_argument("--check-env", action="store_true", help="Check prerequisites without loading a model") + args = parser.parse_args() + if platform.system() != "Darwin" or platform.machine() != "arm64": + parser.error("MLX chat requires native Apple Silicon macOS (not Rosetta).") + try: + root = configure_cache(args.cache_dir, args.volume) + except (OSError, ValueError) as error: + parser.exit(1, f"External cache unavailable: {error}\nNo internal fallback is used.\n") + try: + check_environment() + except (ImportError, OSError, ValueError, PackageNotFoundError) as error: + parser.exit( + 1, f"MLX environment unavailable: {error}\n" + "Run task -g llm:chat to install missing requirements. If setup still " + "fails, repair the venv or move it aside; working packages are not force-upgraded.\n" + ) + if args.check_env: + return + + try: + from huggingface_hub import snapshot_download + from huggingface_hub.errors import LocalEntryNotFoundError + from mlx_lm import load + except ImportError as error: + parser.exit(1, f"Cannot import mlx-lm from this Python environment: {error}\n") + try: + snapshot = snapshot_download( + MODEL, cache_dir=str(root / "hub"), local_files_only=True + ) + except LocalEntryNotFoundError: + print(f"Downloading {MODEL} to {root / 'hub'} (external SSD)...", flush=True) + snapshot = snapshot_download( + MODEL, cache_dir=str(root / "hub"), local_files_only=False + ) + + print(f"Loading {MODEL} from the external cache...", flush=True) + model, tokenizer = load( + snapshot, + tokenizer_config={"trust_remote_code": False}, + trust_remote_code=False, + ) + chat(model, tokenizer) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nBye.") diff --git a/tests/bash/global-mlx-chat.bats b/tests/bash/global-mlx-chat.bats new file mode 100644 index 00000000..7771a167 --- /dev/null +++ b/tests/bash/global-mlx-chat.bats @@ -0,0 +1,361 @@ +#!/usr/bin/env bats + +setup() { + REPO_ROOT="$(cd "$BATS_TEST_DIRNAME/../.." && pwd)" + TEST_HOME="$BATS_TEST_TMPDIR/home with spaces" + mkdir -p "$BATS_TEST_TMPDIR/External Disk Data" + export MLX_TEST_VOLUME + MLX_TEST_VOLUME="$(cd "$BATS_TEST_TMPDIR/External Disk Data" && pwd -P)" + CACHE_DIR="$MLX_TEST_VOLUME/AI/huggingface" + export MOCK_OS=Darwin MOCK_ARCH=arm64 MOCK_MOUNTED=1 MOCK_BASE_PYTHON=1 + export PYTHON_CALLS="$BATS_TEST_TMPDIR/python-calls" + export MLX_TEST_ACTIONS="$BATS_TEST_TMPDIR/actions" + export MLX_TEST_STATE="$BATS_TEST_TMPDIR/state" + export MLX_TEST_PYTHON="$BATS_TEST_TMPDIR/venv-python" + export MLX_TEST_BREW="$BATS_TEST_TMPDIR/homebrew" + printf 'ready\n' >"$MLX_TEST_STATE" + mkdir -p "$TEST_HOME/mlx-qwen/.venv/bin" "$TEST_HOME/.config/mlx" \ + "$CACHE_DIR/hub" "$CACHE_DIR/xet" "$BATS_TEST_TMPDIR/bin" "$BATS_TEST_TMPDIR/work" + cp "$REPO_ROOT/home/Taskfile.yml" "$TEST_HOME/Taskfile.yml" + cp "$REPO_ROOT/home/dot_config/mlx/chat.py" "$TEST_HOME/.config/mlx/chat.py" + cat >"$MLX_TEST_PYTHON" <<'EOF' +#!/bin/sh +case "$*" in + *"--check-env"*) + printf 'check-env\n' >>"$MLX_TEST_ACTIONS" + test "$(cat "$MLX_TEST_STATE")" = ready + exit "$?" + ;; + "-c "*) + printf 'check-venv\n' >>"$MLX_TEST_ACTIONS" + exit "${MOCK_VENV_EXIT:-0}" + ;; + "-m pip --version") + printf 'check-pip\n' >>"$MLX_TEST_ACTIONS" + exit "${MOCK_PIP_EXIT:-0}" + ;; + "-m ensurepip") + printf 'ensurepip\n' >>"$MLX_TEST_ACTIONS" + exit 0 + ;; + "-m pip "*) + printf 'pip: %s\n' "$*" >>"$MLX_TEST_ACTIONS" + printf 'tmp: %s\n' "$TMPDIR" >>"$MLX_TEST_ACTIONS" + [ "${MOCK_INSTALL_EXIT:-0}" -eq 0 ] || exit "$MOCK_INSTALL_EXIT" + [ "${MOCK_STILL_BROKEN:-0}" -eq 1 ] || printf 'ready\n' >"$MLX_TEST_STATE" + exit 0 + ;; +esac +printf 'chat\n' >>"$MLX_TEST_ACTIONS" +printf '<%s>\n' "$@" >"$PYTHON_CALLS" +exit "${MOCK_PYTHON_EXIT:-0}" +EOF + chmod +x "$MLX_TEST_PYTHON" + cp "$MLX_TEST_PYTHON" "$TEST_HOME/mlx-qwen/.venv/bin/python" + cat >"$BATS_TEST_TMPDIR/bin/uname" <<'EOF' +#!/bin/sh +case "$1" in + -s) printf '%s\n' "$MOCK_OS" ;; + -m) printf '%s\n' "$MOCK_ARCH" ;; + *) exit 1 ;; +esac +EOF + cat >"$BATS_TEST_TMPDIR/bin/mount" <<'EOF' +#!/bin/sh +printf 'mount\n' >>"$MLX_TEST_ACTIONS" +[ "$MOCK_MOUNTED" -eq 1 ] || exit 0 +printf '/dev/mock on %s (apfs, local)\n' "$MLX_TEST_VOLUME" +EOF + cat >"$BATS_TEST_TMPDIR/bin/python3" <<'EOF' +#!/bin/sh +case "$1" in + -c) + case "$0" in "$MLX_TEST_BREW"/*) exit 0 ;; esac + [ "$MOCK_BASE_PYTHON" -eq 1 ] || exit 1 + ;; + -m) + [ "$2" = venv ] || exit 1 + printf 'create-venv\n' >>"$MLX_TEST_ACTIONS" + mkdir -p "$3/bin" + cp "$MLX_TEST_PYTHON" "$3/bin/python" + ;; + *) exit 1 ;; +esac +EOF + cat >"$BATS_TEST_TMPDIR/bin/brew" <<'EOF' +#!/bin/sh +case "$*" in + "--prefix python@3") printf '%s\n' "$MLX_TEST_BREW" ;; + "list --versions python@3") test -f "$MLX_TEST_BREW/libexec/bin/python3" ;; + "install python@3") + printf 'install-python\n' >>"$MLX_TEST_ACTIONS" + mkdir -p "$MLX_TEST_BREW/libexec/bin" + cp "$(dirname "$0")/python3" "$MLX_TEST_BREW/libexec/bin/python3" + ;; + *) exit 1 ;; +esac +EOF + chmod +x "$BATS_TEST_TMPDIR/bin/uname" "$BATS_TEST_TMPDIR/bin/mount" \ + "$BATS_TEST_TMPDIR/bin/python3" "$BATS_TEST_TMPDIR/bin/brew" +} + +global_task() { + local task_bin + task_bin="$(command -v task)" || return + cd "$BATS_TEST_TMPDIR/work" || return + HOME="$TEST_HOME" PATH="$BATS_TEST_TMPDIR/bin:${MLX_TEST_PATH:-$PATH}" \ + "$task_bin" --global "$@" "MLX_CACHE_DIR=$CACHE_DIR" "MLX_VOLUME=$MLX_TEST_VOLUME" +} + +@test "global-mlx-chat: Python behavior is validated without MLX or model weights" { + command -v python3 >/dev/null || skip "Python 3 not installed" + run python3 -B -m unittest discover -s "$REPO_ROOT/tests/python" -p test_mlx_chat.py + [ "$status" -eq 0 ] +} + +@test "global-mlx-chat: global task works outside the repository with spaces in paths" { + command -v task >/dev/null || skip "Task not installed" + run global_task llm:chat + [ "$status" -eq 0 ] + [ "$(cat "$PYTHON_CALLS")" = "<$TEST_HOME/.config/mlx/chat.py> +<--cache-dir> +<$CACHE_DIR> +<--volume> +<$MLX_TEST_VOLUME>" ] + [ "$(cat "$MLX_TEST_ACTIONS")" = "mount +check-env +chat" ] +} + +@test "global-mlx-chat: listing tasks does not require the venv or external SSD" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" + rmdir "$CACHE_DIR/hub" "$CACHE_DIR/xet" + export MOCK_OS=Linux MOCK_ARCH=x86_64 + run global_task --list + [ "$status" -eq 0 ] + [[ "$output" == *"llm:chat"* ]] + [ ! -e "$PYTHON_CALLS" ] + [ ! -e "$MLX_TEST_ACTIONS" ] +} + +@test "global-mlx-chat: available Bash, Zsh, and Fish shells use the same global command" { + command -v task >/dev/null || skip "Task not installed" + local shell shell_bin task_bin options + task_bin="$(command -v task)" + cd "$BATS_TEST_TMPDIR/work" + for shell in bash zsh fish; do + shell_bin="$(command -v "$shell")" || continue + case "$shell" in + bash) options=(--noprofile --norc) ;; + zsh) options=(-f) ;; + fish) options=(--no-config) ;; + esac + run env HOME="$TEST_HOME" PATH="$BATS_TEST_TMPDIR/bin:$PATH" \ + MLX_TEST_TASK="$task_bin" MLX_TEST_CACHE_DIR="$CACHE_DIR" \ + "$shell_bin" "${options[@]}" -c \ + '"$MLX_TEST_TASK" --global llm:chat "MLX_CACHE_DIR=$MLX_TEST_CACHE_DIR" "MLX_VOLUME=$MLX_TEST_VOLUME"' + [ "$status" -eq 0 ] + [ "$(cat "$PYTHON_CALLS")" = "<$TEST_HOME/.config/mlx/chat.py> +<--cache-dir> +<$CACHE_DIR> +<--volume> +<$MLX_TEST_VOLUME>" ] + done +} + +@test "global-mlx-chat: unsupported systems fail explicitly before starting Python" { + command -v task >/dev/null || skip "Task not installed" + for MOCK_OS in Linux Darwin Windows; do + export MOCK_OS MOCK_ARCH=x86_64 + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"requires native Apple Silicon"* ]] + [ ! -e "$PYTHON_CALLS" ] + done +} + +@test "global-mlx-chat: broken existing venv is not overwritten" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"Repair it or move it aside"* ]] + [ -d "$TEST_HOME/mlx-qwen/.venv" ] + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: disconnected SSD fails instead of using an internal cache" { + command -v task >/dev/null || skip "Task not installed" + export MOCK_MOUNTED=0 + printf 'missing\n' >"$MLX_TEST_STATE" + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"Connect the SSD"* ]] + [ ! -e "$PYTHON_CALLS" ] + [ "$(cat "$MLX_TEST_ACTIONS")" = mount ] +} + +@test "global-mlx-chat: missing launcher gives selective deployment guidance" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/.config/mlx/chat.py" + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"Apply the global Taskfile and launcher together"* ]] + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: Python errors propagate through Task" { + command -v task >/dev/null || skip "Task not installed" + export MOCK_PYTHON_EXIT=23 + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"exit status 23"* ]] +} + +@test "global-mlx-chat: missing venv is created once and working installs are not upgraded" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" + rmdir "$TEST_HOME/mlx-qwen/.venv/bin" "$TEST_HOME/mlx-qwen/.venv" + printf 'missing\n' >"$MLX_TEST_STATE" + run global_task llm:chat + [ "$status" -eq 0 ] + [[ "$(cat "$MLX_TEST_ACTIONS")" == mount*create-venv*pip:*check-env*chat ]] + grep -F -- "--isolated --require-virtualenv" "$MLX_TEST_ACTIONS" + grep -F -- "--cache-dir $CACHE_DIR/pip install mlx-lm>=0.32.2" "$MLX_TEST_ACTIONS" + grep -F -- "tmp: $CACHE_DIR/tmp" "$MLX_TEST_ACTIONS" + ! grep -F -- "--upgrade" "$MLX_TEST_ACTIONS" + + : >"$MLX_TEST_ACTIONS" + run global_task llm:chat + [ "$status" -eq 0 ] + [ "$(cat "$MLX_TEST_ACTIONS")" = "mount +check-env +chat" ] +} + +@test "global-mlx-chat: missing packages and pip are installed into the existing venv" { + command -v task >/dev/null || skip "Task not installed" + printf 'missing\n' >"$MLX_TEST_STATE" + export MOCK_PIP_EXIT=1 + run global_task llm:chat + [ "$status" -eq 0 ] + [[ "$(cat "$MLX_TEST_ACTIONS")" == *check-pip*ensurepip*pip:*check-env*chat ]] + ! grep -F -- "create-venv" "$MLX_TEST_ACTIONS" + ! grep -F -- "install-python" "$MLX_TEST_ACTIONS" +} + +@test "global-mlx-chat: unusable packages after pip fail the post-install check" { + command -v task >/dev/null || skip "Task not installed" + printf 'broken\n' >"$MLX_TEST_STATE" + export MOCK_STILL_BROKEN=1 + run global_task llm:chat + [ "$status" -ne 0 ] + grep -F -- "pip:" "$MLX_TEST_ACTIONS" + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: pip failures stop before model launch" { + command -v task >/dev/null || skip "Task not installed" + printf 'missing\n' >"$MLX_TEST_STATE" + export MOCK_INSTALL_EXIT=17 + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"exit status 17"* ]] + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: missing base Python uses the repository Homebrew formula" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" + rmdir "$TEST_HOME/mlx-qwen/.venv/bin" "$TEST_HOME/mlx-qwen/.venv" + printf 'missing\n' >"$MLX_TEST_STATE" + export MOCK_BASE_PYTHON=0 + run global_task llm:chat + [ "$status" -eq 0 ] + [[ "$(cat "$MLX_TEST_ACTIONS")" == mount*install-python*create-venv*pip:*chat ]] +} + +@test "global-mlx-chat: absent Python and Homebrew are reported without unsafe bootstrap" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" "$BATS_TEST_TMPDIR/bin/brew" + rmdir "$TEST_HOME/mlx-qwen/.venv/bin" "$TEST_HOME/mlx-qwen/.venv" + printf 'missing\n' >"$MLX_TEST_STATE" + export MOCK_BASE_PYTHON=0 MLX_TEST_PATH=/usr/bin:/bin:/usr/sbin:/sbin + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"cannot safely bootstrap Homebrew"* ]] + ! grep -F -- "install-python" "$MLX_TEST_ACTIONS" + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: existing Homebrew Python is reused when PATH Python is unsuitable" { + command -v task >/dev/null || skip "Task not installed" + rm "$TEST_HOME/mlx-qwen/.venv/bin/python" + rmdir "$TEST_HOME/mlx-qwen/.venv/bin" "$TEST_HOME/mlx-qwen/.venv" + mkdir -p "$MLX_TEST_BREW/libexec/bin" + cp "$BATS_TEST_TMPDIR/bin/python3" "$MLX_TEST_BREW/libexec/bin/python3" + printf 'missing\n' >"$MLX_TEST_STATE" + export MOCK_BASE_PYTHON=0 + run global_task llm:chat + [ "$status" -eq 0 ] + grep -F -- "create-venv" "$MLX_TEST_ACTIONS" + ! grep -F -- "install-python" "$MLX_TEST_ACTIONS" +} + +@test "global-mlx-chat: cache outside the external volume is rejected before setup" { + command -v task >/dev/null || skip "Task not installed" + CACHE_DIR="$BATS_TEST_TMPDIR/internal" + run global_task llm:chat + [ "$status" -ne 0 ] + [[ "$output" == *"must be a directory beneath MLX_VOLUME"* ]] + [ ! -d "$CACHE_DIR" ] + [ ! -e "$PYTHON_CALLS" ] +} + +@test "global-mlx-chat: cache directories are initialized on a mounted SSD" { + command -v task >/dev/null || skip "Task not installed" + rmdir "$CACHE_DIR/hub" "$CACHE_DIR/xet" + run global_task llm:chat + [ "$status" -eq 0 ] + [ -d "$CACHE_DIR/hub" ] + [ -d "$CACHE_DIR/xet" ] + [ -d "$CACHE_DIR/pip" ] + [ -d "$CACHE_DIR/tmp" ] +} + +@test "global-mlx-chat: documented selective Chezmoi deployment manages both files from a fresh home" { + command -v chezmoi >/dev/null || skip "Chezmoi not installed" + local chezmoi_bin deployment_home + chezmoi_bin="$(command -v chezmoi)" + deployment_home="$BATS_TEST_TMPDIR/fresh home" + mkdir "$deployment_home" + printf 'data: {}\n' >"$BATS_TEST_TMPDIR/chezmoi.yaml" + [ ! -e "$deployment_home/.config" ] + [ ! -e "$deployment_home/Taskfile.yml" ] + + mkdir -p "$deployment_home/.config/mlx" + run env HOME="$deployment_home" "$chezmoi_bin" --config="$BATS_TEST_TMPDIR/chezmoi.yaml" \ + --destination="$deployment_home" --source="$REPO_ROOT" apply --exclude=scripts --dry-run --verbose \ + "$deployment_home/Taskfile.yml" "$deployment_home/.config/mlx/chat.py" + [ "$status" -eq 0 ] + [[ "$output" == *"Taskfile.yml"* ]] + [[ "$output" == *".config/mlx/chat.py"* ]] + [[ "$output" == *"llm:chat"* ]] + [ ! -e "$deployment_home/Taskfile.yml" ] + [ ! -e "$deployment_home/.config/mlx/chat.py" ] + + run env HOME="$deployment_home" "$chezmoi_bin" --config="$BATS_TEST_TMPDIR/chezmoi.yaml" \ + --destination="$deployment_home" --source="$REPO_ROOT" apply --exclude=scripts \ + "$deployment_home/Taskfile.yml" "$deployment_home/.config/mlx/chat.py" + [ "$status" -eq 0 ] + cmp "$REPO_ROOT/home/Taskfile.yml" "$deployment_home/Taskfile.yml" + cmp "$REPO_ROOT/home/dot_config/mlx/chat.py" "$deployment_home/.config/mlx/chat.py" + [ ! -e "$deployment_home/.vimrc" ] + [ ! -e "$deployment_home/.config/fish/config.fish" ] + [ ! -e "$deployment_home/mlx-qwen" ] + [ -x "$TEST_HOME/mlx-qwen/.venv/bin/python" ] + [ -d "$CACHE_DIR/hub" ] + [ ! -e "$PYTHON_CALLS" ] + [ ! -e "$MLX_TEST_ACTIONS" ] +} diff --git a/tests/python/test_mlx_chat.py b/tests/python/test_mlx_chat.py new file mode 100644 index 00000000..50d1270e --- /dev/null +++ b/tests/python/test_mlx_chat.py @@ -0,0 +1,384 @@ +"""Model-free tests for the global MLX launcher (standard library only).""" + +from contextlib import redirect_stdout, redirect_stderr +import importlib.util +import io +import os +from pathlib import Path +import sys +import tempfile +from types import ModuleType, SimpleNamespace +import unittest +from unittest.mock import Mock, patch + + +SCRIPT = Path(__file__).resolve().parents[2] / "home/dot_config/mlx/chat.py" +SPEC = importlib.util.spec_from_file_location("mlx_chat", SCRIPT) +launcher = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(launcher) +CHECK_ENVIRONMENT = launcher.check_environment + + +class LocalEntryNotFoundError(Exception): + pass + + +class ChatTests(unittest.TestCase): + def setUp(self): + self.lengths = {} + self.calls = [] + self.responses = [] + self.closed_streams = 0 + self.model = object() + self.tokenizer = Mock() + self.tokenizer.eos_token_ids = {0, 2} + self.tokenizer.encode.side_effect = lambda text, **_: { + "<|im_end|>": [2], "\n": [3] + }[text] + self.tokenizer.apply_chat_template.side_effect = ( + lambda messages, **_: [11] * self.lengths.get(messages[0]["content"], 10) + ) + self.modules = { + name: ModuleType(name) + for name in ( + "mlx", "mlx.core", "mlx_lm", "mlx_lm.models", + "mlx_lm.models.cache", "mlx_lm.sample_utils", + "huggingface_hub", "huggingface_hub.errors", + "packaging", "packaging.version", + ) + } + self.mx = self.modules["mlx.core"] + self.modules["mlx"].core = self.mx + self.mx.clear_cache = Mock() + self.load = self.modules["mlx_lm"].load = Mock( + return_value=(self.model, self.tokenizer) + ) + self.modules["mlx_lm"].stream_generate = self.stream + self.make_cache = self.modules["mlx_lm.models.cache"].make_prompt_cache = Mock( + side_effect=lambda _: [SimpleNamespace(tokens=[])] + ) + self.sampler = self.modules["mlx_lm.sample_utils"].make_sampler = Mock() + self.snapshot = self.modules["huggingface_hub"].snapshot_download = Mock( + return_value="/mock/external/snapshot" + ) + self.modules["huggingface_hub.errors"].LocalEntryNotFoundError = ( + LocalEntryNotFoundError + ) + self.modules["packaging.version"].Version = ( + lambda text: tuple(int(part) for part in text.split(".")) + ) + self.module_patch = patch.dict(sys.modules, self.modules) + self.module_patch.start() + self.addCleanup(self.module_patch.stop) + self.environment_patch = patch.object(launcher, "check_environment") + self.environment_check = self.environment_patch.start() + self.addCleanup(self.environment_patch.stop) + self.output = io.StringIO() + self.redirect = redirect_stdout(self.output) + self.redirect.__enter__() + self.addCleanup(self.redirect.__exit__, None, None, None) + + def stream(self, model, tokenizer, prompt, **kwargs): + self.assertIs(model, self.model) + self.assertIs(tokenizer, self.tokenizer) + cache = kwargs["prompt_cache"] + self.calls.append({ + "prompt": prompt[:], "before": len(cache[0].tokens), **kwargs + }) + cache[0].tokens.extend(prompt) + tokens, reason = self.responses.pop(0) if self.responses else ([42, 2], "stop") + try: + for i, token in enumerate(tokens, 1): + if isinstance(token, BaseException): + raise token + # Match mlx-lm 0.32.2's one-token lookahead, including cached EOS. + cache[0].tokens.append(token) + yield SimpleNamespace( + text="" if token in self.tokenizer.eos_token_ids else "answer", + token=token, generation_tokens=i, generation_tps=12.5, + peak_memory=6.25, + finish_reason=reason if i == len(tokens) else None, + ) + finally: + self.closed_streams += 1 + + def run_chat(self, inputs): + with patch("builtins.input", side_effect=inputs): + launcher.chat(self.model, self.tokenizer) + return self.output.getvalue() + + def test_no_thinking_sampling_and_cache_continuity(self): + output = self.run_chat(["first", "second", "q"]) + self.sampler.assert_called_once_with(temp=0.7, top_p=0.8, top_k=20) + for call in self.tokenizer.apply_chat_template.call_args_list: + self.assertEqual(call.kwargs, { + "tokenize": True, "add_generation_prompt": True, "enable_thinking": False + }) + self.assertEqual(len(call.args[0]), 1) + first, second = self.calls + self.assertIs(first["prompt_cache"], second["prompt_cache"]) + self.assertEqual(first["prompt"], [11] * 10) + self.assertEqual(second["prompt"], [3] + [11] * 10) + self.assertEqual(second["before"], 12) + for call in self.calls: + self.assertEqual(call["max_tokens"], 512) + self.assertEqual(call["prefill_step_size"], 512) + self.assertIs(call["sampler"], self.sampler.return_value) + self.assertEqual(self.tokenizer.detokenizer.reset.call_count, 2) + self.assertIn("12.5 tok/s; peak 6.25 GB", output) + self.assertIn("context 25/8192", output) + self.assertEqual(self.closed_streams, 2) + + def test_reset_releases_history_and_starts_without_separator(self): + output = self.run_chat(["first", "r", "second", "q"]) + self.assertEqual(self.make_cache.call_count, 2) + self.assertIsNot(self.calls[0]["prompt_cache"], self.calls[1]["prompt_cache"]) + self.assertEqual(self.calls[1]["before"], 0) + self.assertEqual(self.calls[1]["prompt"], [11] * 10) + self.assertEqual(self.mx.clear_cache.call_count, 2) + self.assertIn("Conversation reset.", output) + + def test_exact_8192_token_boundary_is_admitted(self): + self.lengths["boundary"] = 8192 - 512 + self.responses = [([42] * 512, "length")] + output = self.run_chat(["boundary", "q"]) + self.assertEqual(len(self.calls), 1) + self.assertIn("context 8192/8192", output) + + def test_oversized_prompt_is_rejected_without_mutating_history(self): + self.lengths["oversized"] = 8192 - 512 + 1 + output = self.run_chat(["oversized", "shorter", "q"]) + self.assertEqual(len(self.calls), 1) + self.assertEqual(self.calls[0]["before"], 0) + self.assertIn("8193/8192", output) + self.assertIn("history has not been dropped", output) + + def test_truncation_boundary_and_cached_history_count_toward_limit(self): + self.lengths.update(first=100, oversized=7067, boundary=7066) + self.responses = [([42] * 512, "length"), ([42, 2], "stop")] + output = self.run_chat(["first", "oversized", "boundary", "q"]) + self.assertEqual(len(self.calls), 2) + self.assertEqual(self.calls[1]["before"], 612) + self.assertEqual(self.calls[1]["prompt"], [2, 3] + [11] * 7066) + self.assertIs(self.calls[0]["prompt_cache"], self.calls[1]["prompt_cache"]) + self.assertIn("8193/8192", output) + self.assertIn("Reply reached the 512-token limit", output) + + def test_alternate_eos_also_closes_the_assistant_turn(self): + self.responses = [([0], "stop")] + self.run_chat(["first", "second", "q"]) + self.assertEqual(self.calls[1]["prompt"], [2, 3] + [11] * 10) + + def test_commands_and_blank_input_do_not_generate(self): + output = self.run_chat([" ", "h", "q"]) + self.assertFalse(self.calls) + self.assertIn(launcher.HELP, output) + + def test_eof_and_keyboard_interrupt_at_input_exit_cleanly(self): + for signal in (EOFError(), KeyboardInterrupt()): + with self.subTest(signal=type(signal).__name__): + output = self.run_chat([signal]) + self.assertIn("Bye.", output) + self.assertFalse(self.calls) + + def test_interrupt_during_generation_closes_stream_and_releases_cache(self): + self.responses = [([42, KeyboardInterrupt()], "stop")] + output = self.run_chat(["first"]) + self.assertEqual(self.closed_streams, 1) + self.assertEqual(self.calls[0]["prompt_cache"], []) + self.mx.clear_cache.assert_called_once_with() + self.assertIn("Bye.", output) + + def test_generation_errors_are_not_hidden(self): + self.responses = [([RuntimeError("GPU failure")], "stop")] + with self.assertRaisesRegex(RuntimeError, "GPU failure"): + self.run_chat(["first"]) + self.assertEqual(self.closed_streams, 1) + self.assertEqual(self.calls[0]["prompt_cache"], []) + self.assertNotIn("Bye.", self.output.getvalue()) + + def test_generation_eof_is_not_mistaken_for_terminal_eof(self): + self.responses = [([EOFError("truncated weights")], "stop")] + with self.assertRaisesRegex(EOFError, "truncated weights"): + self.run_chat(["first"]) + self.assertEqual(self.closed_streams, 1) + self.assertNotIn("Bye.", self.output.getvalue()) + + def test_incomplete_stream_is_an_error(self): + self.responses = [([], "stop")] + with self.assertRaisesRegex(RuntimeError, "without a completed response"): + self.run_chat(["first"]) + + def test_missing_turn_token_is_rejected(self): + self.tokenizer.eos_token_ids = {0} + with self.assertRaisesRegex(ValueError, "end-of-turn token"): + self.run_chat(["q"]) + + def run_main(self, root, system="Darwin", machine="arm64", check=False): + args = ["chat.py", "--cache-dir", str(root), "--volume", str(Path(root).parent)] + if check: + args.append("--check-env") + with ( + patch.object(sys, "argv", args), + patch.object(launcher.platform, "system", return_value=system), + patch.object(launcher.platform, "machine", return_value=machine), + patch.object(Path, "is_mount", return_value=True), + patch("builtins.input", return_value="q"), + ): + launcher.main() + + def test_main_forces_external_cache_and_uses_existing_model_without_download(self): + with tempfile.TemporaryDirectory(prefix="mlx cache ") as directory: + root = Path(directory).resolve() + (root / "hub").mkdir() + (root / "xet").mkdir() + with patch.dict(os.environ, { + "HF_HOME": "/internal", "HF_HUB_CACHE": "/internal/hub", + "HF_XET_CACHE": "/internal/xet", "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + }): + def snapshot(*args, **kwargs): + self.assertEqual(args, (launcher.MODEL,)) + self.assertEqual(kwargs, { + "cache_dir": str(root / "hub"), "local_files_only": True + }) + self.assertEqual(os.environ["HF_HOME"], str(root)) + self.assertEqual(os.environ["HF_HUB_CACHE"], str(root / "hub")) + self.assertEqual(os.environ["HF_XET_CACHE"], str(root / "xet")) + self.assertEqual(os.environ["HF_HUB_OFFLINE"], "0") + self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], "0") + return str(root / "hub/snapshot") + self.snapshot.side_effect = snapshot + self.run_main(root) + self.snapshot.assert_called_once() + self.load.assert_called_once_with( + str(root / "hub/snapshot"), + tokenizer_config={"trust_remote_code": False}, + trust_remote_code=False, + ) + self.assertEqual(sorted(p.name for p in root.iterdir()), ["hub", "xet"]) + + def test_missing_cache_fails_without_loading(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory) + for child in ("", "hub"): + if child: + (root / child).mkdir() + error = io.StringIO() + with redirect_stderr(error), self.assertRaises(SystemExit) as exit_info: + self.run_main(root) + self.assertEqual(exit_info.exception.code, 1) + self.load.assert_not_called() + self.snapshot.assert_not_called() + self.assertIn("No internal fallback", error.getvalue()) + + def test_missing_model_downloads_only_to_external_cache(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory).resolve() + (root / "hub").mkdir() + (root / "xet").mkdir() + self.snapshot.side_effect = [ + LocalEntryNotFoundError("not cached"), str(root / "hub/snapshot") + ] + self.run_main(root) + self.assertEqual(self.snapshot.call_count, 2) + self.assertEqual(self.snapshot.call_args.kwargs, { + "cache_dir": str(root / "hub"), "local_files_only": False + }) + self.assertIn("Downloading", self.output.getvalue()) + self.load.assert_called_once() + + def test_download_failure_is_not_hidden(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory) + (root / "hub").mkdir() + (root / "xet").mkdir() + self.snapshot.side_effect = [ + LocalEntryNotFoundError("not cached"), OSError("download failed") + ] + with self.assertRaisesRegex(OSError, "download failed"): + self.run_main(root) + self.load.assert_not_called() + + def test_status_check_never_resolves_or_loads_a_model(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory) + (root / "hub").mkdir() + (root / "xet").mkdir() + self.run_main(root, check=True) + self.environment_check.assert_called_once() + self.snapshot.assert_not_called() + self.load.assert_not_called() + + def test_status_check_reports_broken_packages(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory) + (root / "hub").mkdir() + (root / "xet").mkdir() + self.environment_check.side_effect = ImportError("missing dependency") + error = io.StringIO() + with redirect_stderr(error), self.assertRaises(SystemExit) as result: + self.run_main(root, check=True) + self.assertEqual(result.exception.code, 1) + self.assertIn("missing dependency", error.getvalue()) + self.snapshot.assert_not_called() + + def test_environment_check_accepts_current_and_newer_without_installing(self): + for release in ("0.32.2", "0.33.0"): + with ( + patch.object(sys, "version_info", (3, 14, 8)), + patch.object(sys, "prefix", "/mock/venv"), + patch.object(launcher, "version", return_value=release), + ): + CHECK_ENVIRONMENT() + self.load.assert_not_called() + self.snapshot.assert_not_called() + + def test_environment_check_rejects_old_packages_and_non_venv_python(self): + with ( + patch.object(sys, "version_info", (3, 14, 8)), + patch.object(sys, "prefix", "/mock/venv"), + patch.object(launcher, "version", return_value="0.31.0"), + ): + with self.assertRaisesRegex(ValueError, "0.32.2"): + CHECK_ENVIRONMENT() + with patch.object(sys, "prefix", sys.base_prefix): + with self.assertRaisesRegex(ValueError, "virtual environment"): + CHECK_ENVIRONMENT() + + def test_unmounted_or_redirected_cache_is_rejected(self): + with tempfile.TemporaryDirectory() as directory: + volume = Path(directory) / "volume" + volume.mkdir() + with patch.object(Path, "is_mount", return_value=False): + with self.assertRaisesRegex(ValueError, "Connect the external SSD"): + launcher.configure_cache(volume / "cache", volume) + outside = Path(directory) / "internal" + outside.mkdir() + with patch.object(Path, "is_mount", return_value=True): + with self.assertRaisesRegex(ValueError, "beneath the external SSD"): + launcher.configure_cache(outside, volume) + + def test_non_apple_platforms_fail_before_cache_lookup(self): + for system, machine in ( + ("Linux", "aarch64"), ("Darwin", "x86_64"), ("Windows", "ARM64") + ): + with self.subTest(system=system, machine=machine): + with redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as error: + self.run_main("/not/accessed", system, machine) + self.assertEqual(error.exception.code, 2) + self.snapshot.assert_not_called() + self.load.assert_not_called() + + def test_model_loading_failure_propagates(self): + with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ): + root = Path(directory) + (root / "hub").mkdir() + (root / "xet").mkdir() + self.load.side_effect = RuntimeError("incomplete model weights") + with self.assertRaisesRegex(RuntimeError, "incomplete model weights"): + self.run_main(root) + + +if __name__ == "__main__": + unittest.main()