Skip to content

feat(subsystembenchmarks): add Ray checkpoint load benchmark - #1047

Merged
zhixiangli merged 7 commits into
fsspec:mainfrom
Yonghui-Lee:ray-subsystembenchmark-checkpointing-load
Sep 17, 2026
Merged

zhixiangli merged 7 commits into
fsspec:mainfrom
Yonghui-Lee:ray-subsystembenchmark-checkpointing-load

Conversation

@Yonghui-Lee

Copy link
Copy Markdown
Collaborator

Depends on #1043

This PR adds the checkpoint load subsystem benchmark for Ray Data and Ray Train PyTorch workloads.

It benchmarks loading single-node and distributed checkpoints across FSDP, DDP, and Tensor Parallel + Data Parallel strategies, including cross-topology resharding.

Key Changes

Ray Checkpoint Load Driver:

  • Implemented RayCheckpointLoadWorker actor executing distributed checkpoint load operations via PyTorch Distributed Checkpoint (DCP) and ray.train.Checkpoint.
  • Implemented RayCheckpointSetupWorker actor that pre-populates the source checkpoint on GCS using the case's setup topology (setup_world_size, setup_tensor_parallel_size, setup_data_parallel_size).
  • Implemented RayCheckpointReadDriver managing cluster initialization, fixture setup, and benchmark execution.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a new subsystem benchmark suite for Ray checkpointing (ray_data), which includes drivers, configurations, and tests for both checkpoint read and write scenarios. It also updates the Cloud Build configuration to support a configurable boot disk size. The review feedback suggests two improvements: logging warnings instead of silently suppressing exceptions when verifying staging directories in common.py, and simplifying the redundant parameter resolution logic in read/driver.py by directly accessing attributes already handled by CheckpointParameters.

Comment thread gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/common.py Outdated
Comment thread gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/read/driver.py Outdated
@codecov

codecov Bot commented Sep 8, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.25%. Comparing base (ef9959c) to head (ee27c8e).

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #1047   +/-   ##
=======================================
  Coverage   90.25%   90.25%           
=======================================
  Files          16       16           
  Lines        3755     3755           
=======================================
  Hits         3389     3389           
  Misses        366      366           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@Yonghui-Lee
Yonghui-Lee force-pushed the ray-subsystembenchmark-checkpointing-load branch from 27c7ee2 to 225597f Compare September 8, 2026 05:09
Comment thread gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/read/driver.py Outdated
… ray_pytorch and scope save timing to GCS

- Rename checkpointing/ray_data -> checkpointing/ray_pytorch; the group uses
  Ray core actors with PyTorch DCP, not Ray Data.
- Time only the staged-to-GCS transfer, so the published write throughput
  measures GCS bandwidth rather than state dict gathering and local staging.
- Reserve one CPU per save worker, matching the Ray macrobenchmark.
@Yonghui-Lee
Yonghui-Lee force-pushed the ray-subsystembenchmark-checkpointing-load branch from 225597f to deb0c8b Compare September 14, 2026 05:32
Comment thread gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_pytorch/common.py Outdated
@zhixiangli zhixiangli changed the title feat(subsystembenchmarks): add Ray Data checkpoint load benchmark feat(subsystembenchmarks): add Ray checkpoint load benchmark Sep 17, 2026
# Conflicts:
#	gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_pytorch/common.py
#	gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_pytorch/configs.yaml
#	gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_pytorch/tests/test_configs.py
#	gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_pytorch/write/driver.py
@zhixiangli
zhixiangli merged commit 765dbe6 into fsspec:main Sep 17, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants