diff --git a/kev/checkpoint.py b/kev/checkpoint.py index 27246247..16aa011f 100644 --- a/kev/checkpoint.py +++ b/kev/checkpoint.py @@ -74,7 +74,7 @@ def to_dict(self): def read_meta(run): - return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu")) + return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu", weights_only=True)) def write_meta(run, meta): diff --git a/kev/mlx_model.py b/kev/mlx_model.py index 0042e36b..7a4dd29b 100644 --- a/kev/mlx_model.py +++ b/kev/mlx_model.py @@ -37,7 +37,14 @@ def load_base(base_dir): - """mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it.""" + """mlx-lm's model of a Hub base snapshot, weights as stored (bf16 for the Qwen3.5 bases); the LoRA path merges into it. + A config.json that names a `model_file` is refused: mlx-lm 0.31's load_model imports and runs that module to build + the model (0.32.2 asks for trust_remote_code first), and the base is whatever the checkpoint's head.pt names, so a + checkpoint could choose code to run. Kev's bases are mlx-lm's built-in architectures; the torch backend ignores the field.""" + config = json.loads((Path(base_dir) / "config.json").read_text(encoding="utf-8")) + if config.get("model_file") is not None: + raise ValueError(f"{base_dir}: config.json names model_file {config['model_file']!r}, model code mlx-lm would import and run; " + "the MLX backend loads only mlx-lm's built-in architectures. Load this checkpoint with backend=torch (KEV_BACKEND=torch), which ignores the field") mx.set_cache_limit(CACHE_LIMIT) return load_model(Path(base_dir))[0] diff --git a/tests/test_mlx.py b/tests/test_mlx.py index 46870a30..2ec3902a 100644 --- a/tests/test_mlx.py +++ b/tests/test_mlx.py @@ -15,7 +15,7 @@ pytestmark = pytest.mark.skipif(platform.system() != "Darwin" or platform.machine() != "arm64", reason="MLX runs on Apple Silicon only") from kev.checkpoint import Checkpoint, LoadOptions, mlx_available # noqa: E402 -from kev.mlx_model import MLXDecisionModel, merge_lora # noqa: E402 +from kev.mlx_model import MLXDecisionModel, load_base, merge_lora # noqa: E402 from kev.model import SCORING_INTERFACE, DecisionModel # noqa: E402 RUN = "jaredpalmer/kev-0.8b@9a45d25eb2ab761841196625383fa1dff0e56c1e" # pinned (round 15, 2026-09-24): bf16 noise is per checkpoint, so a republish must not move these bars @@ -65,6 +65,18 @@ def test_backend_resolution(): Checkpoint("jaredpalmer/kev-4b@qwen3").load("mps", LoadOptions(backend="mlx")) # Qwen3 base: no DeltaNet layers +def test_load_base_refuses_model_code(tmp_path): + """mlx-lm 0.31.x executes config.json's model_file unconditionally; a base that brings one is refused before + load_model gets as far as running it (valid safetensors, so only the guard stands in the way).""" + sentinel = tmp_path / "sentinel" + (tmp_path / "config.json").write_text(json.dumps({"model_type": "qwen3_5", "model_file": "model.py"}), encoding="utf-8") + mx.save_safetensors(str(tmp_path / "model.safetensors"), {"x": mx.zeros((1,))}) + (tmp_path / "model.py").write_text(f"open({str(sentinel)!r}, 'w').write('x')\n", encoding="utf-8") + with pytest.raises(ValueError, match="model_file"): + load_base(tmp_path) + assert not sentinel.exists() + + # --- full-weight checkpoints (kev.train --full_ft; kev.mlx_model.load_full): a 4-layer Qwen3.5 text model with random # weights (norms, conv kernels and A_log included, so a missed layout rule shows), saved as a full-weight run saves it diff --git a/tests/test_unit.py b/tests/test_unit.py index 0557c4df..4b3f74b4 100644 --- a/tests/test_unit.py +++ b/tests/test_unit.py @@ -185,6 +185,22 @@ def test_checkpoint_meta_round_trip_and_defaults(tmp_path): LoadOptions.from_env({"KEV_BACKEND": "metal"}) +def test_read_meta_refuses_a_pickle_payload(tmp_path): + """read_meta loads head.pt with weights_only=True: a checkpoint whose head.pt unpickles into code never runs it + (the Space vendors checkpoint.py onto ZeroGPU's unpinned torch, where the default was not always safe).""" + import pickle + from kev.checkpoint import Meta, read_meta, write_meta + sentinel = tmp_path / "sentinel" + class Payload: + def __reduce__(self): + return (Path.touch, (sentinel,)) + meta = Meta.from_dict({"head": {"w": torch.zeros(1)}, "base": "b"}); meta.extra["payload"] = Payload() + write_meta(tmp_path, meta) + with pytest.raises(pickle.UnpicklingError): + read_meta(tmp_path) + assert not sentinel.exists() + + def test_fused_default_needs_pinned_fla(monkeypatch): """kev.serve's CUDA fused default (checkpoint.fused_available): on only with flash-linear-attention importable at fused_qwen35.FLA_VERSION; it is not in the serve extra, so a plain install serves unfused instead of failing to import fla."""