Repository navigation
feat(subsystembenchmarks): add Ray checkpoint load benchmark - #1047
zhixiangli merged 7 commits into
Conversation
There was a problem hiding this comment.
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.
Codecov Report✅ All modified and coverable lines are covered by tests. 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. 🚀 New features to boost your workflow:
|
27c7ee2 to
225597f
Compare
… 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.
225597f to
deb0c8b
Compare
# Conflicts: # gcsfs/tests/perf/subsystembenchmarks/tests/test_run_groups.py
# 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
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:
RayCheckpointLoadWorkeractor executing distributed checkpoint load operations via PyTorch Distributed Checkpoint (DCP) andray.train.Checkpoint.RayCheckpointSetupWorkeractor 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).RayCheckpointReadDrivermanaging cluster initialization, fixture setup, and benchmark execution.