Repeated `ParameterServer.update()` calls reuse a stale store-based barrier key

[RunningLeon]

Repeated ParameterServer.update() calls reuse a stale store-based barrier key

Summary

ParameterServer.store_based_barrier() calls PyTorch's private _store_based_barrier() with a fixed group_name on the long-lived root TCPStore. The helper is a one-shot process-group-initialization barrier, not a cyclic barrier. Its counter and last_worker keys remain in the store after the first call.

As a result, the barrier at the end of every ParameterServer.update() only synchronizes the first update. Every later update can return before all ParameterServer ranks have finished.

This is especially harmful for targeted P2P updates: a non-target source rank may leave update() and release its P2PStore/Mooncake segment while a target rank is still reading shards owned by that source rank.

Reproduction with the official vLLM example path

This reproduction is based directly on the project's official examples/update.py and uses two vLLM instances. It keeps the official split_tensors, ParameterServer, and request_inference_to_update(.../collective_rpc) path. It only changes the orchestration needed for the failing sequence:

1. launch two independent empty-weight vLLM TP1 instances;
2. map one vLLM endpoint to each of two ParameterServer ranks;
3. Broadcast to both instances;
4. call vLLM sleep(level=2) and wake_up(tags=weights) on both;
5. perform a targeted P2P update with ranks=[0].

Requirements

• two CUDA GPUs;
• a sharded safetensors checkpoint containing model.safetensors.index.json;
• checkpoint-engine with the P2P extra and an RDMA device available;
• vLLM with sleep mode and the checkpoint-engine worker extension enabled.

The failure below was reproduced with checkpoint-engine 0.4.2 and vLLM. The affected barrier implementation is also present on current main at the time of filing.

1. Start two empty-weight vLLM TP1 instances

Set MODEL to a small sharded model. Adapt the GPU IDs, ports, and RDMA devices to the host.

export MODEL=/path/to/sharded-model

CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 \
python -m vllm.entrypoints.openai.api_server \
  --host 127.0.0.1 --port 24201 \
  --model "$MODEL" --served-model-name checkpoint-engine-repro \
  --max-model-len 4096 --enforce-eager --gpu-memory-utilization 0.2 \
  --enable-sleep-mode --load-format dummy \
  --worker-extension-cls checkpoint_engine.worker.VllmColocateWorkerExtension \
  > vllm0.log 2>&1 &

CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 \
python -m vllm.entrypoints.openai.api_server \
  --host 127.0.0.1 --port 24202 \
  --model "$MODEL" --served-model-name checkpoint-engine-repro \
  --max-model-len 4096 --enforce-eager --gpu-memory-utilization 0.2 \
  --enable-sleep-mode --load-format dummy \
  --worker-extension-cls checkpoint_engine.worker.VllmColocateWorkerExtension \
  > vllm1.log 2>&1 &

2. Save this modified official example as examples/repro_repeated_update_vllm.py

#!/usr/bin/env python3
"""Reproduce repeated-update failure with two vLLM TP1 instances.

This is a minimal modification of checkpoint-engine's official
``examples/update.py``.  It keeps the official tensor splitter and the same
``ParameterServer`` / ``request_inference_to_update`` path.  The only material
changes are:

1. map one independent vLLM TP1 endpoint to each ParameterServer rank;
2. run Broadcast, vLLM level-2 sleep/wake(weights), then targeted P2P;
3. target only inference rank 0 in the second update.
"""

import argparse
import json
import os
import time
from collections import defaultdict

import httpx
import torch
from safetensors import safe_open

import checkpoint_engine.distributed as dist
from checkpoint_engine import request_inference_to_update
from checkpoint_engine.ps import ParameterServer


rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])


def log(event: str, **fields) -> None:
    print(
        json.dumps({"time": time.time(), "rank": rank, "event": event, **fields}),
        flush=True,
    )


# Unchanged from the official examples/update.py.
def split_tensors(
    checkpoint_path: str, rank: int, world_size: int
) -> dict[str, torch.Tensor]:
    index_fn = os.path.join(checkpoint_path, "model.safetensors.index.json")
    with open(index_fn) as f:
        weight_map: dict[str, str] = json.load(f)["weight_map"]
    weights_per_rank = (len(weight_map) + world_size - 1) // world_size
    fn_tensors: dict[str, list[str]] = defaultdict(list)
    weight_keys = list(weight_map.items())
    for name, file in weight_keys[
        rank * weights_per_rank : (rank + 1) * weights_per_rank
    ]:
        fn_tensors[file].append(name)
    named_tensors = {}
    for file, names in fn_tensors.items():
        with safe_open(os.path.join(checkpoint_path, file), framework="pt") as f:
            for name in names:
                named_tensors[name] = f.get_tensor(name)
    return named_tensors


def post(endpoint: str, path: str, **kwargs) -> None:
    with httpx.Client(timeout=120.0, trust_env=False) as client:
        response = client.post(endpoint + path, **kwargs)
        response.raise_for_status()


def check_vllm_ready(endpoint: str) -> None:
    while True:
        try:
            with httpx.Client(timeout=10.0, trust_env=False) as client:
                response = client.get(endpoint + "/health")
                response.raise_for_status()
            return
        except (httpx.ConnectError, httpx.HTTPStatusError):
            time.sleep(5)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Official-example-based vLLM repeated-update reproducer"
    )
    parser.add_argument("--checkpoint-path", required=True)
    parser.add_argument(
        "--endpoints",
        required=True,
        help="Comma-separated endpoints, one independent vLLM TP1 server per rank",
    )
    parser.add_argument("--checkpoint-name", default="vllm-repeated-update-repro")
    args = parser.parse_args()

    endpoints = [item.rstrip("/") for item in args.endpoints.split(",")]
    if world_size != 2 or len(endpoints) != world_size:
        raise ValueError("This reproducer requires two PS ranks and two vLLM endpoints")
    endpoint = endpoints[rank]

    ps = ParameterServer(auto_pg=True)
    named_tensors = split_tensors(args.checkpoint_path, rank, world_size)

    # Keep the initialization sequence used by the official example.
    ps.init_process_group()
    dist.barrier()
    ps.register_checkpoint(
        args.checkpoint_name,
        files=[],
        named_tensors=named_tensors,
        use_inplace_pin_memory=False,
    )
    check_vllm_ready(endpoint)
    dist.barrier()
    ps.gather_metas(args.checkpoint_name)

    def request_vllm(socket_paths: list[tuple[str, str]]) -> None:
        # Each PS rank owns one independent TP1 vLLM instance.
        device_uuid, zmq_handle = socket_paths[rank]
        log("callback.begin", endpoint=endpoint, device_uuid=device_uuid)
        request_inference_to_update(
            endpoint + "/collective_rpc",
            {device_uuid: zmq_handle},
            timeout=120.0,
        )
        log("callback.end", endpoint=endpoint)

    log("broadcast.begin")
    ps.update(args.checkpoint_name, request_vllm, ranks=None)
    log("broadcast.end")

    log("sleep.begin", level=2)
    post(endpoint, "/sleep?level=2")
    log("sleep.end")
    log("wake_weights.begin")
    post(endpoint, "/wake_up", params=[("tags", "weights")])
    log("wake_weights.end")

    log("p2p.begin", target_ranks=[0])
    ps.update(args.checkpoint_name, request_vllm, ranks=[0])
    log("p2p.end", target_ranks=[0])

    if rank == 0:
        post(
            endpoint,
            "/wake_up",
            params=[("tags", "kv_cache"), ("tags", "scheduling")],
        )
    log("case.passed")


if __name__ == "__main__":
    main()

3. Run the reproduction

export PS_P2P_STORE_RDMA_DEVICES=mlx5_0,mlx5_1  # adapt to the host
export NO_PROXY=127.0.0.1,localhost

timeout 120s torchrun --standalone --nproc-per-node=2 \
  examples/repro_repeated_update_vllm.py \
  --checkpoint-path "$MODEL" \
  --endpoints http://127.0.0.1:24201,http://127.0.0.1:24202

Actual result

The initial Broadcast and both vLLM sleep/wakeup calls succeed. During the targeted P2P update, ParameterServer rank 1 returns first and destroys its local Mooncake segment. Rank 0 then attempts to read a shard owned by rank 1 and hangs:

rank=1 event=p2p.end target_ranks=[0]
rank=1 event=case.passed
removeSegmentDesc <rank-1-segment> finish
rank=0 event=callback.begin endpoint=http://127.0.0.1:24201
Local segment descriptor not found
Unsupported segment descriptor
# timeout exits with code 124

This reproduction exercises the project's official vLLM integration path.

The sleep/wakeup cycle is not what corrupts the barrier state; it makes the timing window deterministic enough to expose the stale barrier. The underlying requirement is at least two update() calls on the same long-lived ParameterServer/root store, followed by asymmetric work across ranks.

Root cause

The barrier currently uses these keys on every invocation:

store_based_barrier_key:parameter_server_barrier
store_based_barrier_key:parameter_server_barrier:last_worker

For world_size=2, the first call leaves this state in the persistent store:

counter = 2
last_worker = 1

On the second call, the counter becomes 3 and 4, so no rank observes worker_count == world_size. However, last_worker already exists, so both ranks return immediately from store.wait().

The persistent root store is intentional and should remain shared. Process-group rendezvous keys are already isolated with PrefixStore, but the ParameterServer completion barrier has no per-update generation.

PyTorch documents _store_based_barrier() as an initialization helper for init_process_group()/new_group(), not as a generic cyclic replacement for barrier().

Relevant history:

• fix: use tcp store_based_barrier to control p2p update synchronization #51 added the all-ParameterServer-rank completion barrier.
• refactor: use a shared TCPStore in ParameterServer and create ProcessGroup using PrefixStore #82 introduced the shared root TCPStore and per-ProcessGroup PrefixStore namespaces.

Expected behavior

Every invocation of ParameterServer.store_based_barrier() must wait for all ParameterServer ranks participating in that invocation, including when the same ParameterServer and root TCPStore are reused across many updates.

The reproduction above should print case.passed on both ranks and exit successfully.

Suggested fix

Keep the shared root TCPStore, but give every barrier invocation a unique generation, for example:

self._store_barrier_counter += 1
_store_based_barrier(
    store=self._store,
    group_name=f"parameter_server_barrier-{self._store_barrier_counter}",
    ...,
)

The barrier counter should be independent of the ProcessGroup counter because auto_pg=False may reuse an externally managed ProcessGroup without calling ParameterServer.init_process_group().

Deleting the old last_worker key is not sufficient: the arrival counter also persists, and safely resetting both keys would require an additional acknowledgement/cleanup protocol.

Regression coverage

A regression test should:

1. invoke store_based_barrier() twice for two ranks on the same TCPStore;
2. verify that each invocation uses a distinct namespace;
3. verify that the shared root TCPStore itself is reused;
4. cover Broadcast followed by targeted P2P in an end-to-end test, using the vLLM reproducer above.