Skip to content

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

Open
Yonghui-Lee wants to merge 2 commits into
fsspec:mainfrom
Yonghui-Lee:ray-subsystembenchmark-checkpointing-load
Open

feat(subsystembenchmarks): add Ray Data checkpoint load benchmark#1047
Yonghui-Lee wants to merge 2 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.18%. Comparing base (758f35c) to head (225597f).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1047      +/-   ##
==========================================
+ Coverage   90.10%   90.18%   +0.08%     
==========================================
  Files          16       16              
  Lines        3687     3740      +53     
==========================================
+ Hits         3322     3373      +51     
- Misses        365      367       +2     

☔ 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
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.

1 participant