diff --git a/solstone/apps/settings/mlx_bootstrap.py b/solstone/apps/settings/mlx_bootstrap.py index 7e829925f..e550274fd 100644 --- a/solstone/apps/settings/mlx_bootstrap.py +++ b/solstone/apps/settings/mlx_bootstrap.py @@ -23,9 +23,8 @@ from huggingface_hub import constants from huggingface_hub.file_download import repo_folder_name from solstone.think.providers.mlx import ( - MLX_MODEL_REPO, - MLX_MODEL_REVISION, - is_mlx_available, + _MLX_MODEL_REGISTRY, + is_mlx_available_for_model, ) logger = logging.getLogger(__name__) @@ -58,19 +57,23 @@ class MlxVerificationError(RuntimeError): """Raised when a downloaded file fails sha256 verification.""" -_STATE = MlxBootstrapState() -_STATE_LOCK = threading.Lock() +_STATES: dict[str, MlxBootstrapState] = { + name: MlxBootstrapState() for name in _MLX_MODEL_REGISTRY +} +_STATES_LOCK = threading.Lock() -def _snapshot_dir() -> Path: +def _snapshot_dir(model: str) -> Path: + spec = _MLX_MODEL_REGISTRY[model] repo_folder = repo_folder_name( - repo_id=MLX_MODEL_REPO, + repo_id=spec.repo, repo_type="model", ) - return Path(constants.HF_HUB_CACHE) / repo_folder / "snapshots" / MLX_MODEL_REVISION + return Path(constants.HF_HUB_CACHE) / repo_folder / "snapshots" / spec.revision -def _safetensors_paths(snapshot_dir: Path) -> list[str]: +def _safetensors_paths(model: str) -> list[str]: + snapshot_dir = _snapshot_dir(model) index_path = snapshot_dir / "model.safetensors.index.json" data = json.loads(index_path.read_text(encoding="utf-8")) weight_map = data.get("weight_map") @@ -82,14 +85,14 @@ def _safetensors_paths(snapshot_dir: Path) -> list[str]: return paths -def check_model_present() -> bool: +def check_model_present(model: str) -> bool: """Return whether the pinned MLX snapshot is structurally present.""" - snapshot_dir = _snapshot_dir() + snapshot_dir = _snapshot_dir(model) index_path = snapshot_dir / "model.safetensors.index.json" if not snapshot_dir.is_dir() or not index_path.is_file(): return False try: - for rel_path in _safetensors_paths(snapshot_dir): + for rel_path in _safetensors_paths(model): file_path = snapshot_dir / rel_path if not file_path.is_file() or file_path.stat().st_size <= 0: return False @@ -104,10 +107,11 @@ def _is_package_installed(package: str) -> bool: return importlib.util.find_spec(package) is not None -def get_availability_payload() -> dict[str, bool | float | str]: +def get_availability_payload(model: str) -> dict[str, bool | float | int | str]: """Return the MLX availability payload used by Settings.""" - ok, reason = is_mlx_available() - model_present = check_model_present() + spec = _MLX_MODEL_REGISTRY[model] + ok, reason = is_mlx_available_for_model(spec) + model_present = check_model_present(model) available = ok and model_present if ok and not model_present: reason = "model snapshot not present" @@ -116,112 +120,126 @@ def get_availability_payload() -> dict[str, bool | float | str]: total_memory_gb = round(psutil.virtual_memory().total / 1024**3, 1) return { + "model": model, "is_apple_silicon": platform.system() == "Darwin" and platform.machine() == "arm64", "total_memory_gb": total_memory_gb, "mlx_installed": _is_package_installed("mlx_vlm"), + "min_ram_gb": spec.min_ram_bytes // 1024**3, "model_present": model_present, "available": available, "reason": reason, } -def _serialize_state_locked() -> dict[str, int | str | None]: +def _serialize_state_locked(model: str) -> dict[str, int | str | None]: + state = _STATES[model] return { - "state": _STATE.state, - "received_bytes": int(_STATE.received_bytes), - "total_bytes": int(_STATE.total_bytes), - "message": _STATE.message, + "state": state.state, + "received_bytes": int(state.received_bytes), + "total_bytes": int(state.total_bytes), + "message": state.message, } -def _mark_installed_locked() -> None: - _STATE.state = "installed" - _STATE.message = None - _STATE.thread = None - _STATE.last_progress_at = time.monotonic() +def _mark_installed_locked(model: str) -> None: + state = _STATES[model] + state.state = "installed" + state.message = None + state.thread = None + state.last_progress_at = time.monotonic() -def _mark_downloading_locked(thread: threading.Thread, now: float) -> None: - _STATE.state = "downloading" - _STATE.received_bytes = 0 - _STATE.total_bytes = 0 - _STATE.started_at = now - _STATE.last_progress_at = now - _STATE.message = None - _STATE.thread = thread +def _mark_downloading_locked(model: str, thread: threading.Thread, now: float) -> None: + state = _STATES[model] + state.state = "downloading" + state.received_bytes = 0 + state.total_bytes = 0 + state.started_at = now + state.last_progress_at = now + state.message = None + state.thread = thread -def _mark_verifying() -> None: +def _mark_verifying(model: str) -> None: now = time.monotonic() - with _STATE_LOCK: - if _STATE.state != "downloading": + with _STATES_LOCK: + state = _STATES[model] + if state.state != "downloading": return - _STATE.state = "verifying" - _STATE.received_bytes = 0 - _STATE.total_bytes = 0 - _STATE.last_progress_at = now - _STATE.message = None + state.state = "verifying" + state.received_bytes = 0 + state.total_bytes = 0 + state.last_progress_at = now + state.message = None -def _mark_failed_locked(message: str) -> None: - _STATE.state = "failed" - _STATE.message = message - _STATE.last_progress_at = time.monotonic() +def _mark_failed_locked(model: str, message: str) -> None: + state = _STATES[model] + state.state = "failed" + state.message = message + state.last_progress_at = time.monotonic() -def _add_progress(received_delta: int, total: int | None = None) -> None: +def _add_progress(model: str, received_delta: int, total: int | None = None) -> None: if received_delta < 0: received_delta = 0 now = time.monotonic() - with _STATE_LOCK: - _STATE.received_bytes += received_delta + with _STATES_LOCK: + state = _STATES[model] + state.received_bytes += received_delta if total is not None: - _STATE.total_bytes = max(0, int(total)) - _STATE.last_progress_at = now + state.total_bytes = max(0, int(total)) + state.last_progress_at = now -def _set_progress_total(total: int) -> None: - with _STATE_LOCK: - _STATE.total_bytes = max(0, int(total)) - _STATE.last_progress_at = time.monotonic() +def _set_progress_total(model: str, total: int) -> None: + with _STATES_LOCK: + state = _STATES[model] + state.total_bytes = max(0, int(total)) + state.last_progress_at = time.monotonic() -def _set_verify_progress(received_bytes: int, total_bytes: int) -> None: - with _STATE_LOCK: - if _STATE.state != "verifying": +def _set_verify_progress(model: str, received_bytes: int, total_bytes: int) -> None: + with _STATES_LOCK: + state = _STATES[model] + if state.state != "verifying": return - _STATE.received_bytes = max(0, int(received_bytes)) - _STATE.total_bytes = max(0, int(total_bytes)) - _STATE.last_progress_at = time.monotonic() + state.received_bytes = max(0, int(received_bytes)) + state.total_bytes = max(0, int(total_bytes)) + state.last_progress_at = time.monotonic() -def _observe_stall_locked(now: float) -> None: - if _STATE.state not in ("downloading", "verifying"): +def _observe_stall_locked(model: str, now: float) -> None: + state = _STATES[model] + if state.state not in ("downloading", "verifying"): return - thread = _STATE.thread + thread = state.thread if thread is not None and thread.is_alive(): return - last_progress = _STATE.last_progress_at or _STATE.started_at + last_progress = state.last_progress_at or state.started_at if last_progress is None or now - last_progress <= _STALL_SECONDS: return - _mark_failed_locked(f"{_STATE.state} stalled with no progress") + _mark_failed_locked(model, f"{state.state} stalled with no progress") -def get_state() -> dict[str, int | str | None]: +def get_state(model: str) -> dict[str, int | str | None]: """Return the serialized bootstrap state, applying stall detection.""" - with _STATE_LOCK: - _observe_stall_locked(time.monotonic()) - return _serialize_state_locked() + with _STATES_LOCK: + _observe_stall_locked(model, time.monotonic()) + return _serialize_state_locked(model) -def _remote_safetensors_metadata(paths: list[str]) -> dict[str, tuple[str, int]]: +def _remote_safetensors_metadata( + model: str, paths: list[str] +) -> dict[str, tuple[str, int]]: + spec = _MLX_MODEL_REGISTRY[model] wanted = set(paths) found: dict[str, tuple[str, int]] = {} api = huggingface_hub.HfApi() for entry in api.list_repo_tree( - repo_id=MLX_MODEL_REPO, - revision=MLX_MODEL_REVISION, + repo_id=spec.repo, + revision=spec.revision, repo_type="model", recursive=True, ): @@ -236,13 +254,13 @@ def _remote_safetensors_metadata(paths: list[str]) -> dict[str, tuple[str, int]] return found -def _verify_safetensors_sha256_hashes() -> None: - snapshot_dir = _snapshot_dir() - safetensors_paths = _safetensors_paths(snapshot_dir) - metadata = _remote_safetensors_metadata(safetensors_paths) +def _verify_safetensors_sha256_hashes(model: str) -> None: + snapshot_dir = _snapshot_dir(model) + safetensors_paths = _safetensors_paths(model) + metadata = _remote_safetensors_metadata(model, safetensors_paths) total_bytes = sum(size for _sha, size in metadata.values()) hashed_total = 0 - _set_verify_progress(0, total_bytes) + _set_verify_progress(model, 0, total_bytes) for rel_path in safetensors_paths: expected_sha, _expected_size = metadata[rel_path] @@ -253,23 +271,25 @@ def _verify_safetensors_sha256_hashes() -> None: for chunk in iter(lambda: handle.read(_HASH_CHUNK_SIZE), b""): digest.update(chunk) file_received += len(chunk) - _set_verify_progress(hashed_total + file_received, total_bytes) + _set_verify_progress(model, hashed_total + file_received, total_bytes) actual_sha = digest.hexdigest() if actual_sha != expected_sha: raise MlxVerificationError(f"sha256 mismatch for {rel_path}") hashed_total += file_received - _set_verify_progress(hashed_total, total_bytes) + _set_verify_progress(model, hashed_total, total_bytes) class _BootstrapTqdm: + _model = next(iter(_MLX_MODEL_REGISTRY)) + def __init__(self, *args, **kwargs): self._track_bytes = kwargs.get("unit") == "B" self._total = int(kwargs.get("total") or 0) if self._track_bytes: - _set_progress_total(self._total) + _set_progress_total(self._model, self._total) initial = int(kwargs.get("initial") or 0) if initial: - _add_progress(initial) + _add_progress(self._model, initial) @property def total(self) -> int: @@ -279,11 +299,11 @@ class _BootstrapTqdm: def total(self, value: int | float | None) -> None: self._total = int(value or 0) if self._track_bytes: - _set_progress_total(self._total) + _set_progress_total(self._model, self._total) def update(self, n: int | float | None = 1) -> None: if self._track_bytes: - _add_progress(int(n or 0)) + _add_progress(self._model, int(n or 0)) def refresh(self) -> None: return None @@ -302,70 +322,78 @@ class _BootstrapTqdm: return False -def start_bootstrap() -> tuple[dict[str, str], int]: +def start_bootstrap(model: str) -> tuple[dict[str, str], int]: """Start the MLX model bootstrap worker if needed.""" - ok, reason = is_mlx_available() + ok, reason = is_mlx_available_for_model(_MLX_MODEL_REGISTRY[model]) if not ok: raise MlxBootstrapUnavailableError(reason) - present = check_model_present() - with _STATE_LOCK: + present = check_model_present(model) + with _STATES_LOCK: + state = _STATES[model] if ( - _STATE.state in ("downloading", "verifying") - and _STATE.thread - and _STATE.thread.is_alive() + state.state in ("downloading", "verifying") + and state.thread + and state.thread.is_alive() ): - return {"state": _STATE.state}, 200 - retry_after_failed = _STATE.state == "failed" + return {"state": state.state}, 200 + retry_after_failed = state.state == "failed" if present and not retry_after_failed: - _mark_installed_locked() + _mark_installed_locked(model) return {"state": "installed"}, 200 - with _STATE_LOCK: + with _STATES_LOCK: + state = _STATES[model] if ( - _STATE.state in ("downloading", "verifying") - and _STATE.thread - and _STATE.thread.is_alive() + state.state in ("downloading", "verifying") + and state.thread + and state.thread.is_alive() ): - return {"state": _STATE.state}, 200 - retry_after_failed = retry_after_failed or _STATE.state == "failed" - if not retry_after_failed and check_model_present(): - _mark_installed_locked() + return {"state": state.state}, 200 + retry_after_failed = retry_after_failed or state.state == "failed" + if not retry_after_failed and check_model_present(model): + _mark_installed_locked(model) return {"state": "installed"}, 200 try: thread = threading.Thread( target=_run_bootstrap_worker, - name="mlx-model-bootstrap", + args=(model,), + name=f"mlx-model-bootstrap-{model}", daemon=True, ) except Exception as exc: - _mark_failed_locked(str(exc)) + _mark_failed_locked(model, str(exc)) raise MlxBootstrapStartError(str(exc)) from exc - _mark_downloading_locked(thread, time.monotonic()) + _mark_downloading_locked(model, thread, time.monotonic()) try: thread.start() except Exception as exc: - with _STATE_LOCK: - _mark_failed_locked(str(exc)) + with _STATES_LOCK: + _mark_failed_locked(model, str(exc)) raise MlxBootstrapStartError(str(exc)) from exc return {"state": "downloading"}, 202 -def _run_bootstrap_worker() -> None: +def _run_bootstrap_worker(model: str) -> None: + spec = _MLX_MODEL_REGISTRY[model] + + class _ModelBoundTqdm(_BootstrapTqdm): + _model = model + try: # v1.15.0 resumes via .incomplete files automatically; no resume_download kwarg. huggingface_hub.snapshot_download( - repo_id=MLX_MODEL_REPO, - revision=MLX_MODEL_REVISION, - tqdm_class=_BootstrapTqdm, + repo_id=spec.repo, + revision=spec.revision, + tqdm_class=_ModelBoundTqdm, ) - _mark_verifying() - _verify_safetensors_sha256_hashes() - with _STATE_LOCK: - if _STATE.state == "verifying": - _mark_installed_locked() + _mark_verifying(model) + _verify_safetensors_sha256_hashes(model) + with _STATES_LOCK: + if _STATES[model].state == "verifying": + _mark_installed_locked(model) except Exception as exc: logger.exception("MLX model bootstrap failed") - with _STATE_LOCK: - _mark_failed_locked(str(exc)) + with _STATES_LOCK: + _mark_failed_locked(model, str(exc)) diff --git a/solstone/apps/settings/routes.py b/solstone/apps/settings/routes.py index fa1bd521c..956be26b9 100644 --- a/solstone/apps/settings/routes.py +++ b/solstone/apps/settings/routes.py @@ -53,8 +53,10 @@ from solstone.convey.sol_initiated.settings import ( save_settings as save_sol_voice_settings, ) from solstone.convey.utils import error_response +from solstone.think.models import QWEN_35_9B from solstone.think.pairing.config import get_host_url from solstone.think.providers.google import validate_vertex_credentials +from solstone.think.providers.mlx import _MLX_MODEL_REGISTRY from solstone.think.retention import ( _human_bytes, check_storage_health, @@ -592,12 +594,37 @@ def _sol_voice_response(settings: SolVoiceSettings) -> dict[str, Any]: # --------------------------------------------------------------------------- VALID_TIERS = {1, 2, 3} +MLX_MODEL_LABELS = { + QWEN_35_9B: "qwen 3.5 — 16 GB Mac", + "gemma-4-26b-a4b-it-mlx-4bit": "gemma 4 (26B) — 24 GB Mac", +} + + +def _mlx_model_error(model: str) -> Any: + return error_response( + INVALID_REQUEST_VALUE, + detail=( + f"Unknown MLX model: {model}. " + f"Must be one of: {', '.join(_MLX_MODEL_REGISTRY.keys())}" + ), + ) + + +def _mlx_model_from_request() -> tuple[str | None, Any | None]: + model = request.args.get("model") or QWEN_35_9B + if model not in _MLX_MODEL_REGISTRY: + return None, _mlx_model_error(model) + return model, None @settings_bp.route("/api/mlx/availability") def get_mlx_availability() -> Any: try: - return jsonify(mlx_bootstrap.get_availability_payload()) + model, error = _mlx_model_from_request() + if error is not None: + return error + assert model is not None + return jsonify(mlx_bootstrap.get_availability_payload(model)) except Exception: logger.exception("error loading MLX availability") return _settings_operation_failed() @@ -606,7 +633,11 @@ def get_mlx_availability() -> Any: @settings_bp.route("/api/mlx/bootstrap", methods=["POST"]) def start_mlx_bootstrap() -> Any: try: - payload, status = mlx_bootstrap.start_bootstrap() + model, error = _mlx_model_from_request() + if error is not None: + return error + assert model is not None + payload, status = mlx_bootstrap.start_bootstrap(model) return jsonify(payload), status except mlx_bootstrap.MlxBootstrapUnavailableError as exc: return error_response(INVALID_REQUEST_VALUE, detail=str(exc)) @@ -621,12 +652,34 @@ def start_mlx_bootstrap() -> Any: @settings_bp.route("/api/mlx/bootstrap/status") def get_mlx_bootstrap_status() -> Any: try: - return jsonify(mlx_bootstrap.get_state()) + model, error = _mlx_model_from_request() + if error is not None: + return error + assert model is not None + return jsonify(mlx_bootstrap.get_state(model)) except Exception: logger.exception("error loading MLX bootstrap status") return _settings_operation_failed() +@settings_bp.route("/api/mlx/models") +def get_mlx_models() -> Any: + try: + return jsonify( + [ + { + "name": name, + "label": MLX_MODEL_LABELS[name], + "min_ram_gb": spec.min_ram_bytes // 1024**3, + } + for name, spec in _MLX_MODEL_REGISTRY.items() + ] + ) + except Exception: + logger.exception("error loading MLX models") + return _settings_operation_failed() + + @settings_bp.route("/api/providers") def get_providers() -> Any: """Return providers configuration with context defaults and API key status. @@ -731,6 +784,10 @@ def get_providers() -> Any: provider: bundled.get_provider_state(provider) for provider in ("anthropic", "openai") } + mlx_config = providers_config.get("mlx", {}) + mlx_active_model = ( + mlx_config.get("active_model") if isinstance(mlx_config, dict) else None + ) or QWEN_35_9B return jsonify( { @@ -744,6 +801,7 @@ def get_providers() -> Any: "auth": auth, "key_validation": key_validation, "bundled": bundled_status, + "mlx": {"active_model": mlx_active_model}, "google_backend": providers_config.get("google_backend", "auto"), "vertex_credentials_configured": vertex_creds_configured, "vertex_credentials_email": vertex_creds_email, @@ -1117,6 +1175,48 @@ def update_providers() -> Any: } config["providers"]["contexts"][pattern] = ctx_config + # Handle MLX model selection + if "mlx" in request_data: + mlx_data = request_data["mlx"] + if not isinstance(mlx_data, dict): + return error_response( + INVALID_CONFIG_VALUE, + detail="mlx must be an object", + ) + unknown_fields = sorted(set(mlx_data) - {"active_model"}) + if unknown_fields: + return error_response( + INVALID_CONFIG_VALUE, + detail=f"Invalid mlx field: {unknown_fields[0]}", + ) + if "active_model" in mlx_data: + active_model = mlx_data["active_model"] + if not isinstance(active_model, str): + return error_response( + INVALID_CONFIG_VALUE, + detail="mlx.active_model must be a string", + ) + if active_model not in _MLX_MODEL_REGISTRY: + return error_response( + INVALID_CONFIG_VALUE, + detail=( + f"Invalid MLX model: {active_model}. " + f"Must be one of: {', '.join(_MLX_MODEL_REGISTRY.keys())}" + ), + ) + if "mlx" not in config["providers"]: + config["providers"]["mlx"] = {} + old_mlx = old_providers.get("mlx", {}) + old_model = ( + old_mlx.get("active_model") if isinstance(old_mlx, dict) else None + ) + if old_model != active_model: + changed_fields["mlx.active_model"] = { + "old": old_model, + "new": active_model, + } + config["providers"]["mlx"]["active_model"] = active_model + # Handle Google backend settings if "google_backend" in request_data: backend = request_data["google_backend"] diff --git a/solstone/apps/settings/tests/test_mlx_bootstrap_routes.py b/solstone/apps/settings/tests/test_mlx_bootstrap_routes.py index 6be285cc7..ae9b3e004 100644 --- a/solstone/apps/settings/tests/test_mlx_bootstrap_routes.py +++ b/solstone/apps/settings/tests/test_mlx_bootstrap_routes.py @@ -17,7 +17,8 @@ import tomllib from solstone.apps.settings import mlx_bootstrap from solstone.convey import create_app -from solstone.think.providers.mlx import MLX_MODEL_REPO, MLX_MODEL_REVISION +from solstone.think.models import GEMMA4_26B_A4B_4BIT, QWEN_35_9B +from solstone.think.providers.mlx import _MLX_MODEL_REGISTRY def _client(journal_path): @@ -40,14 +41,18 @@ def _settings_config() -> dict: @pytest.fixture(autouse=True) def _reset_mlx_state(monkeypatch, tmp_path): - monkeypatch.setattr(mlx_bootstrap, "_STATE", mlx_bootstrap.MlxBootstrapState()) + monkeypatch.setattr( + mlx_bootstrap, + "_STATES", + {name: mlx_bootstrap.MlxBootstrapState() for name in _MLX_MODEL_REGISTRY}, + ) monkeypatch.setattr(mlx_bootstrap.constants, "HF_HUB_CACHE", str(tmp_path / "hf")) -def _set_state(**updates): - with mlx_bootstrap._STATE_LOCK: +def _set_state(model: str = QWEN_35_9B, **updates): + with mlx_bootstrap._STATES_LOCK: for key, value in updates.items(): - setattr(mlx_bootstrap._STATE, key, value) + setattr(mlx_bootstrap._STATES[model], key, value) class _FakeThread: @@ -75,8 +80,10 @@ class _DeadThread: def test_mlx_availability_payload_exact_shape(settings_env, monkeypatch): journal_path, _config = settings_env(_settings_config()) - monkeypatch.setattr(mlx_bootstrap, "is_mlx_available", lambda: (True, "")) - monkeypatch.setattr(mlx_bootstrap, "check_model_present", lambda: True) + monkeypatch.setattr( + mlx_bootstrap, "is_mlx_available_for_model", lambda _spec: (True, "") + ) + monkeypatch.setattr(mlx_bootstrap, "check_model_present", lambda _model: True) monkeypatch.setattr(mlx_bootstrap.platform, "system", lambda: "Darwin") monkeypatch.setattr(mlx_bootstrap.platform, "machine", lambda: "arm64") monkeypatch.setattr( @@ -92,17 +99,21 @@ def test_mlx_availability_payload_exact_shape(settings_env, monkeypatch): assert response.status_code == 200 payload = response.get_json() assert set(payload) == { + "model", "is_apple_silicon", "total_memory_gb", "mlx_installed", + "min_ram_gb", "model_present", "available", "reason", } assert payload == { + "model": QWEN_35_9B, "is_apple_silicon": True, "total_memory_gb": 32.0, "mlx_installed": True, + "min_ram_gb": 16, "model_present": True, "available": True, "reason": "", @@ -122,12 +133,14 @@ def test_mlx_availability_payload_exact_shape(settings_env, monkeypatch): def test_bootstrap_post_is_idempotent_while_downloading(settings_env, monkeypatch): settings_env(_settings_config()) - monkeypatch.setattr(mlx_bootstrap, "is_mlx_available", lambda: (True, "")) + monkeypatch.setattr( + mlx_bootstrap, "is_mlx_available_for_model", lambda _spec: (True, "") + ) present_barrier = threading.Barrier(2) present_lock = threading.Lock() present_calls = 0 - def fake_check_model_present(): + def fake_check_model_present(_model): nonlocal present_calls with present_lock: present_calls += 1 @@ -149,7 +162,7 @@ def test_bootstrap_post_is_idempotent_while_downloading(settings_env, monkeypatc def call_start_bootstrap(): try: call_barrier.wait(timeout=2) - result = mlx_bootstrap.start_bootstrap() + result = mlx_bootstrap.start_bootstrap(QWEN_35_9B) with results_lock: results.append(result) except Exception as exc: # pragma: no cover - surfaced by assertion below @@ -181,8 +194,10 @@ def test_bootstrap_post_already_installed_returns_installed_without_worker( monkeypatch, ): journal_path, _config = settings_env(_settings_config()) - monkeypatch.setattr(mlx_bootstrap, "is_mlx_available", lambda: (True, "")) - monkeypatch.setattr(mlx_bootstrap, "check_model_present", lambda: True) + monkeypatch.setattr( + mlx_bootstrap, "is_mlx_available_for_model", lambda _spec: (True, "") + ) + monkeypatch.setattr(mlx_bootstrap, "check_model_present", lambda _model: True) monkeypatch.setattr( mlx_bootstrap.threading, "Thread", @@ -207,7 +222,9 @@ def test_bootstrap_post_already_installed_returns_installed_without_worker( ) def test_bootstrap_post_rejects_unqualified_host(settings_env, monkeypatch, reason): journal_path, _config = settings_env(_settings_config()) - monkeypatch.setattr(mlx_bootstrap, "is_mlx_available", lambda: (False, reason)) + monkeypatch.setattr( + mlx_bootstrap, "is_mlx_available_for_model", lambda _spec: (False, reason) + ) client = _client(journal_path) response = client.post("/app/settings/api/mlx/bootstrap") @@ -287,11 +304,199 @@ def test_routes_import_without_mlx_vlm_registers_mlx_endpoints( assert "/app/settings/api/mlx/availability" in registered assert "/app/settings/api/mlx/bootstrap" in registered assert "/app/settings/api/mlx/bootstrap/status" in registered + assert "/app/settings/api/mlx/models" in registered + + +def test_mlx_models_route_returns_settings_shape(settings_env): + journal_path, _config = settings_env(_settings_config()) + client = _client(journal_path) + + response = client.get("/app/settings/api/mlx/models") + + assert response.status_code == 200 + assert response.get_json() == [ + { + "name": QWEN_35_9B, + "label": "qwen 3.5 — 16 GB Mac", + "min_ram_gb": 16, + }, + { + "name": GEMMA4_26B_A4B_4BIT, + "label": "gemma 4 (26B) — 24 GB Mac", + "min_ram_gb": 24, + }, + ] + + +@pytest.mark.parametrize( + ("method", "path"), + [ + ("get", "/app/settings/api/mlx/availability"), + ("post", "/app/settings/api/mlx/bootstrap"), + ("get", "/app/settings/api/mlx/bootstrap/status"), + ], +) +def test_mlx_routes_reject_unknown_model(settings_env, method, path): + journal_path, _config = settings_env(_settings_config()) + client = _client(journal_path) + + response = getattr(client, method)(f"{path}?model=not-real") + + assert response.status_code == 400 + payload = response.get_json() + assert payload["reason_code"] == "invalid_request_value" + assert "not-real" in payload["detail"] + assert QWEN_35_9B in payload["detail"] + assert GEMMA4_26B_A4B_4BIT in payload["detail"] + + +@pytest.mark.parametrize( + ("method", "path", "helper_name", "return_value"), + [ + ( + "get", + "/app/settings/api/mlx/availability", + "get_availability_payload", + {"available": True}, + ), + ( + "post", + "/app/settings/api/mlx/bootstrap", + "start_bootstrap", + ({"state": "installed"}, 200), + ), + ( + "get", + "/app/settings/api/mlx/bootstrap/status", + "get_state", + {"state": "idle"}, + ), + ], +) +def test_mlx_routes_default_to_qwen_model( + settings_env, monkeypatch, method, path, helper_name, return_value +): + journal_path, _config = settings_env(_settings_config()) + calls = [] + + def fake_helper(model): + calls.append(model) + return return_value + monkeypatch.setattr(mlx_bootstrap, helper_name, fake_helper) + client = _client(journal_path) + + response = getattr(client, method)(path) -def _write_snapshot(tmp_path: Path, monkeypatch, files: dict[str, bytes]) -> Path: + assert response.status_code == 200 + assert calls == [QWEN_35_9B] + + +def test_update_providers_mlx_round_trip_persists_active_model( + settings_env, monkeypatch +): + journal_path, _config = settings_env(_settings_config()) + client = _client(journal_path) + + response = client.put( + "/app/settings/api/providers", + json={"mlx": {"active_model": GEMMA4_26B_A4B_4BIT}}, + ) + + assert response.status_code == 200 + assert response.get_json()["mlx"]["active_model"] == GEMMA4_26B_A4B_4BIT + saved = json.loads((journal_path / "config" / "journal.json").read_text()) + assert saved["providers"]["mlx"]["active_model"] == GEMMA4_26B_A4B_4BIT + + from solstone.think.providers import mlx + + monkeypatch.setattr(mlx, "_module_level_cache", {}) + assert mlx._resolve_default_model() == GEMMA4_26B_A4B_4BIT + + +@pytest.mark.parametrize( + "payload", + [ + {"mlx": "string-not-object"}, + {"mlx": {"active_model": 123}}, + {"mlx": {"unknown_field": "x"}}, + {"mlx": {"active_model": "not-a-real-model"}}, + ], +) +def test_update_providers_mlx_rejects_malformed_payload(settings_env, payload): + journal_path, _config = settings_env(_settings_config()) + client = _client(journal_path) + before = (journal_path / "config" / "journal.json").read_text() + + response = client.put("/app/settings/api/providers", json=payload) + + assert response.status_code == 400 + assert response.get_json()["reason_code"] == "invalid_config_value" + assert (journal_path / "config" / "journal.json").read_text() == before + + +def test_bootstrap_state_is_per_model_under_concurrent_access(monkeypatch): + releases = { + QWEN_35_9B: threading.Event(), + GEMMA4_26B_A4B_4BIT: threading.Event(), + } + started = { + QWEN_35_9B: threading.Event(), + GEMMA4_26B_A4B_4BIT: threading.Event(), + } + monkeypatch.setattr( + mlx_bootstrap, + "is_mlx_available_for_model", + lambda _spec: (True, ""), + ) + monkeypatch.setattr(mlx_bootstrap, "check_model_present", lambda _model: False) + + def fake_worker(model): + started[model].set() + releases[model].wait(timeout=2) + with mlx_bootstrap._STATES_LOCK: + mlx_bootstrap._mark_installed_locked(model) + + monkeypatch.setattr(mlx_bootstrap, "_run_bootstrap_worker", fake_worker) + + assert mlx_bootstrap.start_bootstrap(QWEN_35_9B) == ({"state": "downloading"}, 202) + assert started[QWEN_35_9B].wait(timeout=2) + assert mlx_bootstrap.start_bootstrap(GEMMA4_26B_A4B_4BIT) == ( + {"state": "downloading"}, + 202, + ) + assert started[GEMMA4_26B_A4B_4BIT].wait(timeout=2) + assert mlx_bootstrap.get_state(QWEN_35_9B)["state"] == "downloading" + assert mlx_bootstrap.get_state(GEMMA4_26B_A4B_4BIT)["state"] == "downloading" + + releases[QWEN_35_9B].set() + deadline = time.monotonic() + 2 + while ( + mlx_bootstrap.get_state(QWEN_35_9B)["state"] != "installed" + and time.monotonic() < deadline + ): + time.sleep(0.01) + assert mlx_bootstrap.get_state(QWEN_35_9B)["state"] == "installed" + assert mlx_bootstrap.get_state(GEMMA4_26B_A4B_4BIT)["state"] == "downloading" + + releases[GEMMA4_26B_A4B_4BIT].set() + deadline = time.monotonic() + 2 + while ( + mlx_bootstrap.get_state(GEMMA4_26B_A4B_4BIT)["state"] != "installed" + and time.monotonic() < deadline + ): + time.sleep(0.01) + assert mlx_bootstrap.get_state(GEMMA4_26B_A4B_4BIT)["state"] == "installed" + + +def _write_snapshot( + tmp_path: Path, + monkeypatch, + files: dict[str, bytes], + model: str = QWEN_35_9B, +) -> Path: monkeypatch.setattr(mlx_bootstrap.constants, "HF_HUB_CACHE", str(tmp_path / "hf")) - snapshot_dir = mlx_bootstrap._snapshot_dir() + snapshot_dir = mlx_bootstrap._snapshot_dir(model) snapshot_dir.mkdir(parents=True) (snapshot_dir / "model.safetensors.index.json").write_text( json.dumps({"weight_map": {f"w{i}": path for i, path in enumerate(files)}}), @@ -314,10 +519,12 @@ def test_model_present_requires_index_and_all_safetensors(tmp_path, monkeypatch) }, ) - assert mlx_bootstrap.check_model_present() is True + assert mlx_bootstrap.check_model_present(QWEN_35_9B) is True - (mlx_bootstrap._snapshot_dir() / "model-00002-of-00002.safetensors").unlink() - assert mlx_bootstrap.check_model_present() is False + ( + mlx_bootstrap._snapshot_dir(QWEN_35_9B) / "model-00002-of-00002.safetensors" + ).unlink() + assert mlx_bootstrap.check_model_present(QWEN_35_9B) is False def test_snapshot_download_called_without_resume_download(monkeypatch): @@ -331,13 +538,14 @@ def test_snapshot_download_called_without_resume_download(monkeypatch): mlx_bootstrap.huggingface_hub, "snapshot_download", fake_snapshot_download ) monkeypatch.setattr( - mlx_bootstrap, "_verify_safetensors_sha256_hashes", lambda: None + mlx_bootstrap, "_verify_safetensors_sha256_hashes", lambda _model: None ) - mlx_bootstrap._run_bootstrap_worker() + mlx_bootstrap._run_bootstrap_worker(QWEN_35_9B) - assert calls["repo_id"] == MLX_MODEL_REPO - assert calls["revision"] == MLX_MODEL_REVISION + spec = _MLX_MODEL_REGISTRY[QWEN_35_9B] + assert calls["repo_id"] == spec.repo + assert calls["revision"] == spec.revision assert "resume_download" not in calls @@ -361,10 +569,11 @@ def test_hfapi_list_repo_tree_called_with_pinned_revision(tmp_path, monkeypatch) monkeypatch.setattr(mlx_bootstrap.huggingface_hub, "HfApi", lambda: FakeApi()) _set_state(state="verifying") - mlx_bootstrap._verify_safetensors_sha256_hashes() + mlx_bootstrap._verify_safetensors_sha256_hashes(QWEN_35_9B) - assert calls["repo_id"] == MLX_MODEL_REPO - assert calls["revision"] == MLX_MODEL_REVISION + spec = _MLX_MODEL_REGISTRY[QWEN_35_9B] + assert calls["repo_id"] == spec.repo + assert calls["revision"] == spec.revision assert calls["recursive"] is True @@ -372,7 +581,7 @@ def test_worker_enters_verifying_state_between_download_and_install(monkeypatch) entered_verify = threading.Event() release_verify = threading.Event() - def slow_verify(): + def slow_verify(_model): entered_verify.set() release_verify.wait(timeout=2) @@ -381,15 +590,17 @@ def test_worker_enters_verifying_state_between_download_and_install(monkeypatch) mlx_bootstrap.huggingface_hub, "snapshot_download", lambda **_kwargs: None ) monkeypatch.setattr(mlx_bootstrap, "_verify_safetensors_sha256_hashes", slow_verify) - worker = threading.Thread(target=mlx_bootstrap._run_bootstrap_worker) + worker = threading.Thread( + target=mlx_bootstrap._run_bootstrap_worker, args=(QWEN_35_9B,) + ) _set_state(thread=worker) worker.start() assert entered_verify.wait(timeout=2) - assert mlx_bootstrap.get_state()["state"] == "verifying" + assert mlx_bootstrap.get_state(QWEN_35_9B)["state"] == "verifying" release_verify.set() worker.join(timeout=2) - assert mlx_bootstrap.get_state()["state"] == "installed" + assert mlx_bootstrap.get_state(QWEN_35_9B)["state"] == "installed" def test_worker_verify_mismatch_transitions_to_failed_with_filename( @@ -414,8 +625,8 @@ def test_worker_verify_mismatch_transitions_to_failed_with_filename( ) monkeypatch.setattr(mlx_bootstrap.huggingface_hub, "HfApi", lambda: FakeApi()) - mlx_bootstrap._run_bootstrap_worker() - payload = mlx_bootstrap.get_state() + mlx_bootstrap._run_bootstrap_worker(QWEN_35_9B) + payload = mlx_bootstrap.get_state(QWEN_35_9B) assert payload["state"] == "failed" assert "model.safetensors" in payload["message"] @@ -430,9 +641,9 @@ def test_worker_exception_sets_failed_message(monkeypatch): lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("download broke")), ) - mlx_bootstrap._run_bootstrap_worker() + mlx_bootstrap._run_bootstrap_worker(QWEN_35_9B) - payload = mlx_bootstrap.get_state() + payload = mlx_bootstrap.get_state(QWEN_35_9B) assert payload["state"] == "failed" assert "download broke" in payload["message"] diff --git a/solstone/apps/settings/tests/test_workspace_html.py b/solstone/apps/settings/tests/test_workspace_html.py index c3804f759..7a53672a4 100644 --- a/solstone/apps/settings/tests/test_workspace_html.py +++ b/solstone/apps/settings/tests/test_workspace_html.py @@ -54,16 +54,27 @@ def test_workspace_mlx_unavailable_option_disabled_title_and_suffix(): assert reason -def test_workspace_mlx_model_identifier_uses_generate_tier(): +def test_workspace_mlx_model_identifier_uses_active_model(): text = _workspace_text() assert 'id="mlxModelIdentifier"' in text + assert 'id="field-mlx-active-model"' in text assert "function getSelectedMlxModel()" in text - assert "const MLX_PROVIDER_DEFAULTS" in text + assert "providersData?.mlx?.active_model" in text assert "model_present" in text assert "model: ${getSelectedMlxModel()}" in text +def test_workspace_mlx_model_picker_labels_are_pinned(): + text = _workspace_text() + + assert "qwen 3.5 — 16 GB Mac" in text + assert "gemma 4 (26B) — 24 GB Mac" in text + assert text.index('id="field-mlx-active-model"') < text.index( + 'id="mlxBootstrapProgressShell"' + ) + + def test_workspace_mlx_v1_scope_copy_is_generate_gated(): text = _workspace_text() @@ -100,6 +111,7 @@ def test_workspace_mlx_retry_button_is_failed_only(): assert 'id="mlxBootstrapRetry"' in text assert "retry.style.display = state === 'failed' ? '' : 'none'" in text + assert "retry.disabled = !!status?.bootstrap_disabled" in text assert "retry.onclick" in text assert "mountMlxProgress()" in text @@ -109,7 +121,9 @@ def test_workspace_mlx_polling_starts_and_stops_on_terminal_states(): assert "let mlxBootstrapPollTimer = null" in text assert "let mlxBootstrapPostStarted = false" in text - assert "fetch('api/mlx/bootstrap', { method: 'POST' })" in text + assert "api/mlx/bootstrap?model=${model}" in text + assert "api/mlx/bootstrap/status?model=${model}" in text + assert "api/mlx/availability?model=${model}" in text assert "setInterval(pollMlxBootstrap, 1000)" in text assert "clearInterval(mlxBootstrapPollTimer)" in text assert "state === 'installed'" in text diff --git a/solstone/apps/settings/workspace.html b/solstone/apps/settings/workspace.html index d8ffd42ec..f443ad053 100644 --- a/solstone/apps/settings/workspace.html +++ b/solstone/apps/settings/workspace.html @@ -2437,7 +2437,11 @@ button:focus:not(:focus-visible) {