diff --git a/docker/patch/latest/vllm.patch b/docker/patch/latest/vllm.patch index 58093283..15ba63e7 100644 --- a/docker/patch/latest/vllm.patch +++ b/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() diff --git a/tests/utils/test_checkpoint_receiver.py b/tests/utils/test_checkpoint_receiver.py new file mode 100644 index 00000000..3d0ef31d --- /dev/null +++ b/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" diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py index 5f599b39..a4363686 100644 --- a/tests/utils/test_vllm_engine.py +++ b/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( diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py new file mode 100644 index 00000000..ed106881 --- /dev/null +++ b/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), + } diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 82c07e6f..2d9c606f 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/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)