feat: add learned partial mirror planner

README.md

@@ -207,6 +207,26 @@ the hottest ones automatically — colibrì literally gets faster the more you u
 it. On multi-socket hosts, `COLI_NUMA=1` interleaves the resident weights across
 memory controllers ([#82](https://github.com/JustVugg/colibri/issues/82)).
 
+For a second drive that cannot hold the whole model, Colibri can rank a partial
+mirror from the expert history it already learns. Run a few representative
+prompts first so `.coli_usage` reflects the workload, then plan, stage, and
+verify the mirror:
+
+```bash
+./c/coli mirror plan  --model /fast/glm52_i4 --mirror /second/glm52_i4 \
+  --budget-gib 200 --reserve-gib 20
+./c/coli mirror stage --model /fast/glm52_i4 --mirror /second/glm52_i4 \
+  --budget-gib 200 --reserve-gib 20
+./c/coli mirror verify --model /fast/glm52_i4 --mirror /second/glm52_i4
+```
+
+The planner reads safetensors headers directly, follows split-model directories
+from `COLI_MODEL_DIRS`, and prioritizes shards that can serve the hottest routed
+experts. Staging never changes the primary model: it copies through temporary
+files, preserves the requested free-space reserve, verifies every shard with
+SHA-256, never deletes an existing mirror shard, and atomically publishes a
+receipt only after the selected mirror is ready.
+
 ### Never wait for the disk twice
 
 Misses are expensive, so the engine spends most of its cleverness avoiding and

c/coli

@@ -8,6 +8,7 @@ Run GLM-5.2 (744B) locally on CPU with roughly 15-26 GB of RAM.
   coli run "prompt"         one-shot generation
   coli info                 model, RAM, disk, and configuration status
   coli plan                 Disk / RAM / VRAM resource plan
+  coli mirror               Plan, stage, or verify a learned partial mirror
   coli doctor               installation and execution-plan diagnostics
   coli bench [task...]      quality benchmarks (MMLU/HellaSwag/...)
   coli convert              convert GLM-5.2-FP8 to int4, one shard at a time
@@ -1116,6 +1117,27 @@ def cmd_convert(a):
     print(f"  {C.dim}[2/2] int8 MTP head (speculative drafts){C.r}")
     sys.exit(subprocess.call(mtp_cmd+["--mtp"]))
 
+
+def cmd_mirror(a):
+    if not a.mirror:
+        sys.exit("mirror path required: pass --mirror or set COLI_MODEL_MIRROR")
+    command = [sys.executable, os.path.join(TOOLS, "mirror_plan.py"), a.action,
+               "--model", a.model, "--mirror", a.mirror]
+    source_dirs = list(a.source_dir)
+    configured = os.environ.get("COLI_MODEL_DIRS", "")
+    if configured:
+        source_dirs.extend(part.strip() for part in re.split(r"[;,]", configured)
+                           if part.strip())
+    for directory in source_dirs:
+        command += ["--source-dir", directory]
+    if a.usage:
+        command += ["--usage", a.usage]
+    if a.action != "verify":
+        command += ["--budget-gib", str(a.budget_gib),
+                    "--reserve-gib", str(a.reserve_gib)]
+    return subprocess.call(command)
+
+
 def main():
     common=argparse.ArgumentParser(add_help=False)
     common.add_argument("--model", default=DEF_MODEL); common.add_argument("--ram", type=int, default=0)  # 0 = auto (il motore usa l'88% della RAM disponibile)
@@ -1144,6 +1166,14 @@ def main():
     sub.add_parser("build", parents=[common]); sub.add_parser("info", parents=[common])
     pp=sub.add_parser("plan",parents=[common])
     pp.add_argument("--json",action="store_true")
+    pm=sub.add_parser("mirror", parents=[common],
+                      help="plan, stage, or verify a usage-ranked partial model mirror")
+    pm.add_argument("action", choices=("plan", "stage", "verify"))
+    pm.add_argument("--mirror", default=os.environ.get("COLI_MODEL_MIRROR"))
+    pm.add_argument("--source-dir", action="append", default=[])
+    pm.add_argument("--usage")
+    pm.add_argument("--budget-gib", type=float, default=0)
+    pm.add_argument("--reserve-gib", type=float, default=10)
     pd=sub.add_parser("doctor",parents=[common])
     pd.add_argument("--json",action="store_true",help="emit a versioned JSON report")
     pd.add_argument("--deep",action="store_true",
@@ -1207,7 +1237,8 @@ def main():
         help="int4 scale group size: 64 (default, group-scaled quality) or 0 (legacy per-row)")
     pc.add_argument("--no-mtp",action="store_true",help="skip the MTP head (no speculative drafts)")
     a=ap.parse_args()
-    handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"doctor":cmd_doctor,"tune":cmd_tune,
+    handler={"build":cmd_build,"info":cmd_info,"plan":cmd_plan,"mirror":cmd_mirror,
+             "doctor":cmd_doctor,"tune":cmd_tune,
              "run":cmd_run,"chat":cmd_chat,"serve":cmd_serve,"stop":cmd_stop,"bench":cmd_bench,
              "convert":cmd_convert,"web":cmd_web}.get(a.cmd)
     if handler: sys.exit(handler(a) or 0)

c/tests/test_mirror_plan.py

@@ -0,0 +1,148 @@
+import json
+import struct
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+from tools.mirror_plan import (RECEIPT, MirrorError, create_plan, discover_shards,
+                               stage_mirror, usage_counts, verify_mirror)
+
+
+class MirrorPlannerTest(unittest.TestCase):
+    def setUp(self):
+        self.temporary = tempfile.TemporaryDirectory()
+        self.root = Path(self.temporary.name)
+        self.model = self.root / "model"
+        self.mirror = self.root / "mirror"
+        self.split = self.root / "split"
+        self.model.mkdir()
+        self.split.mkdir()
+        self.usage = self.model / ".coli_usage"
+
+    def tearDown(self):
+        self.temporary.cleanup()
+
+    @staticmethod
+    def write_shard(directory, name, tensors):
+        offset = 0
+        header = {}
+        payload = bytearray()
+        for tensor, size in tensors:
+            header[tensor] = {"dtype": "U8", "shape": [size],
+                              "data_offsets": [offset, offset + size]}
+            payload.extend(bytes([len(header) % 251]) * size)
+            offset += size
+        encoded = json.dumps(header, separators=(",", ":")).encode()
+        path = directory / name
+        path.write_bytes(struct.pack("<Q", len(encoded)) + encoded + payload)
+        return path
+
+    def test_usage_parser_ignores_malformed_rows_and_accumulates(self):
+        self.usage.write_text("0 1 7\ninvalid\n0 1 5\n-1 2 9\n1 2 0\n", encoding="utf-8")
+        self.assertEqual(usage_counts(self.usage), {(0, 1): 12})
+
+    def test_plan_activates_hot_gate_shard_before_companion_shard(self):
+        gate = self.write_shard(self.model, "hot-gate.safetensors", [
+            ("model.layers.0.mlp.experts.0.gate_proj.weight", 20),
+        ])
+        companion = self.write_shard(self.model, "hot-down.safetensors", [
+            ("model.layers.0.mlp.experts.0.down_proj.weight", 60),
+        ])
+        self.write_shard(self.model, "cold-gate.safetensors", [
+            ("model.layers.0.mlp.experts.1.gate_proj.weight", 20),
+        ])
+        self.usage.write_text("0 0 100\n0 1 1\n", encoding="utf-8")
+        budget = gate.stat().st_size + companion.stat().st_size
+
+        plan, selected = create_plan(self.model, self.mirror, [], self.usage, budget, 0)
+
+        self.assertTrue(plan["admitted"])
+        self.assertEqual([item["name"] for item in selected],
+                         ["hot-gate.safetensors", "hot-down.safetensors"])
+        self.assertEqual(selected[0]["activated_experts"], 1)
+        self.assertEqual(selected[1]["activated_experts"], 0)
+
+    def test_plan_requires_learned_usage_instead_of_guessing(self):
+        shard = self.write_shard(self.model, "model.safetensors", [
+            ("model.layers.0.mlp.experts.0.gate_proj.weight", 8),
+        ])
+        plan, selected = create_plan(
+            self.model, self.mirror, [], self.usage, shard.stat().st_size, 0)
+        self.assertFalse(plan["admitted"])
+        self.assertEqual(plan["reason"], "usage_history_missing")
+        self.assertEqual(selected, [])
+
+    def test_split_directories_are_searched_and_basenames_are_deduplicated(self):
+        primary = self.write_shard(self.model, "same.safetensors", [
+            ("model.layers.0.mlp.experts.0.gate_proj.weight", 8),
+        ])
+        self.write_shard(self.split, "same.safetensors", [
+            ("model.layers.0.mlp.experts.1.gate_proj.weight", 16),
+        ])
+        extra = self.write_shard(self.split, "extra.safetensors", [
+            ("model.layers.0.mlp.experts.2.gate_proj.weight", 12),
+        ])
+        _directories, candidates = discover_shards(self.model, [self.split])
+        by_name = {item["name"]: item for item in candidates}
+        self.assertEqual(by_name["same.safetensors"]["source"], primary)
+        self.assertEqual(by_name["extra.safetensors"]["source"], extra)
+
+    def test_stage_is_atomic_resumable_and_sha256_verified(self):
+        source = self.write_shard(self.model, "hot.safetensors", [
+            ("model.layers.0.mlp.experts.0.gate_proj.weight", 32),
+        ])
+        self.usage.write_text("0 0 25\n", encoding="utf-8")
+        budget = source.stat().st_size
+
+        result = stage_mirror(self.model, self.mirror, [], self.usage, budget, 0)
+
+        self.assertTrue(result["ready"])
+        target = self.mirror / source.name
+        self.assertEqual(target.read_bytes(), source.read_bytes())
+        self.assertTrue((self.mirror / RECEIPT).is_file())
+        first_mtime = target.stat().st_mtime_ns
+
+        repeated = stage_mirror(self.model, self.mirror, [], self.usage, budget, 0)
+        self.assertTrue(repeated["ready"])
+        self.assertEqual(repeated["plan"]["remaining_copy_bytes"], 0)
+        self.assertEqual(target.stat().st_mtime_ns, first_mtime)
+
+        content = target.read_bytes()
+        target.write_bytes(bytes([content[0] ^ 1]) + content[1:])
+        verification = verify_mirror(self.mirror)
+        self.assertFalse(verification["ready"])
+        self.assertEqual(verification["failures"], ["hot.safetensors (sha256)"])
+
+    def test_reserve_preflight_writes_no_shard_or_receipt(self):
+        source = self.write_shard(self.model, "hot.safetensors", [
+            ("model.layers.0.mlp.experts.0.gate_proj.weight", 32),
+        ])
+        self.usage.write_text("0 0 25\n", encoding="utf-8")
+        disk = mock.Mock(free=source.stat().st_size - 1)
+        with mock.patch("tools.mirror_plan.shutil.disk_usage", return_value=disk):
+            result = stage_mirror(
+                self.model, self.mirror, [], self.usage, source.stat().st_size, 0)
+        self.assertFalse(result["admitted"])
+        self.assertEqual(result["reason"], "free_space_reserve")
+        self.assertFalse((self.mirror / source.name).exists())
+        self.assertFalse((self.mirror / RECEIPT).exists())
+
+    def test_verify_rejects_receipt_path_traversal(self):
+        self.mirror.mkdir()
+        receipt = {"schema": "colibri.partial-mirror.v1", "files": [{
+            "name": "../outside.safetensors", "size": 1, "sha256": "0" * 64,
+        }]}
+        (self.mirror / RECEIPT).write_text(json.dumps(receipt), encoding="utf-8")
+        result = verify_mirror(self.mirror)
+        self.assertFalse(result["ready"])
+        self.assertEqual(result["failures"], ["invalid_receipt_entry"])
+
+    def test_invalid_safetensors_header_fails_closed(self):
+        (self.model / "bad.safetensors").write_bytes(struct.pack("<Q", 999) + b"{}")
+        with self.assertRaises(MirrorError):
+            discover_shards(self.model, [])
+
+
+if __name__ == "__main__":
+    unittest.main()