feat: add full-disk /pull_weights synchronization for vLLM - #360
Conversation
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
766085d to
d98e2a7
Compare
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
d98e2a7 to
bb6f691
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a transactional receiver for full HuggingFace checkpoints published on disk, integrates weight checking and pulling endpoints in the vLLM router, and orchestrates full-disk weight reloads in RayTrainGroup. The review feedback highlights several critical improvement opportunities: optimizing file copying by hashing on-the-fly to avoid redundant disk reads, implementing a quick directory comparison to prevent concurrent hashing storms on shared filesystems, ensuring Python < 3.11 compatibility by safely handling add_note, wrapping directory fsync in a try-except block to handle potential OSErrors, and correcting a misleading error message in RayTrainGroup.
| 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}") |
There was a problem hiding this comment.
Performance Bottleneck: Redundant Disk Reads during Checkpoint Copying\n\nCurrently, _copy_manifest_files copies each file using shutil.copy2 and then reads the entire copied file from disk again to compute its SHA-256 hash via _sha256(destination_file).\n\nFor large LLM checkpoints (often tens of gigabytes), this doubles the disk read I/O on the destination storage, leading to severe performance bottlenecks and potential disk/network contention on shared filesystems.\n\nRecommendation:\nImplement a custom chunked copy loop that updates the SHA-256 hash on the fly as bytes are read from the source and written to the destination. This avoids reading the destination file back from disk entirely.
def _copy_manifest_files(source: Path, destination: Path, files: dict[str, dict[str, Any]]) -> None:\n for relative, metadata in files.items():\n source_file = source / Path(relative)\n destination_file = destination / Path(relative)\n destination_file.parent.mkdir(parents=True, exist_ok=True)\n try:\n digest = hashlib.sha256()\n copied_size = 0\n with source_file.open(\"rb\") as f_src, destination_file.open(\"wb\") as f_dst:\n for chunk in iter(lambda: f_src.read(1024 * 1024), b\"\"):\n f_dst.write(chunk)\n digest.update(chunk)\n copied_size += len(chunk)\n shutil.copystat(source_file, destination_file)\n copied_hash = digest.hexdigest()\n except OSError as exc:\n raise OSError(f\"failed to copy checkpoint file {relative}: {exc}\") from exc\n if copied_size != metadata[\"size\"] or copied_hash != metadata[\"sha256\"]:\n raise OSError(f\"checkpoint file changed while copying: {relative}\")There was a problem hiding this comment.
@aoshen02 The current path uses copy2, then reads the local destination once to verify
its SHA-256 against the manifest. This adds one checkpoint's worth of local
read I/O, but preserves copy2's optimized fast path and verifies the actual
bytes that vLLM will load.
Hashing while copying would avoid the second read, but may lose the optimized
copy path and only verifies the bytes passed to the write operation, not the
completed destination.
I consider this a performance-versus-verification tradeoff, not a correctness
blocker. I prefer to keep the stronger destination verification unless
representative benchmarks show that the extra read is a material bottleneck.
| 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" |
There was a problem hiding this comment.
Performance Bottleneck: Concurrent Hashing Storm on Shared Filesystems\n\nWhen publish_checkpoint_directory is called and destination.exists() is True (which happens for all non-primary writing ranks on a shared filesystem), the code calls _build_checkpoint_manifest on both destination and staging directories.\n\nSince _build_checkpoint_manifest recursively computes the SHA-256 hash of every file, this causes a massive concurrent read storm on the shared filesystem (NFS, Lustre, etc.) from all non-primary ranks, leading to severe I/O degradation and training stalls.\n\nRecommendation:\nIntroduce a fast-path directory comparison helper (_quick_compare_directories) that compares file names and sizes first. If they match, skip the expensive SHA-256 computation entirely.
if destination.exists():\n def _quick_compare_directories(dir1: Path, dir2: Path) -> bool:\n try:\n files1 = {p.relative_to(dir1): p.stat().st_size for p in dir1.rglob(\"*\") if p.is_file() and not p.is_symlink()}\n files2 = {p.relative_to(dir2): p.stat().st_size for p in dir2.rglob(\"*\") if p.is_file() and not p.is_symlink()}\n return files1 == files2\n except OSError:\n return False\n\n if _quick_compare_directories(staging, destination):\n return \"already_published\"\n\n destination_manifest = _build_checkpoint_manifest(destination)\n try:\n staging_manifest = _build_checkpoint_manifest(staging)\n except (CheckpointReceiveError, OSError):\n if not staging.exists() and _published_checkpoint_is_valid(destination):\n return \"published_by_peer\"\n raise\n destination_hash = _manifest_hash(destination_manifest[\"files\"])\n staging_hash = _manifest_hash(staging_manifest[\"files\"])\n if destination_hash != staging_hash:\n raise CheckpointConflictError(\n f\"published checkpoint already exists with different contents: {destination}\"\n )\n return \"already_published\"There was a problem hiding this comment.
@aoshen02 Thanks for flagging this. The performance concern is valid: on a truly
shared filesystem, after one rank atomically publishes the staging directory,
each late publisher may rebuild the manifest and hash the entire checkpoint.
With large checkpoints and many trainer ranks, this can create concurrent
full-file reads and significant shared-storage contention.
I do not think a filename-and-size-only fast path is safe. Different weights
can have identical paths and sizes, so that shortcut would weaken
same-version conflict detection.
Two safer options are:
-
Elect one publisher per physical filesystem view. After all writer ranks
reach the write barrier, the publisher validates checkpoint completeness,
builds the canonical manifest, writes and fsyncs a commit marker, and
atomically publishes the directory. Other ranks wait for the result and
validate only the small committed manifest. -
Compute per-file checksums while writing, gather the metadata from all
writer ranks, and let the publisher construct a canonical, sorted manifest.
Peers can then compare the expected manifest digest without rescanning the
full checkpoint.
The canonical content digest should cover each file's relative path, size, and
SHA-256. Checkpoint version and publisher rank can be stored in the manifest
envelope, but should not affect the content digest or retry idempotency.
The publisher should also verify path uniqueness, index references, expected
shard completeness, and a consistent target version across writers.
For rank-local or object-backed filesystem views, publisher election may need
to happen per host/view rather than using one global rank. Likewise, a rollout
host copying the checkpoint to local disk must still verify the copied bytes
against the canonical manifest; matching manifest metadata alone does not
prove local copy integrity.
I currently lean toward option 1: a single validator and publisher per physical
filesystem view. It gives the publication state a clear owner, avoids
introducing a new distributed checksum-aggregation protocol, and limits the
full checkpoint scan to once per filesystem view. Option 2 could eliminate that
final scan, but it would require integrating checksum generation into every
checkpoint writer and handling missing, failed, or inconsistent writer reports.
I will not adopt the size-only shortcut. I would appreciate maintainer guidance
on whether the single-publisher model matches VIME's intended shared-filesystem
and rank-local topology before implementing the change.
| 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}") |
There was a problem hiding this comment.
Compatibility: Python < 3.11 Compatibility Risk with add_note\n\nException.add_note() was introduced in Python 3.11. If this codebase is run on Python 3.9 or 3.10 (which are still widely used in many deep learning and cluster environments), calling reload_error.add_note(...) will raise an AttributeError, completely masking the original reload_error and replacing it with an AttributeError traceback.\n\nRecommendation:\nCheck if reload_error has the add_note attribute before calling it, or fallback to logging the error context safely.
except Exception as resume_error:\n if reload_error is None:\n raise\n if hasattr(reload_error, \"add_note\"):\n reload_error.add_note(f\"Failed to resume rollout engines after reload failure: {resume_error}\")\n else:\n import logging\n logging.getLogger(__name__).error(f\"Failed to resume rollout engines after reload failure: {resume_error}\")| 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) | ||
|
|
There was a problem hiding this comment.
Robustness: Unhandled OSError on Directory fsync\n\nCalling os.fsync on a directory file descriptor is not universally supported across all filesystems (e.g., some virtualized/shared filesystems, or macOS during local development/testing) and can raise an OSError (such as EINVAL or EBADF).\n\nCurrently, if os.fsync(descriptor) fails, the exception is not caught and will propagate up, failing the entire checkpoint materialization or publication process.\n\nRecommendation:\nWrap os.fsync(descriptor) in a try...except OSError: block to log a warning instead of crashing the entire operation, as directory syncing is a best-effort durability optimization.
def _fsync_directory(path: Path) -> None:\n if not hasattr(os, \"O_DIRECTORY\"):\n return\n try:\n descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)\n except OSError:\n logger.warning(\"Could not open checkpoint directory for fsync: %s\", path, exc_info=True)\n return\n try:\n os.fsync(descriptor)\n except OSError:\n logger.warning(\"Could not fsync checkpoint directory: %s\", path, exc_info=True)\n finally:\n os.close(descriptor)| 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}") |
There was a problem hiding this comment.
Maintainability: Misleading Error Message\n\nThe error message "pull_weights returned one result per engine: ..." is raised when pull_results is not a list or its length does not match len(engines). This is highly misleading and confusing during debugging.\n\nRecommendation:\nUpdate the error message to accurately reflect the failure condition.
if not isinstance(pull_results, list) or len(pull_results) != len(engines):\n raise RuntimeError(f\"Expected one pull_weights result per engine, got: {pull_results!r}\")Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
|
cool, will check it. |
|
Did you raise a pr to vLLM? |
|
Not yet. |
|
gotcha |
|
Hi, could you please try to simplify the code. |
Signed-off-by: 道路自信 <182356120+daoluzixin@users.noreply.github.com>
8d83b50 to
688c2c7
Compare
|
Done. Could you please take another look? |
Thanks, will see it on weekend. |
2 participants
Summary
Implements the minimal full-disk weight refresh and weight-check interfaces discussed in #11.
Full-disk weight refresh
POST /pull_weights.source_dir/weight_v<version>into the rollout-local checkpoint directory through a staging directory.source_dir,local_checkpoint_dir, andtarget_version.Weight checker
weights_checkerendpoint to the patched vLLM server.snapshot,reset_tensors, andcompareacross all vLLM worker ranks.The development endpoints must be enabled with:
Validation
CPU
tests/utils/test_checkpoint_receiver.pycovers:GPU integration
Validated with vLLM
0.25.1andQwen2.5-0.5B-Instruct.snapshot -> reset -> pause -> pull -> reload -> compare -> resume.0.25.1.