fix(raw-block): harden metadata checkpoint recovery

DongDongJu · idevasena · commit e76aa4697e19 · 2026-03-15T22:23:21.000Z
Signed-off-by: DongDongJu &lt;commisori28@gmail.com&gt;
Signed-off-by: Dongjoo Seo &lt;dongjoo.seo1@samsung.com&gt;
diff --git a/lmcache/v1/storage_backend/plugins/rust_raw_block_backend.py b/lmcache/v1/storage_backend/plugins/rust_raw_block_backend.py
@@ -148,6 +148,7 @@ def __init__(
         self.meta_verify_on_load: bool = bool(
             extra.get("rust_raw_block.meta_verify_on_load", True)
         )
+        self._meta_copy_count: int = 2
 
         get_full_chunk_size_bytes = getattr(
             self.local_cpu_backend, "get_full_chunk_size_bytes", None
@@ -192,6 +193,14 @@ def __init__(
                 "rust_raw_block.meta_total_bytes must be > block_align "
                 "(room for metadata header + payload)"
             )
+        self._meta_container_bytes: int = (
+            (self.meta_total_bytes // self._meta_copy_count) // self.block_align
+        ) * self.block_align
+        if self._meta_container_bytes <= self.block_align:
+            raise ValueError(
+                "rust_raw_block.meta_total_bytes must provide room for at least "
+                "two metadata copies (header + payload)"
+            )
 
         self._lock = threading.Lock()
         self._index: dict[CacheEngineKey, _Entry] = {}
@@ -789,7 +798,12 @@ def _checkpoint_loop(self) -> None:
                 logger.warning(f"Periodic metadata checkpoint failed: {e}")
 
     def _meta_payload_capacity(self) -> int:
-        return self.meta_total_bytes - self.block_align
+        return self._meta_container_bytes - self.block_align
+
+    def _meta_container_offsets(self) -> list[int]:
+        return [
+            idx * self._meta_container_bytes for idx in range(self._meta_copy_count)
+        ]
 
     def _read_meta_header(self, container_offset: int) -> Optional[dict[str, int]]:
         raw = self._rawdev()
@@ -832,6 +846,23 @@ def _load_meta_payload(self, header: dict[str, int]) -> Optional[bytes]:
             return None
         return payload
 
+    def _select_latest_checkpoint(
+        self,
+    ) -> tuple[Optional[dict[str, int]], Optional[bytes]]:
+        best_header: Optional[dict[str, int]] = None
+        best_payload: Optional[bytes] = None
+        for offset in self._meta_container_offsets():
+            header = self._read_meta_header(offset)
+            if header is None:
+                continue
+            payload = self._load_meta_payload(header)
+            if payload is None:
+                continue
+            if best_header is None or int(header["seq"]) > int(best_header["seq"]):
+                best_header = header
+                best_payload = payload
+        return best_header, best_payload
+
     def _snapshot_state(self) -> tuple[dict[str, Any], int]:
         with self._lock:
             dirty_total = self._meta_dirty_total
@@ -886,10 +917,9 @@ def _write_checkpoint(self, payload: bytes, dirty_total_snapshot: int) -> bool:
             )
             return False
 
-        target = 0
-        current = self._read_meta_header(target)
-        current_seq = int(current["seq"]) if current is not None else 0
-        next_seq = max(current_seq, self._meta_seq) + 1
+        next_seq = self._meta_seq + 1
+        target_idx = int((next_seq - 1) % self._meta_copy_count)
+        target = self._meta_container_offsets()[target_idx]
 
         payload_len = len(payload)
         payload_total_len = _round_up(payload_len, self.block_align)
@@ -942,6 +972,17 @@ def _checkpoint_once(self, force: bool) -> bool:
             )
         return ok
 
+    def _is_valid_checkpoint_entry(self, offset: int, size: int) -> bool:
+        if offset < self._data_base_offset:
+            return False
+        rel = offset - self._data_base_offset
+        if rel % self.slot_bytes != 0:
+            return False
+        slot = rel // self.slot_bytes
+        if slot < 0 or slot >= self._max_slots:
+            return False
+        return 0 < size <= (self.slot_bytes - self.header_bytes)
+
     def _apply_loaded_state(self, data: dict[str, Any]) -> bool:
         if not isinstance(data, dict):
             return False
@@ -1036,6 +1077,16 @@ def _apply_loaded_state(self, data: dict[str, Any]) -> bool:
                     fmt_name = entry.get("fmt")
                     cached_positions_list = entry.get("cached_positions")
 
+                    if not self._is_valid_checkpoint_entry(offset, size):
+                        logger.warning(
+                            "Skipping invalid checkpoint entry for key '%s': "
+                            "offset=%d size=%d",
+                            k_str,
+                            offset,
+                            size,
+                        )
+                        continue
+
                     shape = (
                         torch.Size(list(shape_list)) if shape_list is not None else None
                     )
@@ -1132,18 +1183,13 @@ def _validate_loaded_entries(self) -> None:
         )
 
     def _load_checkpoint_from_device(self) -> None:
-        header = self._read_meta_header(0)
+        header, payload = self._select_latest_checkpoint()
         if header is None:
             logger.info(
                 "RustRawBlockBackend: no valid on-device metadata checkpoint found"
             )
             return
-        payload = self._load_meta_payload(header)
-        if payload is None:
-            logger.warning(
-                "RustRawBlockBackend: metadata header exists but payload is invalid"
-            )
-            return
+        assert payload is not None
         try:
             data = json.loads(payload.decode("utf-8"))
         except Exception:
diff --git a/tests/v1/storage_backend/test_rust_raw_block_backend.py b/tests/v1/storage_backend/test_rust_raw_block_backend.py
@@ -370,7 +370,7 @@ def test_rust_raw_block_backend_ignores_torn_newer_checkpoint(
     memory_allocator, loop_in_thread
 ):
     """
-    If checkpoint payload CRC is invalid, loader ignores metadata.
+    If a newer checkpoint copy is torn, loader falls back to the older valid copy.
     """
     with tempfile.TemporaryDirectory() as td:
         dev_path = os.path.join(td, "dev.bin")
@@ -423,22 +423,24 @@ def test_rust_raw_block_backend_ignores_torn_newer_checkpoint(
         )
         assert obj is not None and obj.tensor is not None
         obj.tensor.fill_(11)
+        expected = bytes(obj.byte_array)
         try:
             fut = backend1.batched_submit_put_task([key], [obj])[0]
             fut.result(timeout=10)
         finally:
+            torn_offset = backend1._meta_container_offsets()[1]
             backend1.close()
 
-        # Corrupt single checkpoint header/payload with invalid CRC.
+        # Corrupt the newer checkpoint copy with invalid CRC.
         # Header format: <8sIQQI (magic, version, seq, payload_len, crc).
         header = struct.pack(
             "<8sIQQI", _DEFAULT_META_MAGIC, _DEFAULT_META_VERSION, 9999, 2, 0
         )
         padded_header = header + bytes(align - len(header))
         with open(dev_path, "r+b") as f:
-            f.seek(align)
+            f.seek(torn_offset + align)
             f.write(b"{}")
-            f.seek(0)
+            f.seek(torn_offset)
             f.write(padded_header)
 
         backend2 = RustRawBlockBackend(
@@ -449,8 +451,102 @@ def test_rust_raw_block_backend_ignores_torn_newer_checkpoint(
             dst_device="cpu",
         )
         try:
-            assert not backend2.contains(key)
+            assert backend2.contains(key)
             out = backend2.get_blocking(key)
-            assert out is None
+            assert out is not None
+            assert bytes(out.byte_array) == expected
         finally:
             backend2.close()
+
+
+@pytest.mark.skipif(
+    not _has_ext(), reason="lmcache_rust_raw_block_io extension not installed"
+)
+def test_rust_raw_block_backend_skips_invalid_checkpoint_entries(
+    memory_allocator, loop_in_thread
+):
+    """Checkpoint restore should reject invalid offset/size metadata entries."""
+    with tempfile.TemporaryDirectory() as td:
+        dev_path = os.path.join(td, "dev.bin")
+        with open(dev_path, "wb") as f:
+            f.truncate(64 * 1024 * 1024)
+
+        base_cfg = LMCacheEngineConfig.from_defaults(
+            chunk_size=256,
+            local_cpu=True,
+            max_local_cpu_size=0.1,
+            lmcache_instance_id="test_rust_raw_block_backend_invalid_checkpoint",
+        )
+        base_cfg.extra_config = {
+            "rust_raw_block.device_path": dev_path,
+            "rust_raw_block.block_align": 4096,
+            "rust_raw_block.header_bytes": 4096,
+            "rust_raw_block.meta_total_bytes": 4 * 1024 * 1024,
+            "rust_raw_block.meta_enable_periodic": False,
+            "rust_raw_block.meta_verify_on_load": False,
+        }
+        metadata = LMCacheMetadata(
+            model_name="test_model",
+            world_size=1,
+            local_world_size=1,
+            worker_id=0,
+            local_worker_id=0,
+            kv_dtype=torch.bfloat16,
+            kv_shape=(4, 2, 256, 8, 128),
+        )
+
+        local_cpu = LocalCPUBackend(
+            config=base_cfg,
+            metadata=metadata,
+            dst_device="cpu",
+            memory_allocator=memory_allocator,
+        )
+        backend = RustRawBlockBackend(
+            config=base_cfg,
+            metadata=metadata,
+            local_cpu_backend=local_cpu,
+            loop=loop_in_thread,
+            dst_device="cpu",
+        )
+        try:
+            entries = {}
+            for chunk_hash, (offset, size) in {
+                1: (backend._data_base_offset - backend.slot_bytes, 1024),
+                2: (backend._data_base_offset + 1, 1024),
+                3: (
+                    backend._data_base_offset,
+                    backend.slot_bytes - backend.header_bytes + 1,
+                ),
+            }.items():
+                key = CacheEngineKey("test_model", 1, 0, chunk_hash, torch.bfloat16)
+                entries[key.to_string()] = {
+                    "offset": offset,
+                    "size": size,
+                    "shape": [2, 16, 8, 128],
+                    "dtype": "bfloat16",
+                    "fmt": MemoryFormat.KV_T2D.name,
+                    "cached_positions": None,
+                }
+
+            applied = backend._apply_loaded_state(
+                {
+                    "version": 1,
+                    "device_path": dev_path,
+                    "capacity_bytes": backend.capacity_bytes,
+                    "block_align": backend.block_align,
+                    "header_bytes": backend.header_bytes,
+                    "slot_bytes": backend.slot_bytes,
+                    "meta_total_bytes": backend.meta_total_bytes,
+                    "meta_magic": backend.meta_magic_text,
+                    "meta_version": backend.meta_version,
+                    "data_base_offset": backend._data_base_offset,
+                    "next_slot": 0,
+                    "free_slots": [],
+                    "lru_keys": [],
+                    "entries": entries,
+                }
+            )
+            assert applied is True
+            assert backend._index == {}
+        finally:
+            backend.close()
