Loading a checkpoint never runs code it brings: refuse a base's model_file on MLX, read head.pt weights-only

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):

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]
 

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
 

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."""