From 5a5e90f9ab91b53d40301864325b9885abae36a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Sat, 18 Jul 2026 13:36:36 +0800
Subject: [PATCH 1/8] feat: add vLLM weight checks and disk receiver
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 docker/patch/latest/vllm.patch                |  93 ++++++
 tests/utils/test_checkpoint_receiver.py       | 123 +++++++
 tests/utils/test_vllm_engine.py               |  14 +
 .../vllm_utils/checkpoint_receiver.py         | 314 ++++++++++++++++++
 vime/backends/vllm_utils/vllm_engine.py       |   5 +-
 5 files changed, 547 insertions(+), 2 deletions(-)
 create mode 100644 tests/utils/test_checkpoint_receiver.py
 create mode 100644 vime/backends/vllm_utils/checkpoint_receiver.py

diff --git a/docker/patch/latest/vllm.patch b/docker/patch/latest/vllm.patch
index 580932836..15ba63e7a 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 000000000..d289b5444
--- /dev/null
+++ b/tests/utils/test_checkpoint_receiver.py
@@ -0,0 +1,123 @@
+"""CPU tests for the transactional full-disk checkpoint receiver."""
+
+from __future__ import annotations
+
+import json
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+
+from vime.backends.vllm_utils import checkpoint_receiver as receiver
+
+
+def _publish(root: Path, version: int, payload: bytes = b"weights-v1") -> Path:
+    checkpoint = root / f"weight_v{version:06d}"
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "config.json").write_text("{}", encoding="utf-8")
+    (checkpoint / "model-00001.safetensors").write_bytes(payload)
+    (checkpoint / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": len(payload)},
+                "weight_map": {"model.weight": "model-00001.safetensors"},
+            }
+        ),
+        encoding="utf-8",
+    )
+    return checkpoint
+
+
+def test_materialize_is_atomic_and_repeated_pull_is_idempotent(tmp_path: Path):
+    source = tmp_path / "published"
+    _publish(source, 1)
+    local = tmp_path / "local"
+
+    first = receiver.materialize_checkpoint(
+        source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
+    )
+    assert first["status"] == "materialized"
+    assert (local / "model-00001.safetensors").read_bytes() == b"weights-v1"
+
+    repeated = receiver.materialize_checkpoint(
+        source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
+    )
+    assert repeated["status"] == "already_applied"
+    assert repeated["manifest_sha256"] == first["manifest_sha256"]
+
+
+def test_invalid_source_and_malformed_checkpoint_are_rejected(tmp_path: Path):
+    with pytest.raises(receiver.CheckpointReceiveError, match="does not exist"):
+        receiver.materialize_checkpoint(
+            source_dir=str(tmp_path / "missing"),
+            local_checkpoint_dir=str(tmp_path / "local"),
+            target_version=1,
+        )
+
+    source = tmp_path / "published"
+    checkpoint = source / "weight_v000001"
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "model.safetensors.index.json").write_text("not-json", encoding="utf-8")
+    with pytest.raises(receiver.CheckpointReceiveError, match="malformed checkpoint index"):
+        receiver.materialize_checkpoint(
+            source_dir=str(source),
+            local_checkpoint_dir=str(tmp_path / "local"),
+            target_version=1,
+        )
+
+
+def test_stale_and_conflicting_versions_are_rejected(tmp_path: Path):
+    source = tmp_path / "published"
+    _publish(source, 1, b"one")
+    _publish(source, 2, b"two")
+    local = tmp_path / "local"
+    receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
+
+    with pytest.raises(receiver.StaleCheckpointError, match="stale checkpoint"):
+        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
+
+    # Same version with different bytes is a conflict, never an overwrite.
+    (source / "weight_v000002" / "model-00001.safetensors").write_bytes(b"changed")
+    with pytest.raises(receiver.CheckpointConflictError, match="different manifest"):
+        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
+    assert (local / "model-00001.safetensors").read_bytes() == b"two"
+
+
+def test_failed_copy_keeps_old_checkpoint_available(tmp_path: Path, monkeypatch):
+    source = tmp_path / "published"
+    _publish(source, 1, b"old")
+    _publish(source, 2, b"new")
+    local = tmp_path / "local"
+    receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
+
+    original_copy2 = receiver.shutil.copy2
+    calls = 0
+
+    def fail_after_first(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        if calls >= 2:
+            raise OSError("simulated partial copy")
+        return original_copy2(*args, **kwargs)
+
+    monkeypatch.setattr(receiver.shutil, "copy2", fail_after_first)
+    with pytest.raises(OSError, match="simulated partial copy"):
+        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
+    assert json.loads((local / receiver._CHECKPOINT_MARKER).read_text())["version"] == 1
+    assert (local / "model-00001.safetensors").read_bytes() == b"old"
+
+
+def test_concurrent_same_version_pulls_are_serialized(tmp_path: Path):
+    source = tmp_path / "published"
+    _publish(source, 1)
+    local = tmp_path / "local"
+
+    def pull():
+        return receiver.materialize_checkpoint(
+            source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
+        )["status"]
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        statuses = sorted(pool.map(lambda _: pull(), range(2)))
+    assert statuses == ["already_applied", "materialized"]
+    assert (local / "model.safetensors.index.json").is_file()
diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py
index 5f599b397..a43636864 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 000000000..405431889
--- /dev/null
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -0,0 +1,314 @@
+"""Transactional receiver for full HuggingFace checkpoints published on disk.
+
+The vLLM HTTP endpoint is intentionally kept thin.  This module owns the
+filesystem transaction so it can be tested without importing vLLM or starting
+an HTTP server.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import os
+import shutil
+import tempfile
+import threading
+import uuid
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_CHECKPOINT_MARKER = ".vime_checkpoint_manifest.json"
+_INDEX_NAMES = ("model.safetensors.index.json", "pytorch_model.bin.index.json")
+_SINGLE_FILE_WEIGHT_NAMES = ("model.safetensors", "pytorch_model.bin")
+_MATERIALIZE_LOCK = threading.RLock()
+
+
+class CheckpointReceiveError(ValueError):
+    """A malformed request or checkpoint that should be returned as HTTP 400."""
+
+    status_code = 400
+
+
+class StaleCheckpointError(CheckpointReceiveError):
+    """A checkpoint older than the active local checkpoint."""
+
+    status_code = 409
+
+
+class CheckpointConflictError(CheckpointReceiveError):
+    """A version was already applied with different checkpoint contents."""
+
+    status_code = 409
+
+
+def materialize_checkpoint(
+    *,
+    local_checkpoint_dir: str,
+    source_dir: str,
+    target_version: int,
+) -> dict[str, Any]:
+    """Materialize one published checkpoint into a host-local directory.
+
+    The source is validated and checksummed before any destination is touched.
+    Files are copied into a sibling temporary directory, then that directory is
+    swapped into place.  A failed copy or swap therefore leaves the previous
+    local checkpoint intact.
+    """
+
+    version = _parse_version(target_version)
+    source_root = _resolve_existing_directory(source_dir, "source_dir")
+    version_dir = source_root / f"weight_v{version:06d}"
+    version_dir = _resolve_existing_directory(version_dir, "checkpoint version")
+
+    local_dir = _prepare_local_directory(local_checkpoint_dir)
+    if local_dir.exists() and local_dir.resolve() == version_dir:
+        raise CheckpointReceiveError("local_checkpoint_dir must differ from source checkpoint")
+
+    with _MATERIALIZE_LOCK:
+        source_manifest = _build_checkpoint_manifest(version_dir)
+        source_manifest_hash = _manifest_hash(source_manifest["files"])
+        current = _read_marker(local_dir)
+
+        if current is not None:
+            current_version = current["version"]
+            if current_version > version:
+                raise StaleCheckpointError(
+                    f"stale checkpoint version {version}; active version is {current_version}"
+                )
+            if current_version == version:
+                if current["manifest_sha256"] != source_manifest_hash:
+                    raise CheckpointConflictError(
+                        f"checkpoint version {version} is already active with a different manifest"
+                    )
+                if _local_manifest_is_valid(local_dir, current):
+                    return _result("already_applied", version, current)
+
+        staging_dir = Path(tempfile.mkdtemp(prefix=f".{local_dir.name}.", dir=local_dir.parent))
+        try:
+            _copy_manifest_files(version_dir, staging_dir, source_manifest["files"])
+            marker = {
+                "version": version,
+                "manifest_sha256": source_manifest_hash,
+                "files": source_manifest["files"],
+            }
+            _write_json(staging_dir / _CHECKPOINT_MARKER, marker)
+            _atomic_replace_directory(staging_dir, local_dir)
+            staging_dir = None  # ownership transferred to local_dir
+        except Exception:
+            if staging_dir is not None:
+                shutil.rmtree(staging_dir, ignore_errors=True)
+            raise
+
+        return _result("materialized", version, marker)
+
+
+def _parse_version(value: int) -> int:
+    if isinstance(value, bool):
+        raise CheckpointReceiveError("target_version must be a positive integer")
+    try:
+        version = int(value)
+    except (TypeError, ValueError) as exc:
+        raise CheckpointReceiveError("target_version must be a positive integer") from exc
+    if version <= 0 or str(value).strip() != str(version):
+        raise CheckpointReceiveError("target_version must be a positive integer")
+    return version
+
+
+def _resolve_existing_directory(value: str | Path, name: str) -> Path:
+    if not isinstance(value, (str, Path)) or not str(value).strip():
+        raise CheckpointReceiveError(f"{name} must be a non-empty path")
+    path = Path(value).expanduser()
+    try:
+        resolved = path.resolve(strict=True)
+    except OSError as exc:
+        raise CheckpointReceiveError(f"{name} does not exist: {path}") from exc
+    if not resolved.is_dir() or resolved.is_symlink():
+        raise CheckpointReceiveError(f"{name} must be a directory: {path}")
+    return resolved
+
+
+def _prepare_local_directory(value: str | Path) -> Path:
+    if not isinstance(value, (str, Path)) or not str(value).strip():
+        raise CheckpointReceiveError("local_checkpoint_dir must be a non-empty path")
+    path = Path(value).expanduser()
+    if path.exists() and (path.is_symlink() or not path.is_dir()):
+        raise CheckpointReceiveError(f"local_checkpoint_dir must be a directory: {path}")
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        parent = path.parent.resolve(strict=True)
+    except OSError as exc:
+        raise CheckpointReceiveError(f"cannot prepare local checkpoint parent: {path.parent}") from exc
+    return parent / path.name
+
+
+def _build_checkpoint_manifest(root: Path) -> dict[str, Any]:
+    index_candidates = [root / name for name in _INDEX_NAMES if (root / name).exists()]
+    if any(path.is_symlink() for path in index_candidates):
+        raise CheckpointReceiveError("checkpoint index must not be a symlink")
+    indexes = [path for path in index_candidates if path.is_file()]
+    if len(indexes) > 1:
+        raise CheckpointReceiveError(
+            f"checkpoint must contain exactly one supported weight index: {', '.join(_INDEX_NAMES)}"
+        )
+
+    if not indexes:
+        direct_weights = [root / name for name in _SINGLE_FILE_WEIGHT_NAMES if (root / name).is_file()]
+        if not direct_weights or any(path.stat().st_size <= 0 for path in direct_weights):
+            raise CheckpointReceiveError(
+                "checkpoint must contain a supported weight index or a non-empty single-file weight"
+            )
+        return _manifest_files(root)
+
+    index_path = indexes[0]
+    try:
+        with index_path.open("r", encoding="utf-8") as file:
+            index = json.load(file)
+    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise CheckpointReceiveError(f"malformed checkpoint index: {index_path.name}") from exc
+
+    weight_map = index.get("weight_map") if isinstance(index, dict) else None
+    if not isinstance(weight_map, dict) or not weight_map:
+        raise CheckpointReceiveError("checkpoint index must contain a non-empty weight_map")
+    for tensor_name, filename in weight_map.items():
+        if not isinstance(tensor_name, str) or not isinstance(filename, str):
+            raise CheckpointReceiveError("checkpoint index contains a non-string weight mapping")
+        _resolve_checkpoint_file(root, filename)
+
+    return _manifest_files(root)
+
+
+def _manifest_files(root: Path) -> dict[str, Any]:
+    files: dict[str, dict[str, Any]] = {}
+    for path in sorted(root.rglob("*")):
+        if path.is_symlink():
+            raise CheckpointReceiveError(f"checkpoint contains a symlink: {path.relative_to(root)}")
+        if not path.is_file():
+            continue
+        relative = path.relative_to(root).as_posix()
+        try:
+            size = path.stat().st_size
+        except OSError as exc:
+            raise CheckpointReceiveError(f"cannot stat checkpoint file: {relative}") from exc
+        files[relative] = {"size": size, "sha256": _sha256(path)}
+    if not files:
+        raise CheckpointReceiveError("checkpoint contains no files")
+    return {"files": files}
+
+
+def _resolve_checkpoint_file(root: Path, filename: str) -> Path:
+    candidate = root / filename
+    try:
+        resolved = candidate.resolve(strict=True)
+        resolved.relative_to(root)
+    except (OSError, ValueError) as exc:
+        raise CheckpointReceiveError(f"checkpoint index references missing or unsafe file: {filename}") from exc
+    if candidate.is_symlink() or not resolved.is_file() or resolved.stat().st_size <= 0:
+        raise CheckpointReceiveError(f"checkpoint index references invalid file: {filename}")
+    return resolved
+
+
+def _copy_manifest_files(source: Path, destination: Path, files: dict[str, dict[str, Any]]) -> None:
+    for relative, metadata in files.items():
+        source_file = source / Path(relative)
+        destination_file = destination / Path(relative)
+        destination_file.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            shutil.copy2(source_file, destination_file)
+            copied_size = destination_file.stat().st_size
+            copied_hash = _sha256(destination_file)
+        except OSError as exc:
+            raise OSError(f"failed to copy checkpoint file {relative}") from exc
+        if copied_size != metadata["size"] or copied_hash != metadata["sha256"]:
+            raise OSError(f"checkpoint file changed while copying: {relative}")
+
+
+def _read_marker(local_dir: Path) -> dict[str, Any] | None:
+    marker_path = local_dir / _CHECKPOINT_MARKER
+    if not marker_path.exists():
+        return None
+    try:
+        with marker_path.open("r", encoding="utf-8") as file:
+            marker = json.load(file)
+    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
+        raise CheckpointReceiveError("local checkpoint version marker is malformed") from exc
+    if (
+        not isinstance(marker, dict)
+        or isinstance(marker.get("version"), bool)
+        or not isinstance(marker.get("version"), int)
+        or marker["version"] <= 0
+        or not isinstance(marker.get("manifest_sha256"), str)
+        or not isinstance(marker.get("files"), dict)
+    ):
+        raise CheckpointReceiveError("local checkpoint version marker is malformed")
+    return marker
+
+
+def _local_manifest_is_valid(local_dir: Path, marker: dict[str, Any]) -> bool:
+    if _manifest_hash(marker["files"]) != marker["manifest_sha256"]:
+        return False
+    for relative, metadata in marker["files"].items():
+        if not isinstance(relative, str) or not isinstance(metadata, dict):
+            return False
+        path = local_dir / Path(relative)
+        try:
+            path.resolve(strict=True).relative_to(local_dir.resolve(strict=True))
+            if path.is_symlink() or not path.is_file() or path.stat().st_size != metadata.get("size"):
+                return False
+            if _sha256(path) != metadata.get("sha256"):
+                return False
+        except (OSError, ValueError):
+            return False
+    return True
+
+
+def _manifest_hash(files: dict[str, Any]) -> str:
+    encoded = json.dumps(files, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as file:
+        for chunk in iter(lambda: file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _write_json(path: Path, value: dict[str, Any]) -> None:
+    with path.open("w", encoding="utf-8") as file:
+        json.dump(value, file, sort_keys=True, separators=(",", ":"))
+        file.flush()
+        os.fsync(file.fileno())
+
+
+def _atomic_replace_directory(staging: Path, destination: Path) -> None:
+    backup: Path | None = None
+    if destination.exists() or destination.is_symlink():
+        if destination.is_symlink() or not destination.is_dir():
+            raise CheckpointReceiveError(f"local checkpoint destination is not a directory: {destination}")
+        backup = destination.parent / f".{destination.name}.backup-{uuid.uuid4().hex}"
+        os.replace(destination, backup)
+    try:
+        os.replace(staging, destination)
+    except Exception:
+        if backup is not None and backup.exists() and not destination.exists():
+            os.replace(backup, destination)
+        raise
+    if backup is not None:
+        try:
+            shutil.rmtree(backup)
+        except OSError:
+            logger.warning("Could not remove old checkpoint backup %s", backup, exc_info=True)
+
+
+def _result(status: str, version: int, marker: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "success": True,
+        "status": status,
+        "version": version,
+        "manifest_sha256": marker["manifest_sha256"],
+        "files": len(marker["files"]),
+    }
diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py
index 82c07e6ff..ec53a8054 100644
--- a/vime/backends/vllm_utils/vllm_engine.py
+++ b/vime/backends/vllm_utils/vllm_engine.py
@@ -355,8 +355,9 @@ 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}
+        if action not in {"snapshot", "reset_tensors", "compare"}:
+            raise ValueError(f"unsupported weight check action: {action!r}")
+        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)

From 81762c80640806ceab08aaca22344292a04f1328 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Sun, 19 Jul 2026 02:38:02 +0800
Subject: [PATCH 2/8] fix: preserve checkpoint copy failure details
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 vime/backends/vllm_utils/checkpoint_receiver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py
index 405431889..5bf9585aa 100644
--- a/vime/backends/vllm_utils/checkpoint_receiver.py
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -220,7 +220,7 @@ def _copy_manifest_files(source: Path, destination: Path, files: dict[str, dict[
             copied_size = destination_file.stat().st_size
             copied_hash = _sha256(destination_file)
         except OSError as exc:
-            raise OSError(f"failed to copy checkpoint file {relative}") from exc
+            raise OSError(f"failed to copy checkpoint file {relative}: {exc}") from exc
         if copied_size != metadata["size"] or copied_hash != metadata["sha256"]:
             raise OSError(f"checkpoint file changed while copying: {relative}")
 

From 5d9e089c7627e3fa9435e889719dca795f6529b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Sun, 19 Jul 2026 17:54:33 +0800
Subject: [PATCH 3/8] fix: harden full-disk weight synchronization
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_actor_group_disk_update.py   | 137 ++++++++++++++++++
 tests/utils/test_checkpoint_receiver.py       |  72 +++++++++
 tests/utils/test_vllm_engine.py               |  51 ++++++-
 vime/backends/megatron_utils/actor.py         |  11 +-
 .../update_weight/update_weight_from_disk.py  |  31 ++--
 .../vllm_utils/checkpoint_receiver.py         |  85 ++++++++++-
 vime/backends/vllm_utils/vllm_engine.py       |   2 -
 vime/ray/actor_group.py                       |   4 +-
 vime/ray/rollout.py                           |   6 +-
 vime/ray/train_actor.py                       |   2 +-
 vime/utils/arguments.py                       |   3 +
 11 files changed, 382 insertions(+), 22 deletions(-)
 create mode 100644 tests/utils/test_actor_group_disk_update.py

diff --git a/tests/utils/test_actor_group_disk_update.py b/tests/utils/test_actor_group_disk_update.py
new file mode 100644
index 000000000..6306fe2b2
--- /dev/null
+++ b/tests/utils/test_actor_group_disk_update.py
@@ -0,0 +1,137 @@
+"""CPU tests for full-disk version ownership in ``RayTrainGroup``."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+@pytest.fixture
+def actor_group_module(monkeypatch):
+    ray = types.ModuleType("ray")
+    ray.get = lambda refs: refs
+    ray.remote = lambda *args, **kwargs: lambda value: value
+    ray.kill = lambda *args, **kwargs: None
+
+    ray_util = types.ModuleType("ray.util")
+    placement_group = types.ModuleType("ray.util.placement_group")
+    placement_group.PlacementGroup = object
+    scheduling = types.ModuleType("ray.util.scheduling_strategies")
+    scheduling.PlacementGroupSchedulingStrategy = object
+    ray.util = ray_util
+
+    ray_utils = types.ModuleType("vime.ray.utils")
+    ray_utils.NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = []
+    ray_utils.add_default_ray_env_vars = lambda values=None: values or {}
+
+    for name, module in {
+        "ray": ray,
+        "ray.util": ray_util,
+        "ray.util.placement_group": placement_group,
+        "ray.util.scheduling_strategies": scheduling,
+        "vime.ray.utils": ray_utils,
+    }.items():
+        monkeypatch.setitem(sys.modules, name, module)
+
+    sys.modules.pop("vime.ray.actor_group", None)
+    module = importlib.import_module("vime.ray.actor_group")
+    yield module
+    sys.modules.pop("vime.ray.actor_group", None)
+
+
+class _RemoteMethod:
+    def __init__(self, calls: list[dict[str, int]]) -> None:
+        self.calls = calls
+
+    def remote(self, **kwargs):
+        self.calls.append(kwargs)
+        return kwargs
+
+
+class _Actor:
+    def __init__(self, calls: list[dict[str, int]]) -> None:
+        self.update_weights = _RemoteMethod(calls)
+
+
+def _make_group(module, tmp_path: Path):
+    group = module.RayTrainGroup.__new__(module.RayTrainGroup)
+    group.args = SimpleNamespace(
+        update_weight_mode="full",
+        update_weight_transport="disk",
+        update_weight_disk_dir=str(tmp_path),
+        release_train=False,
+    )
+    group.role = "actor"
+    group._disk_weight_version = 0
+    group._actor_handlers = []
+    return group
+
+
+def test_full_disk_group_commits_version_after_reload(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    actor_calls: list[dict[str, int]] = []
+    reload_calls: list[tuple[Path, str]] = []
+    group._actor_handlers = [_Actor(actor_calls)]
+    group._reload_rollout_weights_from_disk = lambda path, version: reload_calls.append((path, version))
+
+    group.update_weights()
+
+    assert actor_calls == [{"weight_version": 1}]
+    assert reload_calls == [(tmp_path / "weight_v000001", "1")]
+    assert group._disk_weight_version == 1
+
+
+def test_full_disk_group_retries_same_version_after_reload_failure(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    actor_calls: list[dict[str, int]] = []
+    group._actor_handlers = [_Actor(actor_calls)]
+    attempts = 0
+
+    def reload_once_then_succeed(path, version):
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise RuntimeError("reload failed")
+
+    group._reload_rollout_weights_from_disk = reload_once_then_succeed
+
+    with pytest.raises(RuntimeError, match="reload failed"):
+        group.update_weights()
+    assert group._disk_weight_version == 0
+
+    group.update_weights()
+
+    assert actor_calls == [{"weight_version": 1}, {"weight_version": 1}]
+    assert group._disk_weight_version == 1
+
+
+def test_full_disk_group_retries_same_version_after_actor_failure(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    actor_calls: list[dict[str, int]] = []
+    group._actor_handlers = [_Actor(actor_calls)]
+    group._reload_rollout_weights_from_disk = lambda path, version: None
+    original_get = actor_group_module.ray.get
+    attempts = 0
+
+    def fail_once(refs):
+        nonlocal attempts
+        attempts += 1
+        if attempts == 1:
+            raise RuntimeError("actor failed")
+        return original_get(refs)
+
+    actor_group_module.ray.get = fail_once
+
+    with pytest.raises(RuntimeError, match="actor failed"):
+        group.update_weights()
+    assert group._disk_weight_version == 0
+
+    group.update_weights()
+
+    assert actor_calls == [{"weight_version": 1}, {"weight_version": 1}]
+    assert group._disk_weight_version == 1
diff --git a/tests/utils/test_checkpoint_receiver.py b/tests/utils/test_checkpoint_receiver.py
index d289b5444..f36a8bfad 100644
--- a/tests/utils/test_checkpoint_receiver.py
+++ b/tests/utils/test_checkpoint_receiver.py
@@ -121,3 +121,75 @@ def pull():
         statuses = sorted(pool.map(lambda _: pull(), range(2)))
     assert statuses == ["already_applied", "materialized"]
     assert (local / "model.safetensors.index.json").is_file()
+
+
+def test_materialize_rejects_paths_outside_configured_roots(tmp_path: Path):
+    configured_source = tmp_path / "configured-source"
+    other_source = tmp_path / "other-source"
+    _publish(configured_source, 1)
+    _publish(other_source, 1)
+    configured_local = tmp_path / "configured-local"
+
+    with pytest.raises(receiver.CheckpointReceiveError, match="configured checkpoint source"):
+        receiver.materialize_checkpoint(
+            source_dir=str(other_source),
+            local_checkpoint_dir=str(configured_local),
+            target_version=1,
+            expected_source_dir=str(configured_source),
+            expected_local_checkpoint_dir=str(configured_local),
+        )
+
+    with pytest.raises(receiver.CheckpointReceiveError, match="configured local checkpoint destination"):
+        receiver.materialize_checkpoint(
+            source_dir=str(configured_source),
+            local_checkpoint_dir=str(tmp_path / "other-local"),
+            target_version=1,
+            expected_source_dir=str(configured_source),
+            expected_local_checkpoint_dir=str(configured_local),
+        )
+
+    result = receiver.materialize_checkpoint(
+        source_dir=str(configured_source),
+        local_checkpoint_dir=str(configured_local),
+        target_version=1,
+        expected_source_dir=str(configured_source),
+        expected_local_checkpoint_dir=str(configured_local),
+    )
+    assert result["status"] == "materialized"
+
+
+def test_materialize_rejects_source_symlink(tmp_path: Path):
+    source = tmp_path / "source"
+    _publish(source, 1)
+    source_link = tmp_path / "source-link"
+    try:
+        source_link.symlink_to(source, target_is_directory=True)
+    except OSError as exc:
+        pytest.skip(f"directory symlinks are unavailable: {exc}")
+
+    with pytest.raises(receiver.CheckpointReceiveError, match="must not be a symlink"):
+        receiver.materialize_checkpoint(
+            source_dir=str(source_link),
+            local_checkpoint_dir=str(tmp_path / "local"),
+            target_version=1,
+        )
+
+
+def test_publish_checkpoint_directory_is_immutable_and_retryable(tmp_path: Path):
+    destination = tmp_path / "weight_v000001"
+    staging = tmp_path / ".weight_v000001.staging"
+    _publish(tmp_path / "first", 1).rename(staging)
+
+    assert receiver.publish_checkpoint_directory(staging, destination) == "published"
+    assert (destination / "model-00001.safetensors").read_bytes() == b"weights-v1"
+
+    retry_staging = tmp_path / ".weight_v000001.retry"
+    _publish(tmp_path / "retry", 1).rename(retry_staging)
+    assert receiver.publish_checkpoint_directory(retry_staging, destination) == "already_published"
+    assert retry_staging.exists()
+
+    conflict_staging = tmp_path / ".weight_v000001.conflict"
+    _publish(tmp_path / "conflict", 1, b"different").rename(conflict_staging)
+    with pytest.raises(receiver.CheckpointConflictError, match="different contents"):
+        receiver.publish_checkpoint_directory(conflict_staging, destination)
+    assert (destination / "model-00001.safetensors").read_bytes() == b"weights-v1"
diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py
index a43636864..9c280e0e0 100644
--- a/tests/utils/test_vllm_engine.py
+++ b/tests/utils/test_vllm_engine.py
@@ -325,6 +325,22 @@ def test_build_vllm_subprocess_env_no_batch_invariant_by_default(vllm_args, monk
     assert "VLLM_BATCH_INVARIANT" not in env
 
 
+@pytest.mark.unit
+def test_build_vllm_subprocess_env_sets_checkpoint_roots(vllm_args):
+    vllm_args.update_weight_disk_dir = "/shared/weights"
+    vllm_args.update_weight_local_checkpoint_dir = "/local/weights"
+
+    env = mod._build_subprocess_env(
+        {
+            "_args": vllm_args,
+            "_visible_devices": "0",
+        }
+    )
+
+    assert env["VIME_WEIGHT_SOURCE_DIR"] == str(Path("/shared/weights"))
+    assert env["VIME_WEIGHT_LOCAL_CHECKPOINT_DIR"] == str(Path("/local/weights"))
+
+
 @pytest.mark.unit
 def test_build_vllm_subprocess_env_sets_disaggregation_side_channel(vllm_args):
     env = mod._build_subprocess_env(
@@ -781,11 +797,11 @@ def test_update_weights_from_disk_posts_collective_rpc(vllm_engine, monkeypatch)
 
     def fake_post(url, *, params=None, timeout=30, json=None):
         seen.append((url, params, timeout, json))
-        return _MockResponse(json_data={"reloaded": True})
+        return _MockResponse(json_data={"results": [None]})
 
     monkeypatch.setattr(mod.requests, "post", fake_post)
 
-    assert vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8") == {"reloaded": True}
+    assert vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8") == {"results": [None]}
     assert seen[0][0] == "http://127.0.0.1:8765/collective_rpc"
     assert seen[0][3]["method"] == "reload_weights"
     assert seen[1][0] == "http://127.0.0.1:8765/update_weight_version"
@@ -861,8 +877,39 @@ def fake_post(url, *, params=None, timeout=30, json=None):
         return _MockResponse(text="boom", status_code=500)
 
     monkeypatch.setattr(mod.requests, "post", fake_post)
+    vllm_engine._weight_version = "7"
     with pytest.raises(requests.exceptions.HTTPError):
         vllm_engine.update_weights_from_disk("/tmp/model")
+    assert vllm_engine._weight_version is None
+
+
+@pytest.mark.unit
+def test_update_weights_from_disk_does_not_commit_version_for_invalid_success_body(vllm_engine, monkeypatch):
+    monkeypatch.setattr(
+        mod.requests,
+        "post",
+        lambda *args, **kwargs: _MockResponse(text="not-json", status_code=200),
+    )
+    vllm_engine._weight_version = "7"
+
+    with pytest.raises(RuntimeError, match="non-JSON success response"):
+        vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8")
+
+    assert vllm_engine._weight_version is None
+
+
+@pytest.mark.unit
+def test_update_weights_from_disk_rejects_empty_rank_results(vllm_engine, monkeypatch):
+    monkeypatch.setattr(
+        mod.requests,
+        "post",
+        lambda *args, **kwargs: _MockResponse(json_data={"results": []}),
+    )
+
+    with pytest.raises(RuntimeError, match="invalid result"):
+        vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8")
+
+    assert vllm_engine._weight_version is None
 
 
 @pytest.mark.unit
diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py
index 13ded4186..dee4403dc 100644
--- a/vime/backends/megatron_utils/actor.py
+++ b/vime/backends/megatron_utils/actor.py
@@ -560,7 +560,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
             self.sleep()
 
     @timer
-    def update_weights(self) -> None:
+    def update_weights(self, weight_version: int | None = None) -> None:
         if self.args.debug_train_only or self.args.debug_rollout_only:
             return
 
@@ -609,7 +609,14 @@ def update_weights(self) -> None:
                     if ".draft_model." in name:
                         param.data = backup[name].to(param.device)
             print_memory("before update_weights")
-            self.weight_updater.update_weights()
+            if self.args.update_weight_mode == "full" and self.args.update_weight_transport == "disk":
+                if weight_version is None:
+                    raise ValueError("full-disk weight update requires an explicit weight_version")
+                self.weight_updater.update_weights(weight_version=weight_version)
+            else:
+                if weight_version is not None:
+                    raise ValueError("weight_version is only valid for full-disk weight updates")
+                self.weight_updater.update_weights()
             print_memory("after update_weights")
 
             if getattr(self.args, "keep_old_actor", False):
diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
index 1f0cc8559..40b4334b9 100644
--- a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
+++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
@@ -9,6 +9,7 @@
 import torch.distributed as dist
 from ray.actor import ActorHandle
 
+from vime.backends.vllm_utils.checkpoint_receiver import publish_checkpoint_directory
 from vime.utils.distributed_utils import get_gloo_group
 
 from ..hf_checkpoint_saver import save_hf_model_to_path
@@ -63,20 +64,27 @@ def pop_metrics(self) -> dict[str, float]:
         return out
 
     @torch.no_grad()
-    def update_weights(self) -> None:
-        self.weight_version += 1
-        version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{self.weight_version:06d}"
-
-        if dist.get_rank() == 0:
-            shutil.rmtree(version_dir, ignore_errors=True)
+    def update_weights(self, *, weight_version: int) -> None:
+        if weight_version < self.weight_version or weight_version > self.weight_version + 1:
+            raise ValueError(
+                f"invalid full-disk weight version {weight_version}; current version is {self.weight_version}"
+            )
+
+        version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{weight_version:06d}"
+        staging_dir = version_dir.with_name(f".{version_dir.name}.staging")
+
+        # Every rank cleans its local view. On a POSIX shared filesystem these
+        # calls converge on the same path; on host-local/object-backed mounts
+        # each writer must clean and create its own staging directory.
+        shutil.rmtree(staging_dir, ignore_errors=True)
         dist.barrier(group=get_gloo_group())
 
-        # every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface
+        # Every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface
         # one rank's mkdir to another until commit
-        version_dir.mkdir(parents=True, exist_ok=True)
+        staging_dir.mkdir(parents=True, exist_ok=True)
         save_hf_model_to_path(
             self.args,
-            version_dir,
+            staging_dir,
             self.model,
             model_name=self.model_name,
             quantization_config=self.quantization_config,
@@ -84,11 +92,16 @@ def update_weights(self) -> None:
         )
         dist.barrier(group=get_gloo_group())
 
+        publish_checkpoint_directory(staging_dir, version_dir)
+        dist.barrier(group=get_gloo_group())
+        shutil.rmtree(staging_dir, ignore_errors=True)
+
         # every rank runs the hook (it gates itself): each container must publish
         # its own writes
         if self._post_write_hook is not None:
             self._post_write_hook(self.args, str(version_dir), list(self.rollout_engines))
         dist.barrier(group=get_gloo_group())
+        self.weight_version = weight_version
 
         # VLLM reload is orchestrated by RayTrainGroup after the checkpoint
         # is fully written, so training-side lifecycle can decide whether
diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py
index 5bf9585aa..4047d4669 100644
--- a/vime/backends/vllm_utils/checkpoint_receiver.py
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -49,6 +49,8 @@ def materialize_checkpoint(
     local_checkpoint_dir: str,
     source_dir: str,
     target_version: int,
+    expected_local_checkpoint_dir: str | None = None,
+    expected_source_dir: str | None = None,
 ) -> dict[str, Any]:
     """Materialize one published checkpoint into a host-local directory.
 
@@ -59,11 +61,25 @@ def materialize_checkpoint(
     """
 
     version = _parse_version(target_version)
-    source_root = _resolve_existing_directory(source_dir, "source_dir")
+    if expected_source_dir is not None:
+        expected_source_root = _resolve_existing_directory(expected_source_dir, "configured source_dir")
+        source_root = _resolve_existing_directory(source_dir, "source_dir")
+        if source_root != expected_source_root:
+            raise CheckpointReceiveError("source_dir does not match the configured checkpoint source")
+    else:
+        source_root = _resolve_existing_directory(source_dir, "source_dir")
     version_dir = source_root / f"weight_v{version:06d}"
     version_dir = _resolve_existing_directory(version_dir, "checkpoint version")
 
-    local_dir = _prepare_local_directory(local_checkpoint_dir)
+    if expected_local_checkpoint_dir is not None:
+        expected_local_dir = _prepare_local_directory(expected_local_checkpoint_dir)
+        local_dir = _resolve_local_directory(local_checkpoint_dir)
+        if local_dir != expected_local_dir:
+            raise CheckpointReceiveError(
+                "local_checkpoint_dir does not match the configured local checkpoint destination"
+            )
+    else:
+        local_dir = _prepare_local_directory(local_checkpoint_dir)
     if local_dir.exists() and local_dir.resolve() == version_dir:
         raise CheckpointReceiveError("local_checkpoint_dir must differ from source checkpoint")
 
@@ -121,11 +137,13 @@ def _resolve_existing_directory(value: str | Path, name: str) -> Path:
     if not isinstance(value, (str, Path)) or not str(value).strip():
         raise CheckpointReceiveError(f"{name} must be a non-empty path")
     path = Path(value).expanduser()
+    if path.is_symlink():
+        raise CheckpointReceiveError(f"{name} must not be a symlink: {path}")
     try:
         resolved = path.resolve(strict=True)
     except OSError as exc:
         raise CheckpointReceiveError(f"{name} does not exist: {path}") from exc
-    if not resolved.is_dir() or resolved.is_symlink():
+    if not resolved.is_dir():
         raise CheckpointReceiveError(f"{name} must be a directory: {path}")
     return resolved
 
@@ -144,6 +162,52 @@ def _prepare_local_directory(value: str | Path) -> Path:
     return parent / path.name
 
 
+def _resolve_local_directory(value: str | Path) -> Path:
+    """Resolve an untrusted destination without creating caller-selected parents."""
+
+    if not isinstance(value, (str, Path)) or not str(value).strip():
+        raise CheckpointReceiveError("local_checkpoint_dir must be a non-empty path")
+    path = Path(value).expanduser()
+    if path.is_symlink() or (path.exists() and not path.is_dir()):
+        raise CheckpointReceiveError(f"local_checkpoint_dir must be a directory: {path}")
+    try:
+        parent = path.parent.resolve(strict=True)
+    except OSError as exc:
+        raise CheckpointReceiveError(f"local checkpoint parent does not exist: {path.parent}") from exc
+    return parent / path.name
+
+
+def publish_checkpoint_directory(staging_dir: str | Path, destination_dir: str | Path) -> str:
+    """Publish an immutable checkpoint directory, allowing an identical retry."""
+
+    staging = _resolve_existing_directory(staging_dir, "checkpoint staging directory")
+    destination = _prepare_local_directory(destination_dir)
+    if staging.parent != destination.parent:
+        raise CheckpointReceiveError("checkpoint staging and destination directories must be siblings")
+
+    if destination.exists():
+        destination_manifest = _build_checkpoint_manifest(destination)
+        staging_manifest = _build_checkpoint_manifest(staging)
+        destination_hash = _manifest_hash(destination_manifest["files"])
+        staging_hash = _manifest_hash(staging_manifest["files"])
+        if destination_hash != staging_hash:
+            raise CheckpointConflictError(
+                f"published checkpoint already exists with different contents: {destination}"
+            )
+        return "already_published"
+
+    try:
+        os.replace(staging, destination)
+    except FileNotFoundError:
+        # On a shared filesystem another writer rank may have renamed the same
+        # staging directory after the distributed write barrier.
+        if destination.is_dir() and not destination.is_symlink():
+            return "published_by_peer"
+        raise
+    _fsync_directory(destination.parent)
+    return "published"
+
+
 def _build_checkpoint_manifest(root: Path) -> dict[str, Any]:
     index_candidates = [root / name for name in _INDEX_NAMES if (root / name).exists()]
     if any(path.is_symlink() for path in index_candidates):
@@ -284,6 +348,20 @@ def _write_json(path: Path, value: dict[str, Any]) -> None:
         os.fsync(file.fileno())
 
 
+def _fsync_directory(path: Path) -> None:
+    if not hasattr(os, "O_DIRECTORY"):
+        return
+    try:
+        descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
+    except OSError:
+        logger.warning("Could not open checkpoint directory for fsync: %s", path, exc_info=True)
+        return
+    try:
+        os.fsync(descriptor)
+    finally:
+        os.close(descriptor)
+
+
 def _atomic_replace_directory(staging: Path, destination: Path) -> None:
     backup: Path | None = None
     if destination.exists() or destination.is_symlink():
@@ -297,6 +375,7 @@ def _atomic_replace_directory(staging: Path, destination: Path) -> None:
         if backup is not None and backup.exists() and not destination.exists():
             os.replace(backup, destination)
         raise
+    _fsync_directory(destination.parent)
     if backup is not None:
         try:
             shutil.rmtree(backup)
diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py
index ec53a8054..2d9c606f1 100644
--- a/vime/backends/vllm_utils/vllm_engine.py
+++ b/vime/backends/vllm_utils/vllm_engine.py
@@ -355,8 +355,6 @@ def resume_memory_occupation(self, tags: list[str] = None):
         return response.json()
 
     def check_weights(self, action: str):
-        if action not in {"snapshot", "reset_tensors", "compare"}:
-            raise ValueError(f"unsupported weight check action: {action!r}")
         return self._make_request("weights_checker", {"action": action})
 
     def init_weight_transfer_engine(self, payload: dict) -> dict:
diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py
index c3725d14c..f7859ac48 100644
--- a/vime/ray/actor_group.py
+++ b/vime/ray/actor_group.py
@@ -166,11 +166,11 @@ def update_weights(self):
 
         weight_version = self._disk_weight_version + 1
         disk_weight_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{weight_version:06d}"
-        ray.get([actor.update_weights.remote() for actor in self._actor_handlers])
-        self._disk_weight_version = weight_version
+        ray.get([actor.update_weights.remote(weight_version=weight_version) for actor in self._actor_handlers])
         if self._release_train_enabled():
             self.release()
         self._reload_rollout_weights_from_disk(disk_weight_dir, str(weight_version))
+        self._disk_weight_version = weight_version
 
     def onload(self):
         return ray.get([actor.wake_up.remote() for actor in self._actor_handlers])
diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py
index 78d21d15f..ed1ecd61f 100644
--- a/vime/ray/rollout.py
+++ b/vime/ray/rollout.py
@@ -271,7 +271,11 @@ def health_monitoring_resume(self) -> None:
             monitor.resume()
 
     def check_weights(self, action: str):
-        return ray.get([engine.check_weights.remote(action=action) for engine in self.rollout_engines])
+        server = self._get_updatable_server()
+        engines = server.engines if server is not None else []
+        if not engines:
+            raise RuntimeError("weight checking requires at least one updatable rollout engine")
+        return ray.get([engine.check_weights.remote(action=action) for engine in engines])
 
     def _get_rollout_data(self, rollout_id):
         if self.args.load_debug_rollout_data:
diff --git a/vime/ray/train_actor.py b/vime/ray/train_actor.py
index 72fea191c..626628310 100644
--- a/vime/ray/train_actor.py
+++ b/vime/ray/train_actor.py
@@ -116,7 +116,7 @@ def save_model(self, rollout_id, force_sync=False):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def update_weights(self):
+    def update_weights(self, weight_version: int | None = None):
         raise NotImplementedError
 
     def set_rollout_manager(self, rollout_manager):
diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py
index 4eefc41fb..62d9e1d94 100644
--- a/vime/utils/arguments.py
+++ b/vime/utils/arguments.py
@@ -239,6 +239,9 @@ def add_train_arguments(parser):
                     "reload from it. Required for --update-weight-mode=delta "
                     "--update-weight-transport=disk; optional for full disk sync (engines then "
                     "pull to local disk instead of reading the shared dir directly). The "
+                    "vLLM server pins these paths through VIME_WEIGHT_SOURCE_DIR and "
+                    "VIME_WEIGHT_LOCAL_CHECKPOINT_DIR; externally launched servers must set "
+                    "both variables to the corresponding rollout-host paths. The "
                     "read-side counterpart of --custom-update-weight-post-write-path is "
                     "--custom-update-weight-pre-read-path."
                 ),

From 828eb967601fa6e42c772fac7653f05a7762c345 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Mon, 20 Jul 2026 00:43:22 +0800
Subject: [PATCH 4/8] fix: make checkpoint publication peer-safe
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_checkpoint_receiver.py       | 10 +++++
 .../vllm_utils/checkpoint_receiver.py         | 41 +++++++++++++++++--
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/tests/utils/test_checkpoint_receiver.py b/tests/utils/test_checkpoint_receiver.py
index f36a8bfad..bc89db9ca 100644
--- a/tests/utils/test_checkpoint_receiver.py
+++ b/tests/utils/test_checkpoint_receiver.py
@@ -193,3 +193,13 @@ def test_publish_checkpoint_directory_is_immutable_and_retryable(tmp_path: Path)
     with pytest.raises(receiver.CheckpointConflictError, match="different contents"):
         receiver.publish_checkpoint_directory(conflict_staging, destination)
     assert (destination / "model-00001.safetensors").read_bytes() == b"weights-v1"
+
+
+def test_publish_checkpoint_directory_accepts_late_shared_fs_publisher(tmp_path: Path):
+    destination = tmp_path / "weight_v000001"
+    staging = tmp_path / ".weight_v000001.staging"
+    _publish(tmp_path / "first", 1).rename(staging)
+
+    assert receiver.publish_checkpoint_directory(staging, destination) == "published"
+    assert not staging.exists()
+    assert receiver.publish_checkpoint_directory(staging, destination) == "published_by_peer"
diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py
index 4047d4669..aa1efcfe7 100644
--- a/vime/backends/vllm_utils/checkpoint_receiver.py
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -180,14 +180,37 @@ def _resolve_local_directory(value: str | Path) -> Path:
 def publish_checkpoint_directory(staging_dir: str | Path, destination_dir: str | Path) -> str:
     """Publish an immutable checkpoint directory, allowing an identical retry."""
 
-    staging = _resolve_existing_directory(staging_dir, "checkpoint staging directory")
+    if not isinstance(staging_dir, (str, Path)) or not str(staging_dir).strip():
+        raise CheckpointReceiveError("checkpoint staging directory must be a non-empty path")
+    staging_path = Path(staging_dir).expanduser()
+    if staging_path.is_symlink():
+        raise CheckpointReceiveError(f"checkpoint staging directory must not be a symlink: {staging_path}")
+    try:
+        staging_parent = staging_path.parent.resolve(strict=True)
+    except OSError as exc:
+        raise CheckpointReceiveError(
+            f"checkpoint staging directory parent does not exist: {staging_path.parent}"
+        ) from exc
     destination = _prepare_local_directory(destination_dir)
-    if staging.parent != destination.parent:
+    if staging_parent != destination.parent:
         raise CheckpointReceiveError("checkpoint staging and destination directories must be siblings")
+    staging_path = staging_parent / staging_path.name
+
+    try:
+        staging = _resolve_existing_directory(staging_path, "checkpoint staging directory")
+    except CheckpointReceiveError:
+        if not staging_path.exists() and _published_checkpoint_is_valid(destination):
+            return "published_by_peer"
+        raise
 
     if destination.exists():
         destination_manifest = _build_checkpoint_manifest(destination)
-        staging_manifest = _build_checkpoint_manifest(staging)
+        try:
+            staging_manifest = _build_checkpoint_manifest(staging)
+        except (CheckpointReceiveError, OSError):
+            if not staging.exists() and _published_checkpoint_is_valid(destination):
+                return "published_by_peer"
+            raise
         destination_hash = _manifest_hash(destination_manifest["files"])
         staging_hash = _manifest_hash(staging_manifest["files"])
         if destination_hash != staging_hash:
@@ -201,13 +224,23 @@ def publish_checkpoint_directory(staging_dir: str | Path, destination_dir: str |
     except FileNotFoundError:
         # On a shared filesystem another writer rank may have renamed the same
         # staging directory after the distributed write barrier.
-        if destination.is_dir() and not destination.is_symlink():
+        if _published_checkpoint_is_valid(destination):
             return "published_by_peer"
         raise
     _fsync_directory(destination.parent)
     return "published"
 
 
+def _published_checkpoint_is_valid(destination: Path) -> bool:
+    if not destination.is_dir() or destination.is_symlink():
+        return False
+    try:
+        _build_checkpoint_manifest(destination)
+    except (CheckpointReceiveError, OSError):
+        return False
+    return True
+
+
 def _build_checkpoint_manifest(root: Path) -> dict[str, Any]:
     index_candidates = [root / name for name in _INDEX_NAMES if (root / name).exists()]
     if any(path.is_symlink() for path in index_candidates):

From e173731ec71dd5e951a8c4312a9fee4dd5d4a05f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Mon, 20 Jul 2026 00:43:42 +0800
Subject: [PATCH 5/8] fix: use vLLM server-local pull paths
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_vllm_engine.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py
index 9c280e0e0..ddf1d4df8 100644
--- a/tests/utils/test_vllm_engine.py
+++ b/tests/utils/test_vllm_engine.py
@@ -341,6 +341,20 @@ def test_build_vllm_subprocess_env_sets_checkpoint_roots(vllm_args):
     assert env["VIME_WEIGHT_LOCAL_CHECKPOINT_DIR"] == str(Path("/local/weights"))
 
 
+@pytest.mark.unit
+def test_pull_weights_leaves_paths_to_server(vllm_engine, monkeypatch):
+    calls: list[tuple[str, dict]] = []
+
+    def fake_make_request(endpoint: str, payload: dict):
+        calls.append((endpoint, payload))
+        return {"success": True, "local_checkpoint_dir": "/remote/local"}
+
+    monkeypatch.setattr(vllm_engine, "_make_request", fake_make_request)
+
+    assert vllm_engine.pull_weights(7)["local_checkpoint_dir"] == "/remote/local"
+    assert calls == [("pull_weights", {"target_version": 7})]
+
+
 @pytest.mark.unit
 def test_build_vllm_subprocess_env_sets_disaggregation_side_channel(vllm_args):
     env = mod._build_subprocess_env(

From 53cae3ee16499e7386c4948e013efd55c05a765c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Mon, 20 Jul 2026 00:44:12 +0800
Subject: [PATCH 6/8] fix: orchestrate per-engine disk reloads
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_actor_group_disk_update.py | 114 ++++++++++++++++++++
 vime/ray/actor_group.py                     |  87 ++++++++++-----
 2 files changed, 172 insertions(+), 29 deletions(-)

diff --git a/tests/utils/test_actor_group_disk_update.py b/tests/utils/test_actor_group_disk_update.py
index 6306fe2b2..890bb3cc1 100644
--- a/tests/utils/test_actor_group_disk_update.py
+++ b/tests/utils/test_actor_group_disk_update.py
@@ -58,6 +58,42 @@ def __init__(self, calls: list[dict[str, int]]) -> None:
         self.update_weights = _RemoteMethod(calls)
 
 
+class _NamedRemoteMethod:
+    def __init__(self, name: str, calls: list[str], result=None) -> None:
+        self.name = name
+        self.calls = calls
+        self.result = result if result is not None else name
+        self.kwargs: list[dict] = []
+
+    def remote(self, *args, **kwargs):
+        self.calls.append(self.name)
+        self.kwargs.append(kwargs)
+        return self.result
+
+
+class _RolloutEngine:
+    def __init__(self, calls: list[str], pull_result=None) -> None:
+        self.pull_weights = _NamedRemoteMethod(
+            "pull",
+            calls,
+            pull_result or {"success": True, "local_checkpoint_dir": "/remote/local"},
+        )
+        self.pause_generation = _NamedRemoteMethod("pause", calls)
+        self.flush_cache = _NamedRemoteMethod("flush", calls)
+        self.update_weights_from_disk = _NamedRemoteMethod("reload", calls)
+        self.continue_generation = _NamedRemoteMethod("continue", calls)
+
+
+class _RolloutManager:
+    def __init__(self, engine) -> None:
+        engines = engine if isinstance(engine, list) else [engine]
+        self.get_updatable_engines_and_lock = _NamedRemoteMethod(
+            "get_engines",
+            [],
+            (engines, None, 0, [], []),
+        )
+
+
 def _make_group(module, tmp_path: Path):
     group = module.RayTrainGroup.__new__(module.RayTrainGroup)
     group.args = SimpleNamespace(
@@ -135,3 +171,81 @@ def fail_once(refs):
 
     assert actor_calls == [{"weight_version": 1}, {"weight_version": 1}]
     assert group._disk_weight_version == 1
+
+
+@pytest.mark.parametrize(
+    ("failed_call", "expected_calls"),
+    [
+        ("pause", ["pause", "continue"]),
+        ("flush", ["pause", "flush", "continue"]),
+        ("reload", ["pause", "flush", "reload", "continue"]),
+    ],
+)
+def test_disk_reload_resumes_engines_after_failure(
+    actor_group_module,
+    tmp_path: Path,
+    failed_call: str,
+    expected_calls: list[str],
+):
+    group = _make_group(actor_group_module, tmp_path)
+    group.args.offload_rollout = False
+    group.args.update_weight_local_checkpoint_dir = None
+    group.args.update_weight_disk_keep_files = True
+    group.args.ci_test = False
+    calls: list[str] = []
+    engine = _RolloutEngine(calls)
+    group._rollout_manager = _RolloutManager(engine)
+
+    def ray_get(refs):
+        if refs == [failed_call]:
+            raise RuntimeError(f"{failed_call} failed")
+        return refs
+
+    actor_group_module.ray.get = ray_get
+
+    with pytest.raises(RuntimeError, match=rf"{failed_call} failed"):
+        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
+
+    assert calls == expected_calls
+
+
+def test_disk_reload_uses_server_returned_local_path(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    group.args.offload_rollout = False
+    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
+    group.args.update_weight_disk_keep_files = True
+    group.args.ci_test = False
+    calls: list[str] = []
+    engine = _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/local"})
+    group._rollout_manager = _RolloutManager(engine)
+
+    group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
+
+    assert calls == ["pull", "pause", "flush", "reload", "continue"]
+    assert engine.pull_weights.kwargs == [{"target_version": 1}]
+    assert engine.update_weights_from_disk.kwargs[0]["model_path"] == "/remote/local"
+
+
+def test_disk_reload_fans_out_server_returned_local_paths(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    group.args.offload_rollout = False
+    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
+    group.args.update_weight_disk_keep_files = True
+    group.args.ci_test = False
+    calls: list[str] = []
+    engines = [
+        _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/engine-0"}),
+        _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/engine-1"}),
+    ]
+    group._rollout_manager = _RolloutManager(engines)
+
+    group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
+
+    assert [engine.pull_weights.kwargs for engine in engines] == [
+        [{"target_version": 1}],
+        [{"target_version": 1}],
+    ]
+    assert [engine.update_weights_from_disk.kwargs[0]["model_path"] for engine in engines] == [
+        "/remote/engine-0",
+        "/remote/engine-1",
+    ]
diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py
index f7859ac48..f05cc13fd 100644
--- a/vime/ray/actor_group.py
+++ b/vime/ray/actor_group.py
@@ -237,33 +237,62 @@ def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version):
             # each host pulls the published checkpoint onto local disk (e.g. NVMe) and
             # the engines reload from there; the pull is disk-only, so it runs before
             # pause and overlaps generation
-            ray.get([engine.pull_weights.remote(int(weight_version)) for engine in engines])
-            model_path = self.args.update_weight_local_checkpoint_dir
+            pull_results = ray.get(
+                [engine.pull_weights.remote(target_version=int(weight_version)) for engine in engines]
+            )
+            if not isinstance(pull_results, list) or len(pull_results) != len(engines):
+                raise RuntimeError(f"pull_weights returned one result per engine: {pull_results!r}")
+            model_paths = []
+            for index, result in enumerate(pull_results):
+                if (
+                    not isinstance(result, dict)
+                    or result.get("success") is not True
+                    or not isinstance(result.get("local_checkpoint_dir"), str)
+                    or not result["local_checkpoint_dir"]
+                ):
+                    raise RuntimeError(f"pull_weights returned an invalid result for engine {index}: {result!r}")
+                model_paths.append(result["local_checkpoint_dir"])
         else:
-            model_path = str(disk_weight_dir)
-        ray.get([engine.pause_generation.remote() for engine in engines])
-        ray.get([engine.flush_cache.remote() for engine in engines])
-        ray.get(
-            [
-                engine.update_weights_from_disk.remote(
-                    model_path=model_path,
-                    weight_version=weight_version,
-                )
-                for engine in engines
-            ]
-        )
-        if self.args.ci_test:
-            engine_versions = ray.get([engine.get_weight_version.remote() for engine in engines])
-            mismatches = [
-                f"engine {idx}: {engine_version}"
-                for idx, engine_version in enumerate(engine_versions)
-                if str(engine_version) != str(weight_version)
-            ]
-            if mismatches:
-                raise RuntimeError(
-                    "Weight version mismatch after disk reload! "
-                    f"Expected: {weight_version}; " + ", ".join(mismatches)
-                )
-        if not self.args.update_weight_disk_keep_files:
-            shutil.rmtree(disk_weight_dir, ignore_errors=True)
-        ray.get([engine.continue_generation.remote() for engine in engines])
+            model_paths = [str(disk_weight_dir)] * len(engines)
+        reload_error: Exception | None = None
+        pause_attempted = False
+        try:
+            # Resume every engine even if only a subset completed the batched
+            # pause call before Ray surfaced an error.
+            pause_attempted = True
+            ray.get([engine.pause_generation.remote() for engine in engines])
+            ray.get([engine.flush_cache.remote() for engine in engines])
+            ray.get(
+                [
+                    engine.update_weights_from_disk.remote(
+                        model_path=model_path,
+                        weight_version=weight_version,
+                    )
+                    for engine, model_path in zip(engines, model_paths, strict=True)
+                ]
+            )
+            if self.args.ci_test:
+                engine_versions = ray.get([engine.get_weight_version.remote() for engine in engines])
+                mismatches = [
+                    f"engine {idx}: {engine_version}"
+                    for idx, engine_version in enumerate(engine_versions)
+                    if str(engine_version) != str(weight_version)
+                ]
+                if mismatches:
+                    raise RuntimeError(
+                        "Weight version mismatch after disk reload! "
+                        f"Expected: {weight_version}; " + ", ".join(mismatches)
+                    )
+            if not self.args.update_weight_disk_keep_files:
+                shutil.rmtree(disk_weight_dir, ignore_errors=True)
+        except Exception as exc:
+            reload_error = exc
+            raise
+        finally:
+            if pause_attempted:
+                try:
+                    ray.get([engine.continue_generation.remote() for engine in engines])
+                except Exception as resume_error:
+                    if reload_error is None:
+                        raise
+                    reload_error.add_note(f"Failed to resume rollout engines after reload failure: {resume_error}")

From ecefe764e4246d367bebcec19d116a9b49a0539c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Mon, 20 Jul 2026 15:20:14 +0800
Subject: [PATCH 7/8] fix: harden disk sync compatibility diagnostics
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_actor_group_disk_update.py   | 48 +++++++++++++++++++
 tests/utils/test_checkpoint_receiver.py       | 26 ++++++++--
 .../vllm_utils/checkpoint_receiver.py         |  6 +--
 vime/ray/actor_group.py                       | 11 ++++-
 4 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/tests/utils/test_actor_group_disk_update.py b/tests/utils/test_actor_group_disk_update.py
index 890bb3cc1..3053b2d6c 100644
--- a/tests/utils/test_actor_group_disk_update.py
+++ b/tests/utils/test_actor_group_disk_update.py
@@ -209,6 +209,39 @@ def ray_get(refs):
     assert calls == expected_calls
 
 
+def test_disk_reload_preserves_primary_failure_without_exception_notes(
+    actor_group_module, tmp_path: Path, monkeypatch
+):
+    class LegacyRuntimeError(RuntimeError):
+        add_note = None
+
+    group = _make_group(actor_group_module, tmp_path)
+    group.args.offload_rollout = False
+    group.args.update_weight_local_checkpoint_dir = None
+    group.args.update_weight_disk_keep_files = True
+    group.args.ci_test = False
+    calls: list[str] = []
+    engine = _RolloutEngine(calls)
+    group._rollout_manager = _RolloutManager(engine)
+    logged = []
+    monkeypatch.setattr(actor_group_module.logger, "exception", logged.append)
+
+    def ray_get(refs):
+        if refs == ["reload"]:
+            raise LegacyRuntimeError("reload failed")
+        if refs == ["continue"]:
+            raise RuntimeError("continue failed")
+        return refs
+
+    actor_group_module.ray.get = ray_get
+
+    with pytest.raises(LegacyRuntimeError, match="reload failed"):
+        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
+
+    assert calls == ["pause", "flush", "reload", "continue"]
+    assert logged == ["Failed to resume rollout engines after reload failure"]
+
+
 def test_disk_reload_uses_server_returned_local_path(actor_group_module, tmp_path: Path):
     group = _make_group(actor_group_module, tmp_path)
     group.args.offload_rollout = False
@@ -226,6 +259,21 @@ def test_disk_reload_uses_server_returned_local_path(actor_group_module, tmp_pat
     assert engine.update_weights_from_disk.kwargs[0]["model_path"] == "/remote/local"
 
 
+def test_disk_reload_rejects_missing_pull_result(actor_group_module, tmp_path: Path):
+    group = _make_group(actor_group_module, tmp_path)
+    group.args.offload_rollout = False
+    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
+    group.args.update_weight_disk_keep_files = True
+    group.args.ci_test = False
+    engine = _RolloutEngine([])
+    group._rollout_manager = _RolloutManager(engine)
+
+    actor_group_module.ray.get = lambda refs: [] if refs == [engine.pull_weights.result] else refs
+
+    with pytest.raises(RuntimeError, match="Expected one pull_weights result per engine, got: \\[\\]"):
+        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
+
+
 def test_disk_reload_fans_out_server_returned_local_paths(actor_group_module, tmp_path: Path):
     group = _make_group(actor_group_module, tmp_path)
     group.args.offload_rollout = False
diff --git a/tests/utils/test_checkpoint_receiver.py b/tests/utils/test_checkpoint_receiver.py
index bc89db9ca..8361c7553 100644
--- a/tests/utils/test_checkpoint_receiver.py
+++ b/tests/utils/test_checkpoint_receiver.py
@@ -33,9 +33,7 @@ def test_materialize_is_atomic_and_repeated_pull_is_idempotent(tmp_path: Path):
     _publish(source, 1)
     local = tmp_path / "local"
 
-    first = receiver.materialize_checkpoint(
-        source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
-    )
+    first = receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
     assert first["status"] == "materialized"
     assert (local / "model-00001.safetensors").read_bytes() == b"weights-v1"
 
@@ -203,3 +201,25 @@ def test_publish_checkpoint_directory_accepts_late_shared_fs_publisher(tmp_path:
     assert receiver.publish_checkpoint_directory(staging, destination) == "published"
     assert not staging.exists()
     assert receiver.publish_checkpoint_directory(staging, destination) == "published_by_peer"
+
+
+def test_fsync_directory_is_best_effort(tmp_path: Path, monkeypatch):
+    warnings = []
+    closed = []
+
+    def unsupported_fsync(descriptor):
+        assert descriptor == 7
+        raise OSError("unsupported")
+
+    monkeypatch.setattr(receiver.os, "O_DIRECTORY", 0, raising=False)
+    monkeypatch.setattr(receiver.os, "open", lambda *args: 7)
+    monkeypatch.setattr(receiver.os, "fsync", unsupported_fsync)
+    monkeypatch.setattr(receiver.os, "close", closed.append)
+    monkeypatch.setattr(receiver.logger, "warning", lambda *args, **kwargs: warnings.append((args, kwargs)))
+
+    receiver._fsync_directory(tmp_path)
+
+    assert closed == [7]
+    assert warnings[0][0][0] == "Could not fsync checkpoint directory: %s"
+    assert warnings[0][0][1] == tmp_path
+    assert warnings[0][1] == {"exc_info": True}
diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py
index aa1efcfe7..19090899b 100644
--- a/vime/backends/vllm_utils/checkpoint_receiver.py
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -91,9 +91,7 @@ def materialize_checkpoint(
         if current is not None:
             current_version = current["version"]
             if current_version > version:
-                raise StaleCheckpointError(
-                    f"stale checkpoint version {version}; active version is {current_version}"
-                )
+                raise StaleCheckpointError(f"stale checkpoint version {version}; active version is {current_version}")
             if current_version == version:
                 if current["manifest_sha256"] != source_manifest_hash:
                     raise CheckpointConflictError(
@@ -391,6 +389,8 @@ def _fsync_directory(path: Path) -> None:
         return
     try:
         os.fsync(descriptor)
+    except OSError:
+        logger.warning("Could not fsync checkpoint directory: %s", path, exc_info=True)
     finally:
         os.close(descriptor)
 
diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py
index f05cc13fd..0284c5811 100644
--- a/vime/ray/actor_group.py
+++ b/vime/ray/actor_group.py
@@ -1,3 +1,4 @@
+import logging
 import os
 import shutil
 import time
@@ -9,6 +10,8 @@
 
 from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars
 
+logger = logging.getLogger(__name__)
+
 
 class RayTrainGroup:
     """
@@ -241,7 +244,7 @@ def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version):
                 [engine.pull_weights.remote(target_version=int(weight_version)) for engine in engines]
             )
             if not isinstance(pull_results, list) or len(pull_results) != len(engines):
-                raise RuntimeError(f"pull_weights returned one result per engine: {pull_results!r}")
+                raise RuntimeError(f"Expected one pull_weights result per engine, got: {pull_results!r}")
             model_paths = []
             for index, result in enumerate(pull_results):
                 if (
@@ -295,4 +298,8 @@ def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version):
                 except Exception as resume_error:
                     if reload_error is None:
                         raise
-                    reload_error.add_note(f"Failed to resume rollout engines after reload failure: {resume_error}")
+                    add_note = getattr(reload_error, "add_note", None)
+                    if callable(add_note):
+                        add_note(f"Failed to resume rollout engines after reload failure: {resume_error}")
+                    else:
+                        logger.exception("Failed to resume rollout engines after reload failure")

From 9d4dd00ed11221f3befff9469b4adaa81b67f300 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=81=93=E8=B7=AF=E8=87=AA=E4=BF=A1?=
 <182356120+daoluzixin@users.noreply.github.com>
Date: Tue, 4 Aug 2026 00:30:45 +0800
Subject: [PATCH 8/8] feat: simplify vLLM full-disk weight refresh
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
---
 tests/utils/test_actor_group_disk_update.py   | 299 ------------
 tests/utils/test_checkpoint_receiver.py       | 243 ++--------
 tests/utils/test_vllm_engine.py               |  65 +--
 vime/backends/megatron_utils/actor.py         |  11 +-
 .../update_weight/update_weight_from_disk.py  |  31 +-
 .../vllm_utils/checkpoint_receiver.py         | 427 ++----------------
 vime/ray/actor_group.py                       |  98 ++--
 vime/ray/rollout.py                           |   6 +-
 vime/ray/train_actor.py                       |   2 +-
 vime/utils/arguments.py                       |   3 -
 10 files changed, 115 insertions(+), 1070 deletions(-)
 delete mode 100644 tests/utils/test_actor_group_disk_update.py

diff --git a/tests/utils/test_actor_group_disk_update.py b/tests/utils/test_actor_group_disk_update.py
deleted file mode 100644
index 3053b2d6c..000000000
--- a/tests/utils/test_actor_group_disk_update.py
+++ /dev/null
@@ -1,299 +0,0 @@
-"""CPU tests for full-disk version ownership in ``RayTrainGroup``."""
-
-from __future__ import annotations
-
-import importlib
-import sys
-import types
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-
-
-@pytest.fixture
-def actor_group_module(monkeypatch):
-    ray = types.ModuleType("ray")
-    ray.get = lambda refs: refs
-    ray.remote = lambda *args, **kwargs: lambda value: value
-    ray.kill = lambda *args, **kwargs: None
-
-    ray_util = types.ModuleType("ray.util")
-    placement_group = types.ModuleType("ray.util.placement_group")
-    placement_group.PlacementGroup = object
-    scheduling = types.ModuleType("ray.util.scheduling_strategies")
-    scheduling.PlacementGroupSchedulingStrategy = object
-    ray.util = ray_util
-
-    ray_utils = types.ModuleType("vime.ray.utils")
-    ray_utils.NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = []
-    ray_utils.add_default_ray_env_vars = lambda values=None: values or {}
-
-    for name, module in {
-        "ray": ray,
-        "ray.util": ray_util,
-        "ray.util.placement_group": placement_group,
-        "ray.util.scheduling_strategies": scheduling,
-        "vime.ray.utils": ray_utils,
-    }.items():
-        monkeypatch.setitem(sys.modules, name, module)
-
-    sys.modules.pop("vime.ray.actor_group", None)
-    module = importlib.import_module("vime.ray.actor_group")
-    yield module
-    sys.modules.pop("vime.ray.actor_group", None)
-
-
-class _RemoteMethod:
-    def __init__(self, calls: list[dict[str, int]]) -> None:
-        self.calls = calls
-
-    def remote(self, **kwargs):
-        self.calls.append(kwargs)
-        return kwargs
-
-
-class _Actor:
-    def __init__(self, calls: list[dict[str, int]]) -> None:
-        self.update_weights = _RemoteMethod(calls)
-
-
-class _NamedRemoteMethod:
-    def __init__(self, name: str, calls: list[str], result=None) -> None:
-        self.name = name
-        self.calls = calls
-        self.result = result if result is not None else name
-        self.kwargs: list[dict] = []
-
-    def remote(self, *args, **kwargs):
-        self.calls.append(self.name)
-        self.kwargs.append(kwargs)
-        return self.result
-
-
-class _RolloutEngine:
-    def __init__(self, calls: list[str], pull_result=None) -> None:
-        self.pull_weights = _NamedRemoteMethod(
-            "pull",
-            calls,
-            pull_result or {"success": True, "local_checkpoint_dir": "/remote/local"},
-        )
-        self.pause_generation = _NamedRemoteMethod("pause", calls)
-        self.flush_cache = _NamedRemoteMethod("flush", calls)
-        self.update_weights_from_disk = _NamedRemoteMethod("reload", calls)
-        self.continue_generation = _NamedRemoteMethod("continue", calls)
-
-
-class _RolloutManager:
-    def __init__(self, engine) -> None:
-        engines = engine if isinstance(engine, list) else [engine]
-        self.get_updatable_engines_and_lock = _NamedRemoteMethod(
-            "get_engines",
-            [],
-            (engines, None, 0, [], []),
-        )
-
-
-def _make_group(module, tmp_path: Path):
-    group = module.RayTrainGroup.__new__(module.RayTrainGroup)
-    group.args = SimpleNamespace(
-        update_weight_mode="full",
-        update_weight_transport="disk",
-        update_weight_disk_dir=str(tmp_path),
-        release_train=False,
-    )
-    group.role = "actor"
-    group._disk_weight_version = 0
-    group._actor_handlers = []
-    return group
-
-
-def test_full_disk_group_commits_version_after_reload(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    actor_calls: list[dict[str, int]] = []
-    reload_calls: list[tuple[Path, str]] = []
-    group._actor_handlers = [_Actor(actor_calls)]
-    group._reload_rollout_weights_from_disk = lambda path, version: reload_calls.append((path, version))
-
-    group.update_weights()
-
-    assert actor_calls == [{"weight_version": 1}]
-    assert reload_calls == [(tmp_path / "weight_v000001", "1")]
-    assert group._disk_weight_version == 1
-
-
-def test_full_disk_group_retries_same_version_after_reload_failure(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    actor_calls: list[dict[str, int]] = []
-    group._actor_handlers = [_Actor(actor_calls)]
-    attempts = 0
-
-    def reload_once_then_succeed(path, version):
-        nonlocal attempts
-        attempts += 1
-        if attempts == 1:
-            raise RuntimeError("reload failed")
-
-    group._reload_rollout_weights_from_disk = reload_once_then_succeed
-
-    with pytest.raises(RuntimeError, match="reload failed"):
-        group.update_weights()
-    assert group._disk_weight_version == 0
-
-    group.update_weights()
-
-    assert actor_calls == [{"weight_version": 1}, {"weight_version": 1}]
-    assert group._disk_weight_version == 1
-
-
-def test_full_disk_group_retries_same_version_after_actor_failure(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    actor_calls: list[dict[str, int]] = []
-    group._actor_handlers = [_Actor(actor_calls)]
-    group._reload_rollout_weights_from_disk = lambda path, version: None
-    original_get = actor_group_module.ray.get
-    attempts = 0
-
-    def fail_once(refs):
-        nonlocal attempts
-        attempts += 1
-        if attempts == 1:
-            raise RuntimeError("actor failed")
-        return original_get(refs)
-
-    actor_group_module.ray.get = fail_once
-
-    with pytest.raises(RuntimeError, match="actor failed"):
-        group.update_weights()
-    assert group._disk_weight_version == 0
-
-    group.update_weights()
-
-    assert actor_calls == [{"weight_version": 1}, {"weight_version": 1}]
-    assert group._disk_weight_version == 1
-
-
-@pytest.mark.parametrize(
-    ("failed_call", "expected_calls"),
-    [
-        ("pause", ["pause", "continue"]),
-        ("flush", ["pause", "flush", "continue"]),
-        ("reload", ["pause", "flush", "reload", "continue"]),
-    ],
-)
-def test_disk_reload_resumes_engines_after_failure(
-    actor_group_module,
-    tmp_path: Path,
-    failed_call: str,
-    expected_calls: list[str],
-):
-    group = _make_group(actor_group_module, tmp_path)
-    group.args.offload_rollout = False
-    group.args.update_weight_local_checkpoint_dir = None
-    group.args.update_weight_disk_keep_files = True
-    group.args.ci_test = False
-    calls: list[str] = []
-    engine = _RolloutEngine(calls)
-    group._rollout_manager = _RolloutManager(engine)
-
-    def ray_get(refs):
-        if refs == [failed_call]:
-            raise RuntimeError(f"{failed_call} failed")
-        return refs
-
-    actor_group_module.ray.get = ray_get
-
-    with pytest.raises(RuntimeError, match=rf"{failed_call} failed"):
-        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
-
-    assert calls == expected_calls
-
-
-def test_disk_reload_preserves_primary_failure_without_exception_notes(
-    actor_group_module, tmp_path: Path, monkeypatch
-):
-    class LegacyRuntimeError(RuntimeError):
-        add_note = None
-
-    group = _make_group(actor_group_module, tmp_path)
-    group.args.offload_rollout = False
-    group.args.update_weight_local_checkpoint_dir = None
-    group.args.update_weight_disk_keep_files = True
-    group.args.ci_test = False
-    calls: list[str] = []
-    engine = _RolloutEngine(calls)
-    group._rollout_manager = _RolloutManager(engine)
-    logged = []
-    monkeypatch.setattr(actor_group_module.logger, "exception", logged.append)
-
-    def ray_get(refs):
-        if refs == ["reload"]:
-            raise LegacyRuntimeError("reload failed")
-        if refs == ["continue"]:
-            raise RuntimeError("continue failed")
-        return refs
-
-    actor_group_module.ray.get = ray_get
-
-    with pytest.raises(LegacyRuntimeError, match="reload failed"):
-        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
-
-    assert calls == ["pause", "flush", "reload", "continue"]
-    assert logged == ["Failed to resume rollout engines after reload failure"]
-
-
-def test_disk_reload_uses_server_returned_local_path(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    group.args.offload_rollout = False
-    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
-    group.args.update_weight_disk_keep_files = True
-    group.args.ci_test = False
-    calls: list[str] = []
-    engine = _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/local"})
-    group._rollout_manager = _RolloutManager(engine)
-
-    group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
-
-    assert calls == ["pull", "pause", "flush", "reload", "continue"]
-    assert engine.pull_weights.kwargs == [{"target_version": 1}]
-    assert engine.update_weights_from_disk.kwargs[0]["model_path"] == "/remote/local"
-
-
-def test_disk_reload_rejects_missing_pull_result(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    group.args.offload_rollout = False
-    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
-    group.args.update_weight_disk_keep_files = True
-    group.args.ci_test = False
-    engine = _RolloutEngine([])
-    group._rollout_manager = _RolloutManager(engine)
-
-    actor_group_module.ray.get = lambda refs: [] if refs == [engine.pull_weights.result] else refs
-
-    with pytest.raises(RuntimeError, match="Expected one pull_weights result per engine, got: \\[\\]"):
-        group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
-
-
-def test_disk_reload_fans_out_server_returned_local_paths(actor_group_module, tmp_path: Path):
-    group = _make_group(actor_group_module, tmp_path)
-    group.args.offload_rollout = False
-    group.args.update_weight_local_checkpoint_dir = "/trainer/local"
-    group.args.update_weight_disk_keep_files = True
-    group.args.ci_test = False
-    calls: list[str] = []
-    engines = [
-        _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/engine-0"}),
-        _RolloutEngine(calls, {"success": True, "local_checkpoint_dir": "/remote/engine-1"}),
-    ]
-    group._rollout_manager = _RolloutManager(engines)
-
-    group._reload_rollout_weights_from_disk(tmp_path / "weight_v000001", "1")
-
-    assert [engine.pull_weights.kwargs for engine in engines] == [
-        [{"target_version": 1}],
-        [{"target_version": 1}],
-    ]
-    assert [engine.update_weights_from_disk.kwargs[0]["model_path"] for engine in engines] == [
-        "/remote/engine-0",
-        "/remote/engine-1",
-    ]
diff --git a/tests/utils/test_checkpoint_receiver.py b/tests/utils/test_checkpoint_receiver.py
index 8361c7553..3d0ef31d2 100644
--- a/tests/utils/test_checkpoint_receiver.py
+++ b/tests/utils/test_checkpoint_receiver.py
@@ -1,225 +1,68 @@
-"""CPU tests for the transactional full-disk checkpoint receiver."""
-
-from __future__ import annotations
-
-import json
-from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import pytest
 
-from vime.backends.vllm_utils import checkpoint_receiver as receiver
+from vime.backends.vllm_utils import checkpoint_receiver
 
 
-def _publish(root: Path, version: int, payload: bytes = b"weights-v1") -> Path:
-    checkpoint = root / f"weight_v{version:06d}"
+def _checkpoint(source_dir: Path, version: int, content: bytes) -> Path:
+    checkpoint = source_dir / f"weight_v{version:06d}"
     checkpoint.mkdir(parents=True)
-    (checkpoint / "config.json").write_text("{}", encoding="utf-8")
-    (checkpoint / "model-00001.safetensors").write_bytes(payload)
-    (checkpoint / "model.safetensors.index.json").write_text(
-        json.dumps(
-            {
-                "metadata": {"total_size": len(payload)},
-                "weight_map": {"model.weight": "model-00001.safetensors"},
-            }
-        ),
-        encoding="utf-8",
-    )
+    (checkpoint / "model.safetensors").write_bytes(content)
     return checkpoint
 
 
-def test_materialize_is_atomic_and_repeated_pull_is_idempotent(tmp_path: Path):
-    source = tmp_path / "published"
-    _publish(source, 1)
-    local = tmp_path / "local"
-
-    first = receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
-    assert first["status"] == "materialized"
-    assert (local / "model-00001.safetensors").read_bytes() == b"weights-v1"
+@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")
 
-    repeated = receiver.materialize_checkpoint(
-        source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
+    result = checkpoint_receiver.materialize_checkpoint(
+        source_dir=str(source_dir),
+        local_checkpoint_dir=str(local_dir),
+        target_version=1,
     )
-    assert repeated["status"] == "already_applied"
-    assert repeated["manifest_sha256"] == first["manifest_sha256"]
 
+    assert result == {
+        "success": True,
+        "version": 1,
+        "local_checkpoint_dir": str(local_dir),
+    }
+    assert (local_dir / "model.safetensors").read_bytes() == b"new"
 
-def test_invalid_source_and_malformed_checkpoint_are_rejected(tmp_path: Path):
-    with pytest.raises(receiver.CheckpointReceiveError, match="does not exist"):
-        receiver.materialize_checkpoint(
-            source_dir=str(tmp_path / "missing"),
-            local_checkpoint_dir=str(tmp_path / "local"),
-            target_version=1,
-        )
 
-    source = tmp_path / "published"
-    checkpoint = source / "weight_v000001"
-    checkpoint.mkdir(parents=True)
-    (checkpoint / "model.safetensors.index.json").write_text("not-json", encoding="utf-8")
-    with pytest.raises(receiver.CheckpointReceiveError, match="malformed checkpoint index"):
-        receiver.materialize_checkpoint(
-            source_dir=str(source),
+@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=1,
-        )
-
-
-def test_stale_and_conflicting_versions_are_rejected(tmp_path: Path):
-    source = tmp_path / "published"
-    _publish(source, 1, b"one")
-    _publish(source, 2, b"two")
-    local = tmp_path / "local"
-    receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
-
-    with pytest.raises(receiver.StaleCheckpointError, match="stale checkpoint"):
-        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
-
-    # Same version with different bytes is a conflict, never an overwrite.
-    (source / "weight_v000002" / "model-00001.safetensors").write_bytes(b"changed")
-    with pytest.raises(receiver.CheckpointConflictError, match="different manifest"):
-        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
-    assert (local / "model-00001.safetensors").read_bytes() == b"two"
-
-
-def test_failed_copy_keeps_old_checkpoint_available(tmp_path: Path, monkeypatch):
-    source = tmp_path / "published"
-    _publish(source, 1, b"old")
-    _publish(source, 2, b"new")
-    local = tmp_path / "local"
-    receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=1)
-
-    original_copy2 = receiver.shutil.copy2
-    calls = 0
-
-    def fail_after_first(*args, **kwargs):
-        nonlocal calls
-        calls += 1
-        if calls >= 2:
-            raise OSError("simulated partial copy")
-        return original_copy2(*args, **kwargs)
-
-    monkeypatch.setattr(receiver.shutil, "copy2", fail_after_first)
-    with pytest.raises(OSError, match="simulated partial copy"):
-        receiver.materialize_checkpoint(source_dir=str(source), local_checkpoint_dir=str(local), target_version=2)
-    assert json.loads((local / receiver._CHECKPOINT_MARKER).read_text())["version"] == 1
-    assert (local / "model-00001.safetensors").read_bytes() == b"old"
-
-
-def test_concurrent_same_version_pulls_are_serialized(tmp_path: Path):
-    source = tmp_path / "published"
-    _publish(source, 1)
-    local = tmp_path / "local"
-
-    def pull():
-        return receiver.materialize_checkpoint(
-            source_dir=str(source), local_checkpoint_dir=str(local), target_version=1
-        )["status"]
-
-    with ThreadPoolExecutor(max_workers=2) as pool:
-        statuses = sorted(pool.map(lambda _: pull(), range(2)))
-    assert statuses == ["already_applied", "materialized"]
-    assert (local / "model.safetensors.index.json").is_file()
-
-
-def test_materialize_rejects_paths_outside_configured_roots(tmp_path: Path):
-    configured_source = tmp_path / "configured-source"
-    other_source = tmp_path / "other-source"
-    _publish(configured_source, 1)
-    _publish(other_source, 1)
-    configured_local = tmp_path / "configured-local"
-
-    with pytest.raises(receiver.CheckpointReceiveError, match="configured checkpoint source"):
-        receiver.materialize_checkpoint(
-            source_dir=str(other_source),
-            local_checkpoint_dir=str(configured_local),
-            target_version=1,
-            expected_source_dir=str(configured_source),
-            expected_local_checkpoint_dir=str(configured_local),
+            target_version=version,  # type: ignore[arg-type]
         )
 
-    with pytest.raises(receiver.CheckpointReceiveError, match="configured local checkpoint destination"):
-        receiver.materialize_checkpoint(
-            source_dir=str(configured_source),
-            local_checkpoint_dir=str(tmp_path / "other-local"),
-            target_version=1,
-            expected_source_dir=str(configured_source),
-            expected_local_checkpoint_dir=str(configured_local),
-        )
 
-    result = receiver.materialize_checkpoint(
-        source_dir=str(configured_source),
-        local_checkpoint_dir=str(configured_local),
-        target_version=1,
-        expected_source_dir=str(configured_source),
-        expected_local_checkpoint_dir=str(configured_local),
-    )
-    assert result["status"] == "materialized"
+@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")
 
-def test_materialize_rejects_source_symlink(tmp_path: Path):
-    source = tmp_path / "source"
-    _publish(source, 1)
-    source_link = tmp_path / "source-link"
-    try:
-        source_link.symlink_to(source, target_is_directory=True)
-    except OSError as exc:
-        pytest.skip(f"directory symlinks are unavailable: {exc}")
+    monkeypatch.setattr(checkpoint_receiver.shutil, "copytree", fail_copy)
 
-    with pytest.raises(receiver.CheckpointReceiveError, match="must not be a symlink"):
-        receiver.materialize_checkpoint(
-            source_dir=str(source_link),
-            local_checkpoint_dir=str(tmp_path / "local"),
+    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,
         )
 
-
-def test_publish_checkpoint_directory_is_immutable_and_retryable(tmp_path: Path):
-    destination = tmp_path / "weight_v000001"
-    staging = tmp_path / ".weight_v000001.staging"
-    _publish(tmp_path / "first", 1).rename(staging)
-
-    assert receiver.publish_checkpoint_directory(staging, destination) == "published"
-    assert (destination / "model-00001.safetensors").read_bytes() == b"weights-v1"
-
-    retry_staging = tmp_path / ".weight_v000001.retry"
-    _publish(tmp_path / "retry", 1).rename(retry_staging)
-    assert receiver.publish_checkpoint_directory(retry_staging, destination) == "already_published"
-    assert retry_staging.exists()
-
-    conflict_staging = tmp_path / ".weight_v000001.conflict"
-    _publish(tmp_path / "conflict", 1, b"different").rename(conflict_staging)
-    with pytest.raises(receiver.CheckpointConflictError, match="different contents"):
-        receiver.publish_checkpoint_directory(conflict_staging, destination)
-    assert (destination / "model-00001.safetensors").read_bytes() == b"weights-v1"
-
-
-def test_publish_checkpoint_directory_accepts_late_shared_fs_publisher(tmp_path: Path):
-    destination = tmp_path / "weight_v000001"
-    staging = tmp_path / ".weight_v000001.staging"
-    _publish(tmp_path / "first", 1).rename(staging)
-
-    assert receiver.publish_checkpoint_directory(staging, destination) == "published"
-    assert not staging.exists()
-    assert receiver.publish_checkpoint_directory(staging, destination) == "published_by_peer"
-
-
-def test_fsync_directory_is_best_effort(tmp_path: Path, monkeypatch):
-    warnings = []
-    closed = []
-
-    def unsupported_fsync(descriptor):
-        assert descriptor == 7
-        raise OSError("unsupported")
-
-    monkeypatch.setattr(receiver.os, "O_DIRECTORY", 0, raising=False)
-    monkeypatch.setattr(receiver.os, "open", lambda *args: 7)
-    monkeypatch.setattr(receiver.os, "fsync", unsupported_fsync)
-    monkeypatch.setattr(receiver.os, "close", closed.append)
-    monkeypatch.setattr(receiver.logger, "warning", lambda *args, **kwargs: warnings.append((args, kwargs)))
-
-    receiver._fsync_directory(tmp_path)
-
-    assert closed == [7]
-    assert warnings[0][0][0] == "Could not fsync checkpoint directory: %s"
-    assert warnings[0][0][1] == tmp_path
-    assert warnings[0][1] == {"exc_info": True}
+    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 ddf1d4df8..a43636864 100644
--- a/tests/utils/test_vllm_engine.py
+++ b/tests/utils/test_vllm_engine.py
@@ -325,36 +325,6 @@ def test_build_vllm_subprocess_env_no_batch_invariant_by_default(vllm_args, monk
     assert "VLLM_BATCH_INVARIANT" not in env
 
 
-@pytest.mark.unit
-def test_build_vllm_subprocess_env_sets_checkpoint_roots(vllm_args):
-    vllm_args.update_weight_disk_dir = "/shared/weights"
-    vllm_args.update_weight_local_checkpoint_dir = "/local/weights"
-
-    env = mod._build_subprocess_env(
-        {
-            "_args": vllm_args,
-            "_visible_devices": "0",
-        }
-    )
-
-    assert env["VIME_WEIGHT_SOURCE_DIR"] == str(Path("/shared/weights"))
-    assert env["VIME_WEIGHT_LOCAL_CHECKPOINT_DIR"] == str(Path("/local/weights"))
-
-
-@pytest.mark.unit
-def test_pull_weights_leaves_paths_to_server(vllm_engine, monkeypatch):
-    calls: list[tuple[str, dict]] = []
-
-    def fake_make_request(endpoint: str, payload: dict):
-        calls.append((endpoint, payload))
-        return {"success": True, "local_checkpoint_dir": "/remote/local"}
-
-    monkeypatch.setattr(vllm_engine, "_make_request", fake_make_request)
-
-    assert vllm_engine.pull_weights(7)["local_checkpoint_dir"] == "/remote/local"
-    assert calls == [("pull_weights", {"target_version": 7})]
-
-
 @pytest.mark.unit
 def test_build_vllm_subprocess_env_sets_disaggregation_side_channel(vllm_args):
     env = mod._build_subprocess_env(
@@ -811,11 +781,11 @@ def test_update_weights_from_disk_posts_collective_rpc(vllm_engine, monkeypatch)
 
     def fake_post(url, *, params=None, timeout=30, json=None):
         seen.append((url, params, timeout, json))
-        return _MockResponse(json_data={"results": [None]})
+        return _MockResponse(json_data={"reloaded": True})
 
     monkeypatch.setattr(mod.requests, "post", fake_post)
 
-    assert vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8") == {"results": [None]}
+    assert vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8") == {"reloaded": True}
     assert seen[0][0] == "http://127.0.0.1:8765/collective_rpc"
     assert seen[0][3]["method"] == "reload_weights"
     assert seen[1][0] == "http://127.0.0.1:8765/update_weight_version"
@@ -891,39 +861,8 @@ def fake_post(url, *, params=None, timeout=30, json=None):
         return _MockResponse(text="boom", status_code=500)
 
     monkeypatch.setattr(mod.requests, "post", fake_post)
-    vllm_engine._weight_version = "7"
     with pytest.raises(requests.exceptions.HTTPError):
         vllm_engine.update_weights_from_disk("/tmp/model")
-    assert vllm_engine._weight_version is None
-
-
-@pytest.mark.unit
-def test_update_weights_from_disk_does_not_commit_version_for_invalid_success_body(vllm_engine, monkeypatch):
-    monkeypatch.setattr(
-        mod.requests,
-        "post",
-        lambda *args, **kwargs: _MockResponse(text="not-json", status_code=200),
-    )
-    vllm_engine._weight_version = "7"
-
-    with pytest.raises(RuntimeError, match="non-JSON success response"):
-        vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8")
-
-    assert vllm_engine._weight_version is None
-
-
-@pytest.mark.unit
-def test_update_weights_from_disk_rejects_empty_rank_results(vllm_engine, monkeypatch):
-    monkeypatch.setattr(
-        mod.requests,
-        "post",
-        lambda *args, **kwargs: _MockResponse(json_data={"results": []}),
-    )
-
-    with pytest.raises(RuntimeError, match="invalid result"):
-        vllm_engine.update_weights_from_disk("/tmp/model", weight_version="8")
-
-    assert vllm_engine._weight_version is None
 
 
 @pytest.mark.unit
diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py
index dee4403dc..13ded4186 100644
--- a/vime/backends/megatron_utils/actor.py
+++ b/vime/backends/megatron_utils/actor.py
@@ -560,7 +560,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
             self.sleep()
 
     @timer
-    def update_weights(self, weight_version: int | None = None) -> None:
+    def update_weights(self) -> None:
         if self.args.debug_train_only or self.args.debug_rollout_only:
             return
 
@@ -609,14 +609,7 @@ def update_weights(self, weight_version: int | None = None) -> None:
                     if ".draft_model." in name:
                         param.data = backup[name].to(param.device)
             print_memory("before update_weights")
-            if self.args.update_weight_mode == "full" and self.args.update_weight_transport == "disk":
-                if weight_version is None:
-                    raise ValueError("full-disk weight update requires an explicit weight_version")
-                self.weight_updater.update_weights(weight_version=weight_version)
-            else:
-                if weight_version is not None:
-                    raise ValueError("weight_version is only valid for full-disk weight updates")
-                self.weight_updater.update_weights()
+            self.weight_updater.update_weights()
             print_memory("after update_weights")
 
             if getattr(self.args, "keep_old_actor", False):
diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
index 40b4334b9..1f0cc8559 100644
--- a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
+++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py
@@ -9,7 +9,6 @@
 import torch.distributed as dist
 from ray.actor import ActorHandle
 
-from vime.backends.vllm_utils.checkpoint_receiver import publish_checkpoint_directory
 from vime.utils.distributed_utils import get_gloo_group
 
 from ..hf_checkpoint_saver import save_hf_model_to_path
@@ -64,27 +63,20 @@ def pop_metrics(self) -> dict[str, float]:
         return out
 
     @torch.no_grad()
-    def update_weights(self, *, weight_version: int) -> None:
-        if weight_version < self.weight_version or weight_version > self.weight_version + 1:
-            raise ValueError(
-                f"invalid full-disk weight version {weight_version}; current version is {self.weight_version}"
-            )
-
-        version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{weight_version:06d}"
-        staging_dir = version_dir.with_name(f".{version_dir.name}.staging")
-
-        # Every rank cleans its local view. On a POSIX shared filesystem these
-        # calls converge on the same path; on host-local/object-backed mounts
-        # each writer must clean and create its own staging directory.
-        shutil.rmtree(staging_dir, ignore_errors=True)
+    def update_weights(self) -> None:
+        self.weight_version += 1
+        version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{self.weight_version:06d}"
+
+        if dist.get_rank() == 0:
+            shutil.rmtree(version_dir, ignore_errors=True)
         dist.barrier(group=get_gloo_group())
 
-        # Every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface
+        # every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface
         # one rank's mkdir to another until commit
-        staging_dir.mkdir(parents=True, exist_ok=True)
+        version_dir.mkdir(parents=True, exist_ok=True)
         save_hf_model_to_path(
             self.args,
-            staging_dir,
+            version_dir,
             self.model,
             model_name=self.model_name,
             quantization_config=self.quantization_config,
@@ -92,16 +84,11 @@ def update_weights(self, *, weight_version: int) -> None:
         )
         dist.barrier(group=get_gloo_group())
 
-        publish_checkpoint_directory(staging_dir, version_dir)
-        dist.barrier(group=get_gloo_group())
-        shutil.rmtree(staging_dir, ignore_errors=True)
-
         # every rank runs the hook (it gates itself): each container must publish
         # its own writes
         if self._post_write_hook is not None:
             self._post_write_hook(self.args, str(version_dir), list(self.rollout_engines))
         dist.barrier(group=get_gloo_group())
-        self.weight_version = weight_version
 
         # VLLM reload is orchestrated by RayTrainGroup after the checkpoint
         # is fully written, so training-side lifecycle can decide whether
diff --git a/vime/backends/vllm_utils/checkpoint_receiver.py b/vime/backends/vllm_utils/checkpoint_receiver.py
index 19090899b..ed106881c 100644
--- a/vime/backends/vllm_utils/checkpoint_receiver.py
+++ b/vime/backends/vllm_utils/checkpoint_receiver.py
@@ -1,47 +1,10 @@
-"""Transactional receiver for full HuggingFace checkpoints published on disk.
-
-The vLLM HTTP endpoint is intentionally kept thin.  This module owns the
-filesystem transaction so it can be tested without importing vLLM or starting
-an HTTP server.
-"""
+"""Copy a published checkpoint to rollout-host-local storage."""
 
 from __future__ import annotations
 
-import hashlib
-import json
-import logging
-import os
 import shutil
-import tempfile
-import threading
 import uuid
 from pathlib import Path
-from typing import Any
-
-logger = logging.getLogger(__name__)
-
-_CHECKPOINT_MARKER = ".vime_checkpoint_manifest.json"
-_INDEX_NAMES = ("model.safetensors.index.json", "pytorch_model.bin.index.json")
-_SINGLE_FILE_WEIGHT_NAMES = ("model.safetensors", "pytorch_model.bin")
-_MATERIALIZE_LOCK = threading.RLock()
-
-
-class CheckpointReceiveError(ValueError):
-    """A malformed request or checkpoint that should be returned as HTTP 400."""
-
-    status_code = 400
-
-
-class StaleCheckpointError(CheckpointReceiveError):
-    """A checkpoint older than the active local checkpoint."""
-
-    status_code = 409
-
-
-class CheckpointConflictError(CheckpointReceiveError):
-    """A version was already applied with different checkpoint contents."""
-
-    status_code = 409
 
 
 def materialize_checkpoint(
@@ -49,378 +12,40 @@ def materialize_checkpoint(
     local_checkpoint_dir: str,
     source_dir: str,
     target_version: int,
-    expected_local_checkpoint_dir: str | None = None,
-    expected_source_dir: str | None = None,
-) -> dict[str, Any]:
-    """Materialize one published checkpoint into a host-local directory.
-
-    The source is validated and checksummed before any destination is touched.
-    Files are copied into a sibling temporary directory, then that directory is
-    swapped into place.  A failed copy or swap therefore leaves the previous
-    local checkpoint intact.
-    """
-
-    version = _parse_version(target_version)
-    if expected_source_dir is not None:
-        expected_source_root = _resolve_existing_directory(expected_source_dir, "configured source_dir")
-        source_root = _resolve_existing_directory(source_dir, "source_dir")
-        if source_root != expected_source_root:
-            raise CheckpointReceiveError("source_dir does not match the configured checkpoint source")
-    else:
-        source_root = _resolve_existing_directory(source_dir, "source_dir")
-    version_dir = source_root / f"weight_v{version:06d}"
-    version_dir = _resolve_existing_directory(version_dir, "checkpoint version")
-
-    if expected_local_checkpoint_dir is not None:
-        expected_local_dir = _prepare_local_directory(expected_local_checkpoint_dir)
-        local_dir = _resolve_local_directory(local_checkpoint_dir)
-        if local_dir != expected_local_dir:
-            raise CheckpointReceiveError(
-                "local_checkpoint_dir does not match the configured local checkpoint destination"
-            )
-    else:
-        local_dir = _prepare_local_directory(local_checkpoint_dir)
-    if local_dir.exists() and local_dir.resolve() == version_dir:
-        raise CheckpointReceiveError("local_checkpoint_dir must differ from source checkpoint")
-
-    with _MATERIALIZE_LOCK:
-        source_manifest = _build_checkpoint_manifest(version_dir)
-        source_manifest_hash = _manifest_hash(source_manifest["files"])
-        current = _read_marker(local_dir)
-
-        if current is not None:
-            current_version = current["version"]
-            if current_version > version:
-                raise StaleCheckpointError(f"stale checkpoint version {version}; active version is {current_version}")
-            if current_version == version:
-                if current["manifest_sha256"] != source_manifest_hash:
-                    raise CheckpointConflictError(
-                        f"checkpoint version {version} is already active with a different manifest"
-                    )
-                if _local_manifest_is_valid(local_dir, current):
-                    return _result("already_applied", version, current)
-
-        staging_dir = Path(tempfile.mkdtemp(prefix=f".{local_dir.name}.", dir=local_dir.parent))
-        try:
-            _copy_manifest_files(version_dir, staging_dir, source_manifest["files"])
-            marker = {
-                "version": version,
-                "manifest_sha256": source_manifest_hash,
-                "files": source_manifest["files"],
-            }
-            _write_json(staging_dir / _CHECKPOINT_MARKER, marker)
-            _atomic_replace_directory(staging_dir, local_dir)
-            staging_dir = None  # ownership transferred to local_dir
-        except Exception:
-            if staging_dir is not None:
-                shutil.rmtree(staging_dir, ignore_errors=True)
-            raise
-
-        return _result("materialized", version, marker)
+) -> 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")
 
-def _parse_version(value: int) -> int:
-    if isinstance(value, bool):
-        raise CheckpointReceiveError("target_version must be a positive integer")
-    try:
-        version = int(value)
-    except (TypeError, ValueError) as exc:
-        raise CheckpointReceiveError("target_version must be a positive integer") from exc
-    if version <= 0 or str(value).strip() != str(version):
-        raise CheckpointReceiveError("target_version must be a positive integer")
-    return version
-
-
-def _resolve_existing_directory(value: str | Path, name: str) -> Path:
-    if not isinstance(value, (str, Path)) or not str(value).strip():
-        raise CheckpointReceiveError(f"{name} must be a non-empty path")
-    path = Path(value).expanduser()
-    if path.is_symlink():
-        raise CheckpointReceiveError(f"{name} must not be a symlink: {path}")
-    try:
-        resolved = path.resolve(strict=True)
-    except OSError as exc:
-        raise CheckpointReceiveError(f"{name} does not exist: {path}") from exc
-    if not resolved.is_dir():
-        raise CheckpointReceiveError(f"{name} must be a directory: {path}")
-    return resolved
-
-
-def _prepare_local_directory(value: str | Path) -> Path:
-    if not isinstance(value, (str, Path)) or not str(value).strip():
-        raise CheckpointReceiveError("local_checkpoint_dir must be a non-empty path")
-    path = Path(value).expanduser()
-    if path.exists() and (path.is_symlink() or not path.is_dir()):
-        raise CheckpointReceiveError(f"local_checkpoint_dir must be a directory: {path}")
-    try:
-        path.parent.mkdir(parents=True, exist_ok=True)
-        parent = path.parent.resolve(strict=True)
-    except OSError as exc:
-        raise CheckpointReceiveError(f"cannot prepare local checkpoint parent: {path.parent}") from exc
-    return parent / path.name
+    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")
 
-def _resolve_local_directory(value: str | Path) -> Path:
-    """Resolve an untrusted destination without creating caller-selected parents."""
+    suffix = uuid.uuid4().hex
+    staging = destination.parent / f".{destination.name}.staging-{suffix}"
+    backup = destination.parent / f".{destination.name}.backup-{suffix}"
 
-    if not isinstance(value, (str, Path)) or not str(value).strip():
-        raise CheckpointReceiveError("local_checkpoint_dir must be a non-empty path")
-    path = Path(value).expanduser()
-    if path.is_symlink() or (path.exists() and not path.is_dir()):
-        raise CheckpointReceiveError(f"local_checkpoint_dir must be a directory: {path}")
     try:
-        parent = path.parent.resolve(strict=True)
-    except OSError as exc:
-        raise CheckpointReceiveError(f"local checkpoint parent does not exist: {path.parent}") from exc
-    return parent / path.name
-
-
-def publish_checkpoint_directory(staging_dir: str | Path, destination_dir: str | Path) -> str:
-    """Publish an immutable checkpoint directory, allowing an identical retry."""
-
-    if not isinstance(staging_dir, (str, Path)) or not str(staging_dir).strip():
-        raise CheckpointReceiveError("checkpoint staging directory must be a non-empty path")
-    staging_path = Path(staging_dir).expanduser()
-    if staging_path.is_symlink():
-        raise CheckpointReceiveError(f"checkpoint staging directory must not be a symlink: {staging_path}")
-    try:
-        staging_parent = staging_path.parent.resolve(strict=True)
-    except OSError as exc:
-        raise CheckpointReceiveError(
-            f"checkpoint staging directory parent does not exist: {staging_path.parent}"
-        ) from exc
-    destination = _prepare_local_directory(destination_dir)
-    if staging_parent != destination.parent:
-        raise CheckpointReceiveError("checkpoint staging and destination directories must be siblings")
-    staging_path = staging_parent / staging_path.name
-
-    try:
-        staging = _resolve_existing_directory(staging_path, "checkpoint staging directory")
-    except CheckpointReceiveError:
-        if not staging_path.exists() and _published_checkpoint_is_valid(destination):
-            return "published_by_peer"
-        raise
-
-    if destination.exists():
-        destination_manifest = _build_checkpoint_manifest(destination)
-        try:
-            staging_manifest = _build_checkpoint_manifest(staging)
-        except (CheckpointReceiveError, OSError):
-            if not staging.exists() and _published_checkpoint_is_valid(destination):
-                return "published_by_peer"
-            raise
-        destination_hash = _manifest_hash(destination_manifest["files"])
-        staging_hash = _manifest_hash(staging_manifest["files"])
-        if destination_hash != staging_hash:
-            raise CheckpointConflictError(
-                f"published checkpoint already exists with different contents: {destination}"
-            )
-        return "already_published"
-
-    try:
-        os.replace(staging, destination)
-    except FileNotFoundError:
-        # On a shared filesystem another writer rank may have renamed the same
-        # staging directory after the distributed write barrier.
-        if _published_checkpoint_is_valid(destination):
-            return "published_by_peer"
-        raise
-    _fsync_directory(destination.parent)
-    return "published"
-
-
-def _published_checkpoint_is_valid(destination: Path) -> bool:
-    if not destination.is_dir() or destination.is_symlink():
-        return False
-    try:
-        _build_checkpoint_manifest(destination)
-    except (CheckpointReceiveError, OSError):
-        return False
-    return True
-
-
-def _build_checkpoint_manifest(root: Path) -> dict[str, Any]:
-    index_candidates = [root / name for name in _INDEX_NAMES if (root / name).exists()]
-    if any(path.is_symlink() for path in index_candidates):
-        raise CheckpointReceiveError("checkpoint index must not be a symlink")
-    indexes = [path for path in index_candidates if path.is_file()]
-    if len(indexes) > 1:
-        raise CheckpointReceiveError(
-            f"checkpoint must contain exactly one supported weight index: {', '.join(_INDEX_NAMES)}"
-        )
-
-    if not indexes:
-        direct_weights = [root / name for name in _SINGLE_FILE_WEIGHT_NAMES if (root / name).is_file()]
-        if not direct_weights or any(path.stat().st_size <= 0 for path in direct_weights):
-            raise CheckpointReceiveError(
-                "checkpoint must contain a supported weight index or a non-empty single-file weight"
-            )
-        return _manifest_files(root)
-
-    index_path = indexes[0]
-    try:
-        with index_path.open("r", encoding="utf-8") as file:
-            index = json.load(file)
-    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
-        raise CheckpointReceiveError(f"malformed checkpoint index: {index_path.name}") from exc
-
-    weight_map = index.get("weight_map") if isinstance(index, dict) else None
-    if not isinstance(weight_map, dict) or not weight_map:
-        raise CheckpointReceiveError("checkpoint index must contain a non-empty weight_map")
-    for tensor_name, filename in weight_map.items():
-        if not isinstance(tensor_name, str) or not isinstance(filename, str):
-            raise CheckpointReceiveError("checkpoint index contains a non-string weight mapping")
-        _resolve_checkpoint_file(root, filename)
-
-    return _manifest_files(root)
-
-
-def _manifest_files(root: Path) -> dict[str, Any]:
-    files: dict[str, dict[str, Any]] = {}
-    for path in sorted(root.rglob("*")):
-        if path.is_symlink():
-            raise CheckpointReceiveError(f"checkpoint contains a symlink: {path.relative_to(root)}")
-        if not path.is_file():
-            continue
-        relative = path.relative_to(root).as_posix()
-        try:
-            size = path.stat().st_size
-        except OSError as exc:
-            raise CheckpointReceiveError(f"cannot stat checkpoint file: {relative}") from exc
-        files[relative] = {"size": size, "sha256": _sha256(path)}
-    if not files:
-        raise CheckpointReceiveError("checkpoint contains no files")
-    return {"files": files}
-
-
-def _resolve_checkpoint_file(root: Path, filename: str) -> Path:
-    candidate = root / filename
-    try:
-        resolved = candidate.resolve(strict=True)
-        resolved.relative_to(root)
-    except (OSError, ValueError) as exc:
-        raise CheckpointReceiveError(f"checkpoint index references missing or unsafe file: {filename}") from exc
-    if candidate.is_symlink() or not resolved.is_file() or resolved.stat().st_size <= 0:
-        raise CheckpointReceiveError(f"checkpoint index references invalid file: {filename}")
-    return resolved
-
-
-def _copy_manifest_files(source: Path, destination: Path, files: dict[str, dict[str, Any]]) -> None:
-    for relative, metadata in files.items():
-        source_file = source / Path(relative)
-        destination_file = destination / Path(relative)
-        destination_file.parent.mkdir(parents=True, exist_ok=True)
-        try:
-            shutil.copy2(source_file, destination_file)
-            copied_size = destination_file.stat().st_size
-            copied_hash = _sha256(destination_file)
-        except OSError as exc:
-            raise OSError(f"failed to copy checkpoint file {relative}: {exc}") from exc
-        if copied_size != metadata["size"] or copied_hash != metadata["sha256"]:
-            raise OSError(f"checkpoint file changed while copying: {relative}")
-
-
-def _read_marker(local_dir: Path) -> dict[str, Any] | None:
-    marker_path = local_dir / _CHECKPOINT_MARKER
-    if not marker_path.exists():
-        return None
-    try:
-        with marker_path.open("r", encoding="utf-8") as file:
-            marker = json.load(file)
-    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
-        raise CheckpointReceiveError("local checkpoint version marker is malformed") from exc
-    if (
-        not isinstance(marker, dict)
-        or isinstance(marker.get("version"), bool)
-        or not isinstance(marker.get("version"), int)
-        or marker["version"] <= 0
-        or not isinstance(marker.get("manifest_sha256"), str)
-        or not isinstance(marker.get("files"), dict)
-    ):
-        raise CheckpointReceiveError("local checkpoint version marker is malformed")
-    return marker
-
-
-def _local_manifest_is_valid(local_dir: Path, marker: dict[str, Any]) -> bool:
-    if _manifest_hash(marker["files"]) != marker["manifest_sha256"]:
-        return False
-    for relative, metadata in marker["files"].items():
-        if not isinstance(relative, str) or not isinstance(metadata, dict):
-            return False
-        path = local_dir / Path(relative)
-        try:
-            path.resolve(strict=True).relative_to(local_dir.resolve(strict=True))
-            if path.is_symlink() or not path.is_file() or path.stat().st_size != metadata.get("size"):
-                return False
-            if _sha256(path) != metadata.get("sha256"):
-                return False
-        except (OSError, ValueError):
-            return False
-    return True
-
-
-def _manifest_hash(files: dict[str, Any]) -> str:
-    encoded = json.dumps(files, sort_keys=True, separators=(",", ":")).encode("utf-8")
-    return hashlib.sha256(encoded).hexdigest()
-
-
-def _sha256(path: Path) -> str:
-    digest = hashlib.sha256()
-    with path.open("rb") as file:
-        for chunk in iter(lambda: file.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _write_json(path: Path, value: dict[str, Any]) -> None:
-    with path.open("w", encoding="utf-8") as file:
-        json.dump(value, file, sort_keys=True, separators=(",", ":"))
-        file.flush()
-        os.fsync(file.fileno())
-
-
-def _fsync_directory(path: Path) -> None:
-    if not hasattr(os, "O_DIRECTORY"):
-        return
-    try:
-        descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
-    except OSError:
-        logger.warning("Could not open checkpoint directory for fsync: %s", path, exc_info=True)
-        return
-    try:
-        os.fsync(descriptor)
-    except OSError:
-        logger.warning("Could not fsync checkpoint directory: %s", path, exc_info=True)
-    finally:
-        os.close(descriptor)
-
-
-def _atomic_replace_directory(staging: Path, destination: Path) -> None:
-    backup: Path | None = None
-    if destination.exists() or destination.is_symlink():
-        if destination.is_symlink() or not destination.is_dir():
-            raise CheckpointReceiveError(f"local checkpoint destination is not a directory: {destination}")
-        backup = destination.parent / f".{destination.name}.backup-{uuid.uuid4().hex}"
-        os.replace(destination, backup)
-    try:
-        os.replace(staging, destination)
+        shutil.copytree(source, staging)
+        if destination.exists():
+            destination.replace(backup)
+        staging.replace(destination)
     except Exception:
-        if backup is not None and backup.exists() and not destination.exists():
-            os.replace(backup, destination)
+        if backup.exists() and not destination.exists():
+            backup.replace(destination)
         raise
-    _fsync_directory(destination.parent)
-    if backup is not None:
-        try:
-            shutil.rmtree(backup)
-        except OSError:
-            logger.warning("Could not remove old checkpoint backup %s", backup, exc_info=True)
-
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
 
-def _result(status: str, version: int, marker: dict[str, Any]) -> dict[str, Any]:
+    shutil.rmtree(backup, ignore_errors=True)
     return {
         "success": True,
-        "status": status,
-        "version": version,
-        "manifest_sha256": marker["manifest_sha256"],
-        "files": len(marker["files"]),
+        "version": target_version,
+        "local_checkpoint_dir": str(destination),
     }
diff --git a/vime/ray/actor_group.py b/vime/ray/actor_group.py
index 0284c5811..c3725d14c 100644
--- a/vime/ray/actor_group.py
+++ b/vime/ray/actor_group.py
@@ -1,4 +1,3 @@
-import logging
 import os
 import shutil
 import time
@@ -10,8 +9,6 @@
 
 from vime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, add_default_ray_env_vars
 
-logger = logging.getLogger(__name__)
-
 
 class RayTrainGroup:
     """
@@ -169,11 +166,11 @@ def update_weights(self):
 
         weight_version = self._disk_weight_version + 1
         disk_weight_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{weight_version:06d}"
-        ray.get([actor.update_weights.remote(weight_version=weight_version) for actor in self._actor_handlers])
+        ray.get([actor.update_weights.remote() for actor in self._actor_handlers])
+        self._disk_weight_version = weight_version
         if self._release_train_enabled():
             self.release()
         self._reload_rollout_weights_from_disk(disk_weight_dir, str(weight_version))
-        self._disk_weight_version = weight_version
 
     def onload(self):
         return ray.get([actor.wake_up.remote() for actor in self._actor_handlers])
@@ -240,66 +237,33 @@ def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version):
             # each host pulls the published checkpoint onto local disk (e.g. NVMe) and
             # the engines reload from there; the pull is disk-only, so it runs before
             # pause and overlaps generation
-            pull_results = ray.get(
-                [engine.pull_weights.remote(target_version=int(weight_version)) for engine in engines]
-            )
-            if not isinstance(pull_results, list) or len(pull_results) != len(engines):
-                raise RuntimeError(f"Expected one pull_weights result per engine, got: {pull_results!r}")
-            model_paths = []
-            for index, result in enumerate(pull_results):
-                if (
-                    not isinstance(result, dict)
-                    or result.get("success") is not True
-                    or not isinstance(result.get("local_checkpoint_dir"), str)
-                    or not result["local_checkpoint_dir"]
-                ):
-                    raise RuntimeError(f"pull_weights returned an invalid result for engine {index}: {result!r}")
-                model_paths.append(result["local_checkpoint_dir"])
+            ray.get([engine.pull_weights.remote(int(weight_version)) for engine in engines])
+            model_path = self.args.update_weight_local_checkpoint_dir
         else:
-            model_paths = [str(disk_weight_dir)] * len(engines)
-        reload_error: Exception | None = None
-        pause_attempted = False
-        try:
-            # Resume every engine even if only a subset completed the batched
-            # pause call before Ray surfaced an error.
-            pause_attempted = True
-            ray.get([engine.pause_generation.remote() for engine in engines])
-            ray.get([engine.flush_cache.remote() for engine in engines])
-            ray.get(
-                [
-                    engine.update_weights_from_disk.remote(
-                        model_path=model_path,
-                        weight_version=weight_version,
-                    )
-                    for engine, model_path in zip(engines, model_paths, strict=True)
-                ]
-            )
-            if self.args.ci_test:
-                engine_versions = ray.get([engine.get_weight_version.remote() for engine in engines])
-                mismatches = [
-                    f"engine {idx}: {engine_version}"
-                    for idx, engine_version in enumerate(engine_versions)
-                    if str(engine_version) != str(weight_version)
-                ]
-                if mismatches:
-                    raise RuntimeError(
-                        "Weight version mismatch after disk reload! "
-                        f"Expected: {weight_version}; " + ", ".join(mismatches)
-                    )
-            if not self.args.update_weight_disk_keep_files:
-                shutil.rmtree(disk_weight_dir, ignore_errors=True)
-        except Exception as exc:
-            reload_error = exc
-            raise
-        finally:
-            if pause_attempted:
-                try:
-                    ray.get([engine.continue_generation.remote() for engine in engines])
-                except Exception as resume_error:
-                    if reload_error is None:
-                        raise
-                    add_note = getattr(reload_error, "add_note", None)
-                    if callable(add_note):
-                        add_note(f"Failed to resume rollout engines after reload failure: {resume_error}")
-                    else:
-                        logger.exception("Failed to resume rollout engines after reload failure")
+            model_path = str(disk_weight_dir)
+        ray.get([engine.pause_generation.remote() for engine in engines])
+        ray.get([engine.flush_cache.remote() for engine in engines])
+        ray.get(
+            [
+                engine.update_weights_from_disk.remote(
+                    model_path=model_path,
+                    weight_version=weight_version,
+                )
+                for engine in engines
+            ]
+        )
+        if self.args.ci_test:
+            engine_versions = ray.get([engine.get_weight_version.remote() for engine in engines])
+            mismatches = [
+                f"engine {idx}: {engine_version}"
+                for idx, engine_version in enumerate(engine_versions)
+                if str(engine_version) != str(weight_version)
+            ]
+            if mismatches:
+                raise RuntimeError(
+                    "Weight version mismatch after disk reload! "
+                    f"Expected: {weight_version}; " + ", ".join(mismatches)
+                )
+        if not self.args.update_weight_disk_keep_files:
+            shutil.rmtree(disk_weight_dir, ignore_errors=True)
+        ray.get([engine.continue_generation.remote() for engine in engines])
diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py
index ed1ecd61f..78d21d15f 100644
--- a/vime/ray/rollout.py
+++ b/vime/ray/rollout.py
@@ -271,11 +271,7 @@ def health_monitoring_resume(self) -> None:
             monitor.resume()
 
     def check_weights(self, action: str):
-        server = self._get_updatable_server()
-        engines = server.engines if server is not None else []
-        if not engines:
-            raise RuntimeError("weight checking requires at least one updatable rollout engine")
-        return ray.get([engine.check_weights.remote(action=action) for engine in engines])
+        return ray.get([engine.check_weights.remote(action=action) for engine in self.rollout_engines])
 
     def _get_rollout_data(self, rollout_id):
         if self.args.load_debug_rollout_data:
diff --git a/vime/ray/train_actor.py b/vime/ray/train_actor.py
index 626628310..72fea191c 100644
--- a/vime/ray/train_actor.py
+++ b/vime/ray/train_actor.py
@@ -116,7 +116,7 @@ def save_model(self, rollout_id, force_sync=False):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def update_weights(self, weight_version: int | None = None):
+    def update_weights(self):
         raise NotImplementedError
 
     def set_rollout_manager(self, rollout_manager):
diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py
index 62d9e1d94..4eefc41fb 100644
--- a/vime/utils/arguments.py
+++ b/vime/utils/arguments.py
@@ -239,9 +239,6 @@ def add_train_arguments(parser):
                     "reload from it. Required for --update-weight-mode=delta "
                     "--update-weight-transport=disk; optional for full disk sync (engines then "
                     "pull to local disk instead of reading the shared dir directly). The "
-                    "vLLM server pins these paths through VIME_WEIGHT_SOURCE_DIR and "
-                    "VIME_WEIGHT_LOCAL_CHECKPOINT_DIR; externally launched servers must set "
-                    "both variables to the corresponding rollout-host paths. The "
                     "read-side counterpart of --custom-update-weight-post-write-path is "
                     "--custom-update-weight-pre-read-path."
                 ),