feat: expose ParameterServer metas via HTTP for cross-process join (MoonshotAI#90)

* feat: expose ParameterServer metas via HTTP for cross-process join

Add two HTTP endpoints (GET/POST /v1/checkpoints/{name}/{metas,load-metas})
and a standalone `python -m checkpoint_engine.join_cli` entrypoint, so a
new ParameterServer instance can join an existing P2P weight world over
mooncake RDMA without re-reading checkpoints from disk.

Motivation: in elastic-rollout setups (e.g. mshrl), a long-running training
job already holds pinned CPU weight buffers registered with the mooncake
P2PStore. Newly-started inference replicas should be able to pull these
weights over RDMA instead of re-converting the checkpoint from disk.

Changes:
* api.py: GET /v1/checkpoints/{name}/metas returns pickle.dumps(ps.get_metas())
  as application/octet-stream; POST /v1/checkpoints/{name}/load-metas accepts
  the same bytes and feeds them into ps.load_metas(). Bad pickle is rejected
  with 400; PS errors are surfaced as 500.
* join_cli.py: `python -m checkpoint_engine.join_cli` -- the join() flow from
  examples/update.py, packaged as a first-class CLI under the published
  package so consumers can invoke it without checking out the source tree.
  Accepts metas from either a local pickle file or a remote HTTP URL.
* tests/test_api.py: 6 CPU-only tests covering pickle round-trip, ps-error
  propagation, bad-input rejection, and a GET-then-POST chain that validates
  the new endpoints are mutually consistent.

Verified end-to-end on a 2-node launchpad job: 14.5 GiB Qwen2.5-7B weights
transferred from main to elastic in 1.49s over real RDMA (4 mlx5_bond HCAs)
vs 6.49s over TCP fallback in environments without RDMA passthrough.

* 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).

* refactor: merge join_cli into examples/update.py

join_cli was a duplicate of examples/update.py:join() with an extra
--metas-url flag. Add the flag to update.py, remove join_cli.

* refactor: simplify metas endpoints to /v1/metas

The checkpoint_name path param was never read — ps.get_metas() and
ps.load_metas() act on a single global field. Drop it from the URL.

* refactor: drop _METAS_ADAPTER from api.py

* test: set Content-Type on POST /v1/metas requests

---------

Co-authored-by: liuweijian <liuweijian@msh.team>
Co-authored-by: 刘伟健 <liuweijian@moonshot.ai>

Author: 3 people

checkpoint_engine/api.py

@@ -3,11 +3,12 @@
 
 import fastapi
 import httpx
-from fastapi import Request
+from fastapi import HTTPException, Request
 from fastapi.responses import JSONResponse, Response
 from loguru import logger
 from pydantic import BaseModel
 
+from checkpoint_engine.data_types import MemoryBufferMetaList
 from checkpoint_engine.ps import ParameterServer
 
 
@@ -79,6 +80,18 @@ async def healthz() -> Response:
     async def gather_metas(checkpoint_name: str) -> Response:
         return wrap_exception(lambda: ps.gather_metas(checkpoint_name))
 
+    @app.get("/v1/metas")
+    async def get_metas() -> dict[int, MemoryBufferMetaList]:
+        try:
+            return ps.get_metas()
+        except Exception as e:
+            logger.exception("get_metas failed")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+    @app.post("/v1/metas")
+    async def load_metas(metas: dict[int, MemoryBufferMetaList]) -> Response:
+        return wrap_exception(lambda: ps.load_metas(metas))
+
     @app.post("/v1/checkpoints/{checkpoint_name}/update")
     async def update(checkpoint_name: str, req: UpdateRequest) -> Response:
         def update_func(socket_paths: list[tuple[str, str]]):

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"):
@@ -127,15 +131,22 @@ def update_weights(
 def join(
     ps: ParameterServer,
     checkpoint_name: str,
-    load_metas_file: str,
+    load_metas_file: str | None,
+    metas_url: str | None,
     req_func: Callable[[list[tuple[str, str]]], None],
     inference_parallel_size: int,
     endpoint: str,
     uds: str | None = None,
 ):
-    assert load_metas_file, "load_metas_file is required"
-    with open(load_metas_file, "rb") as f:
-        metas = pickle.load(f)
+    if load_metas_file:
+        with open(load_metas_file, "rb") as f:
+            metas = _METAS_ADAPTER.validate_json(f.read())
+    elif metas_url:
+        resp = httpx.get(metas_url, timeout=300.0)
+        resp.raise_for_status()
+        metas = _METAS_ADAPTER.validate_json(resp.content)
+    else:
+        raise ValueError("either load_metas_file or metas_url is required")
     ps.init_process_group()
     check_vllm_ready(endpoint, inference_parallel_size, uds)
     dist.barrier()
@@ -152,7 +163,19 @@ def join(
     parser = argparse.ArgumentParser(description="Update weights example")
     parser.add_argument("--checkpoint-path", type=str, default=None)
     parser.add_argument("--save-metas-file", type=str, default=None)
-    parser.add_argument("--load-metas-file", type=str, default=None)
+    metas_src = parser.add_mutually_exclusive_group()
+    metas_src.add_argument(
+        "--load-metas-file",
+        type=str,
+        default=None,
+        help="Path to a metas JSON file (triggers join mode)",
+    )
+    metas_src.add_argument(
+        "--metas-url",
+        type=str,
+        default=None,
+        help="HTTP URL returning a metas JSON (triggers join mode)",
+    )
     parser.add_argument("--sleep-time", type=int, default=0)
     parser.add_argument("--endpoint", type=str, default="http://localhost:19730")
     parser.add_argument("--inference-parallel-size", type=int, default=8)
@@ -167,11 +190,12 @@ def join(
     req_func = req_inference(args.endpoint, args.inference_parallel_size, args.uds)
     dist.use_backend(args.custom_dist)
     ps = ParameterServer(auto_pg=True)
-    if args.load_metas_file:
+    if args.load_metas_file or args.metas_url:
         join(
             ps,
             args.checkpoint_name,
             args.load_metas_file,
+            args.metas_url,
             req_func,
             args.inference_parallel_size,
             args.endpoint,

tests/test_api.py

@@ -0,0 +1,139 @@
+"""CPU-only tests for the metas endpoints in api.py."""
+
+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[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[int, MemoryBufferMetaList]) -> MagicMock:
+    ps = MagicMock()
+    ps.get_metas.return_value = fake_metas
+    return ps
+
+
+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/metas")
+    assert resp.status_code == 200
+    assert resp.headers["content-type"] == "application/json"
+    assert _METAS_ADAPTER.validate_json(resp.content) == fake_metas
+    ps_mock.get_metas.assert_called_once_with()
+
+
+def test_get_metas_propagates_ps_error(ps_mock: MagicMock) -> None:
+    ps_mock.get_metas.side_effect = RuntimeError("metas not gathered yet")
+    client = TestClient(_init_api(ps_mock))
+    resp = client.get("/v1/metas")
+    assert resp.status_code == 500
+    assert "metas not gathered yet" in resp.text
+
+
+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/metas",
+        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_json(ps_mock: MagicMock) -> None:
+    client = TestClient(_init_api(ps_mock))
+    resp = client.post(
+        "/v1/metas",
+        content=b"not a valid json",
+        headers={"content-type": "application/json"},
+    )
+    assert resp.status_code == 422
+    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 -> 422."""
+    client = TestClient(_init_api(ps_mock))
+    resp = client.post(
+        "/v1/metas",
+        content=b'{"0": {"foo": "bar"}}',
+        headers={"content-type": "application/json"},
+    )
+    assert resp.status_code == 422
+    ps_mock.load_metas.assert_not_called()
+
+
+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/metas",
+        content=_METAS_ADAPTER.dump_json(fake_metas),
+        headers={"content-type": "application/json"},
+    )
+    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[int, MemoryBufferMetaList]
+) -> None:
+    """JSON bytes returned by GET /v1/metas must be accepted by POST /v1/metas."""
+    client = TestClient(_init_api(ps_mock))
+    get_resp = client.get("/v1/metas")
+    assert get_resp.status_code == 200
+    load_resp = client.post(
+        "/v1/metas",
+        content=get_resp.content,
+        headers={"content-type": "application/json"},
+    )
+    assert load_resp.status_code == 200
+    ps_mock.load_metas.assert_called_once_with(fake_metas)