refactor: serialize metas as JSON instead of pickle

Replace pickle with pydantic TypeAdapter(dict[int, MemoryBufferMetaList])
for the metas wire format across the HTTP endpoints, join_cli, and
examples/update.py. This reuses the existing pydantic schema (torch.dtype
/ torch.Size already have serializers in data_types.py), removes the
arbitrary-code-execution risk of pickle.loads on request bodies, and makes
the metas self-describing for cross-language consumers.

- api.py: GET /metas returns application/json; POST /load-metas validates
  via validate_json and returns 400 on ValidationError (was broad except).
- join_cli.py / examples/update.py: read/write metas as JSON; document the
  --metas-url HTTP path alongside --load-metas-file.
- tests/test_api.py: use real MemoryBufferMetaList fixtures; add a
  schema-mismatch case (valid JSON, wrong shape -> 400).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 刘伟健

checkpoint_engine/api.py

@@ -1,4 +1,3 @@
-import pickle
 from collections.abc import Callable
 from typing import Any
 
@@ -7,11 +6,15 @@
 from fastapi import Request
 from fastapi.responses import JSONResponse, Response
 from loguru import logger
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter, ValidationError
 
+from checkpoint_engine.data_types import MemoryBufferMetaList
 from checkpoint_engine.ps import ParameterServer
 
 
+_METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList])
+
+
 def request_inference_to_update(
     url: str,
     socket_paths: dict[str, str],
@@ -87,15 +90,18 @@ async def get_metas(checkpoint_name: str) -> Response:
         except Exception as e:  # noqa: BLE001
             logger.exception(f"get_metas for {checkpoint_name} failed")
             return JSONResponse(content=str(e), status_code=500)
-        return Response(content=pickle.dumps(metas), media_type="application/octet-stream")
+        return Response(
+            content=_METAS_ADAPTER.dump_json(metas),
+            media_type="application/json",
+        )
 
     @app.post("/v1/checkpoints/{checkpoint_name}/load-metas")
     async def load_metas(checkpoint_name: str, raw: Request) -> Response:
         body = await raw.body()
         try:
-            metas = pickle.loads(body)
-        except Exception as e:  # noqa: BLE001
-            logger.exception(f"load_metas pickle decode for {checkpoint_name} failed")
+            metas = _METAS_ADAPTER.validate_json(body)
+        except ValidationError as e:
+            logger.exception(f"load_metas json validation for {checkpoint_name} failed")
             return JSONResponse(content=str(e), status_code=400)
         return wrap_exception(lambda: ps.load_metas(metas))
 

checkpoint_engine/join_cli.py

@@ -6,36 +6,47 @@
 
 The remote side must already have done ``gather_metas`` so that its
 ``ps.get_metas()`` returns a usable ``dict[int, MemoryBufferMetaList]``,
-and the metas pickle bytes must be reachable via either a file path or
+and the metas JSON bytes must be reachable via either a file path or
 an HTTP URL.
 
 Usage (one process per local GPU, e.g. via torchrun):
 
+    # From a local file (e.g. shared moonfs):
     torchrun --nproc-per-node N -m checkpoint_engine.join_cli \\
-        --load-metas-file /path/to/metas.pkl \\
+        --load-metas-file /path/to/metas.json \\
         --endpoint http://localhost:19730 \\
         --inference-parallel-size N \\
         [--checkpoint-name <name>]
 
+    # Or directly from the source ParameterServer's HTTP endpoint:
+    torchrun --nproc-per-node N -m checkpoint_engine.join_cli \\
+        --metas-url http://main-ps-host:19710/v1/checkpoints/<name>/metas \\
+        --endpoint http://localhost:19730 \\
+        --inference-parallel-size N
+
 Environment variables (same as ``torchrun`` sets): ``RANK``, ``WORLD_SIZE``,
 ``LOCAL_RANK``, ``MASTER_ADDR``, ``MASTER_PORT``.
 """
 
 import argparse
 import os
-import pickle
 import time
 from collections.abc import Callable
 from contextlib import contextmanager
 
 import httpx
 from loguru import logger
+from pydantic import TypeAdapter
 
 import checkpoint_engine.distributed as dist
 from checkpoint_engine import request_inference_to_update
+from checkpoint_engine.data_types import MemoryBufferMetaList
 from checkpoint_engine.ps import ParameterServer
 
 
+_METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList])
+
+
 @contextmanager
 def _timer(msg: str):
     start = time.perf_counter()
@@ -78,14 +89,14 @@ def req(socket_paths: list[tuple[str, str]]) -> None:
     return req
 
 
-def _load_metas(args: argparse.Namespace) -> dict:
+def _load_metas(args: argparse.Namespace) -> dict[int, MemoryBufferMetaList]:
     if args.load_metas_file:
         with open(args.load_metas_file, "rb") as f:
-            return pickle.load(f)
+            return _METAS_ADAPTER.validate_json(f.read())
     if args.metas_url:
         resp = httpx.get(args.metas_url, timeout=300.0)
         resp.raise_for_status()
-        return pickle.loads(resp.content)
+        return _METAS_ADAPTER.validate_json(resp.content)
     raise ValueError("either --load-metas-file or --metas-url is required")
 
 
@@ -116,11 +127,11 @@ def main() -> None:
         description="Join an existing P2P weight world via mooncake RDMA"
     )
     src = parser.add_mutually_exclusive_group(required=True)
-    src.add_argument("--load-metas-file", type=str, help="Path to a metas pickle file")
+    src.add_argument("--load-metas-file", type=str, help="Path to a metas JSON file")
     src.add_argument(
         "--metas-url",
         type=str,
-        help="HTTP URL returning a metas pickle (application/octet-stream)",
+        help="HTTP URL returning a metas JSON (application/json)",
     )
     parser.add_argument(
         "--endpoint",

examples/update.py

@@ -1,7 +1,6 @@
 import argparse
 import json
 import os
-import pickle
 import time
 from collections import defaultdict
 from collections.abc import Callable
@@ -11,13 +10,18 @@
 import httpx
 import torch
 from loguru import logger
+from pydantic import TypeAdapter
 from safetensors import safe_open
 
 import checkpoint_engine.distributed as dist
 from checkpoint_engine import request_inference_to_update
+from checkpoint_engine.data_types import MemoryBufferMetaList
 from checkpoint_engine.ps import ParameterServer
 
 
+_METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList])
+
+
 @contextmanager
 def timer(msg: str):
     start = time.perf_counter()
@@ -110,7 +114,7 @@ def update_weights(
         ps.gather_metas(checkpoint_name)
     if save_metas_file and int(os.getenv("RANK")) == 0:
         with open(save_metas_file, "wb") as f:
-            pickle.dump(ps.get_metas(), f)
+            f.write(_METAS_ADAPTER.dump_json(ps.get_metas()))
 
     if update_method == "broadcast" or update_method == "all":
         with timer("Update weights without setting ranks"):
@@ -135,7 +139,7 @@ def join(
 ):
     assert load_metas_file, "load_metas_file is required"
     with open(load_metas_file, "rb") as f:
-        metas = pickle.load(f)
+        metas = _METAS_ADAPTER.validate_json(f.read())
     ps.init_process_group()
     check_vllm_ready(endpoint, inference_parallel_size, uds)
     dist.barrier()

tests/test_api.py

@@ -1,34 +1,67 @@
 """CPU-only tests for the metas endpoints in api.py."""
 
-import pickle
 from unittest.mock import MagicMock
 
 import pytest
+import torch
 from fastapi.testclient import TestClient
+from pydantic import TypeAdapter
 
 from checkpoint_engine.api import _init_api
+from checkpoint_engine.data_types import (
+    MemoryBufferMetaList,
+    MemoryBufferMetas,
+    ParameterMeta,
+)
+
+
+_METAS_ADAPTER = TypeAdapter(dict[int, MemoryBufferMetaList])
+
+
+def _make_meta(rdma_device: str, ip: str) -> MemoryBufferMetaList:
+    return MemoryBufferMetaList(
+        p2p_store_addr=f"{ip}:12345",
+        rdma_device=rdma_device,
+        memory_buffer_metas_list=[
+            MemoryBufferMetas(
+                metas=[
+                    ParameterMeta(
+                        name="w",
+                        dtype=torch.float16,
+                        shape=torch.Size([2, 3]),
+                        aligned_size=12,
+                    )
+                ],
+                ptr=0x12345678,
+                size=1024,
+            )
+        ],
+    )
 
 
 @pytest.fixture
-def fake_metas() -> dict:
-    # Mimic ParameterServer.get_metas() return shape (dict[int, ...]),
-    # exact value type doesn't matter for the round-trip test.
-    return {0: {"foo": [1, 2, 3]}, 1: {"bar": "baz"}}
+def fake_metas() -> dict[int, MemoryBufferMetaList]:
+    return {
+        0: _make_meta("mlx5_0", "192.168.1.1"),
+        1: _make_meta("mlx5_1", "192.168.1.1"),
+    }
 
 
 @pytest.fixture
-def ps_mock(fake_metas: dict) -> MagicMock:
+def ps_mock(fake_metas: dict[int, MemoryBufferMetaList]) -> MagicMock:
     ps = MagicMock()
     ps.get_metas.return_value = fake_metas
     return ps
 
 
-def test_get_metas_returns_pickle_bytes(ps_mock: MagicMock, fake_metas: dict) -> None:
+def test_get_metas_returns_json(
+    ps_mock: MagicMock, fake_metas: dict[int, MemoryBufferMetaList]
+) -> None:
     client = TestClient(_init_api(ps_mock))
     resp = client.get("/v1/checkpoints/my-ckpt/metas")
     assert resp.status_code == 200
-    assert resp.headers["content-type"] == "application/octet-stream"
-    assert pickle.loads(resp.content) == fake_metas
+    assert resp.headers["content-type"] == "application/json"
+    assert _METAS_ADAPTER.validate_json(resp.content) == fake_metas
     ps_mock.get_metas.assert_called_once_with()
 
 
@@ -40,40 +73,57 @@ def test_get_metas_propagates_ps_error(ps_mock: MagicMock) -> None:
     assert "metas not gathered yet" in resp.text
 
 
-def test_load_metas_decodes_and_calls_ps(ps_mock: MagicMock, fake_metas: dict) -> None:
+def test_load_metas_decodes_and_calls_ps(
+    ps_mock: MagicMock, fake_metas: dict[int, MemoryBufferMetaList]
+) -> None:
     client = TestClient(_init_api(ps_mock))
     resp = client.post(
         "/v1/checkpoints/my-ckpt/load-metas",
-        content=pickle.dumps(fake_metas),
-        headers={"content-type": "application/octet-stream"},
+        content=_METAS_ADAPTER.dump_json(fake_metas),
+        headers={"content-type": "application/json"},
     )
     assert resp.status_code == 200
     ps_mock.load_metas.assert_called_once_with(fake_metas)
 
 
-def test_load_metas_rejects_bad_pickle(ps_mock: MagicMock) -> None:
+def test_load_metas_rejects_bad_json(ps_mock: MagicMock) -> None:
+    client = TestClient(_init_api(ps_mock))
+    resp = client.post(
+        "/v1/checkpoints/my-ckpt/load-metas",
+        content=b"not a valid json",
+    )
+    assert resp.status_code == 400
+    ps_mock.load_metas.assert_not_called()
+
+
+def test_load_metas_rejects_schema_mismatch(ps_mock: MagicMock) -> None:
+    """JSON that parses but doesn't match MemoryBufferMetaList shape -> 400."""
     client = TestClient(_init_api(ps_mock))
     resp = client.post(
         "/v1/checkpoints/my-ckpt/load-metas",
-        content=b"not a valid pickle",
+        content=b'{"0": {"foo": "bar"}}',
     )
     assert resp.status_code == 400
     ps_mock.load_metas.assert_not_called()
 
 
-def test_load_metas_propagates_ps_error(ps_mock: MagicMock, fake_metas: dict) -> None:
+def test_load_metas_propagates_ps_error(
+    ps_mock: MagicMock, fake_metas: dict[int, MemoryBufferMetaList]
+) -> None:
     ps_mock.load_metas.side_effect = RuntimeError("rdma device mismatch")
     client = TestClient(_init_api(ps_mock))
     resp = client.post(
         "/v1/checkpoints/my-ckpt/load-metas",
-        content=pickle.dumps(fake_metas),
+        content=_METAS_ADAPTER.dump_json(fake_metas),
     )
     assert resp.status_code == 500
     assert "rdma device mismatch" in resp.text
 
 
-def test_round_trip_get_then_load(ps_mock: MagicMock, fake_metas: dict) -> None:
-    """Pickle bytes returned by GET /metas must be accepted by POST /load-metas."""
+def test_round_trip_get_then_load(
+    ps_mock: MagicMock, fake_metas: dict[int, MemoryBufferMetaList]
+) -> None:
+    """JSON bytes returned by GET /metas must be accepted by POST /load-metas."""
     client = TestClient(_init_api(ps_mock))
     get_resp = client.get("/v1/checkpoints/source/metas")
     assert get_resp.status_code == 200