feat: add header format check and key __metadata__ ignored

Author: specture724

checkpoint_engine/ps.py

@@ -19,7 +19,7 @@
 import zmq
 from loguru import logger
 from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
-from safetensors.torch import _getdtype, safe_open
+from safetensors.torch import _TYPES, _getdtype, safe_open
 from torch.multiprocessing.reductions import reduce_tensor
 
 from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
@@ -499,6 +499,23 @@ def _inplace_pin_memory(file_path: str) -> MemoryBuffer:
             time.sleep(3)
             header_tensor = t[flag_size:start_pos]
             header = json.loads(header_tensor.numpy().tobytes())
+            if "__metadata__" in header:
+                header.pop("__metadata__")
+
+            # check the header format
+            for name, meta in header.items():
+                assert isinstance(name, str), f"parameter name {name} should be str"
+                assert isinstance(meta, dict), f"parameter meta {meta} should be dict"
+                assert "data_offsets" in meta, f"data_offsets not in meta for parameter {name}"
+                assert "shape" in meta, f"shape not in meta for parameter {name}"
+                assert "dtype" in meta, f"dtype not in meta for parameter {name}"
+                assert meta["dtype"] in _TYPES, (
+                    f"dtype {meta['dtype']} not a valid safetensors dtype, supported dtypes: {_TYPES.keys()}"
+                )
+                assert isinstance(meta["shape"], list), f"shape {meta['shape']} should be list"
+                assert isinstance(meta["data_offsets"], list) and len(meta["data_offsets"]) == 2, (
+                    f"data_offsets {meta['data_offsets']} should be list of length 2"
+                )
 
             metas: list[ParameterMeta] = []
             offset = 0