feat(subsystembenchmarks): add Ray Data checkpoint load benchmark

Author: Yonghui-Lee

gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/common.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import shutil
 import socket
 import tempfile
 
@@ -10,6 +11,8 @@
 import ray
 import torch
 import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict
 
 from gcsfs.tests.perf.subsystembenchmarks.dataloading.driver import assert_fsspec_gcsfs
 
@@ -163,6 +166,15 @@ def parallelize_model(model, params):
     raise ValueError(f"Unknown strategy: {strategy}")
 
 
+def setup_model_and_optimizer(params):
+    """Loads, parallelizes the benchmark model and materializes AdamW optimizer states."""
+    model = load_benchmark_model(params)
+    model = parallelize_model(model, params)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+    materialize_adamw_states(optimizer)
+    return model, optimizer
+
+
 def resolve_storage(prefix: str):
     """Resolves fsspec and PyArrow filesystems for checkpoint storage."""
     fs, base_path = fsspec.core.url_to_fs(prefix)
@@ -201,3 +213,64 @@ def _get_staging_dir(prefix: str, min_free_gb: float = 50.0) -> str:
             except Exception:
                 pass
     return tempfile.mkdtemp(prefix=prefix)
+
+
+def save_checkpoint_step(
+    model,
+    optimizer,
+    params,
+    rank: int,
+    arrow_fs,
+    fs,
+    destination_ckpt: str,
+    staging_prefix: str = "ray-ckpt",
+):
+    """Performs a single distributed or single-node checkpoint save to storage."""
+    is_sharded = params.strategy in (
+        "fsdp_sharded",
+        "model_parallel_sharded",
+    )
+    options = StateDictOptions(
+        full_state_dict=not is_sharded,
+        cpu_offload=not is_sharded,
+    )
+
+    local_dir = None
+    try:
+        model_state, opt_state = get_state_dict(
+            model, optimizer, options=options
+        )
+        app_state = {"model": model_state, "optimizer": opt_state}
+
+        if is_sharded:
+            local_dir = _get_staging_dir(f"{staging_prefix}-rank{rank}-")
+            dcp.save(
+                {"app": app_state},
+                storage_writer=dcp.FileSystemWriter(local_dir),
+            )
+            arrow_fs.create_dir(destination_ckpt)
+            _pyarrow_fs_copy_files(
+                local_dir,
+                destination_ckpt,
+                destination_filesystem=arrow_fs,
+            )
+        else:
+            if rank == 0:
+                local_dir = _get_staging_dir(f"{staging_prefix}-rank0-")
+                ckpt_file = os.path.join(local_dir, "checkpoint.pt")
+                torch.save(app_state, ckpt_file)
+                arrow_fs.create_dir(destination_ckpt)
+                _pyarrow_fs_copy_files(
+                    local_dir,
+                    destination_ckpt,
+                    destination_filesystem=arrow_fs,
+                )
+
+        del app_state, model_state, opt_state
+        dist.barrier()
+        if is_sharded and rank == 0:
+            fs.touch(f"{destination_ckpt.rstrip('/')}/_SUCCESS")
+    finally:
+        if local_dir and os.path.exists(local_dir):
+            shutil.rmtree(local_dir, ignore_errors=True)
+

gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/configs.yaml

@@ -15,3 +15,16 @@ scenarios:
       - {axis: "strategy", strategy: "fsdp_full", world_size: 8}
       - {axis: "strategy", strategy: "model_parallel_full", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8}
       - {axis: "strategy", strategy: "model_parallel_sharded", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8}
+  - name: "checkpoint_read"
+    scenario: "checkpoint_read"
+    variants:
+      # Strategy variants
+      - {axis: "strategy", strategy: "ddp", world_size: 8}
+      - {axis: "strategy", strategy: "fsdp_sharded", world_size: 8}
+      - {axis: "strategy", strategy: "fsdp_full", world_size: 8}
+      - {axis: "strategy", strategy: "model_parallel_full", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8}
+      - {axis: "strategy", strategy: "model_parallel_sharded", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8}
+      # Cross-topology restore variants
+      - {axis: "cross_size", strategy: "fsdp_sharded", setup_world_size: 4, world_size: 8}
+      - {axis: "cross_size", strategy: "fsdp_sharded", setup_world_size: 8, world_size: 4}
+      - {axis: "cross_size", strategy: "model_parallel_sharded", setup_tensor_parallel_size: 4, setup_data_parallel_size: 2, setup_world_size: 8, tensor_parallel_size: 2, data_parallel_size: 2, world_size: 4}

gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/read/__init__.py

@@ -0,0 +1 @@
+"""Ray checkpoint load benchmark."""

gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/read/driver.py

@@ -0,0 +1,237 @@
+"""Driver for Ray checkpoint load benchmark."""
+
+import contextlib
+import dataclasses
+import os
+import time
+import uuid
+
+import ray
+import ray.train
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_state_dict,
+    set_model_state_dict,
+    set_optimizer_state_dict,
+    set_state_dict,
+)
+
+from gcsfs.tests.perf.subsystembenchmarks.checkpointing.driver import (
+    CheckpointDriver,
+    CheckpointResult,
+)
+from gcsfs.tests.perf.subsystembenchmarks.checkpointing.ray_data.common import (
+    ensure_ray_initialized,
+    find_free_port,
+    resolve_storage,
+    save_checkpoint_step,
+    setup_distributed_env,
+    setup_model_and_optimizer,
+)
+
+
+@ray.remote
+class RayCheckpointSetupWorker:
+    """Ray Actor that creates the initial checkpoint on storage for load benchmarks."""
+
+    def __init__(self, rank, world_size, port, prefix, params):
+        self.rank = rank
+        self.world_size = world_size
+        self.port = port
+        self.prefix = prefix
+        self.params = params
+
+    def setup_and_save(self):
+        try:
+            setup_distributed_env(self.rank, self.world_size, self.port)
+            model, optimizer = setup_model_and_optimizer(self.params)
+            fs, arrow_fs, base_path = resolve_storage(self.prefix)
+            destination_ckpt = f"{base_path.rstrip('/')}/model.ckpt"
+            save_checkpoint_step(
+                model,
+                optimizer,
+                self.params,
+                self.rank,
+                arrow_fs,
+                fs,
+                destination_ckpt,
+                staging_prefix="ray-setup-ckpt",
+            )
+        finally:
+            dist.destroy_process_group()
+
+
+@ray.remote
+class RayCheckpointLoadWorker:
+    """Ray Actor executing distributed checkpoint load operations on CPU."""
+
+    def __init__(self, rank, world_size, port, prefix, params):
+        self.rank = rank
+        self.world_size = world_size
+        self.port = port
+        self.prefix = prefix
+        self.params = params
+
+    def setup(self):
+        setup_distributed_env(self.rank, self.world_size, self.port)
+        self.model, self.optimizer = setup_model_and_optimizer(self.params)
+        self.fs, self.arrow_fs, self.base_path = resolve_storage(self.prefix)
+        self.destination_ckpt = f"{self.base_path.rstrip('/')}/model.ckpt"
+
+    def load_rounds(self):
+        try:
+            durations = []
+            is_sharded = self.params.strategy in (
+                "fsdp_sharded",
+                "model_parallel_sharded",
+            )
+            for round_idx in range(self.params.rounds):
+                dist.barrier()
+                t_start = time.perf_counter()
+
+                checkpoint = ray.train.Checkpoint(
+                    path=self.destination_ckpt, filesystem=self.arrow_fs
+                )
+                # Ensure all workers on the host share the same UUID for this round
+                # so Ray's built-in file locking deduplicates the download across workers
+                # and cleans up the shared temporary directory after all workers exit.
+                checkpoint._uuid = uuid.uuid5(
+                    uuid.NAMESPACE_URL, f"{self.destination_ckpt}-round-{round_idx}"
+                )
+
+                directory_context = (
+                    checkpoint.as_directory()
+                    if is_sharded or self.rank == 0
+                    else contextlib.nullcontext(None)
+                )
+                with directory_context as local_dir:
+                    if is_sharded:
+                        options = StateDictOptions(
+                            full_state_dict=False,
+                            cpu_offload=False,
+                        )
+                        model_state, opt_state = get_state_dict(
+                            self.model, self.optimizer, options=options
+                        )
+                        app_state = {"model": model_state, "optimizer": opt_state}
+                        dcp.load(
+                            {"app": app_state},
+                            storage_reader=dcp.FileSystemReader(local_dir),
+                        )
+                        set_state_dict(
+                            self.model,
+                            self.optimizer,
+                            model_state_dict=app_state["model"],
+                            optim_state_dict=app_state["optimizer"],
+                            options=options,
+                        )
+                        del app_state, model_state, opt_state
+                    else:
+                        options = StateDictOptions(
+                            full_state_dict=True,
+                            cpu_offload=False,
+                            broadcast_from_rank0=True,
+                        )
+                        if self.rank == 0:
+                            ckpt_file = os.path.join(local_dir, "checkpoint.pt")
+                            state = torch.load(
+                                ckpt_file, map_location="cpu", weights_only=False
+                            )
+                            model_state = state["model"]
+                            opt_state = state["optimizer"]
+                        else:
+                            model_state = {}
+                            opt_state = {}
+
+                        set_model_state_dict(
+                            self.model,
+                            model_state,
+                            options=options,
+                        )
+                        set_optimizer_state_dict(
+                            self.model,
+                            self.optimizer,
+                            opt_state,
+                            options=options,
+                        )
+                        if self.rank == 0:
+                            del state, model_state, opt_state
+
+                    dist.barrier()
+                    t_end = time.perf_counter()
+                    durations.append((t_start, t_end))
+
+            return durations
+        finally:
+            dist.destroy_process_group()
+
+
+def run_ray_load(prefix, params):
+    """Runs single or distributed checkpoint load benchmark across Ray actor workers."""
+    ensure_ray_initialized()
+    world_size = params.world_size
+    port = find_free_port()
+
+    workers = [
+        RayCheckpointLoadWorker.remote(rank, world_size, port, prefix, params)
+        for rank in range(world_size)
+    ]
+    ray.get([w.setup.remote() for w in workers])
+    results = ray.get([w.load_rounds.remote() for w in workers])
+
+    durations = []
+    for r in range(params.rounds):
+        begins = [results[rank][r][0] for rank in range(world_size)]
+        ends = [results[rank][r][1] for rank in range(world_size)]
+        durations.append(max(ends) - min(begins))
+    return durations
+
+
+class RayCheckpointReadDriver(CheckpointDriver):
+    """Driver for Ray checkpoint load benchmarks."""
+
+    def setup(self, prefix: str, params):
+        """Generates the source checkpoint on storage using setup topology."""
+        ensure_ray_initialized()
+        setup_world_size = (
+            getattr(params, "setup_world_size", None) or params.world_size
+        )
+        setup_tp = (
+            getattr(params, "setup_tensor_parallel_size", None)
+            or params.tensor_parallel_size
+        )
+        setup_dp = (
+            getattr(params, "setup_data_parallel_size", None)
+            or params.data_parallel_size
+        )
+
+        setup_params = dataclasses.replace(
+            params,
+            world_size=setup_world_size,
+            tensor_parallel_size=setup_tp,
+            data_parallel_size=setup_dp,
+        )
+        port = find_free_port()
+        try:
+            workers = [
+                RayCheckpointSetupWorker.remote(
+                    rank, setup_world_size, port, prefix, setup_params
+                )
+                for rank in range(setup_world_size)
+            ]
+            ray.get([w.setup_and_save.remote() for w in workers])
+        finally:
+            if ray.is_initialized():
+                ray.shutdown()
+
+    def run(self, prefix: str, params) -> CheckpointResult:
+        """Executes the checkpoint load benchmark across Ray actor workers."""
+        try:
+            durations = run_ray_load(prefix, params)
+            return CheckpointResult(durations=durations)
+        finally:
+            if ray.is_initialized():
+                ray.shutdown()

gcsfs/tests/perf/subsystembenchmarks/checkpointing/ray_data/read/test_checkpoint.py

@@ -0,0 +1,35 @@
+import os
+
+import pytest
+
+from gcsfs.tests.perf.subsystembenchmarks.checkpointing.ray_data import configs
+from gcsfs.tests.perf.subsystembenchmarks.checkpointing.ray_data.configs import (
+    RayCheckpointConfigurator,
+)
+
+pytest.importorskip("ray")
+pytest.importorskip("torch")
+
+pytestmark = pytest.mark.skipif(
+    not os.environ.get("GCSFS_SUBSYSTEM_BUCKET_PREFIX"),
+    reason="the checkpoint benchmarks create a bucket per case; CI-only (run.py exports the prefix)",
+)
+
+CASES = [
+    c
+    for c in RayCheckpointConfigurator(configs.__file__).generate_cases()
+    if c.scenario == "checkpoint_read"
+]
+
+
+@pytest.mark.timeout(7200)
+@pytest.mark.parametrize("params", CASES, ids=lambda p: p.name)
+def test_checkpoint_load(benchmark, params, monitor):
+    from gcsfs.tests.perf.subsystembenchmarks.checkpointing.checkpoint_case import (
+        run_checkpoint_case,
+    )
+    from gcsfs.tests.perf.subsystembenchmarks.checkpointing.ray_data.read.driver import (
+        RayCheckpointReadDriver,
+    )
+
+    run_checkpoint_case(benchmark, monitor, params, RayCheckpointReadDriver())