feat: add full-disk /pull_weights synchronization for vLLM

docker/patch/latest/vllm.patch

@@ -166,3 +166,96 @@ diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_e
          )
 
          self.attn = MMEncoderAttention(
+diff --git a/vllm/entrypoints/serve/dev/rlhf/api_router.py b/vllm/entrypoints/serve/dev/rlhf/api_router.py
+--- a/vllm/entrypoints/serve/dev/rlhf/api_router.py
++++ b/vllm/entrypoints/serve/dev/rlhf/api_router.py
+@@ -220,6 +220,28 @@ async def update_weight_version(
+     return JSONResponse(content={"success": True, "new_version": new_version})
+
+
++@router.post("/weights_checker")
++async def check_weights(raw_request: Request):
++    body = await raw_request.json()
++    action = body.get("action")
++    if action not in {"snapshot", "reset_tensors", "compare"}:
++        raise HTTPException(
++            status_code=HTTPStatus.BAD_REQUEST.value, detail="unsupported action"
++        )
++    results = await engine_client(raw_request).collective_rpc(
++        method="check_weights", kwargs={"action": action}
++    )
++    success = all(
++        isinstance(result, dict) and result.get("success") for result in results
++    )
++    return JSONResponse(
++        content={"success": success, "results": results},
++        status_code=(
++            HTTPStatus.OK.value if success else HTTPStatus.BAD_REQUEST.value
++        ),
++    )
++
++
+ @router.get("/weight_info")
+ async def weight_info(raw_request: Request):
+     weight_version = await engine_client(raw_request).get_weight_version()
+diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
+--- a/vllm/v1/worker/gpu/model_runner.py
++++ b/vllm/v1/worker/gpu/model_runner.py
+@@ -492,6 +492,44 @@ class GPUModelRunner(LoRAModelRunnerMixin):
+         if not isinstance(speculator, DraftModelSpeculator):
+             return None
+         return speculator.model
++
++    def check_weights(self, action: str) -> dict[str, Any]:
++        tensors = dict(self._weight_check_state())
++        if action == "snapshot":
++            self._weight_check_snapshot = {
++                name: tensor.detach().cpu().clone() for name, tensor in tensors.items()
++            }
++            return {"success": True}
++        if action == "reset_tensors":
++            ignored = ("cos_sin_cache", "inv_freq", "freqs_cis", "_weight_fp32")
++            with torch.inference_mode():
++                for name, tensor in tensors.items():
++                    if any(pattern in name for pattern in ignored):
++                        continue
++                    if tensor.is_floating_point():
++                        tensor.normal_()
++                    else:
++                        tensor.random_()
++            return {"success": True}
++        if action == "compare":
++            snapshot = getattr(self, "_weight_check_snapshot", None)
++            if snapshot is None:
++                return {
++                    "success": False,
++                    "error": "compare requires a prior snapshot",
++                }
++            mismatches = sorted(tensors.keys() ^ snapshot.keys())
++            mismatches.extend(
++                name
++                for name in tensors.keys() & snapshot.keys()
++                if not torch.equal(tensors[name].detach().cpu(), snapshot[name])
++            )
++            return {"success": not mismatches, "mismatches": mismatches[:32]}
++        return {"success": False, "error": f"unsupported action: {action!r}"}
++
++    def _weight_check_state(self):
++        yield from self.model.named_parameters()
++        yield from self.model.named_buffers()
+
+     def reload_weights(self, *args, **kwargs) -> None:
+         # TODO(Wentao): Use full version instead of import when fully migrated to v2
+diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
+--- a/vllm/v1/worker/gpu_worker.py
++++ b/vllm/v1/worker/gpu_worker.py
+@@ -893,6 +893,9 @@ class Worker(WorkerBase):
+
+     def get_model(self) -> nn.Module:
+         return self.model_runner.get_model()
++
++    def check_weights(self, action: str) -> dict:
++        return self.model_runner.check_weights(action)
+
+     def get_draft_model(self) -> nn.Module | None:
+         return self.model_runner.get_draft_model()

tests/utils/test_checkpoint_receiver.py

@@ -0,0 +1,68 @@
+from pathlib import Path
+
+import pytest
+
+from vime.backends.vllm_utils import checkpoint_receiver
+
+
+def _checkpoint(source_dir: Path, version: int, content: bytes) -> Path:
+    checkpoint = source_dir / f"weight_v{version:06d}"
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "model.safetensors").write_bytes(content)
+    return checkpoint
+
+
+@pytest.mark.unit
+def test_materialize_checkpoint_replaces_local_copy(tmp_path: Path) -> None:
+    source_dir = tmp_path / "published"
+    _checkpoint(source_dir, 1, b"new")
+    local_dir = tmp_path / "local"
+    local_dir.mkdir()
+    (local_dir / "model.safetensors").write_bytes(b"old")
+
+    result = checkpoint_receiver.materialize_checkpoint(
+        source_dir=str(source_dir),
+        local_checkpoint_dir=str(local_dir),
+        target_version=1,
+    )
+
+    assert result == {
+        "success": True,
+        "version": 1,
+        "local_checkpoint_dir": str(local_dir),
+    }
+    assert (local_dir / "model.safetensors").read_bytes() == b"new"
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("version", [0, -1, True, "1"])
+def test_materialize_checkpoint_rejects_invalid_version(tmp_path: Path, version: object) -> None:
+    with pytest.raises(ValueError, match="positive integer"):
+        checkpoint_receiver.materialize_checkpoint(
+            source_dir=str(tmp_path / "published"),
+            local_checkpoint_dir=str(tmp_path / "local"),
+            target_version=version,  # type: ignore[arg-type]
+        )
+
+
+@pytest.mark.unit
+def test_failed_copy_preserves_local_checkpoint(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    source_dir = tmp_path / "published"
+    _checkpoint(source_dir, 1, b"new")
+    local_dir = tmp_path / "local"
+    local_dir.mkdir()
+    (local_dir / "model.safetensors").write_bytes(b"old")
+
+    def fail_copy(*args: object, **kwargs: object) -> None:
+        raise OSError("copy failed")
+
+    monkeypatch.setattr(checkpoint_receiver.shutil, "copytree", fail_copy)
+
+    with pytest.raises(OSError, match="copy failed"):
+        checkpoint_receiver.materialize_checkpoint(
+            source_dir=str(source_dir),
+            local_checkpoint_dir=str(local_dir),
+            target_version=1,
+        )
+
+    assert (local_dir / "model.safetensors").read_bytes() == b"old"

tests/utils/test_vllm_engine.py

@@ -515,6 +515,20 @@ def fake_post_fail(endpoint: str, payload: dict) -> dict:
     assert vllm_engine._weight_version == "old"
 
 
+@pytest.mark.unit
+@pytest.mark.parametrize("action", ["snapshot", "reset_tensors", "compare"])
+def test_check_weights_posts_weights_checker_action(vllm_engine, monkeypatch, action):
+    calls: list[tuple[str, dict]] = []
+
+    def fake_make_request(endpoint: str, payload: dict):
+        calls.append((endpoint, payload))
+        return {"action": action, "results": []}
+
+    monkeypatch.setattr(vllm_engine, "_make_request", fake_make_request)
+    assert vllm_engine.check_weights(action) == {"action": action, "results": []}
+    assert calls == [("weights_checker", {"action": action})]
+
+
 @pytest.mark.unit
 def test_get_weight_version_reads_vllm_weight_info(vllm_engine, monkeypatch):
     monkeypatch.setattr(

vime/backends/vllm_utils/checkpoint_receiver.py

@@ -0,0 +1,51 @@
+"""Copy a published checkpoint to rollout-host-local storage."""
+
+from __future__ import annotations
+
+import shutil
+import uuid
+from pathlib import Path
+
+
+def materialize_checkpoint(
+    *,
+    local_checkpoint_dir: str,
+    source_dir: str,
+    target_version: int,
+) -> dict[str, object]:
+    """Copy one complete checkpoint and atomically make it active."""
+
+    if isinstance(target_version, bool) or not isinstance(target_version, int) or target_version <= 0:
+        raise ValueError("target_version must be a positive integer")
+
+    source = Path(source_dir).expanduser() / f"weight_v{target_version:06d}"
+    if not source.is_dir():
+        raise FileNotFoundError(f"checkpoint does not exist: {source}")
+
+    destination = Path(local_checkpoint_dir).expanduser()
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    if destination.exists() and destination.resolve() == source.resolve():
+        raise ValueError("local checkpoint directory must differ from the source")
+
+    suffix = uuid.uuid4().hex
+    staging = destination.parent / f".{destination.name}.staging-{suffix}"
+    backup = destination.parent / f".{destination.name}.backup-{suffix}"
+
+    try:
+        shutil.copytree(source, staging)
+        if destination.exists():
+            destination.replace(backup)
+        staging.replace(destination)
+    except Exception:
+        if backup.exists() and not destination.exists():
+            backup.replace(destination)
+        raise
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+    shutil.rmtree(backup, ignore_errors=True)
+    return {
+        "success": True,
+        "version": target_version,
+        "local_checkpoint_dir": str(destination),
+    }

vime/backends/vllm_utils/vllm_engine.py

@@ -355,8 +355,7 @@ def resume_memory_occupation(self, tags: list[str] = None):
         return response.json()
 
     def check_weights(self, action: str):
-        del action
-        return {"ok": True, "supported": False}
+        return self._make_request("weights_checker", {"action": action})
 
     def init_weight_transfer_engine(self, payload: dict) -> dict:
         return self._make_request("init_weight_transfer_engine", payload)