chore: import upstream snapshot with attribution

release/train_tests/async_checkpointing_validation_benchmark/compute_aws.yaml

@@ -0,0 +1,27 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - name: train-worker-node
+      instance_type: g4dn.xlarge
+      min_nodes: 2
+      max_nodes: 2
+      market_type: ON_DEMAND
+      labels:
+        ray-subcluster: train
+    - name: validation-worker-node
+      instance_type: g4dn.xlarge
+      min_nodes: 0
+      max_nodes: 2
+      market_type: ON_DEMAND
+      labels:
+        ray-subcluster: validation

release/train_tests/async_checkpointing_validation_benchmark/test_async_checkpointing_validation_benchmark.py

@@ -0,0 +1,723 @@
+from enum import Enum
+import logging
+import os
+import tempfile
+import time
+
+import torch
+import torch.distributed.checkpoint as dist_cp
+import torchmetrics
+from torch.distributed.checkpoint.state_dict import get_state_dict
+from torch.distributed.checkpoint.state_dict_saver import async_save
+from torch.nn import CrossEntropyLoss
+from torch.optim import Adam
+from torchvision import transforms
+from torchvision.models import VisionTransformer
+from torchvision.transforms import ToTensor, Normalize
+import ray
+import ray.train
+import ray.train.torch
+from ray.data import ExecutionOptions
+from ray.train import CheckpointUploadMode, ValidationConfig, ValidationTaskConfig
+from ray._private.test_utils import safe_write_to_results_json
+
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class ValidationType(Enum):
+    # run synchronously with the training loop
+    INLINE = "inline"
+    # run asynchronously with a torch trainer
+    TORCH_TRAINER = "torch_trainer"
+    # run asynchronously with a map batches function
+    MAP_BATCHES = "map_batches"
+
+
+class CheckpointSaveMode(Enum):
+    # save to disk with torch.save
+    TORCH_SAVE = "torch_save"
+    # synchronous save via Torch DCP
+    TORCH_DCP_SYNC = "torch_dcp_sync"
+    # asynchronous save, Ray Train's background thread waits for completion.
+    TORCH_DCP_ASYNC = "torch_dcp_async"
+
+
+MAXIMUM_ALLOWED_ACCURACY_DIFF = 0.2
+MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER = 1.1
+
+# ==== Start dataset and model creation ======
+
+STORAGE_PATH_PREFIX = os.environ.get("ANYSCALE_ARTIFACT_STORAGE", "artifact_storage")
+STORAGE_PATH = f"{STORAGE_PATH_PREFIX}/ray_summit_24_train_demo"
+
+
+def transform_cifar(row: dict):
+    transform = transforms.Compose(
+        [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    )
+    row["image"] = transform(row["image"])
+    return row
+
+
+validation_dataset = ray.data.read_parquet(f"{STORAGE_PATH}/cifar10-parquet/test").map(
+    transform_cifar
+)
+
+
+def create_model():
+    return VisionTransformer(
+        image_size=32,  # CIFAR-10 image size is 32x32
+        patch_size=4,  # Patch size is 4x4
+        num_layers=24,  # Number of transformer layers
+        num_heads=8,  # Number of attention heads
+        hidden_dim=384,  # Hidden size (can be adjusted)
+        mlp_dim=768,  # MLP dimension (can be adjusted)
+        num_classes=10,  # CIFAR-10 has 10 classes
+    )
+
+
+# ==== End dataset and model creation ======
+
+# ==== Start map_batches approach ======
+
+
+class Predictor:
+    def __init__(self, checkpoint):
+        self.model = create_model()
+
+        with checkpoint.as_directory() as checkpoint_dir:
+            model_pt = os.path.join(checkpoint_dir, "model.pt")
+            if os.path.exists(model_pt):
+                self.model.load_state_dict(torch.load(model_pt))
+            else:
+                state_dict = {"model": self.model.state_dict()}
+                dist_cp.load(
+                    state_dict,
+                    storage_reader=dist_cp.FileSystemReader(checkpoint_dir),
+                )
+                self.model.load_state_dict(state_dict["model"])
+
+        self.model.cuda().eval()
+
+    def __call__(self, batch):
+        image = torch.as_tensor(batch["image"], dtype=torch.float32, device="cuda")
+        label = torch.as_tensor(batch["label"], dtype=torch.int8, device="cuda")
+        pred = self.model(image)
+        return {"res": (pred.argmax(1) == label).cpu().numpy()}
+
+
+def validate_with_map_batches(checkpoint):
+    validation_dataset.set_name("async_val_map_batches")
+    validation_dataset.context.execution_options.label_selector = {
+        "ray-subcluster": "validation"
+    }
+    start_time = time.time()
+    eval_res = validation_dataset.map_batches(
+        Predictor,
+        batch_size=128,
+        num_gpus=1,
+        fn_constructor_kwargs={"checkpoint": checkpoint},
+        concurrency=2,
+    )
+    mean = eval_res.mean(["res"])
+    return {
+        "score": mean,
+        "validation_time": time.time() - start_time,
+    }
+
+
+# ==== End map_batches approach ======
+
+# ==== Start TorchTrainer approach ======
+
+
+def eval_only_train_func(config_dict):
+    # Load the checkpoint
+    model = create_model()
+
+    checkpoint = config_dict["checkpoint"]
+    with checkpoint.as_directory() as checkpoint_dir:
+        model_pt = os.path.join(checkpoint_dir, "model.pt")
+        if os.path.exists(model_pt):
+            model.load_state_dict(torch.load(model_pt))
+        else:
+            state_dict = {"model": model.state_dict()}
+            dist_cp.load(
+                state_dict,
+                storage_reader=dist_cp.FileSystemReader(checkpoint_dir),
+            )
+            model.load_state_dict(state_dict["model"])
+
+    model.cuda().eval()
+
+    # Get the data
+    test_data_shard = ray.train.get_dataset_shard("async_val_torch_trainer")
+    test_dataloader = test_data_shard.iter_torch_batches(batch_size=128)
+
+    # Report metrics
+    mean_acc = torchmetrics.Accuracy(task="multiclass", num_classes=10, top_k=1).cuda()
+    with torch.no_grad():
+        for batch in test_dataloader:
+            images, labels = batch["image"], batch["label"]
+            outputs = model(images)
+            mean_acc(outputs.argmax(1), labels)
+    return {"score": mean_acc.compute().item()}
+
+
+def validate_with_torch_trainer(checkpoint, parent_run_name, epoch, batch_idx):
+    start_time = time.time()
+    trainer = ray.train.torch.TorchTrainer(
+        eval_only_train_func,
+        train_loop_config={"checkpoint": checkpoint},
+        scaling_config=ray.train.ScalingConfig(num_workers=2, use_gpu=True),
+        datasets={"async_val_torch_trainer": validation_dataset},
+        run_config=ray.train.RunConfig(
+            name=f"{parent_run_name}-validation_epoch={epoch}_batch_idx={batch_idx}"
+        ),
+        dataset_config=ray.train.DataConfig(
+            execution_options={
+                "async_val_torch_trainer": ExecutionOptions(
+                    label_selector={"ray-subcluster": "validation"}
+                ),
+            },
+        ),
+    )
+    result = trainer.fit()
+    return {
+        "score": result.return_value["score"],
+        "validation_time": time.time() - start_time,
+    }
+
+
+# ==== End TorchTrainer approach ======
+
+
+def validate_and_report(
+    model,
+    epoch,
+    batch_idx,
+    blocked_times,
+    config,
+    loss,
+):
+    validate_within_trainer = config["validate_within_trainer"]
+    num_epochs = config["num_epochs"]
+    checkpoint_upload_mode = config["checkpoint_upload_mode"]
+    validation_type = config["validation_type"]
+    checkpoint_save_mode = config["checkpoint_save_mode"]
+
+    if validate_within_trainer:
+        test_dataloader = ray.train.get_dataset_shard("inline_val").iter_torch_batches(
+            batch_size=128
+        )
+
+    # Validate model within training loop
+    val_elapsed_time = None
+    if validate_within_trainer:
+        val_start_time = time.time()
+        mean_acc = torchmetrics.Accuracy(
+            task="multiclass", num_classes=10, top_k=1
+        ).cuda()
+        model.eval()
+        with torch.no_grad():
+            for batch in test_dataloader:
+                X, y = batch["image"], batch["label"]
+                outputs = model(X)
+                mean_acc(outputs.argmax(1), y)
+        val_elapsed_time = time.time() - val_start_time
+
+    # Report metrics + checkpoint + validate
+    metrics = {"loss": loss.item(), "epoch": epoch}
+    if validate_within_trainer and epoch == num_epochs - 1:
+        metrics["score"] = mean_acc.compute().item()
+
+    # Record how long the upload process takes
+    start_time = time.time()
+
+    # DCP save is a distributed collective so all ranks must call it together.
+    ckpt_ref = None  # Only used by TORCH_DCP_ASYNC
+    iteration_checkpoint_dir = None  # Not used by TORCH_SAVE
+    if checkpoint_save_mode in (
+        CheckpointSaveMode.TORCH_DCP_SYNC,
+        CheckpointSaveMode.TORCH_DCP_ASYNC,
+    ):
+        # For DCP, all workers write shards to the same shared storage path so that
+        # the full checkpoint is available without any upload step.
+        iteration_checkpoint_dir = (
+            ray.train.get_context()
+            .get_storage()
+            .build_checkpoint_path_from_name(f"dcp_epoch_{epoch}_batch_{batch_idx}")
+        )
+
+        storage_writer = dist_cp.FileSystemWriter(iteration_checkpoint_dir)
+        model_dict, _ = get_state_dict(model=model, optimizers=())
+
+        if checkpoint_save_mode == CheckpointSaveMode.TORCH_DCP_SYNC:
+            # Save via Torch DCP
+            dist_cp.save({"model": model_dict}, storage_writer=storage_writer)
+        elif checkpoint_save_mode == CheckpointSaveMode.TORCH_DCP_ASYNC:
+            # Initiate async save; rank 0 will wait via checkpoint_upload_fn
+            ckpt_ref = async_save({"model": model_dict}, storage_writer=storage_writer)
+        else:
+            raise NotImplementedError
+
+    if ray.train.get_context().get_world_rank() == 0:
+        if val_elapsed_time:
+            metrics["validation_time"] = val_elapsed_time
+
+        if validation_type == ValidationType.TORCH_TRAINER:
+            validation = ValidationTaskConfig(
+                fn_kwargs={
+                    "parent_run_name": ray.train.get_context().get_experiment_name(),
+                    "epoch": epoch,
+                    "batch_idx": batch_idx,
+                }
+            )
+        elif validation_type == ValidationType.MAP_BATCHES:
+            validation = True
+        else:
+            validation = False
+
+        if checkpoint_save_mode == CheckpointSaveMode.TORCH_SAVE:
+            # We can't use `tempfile.TemporaryDirectory()` due to CheckpointUploadMode.ASYNC
+            iteration_checkpoint_dir = tempfile.mkdtemp()
+            torch.save(
+                model.module.state_dict(),
+                os.path.join(iteration_checkpoint_dir, "model.pt"),
+            )
+
+            ray.train.report(
+                metrics,
+                checkpoint=ray.train.Checkpoint.from_directory(
+                    iteration_checkpoint_dir
+                ),
+                checkpoint_upload_mode=checkpoint_upload_mode,
+                delete_local_checkpoint_after_upload=True,
+                validation=validation,
+            )
+        elif checkpoint_save_mode == CheckpointSaveMode.TORCH_DCP_SYNC:
+            # Shards are already in shared storage; no upload needed.
+            ray.train.report(
+                metrics,
+                checkpoint=ray.train.Checkpoint.from_directory(
+                    iteration_checkpoint_dir
+                ),
+                checkpoint_upload_mode=CheckpointUploadMode.NO_UPLOAD,
+                validation=validation,
+            )
+        elif checkpoint_save_mode == CheckpointSaveMode.TORCH_DCP_ASYNC:
+            # Shards are written directly to shared storage. The `async_save`
+            # returns a future that will wait until all workers are complete.
+            # Internally it has a barrier before `future.result()` is returned.
+            def wait_async_save(
+                checkpoint, checkpoint_dir_name, upload_complete_ref=ckpt_ref
+            ):
+                upload_complete_ref.result()
+                return checkpoint
+
+            ray.train.report(
+                metrics,
+                checkpoint=ray.train.Checkpoint.from_directory(
+                    iteration_checkpoint_dir
+                ),
+                checkpoint_upload_fn=wait_async_save,
+                checkpoint_dir_name=f"dcp_epoch_{epoch}_batch_{batch_idx}",
+                checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
+                # iteration_checkpoint_dir is already in shared storage so don't delete it.
+                delete_local_checkpoint_after_upload=False,
+                validation=validation,
+            )
+        else:
+            raise NotImplementedError
+
+        blocked_times.append(time.time() - start_time)
+    else:
+        ray.train.report({}, None)
+
+
+def train_func(config):
+    batch_size = 256
+    num_epochs = config["num_epochs"]
+    midpoint_batch = int(config["rows_per_worker"] / batch_size / 2)
+
+    # Prepare model, dataloader, and possibly metrics
+    model = create_model()
+    model = ray.train.torch.prepare_model(model)
+    criterion = CrossEntropyLoss()
+    optimizer = Adam(model.parameters(), lr=0.001)
+    train_data_shard = ray.train.get_dataset_shard("train")
+    train_dataloader = train_data_shard.iter_torch_batches(batch_size=batch_size)
+
+    # Train / eval / report loop
+    blocked_times = []
+    for epoch in range(num_epochs):
+
+        # Train model, then validate/report at midpoint and end of epoch
+        model.train()
+        i = 0
+        for i, batch in enumerate(train_dataloader):
+            images, labels = batch["image"], batch["label"]
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if i == midpoint_batch:
+                validate_and_report(model, epoch, i, blocked_times, config, loss)
+        validate_and_report(model, epoch, i, blocked_times, config, loss)
+
+    # Return train_func metrics
+    return {
+        "report_blocked_times": blocked_times,
+        "train_func_return_time": time.time(),
+    }
+
+
+def run_training_with_validation(
+    checkpoint_upload_mode: CheckpointUploadMode,
+    validation_type: ValidationType,
+    validate_within_trainer: bool,
+    num_epochs: int,
+    train_dataset: ray.data.Dataset,
+    training_rows: int,
+    checkpoint_save_mode: CheckpointSaveMode,
+):
+    # Launch distributed training job.
+    start_time = time.time()
+    scaling_config = ray.train.ScalingConfig(num_workers=2, use_gpu=True)
+
+    if validation_type == ValidationType.INLINE:
+        validation_config = None
+    elif validation_type == ValidationType.TORCH_TRAINER:
+        validation_config = ValidationConfig(validate_with_torch_trainer)
+    elif validation_type == ValidationType.MAP_BATCHES:
+        validation_config = ValidationConfig(validate_with_map_batches)
+    else:
+        raise NotImplementedError
+
+    datasets = {"train": train_dataset}
+    train_loop_config = {
+        "validate_within_trainer": validate_within_trainer,
+        "num_epochs": num_epochs,
+        "checkpoint_upload_mode": checkpoint_upload_mode,
+        "rows_per_worker": training_rows / 2,
+        "validation_type": validation_type,
+        "checkpoint_save_mode": checkpoint_save_mode,
+    }
+    if validate_within_trainer:
+        datasets["inline_val"] = validation_dataset
+        # Sync validation: train workers iterate both datasets, so split each
+        # across the train subcluster and the validation subcluster respectively.
+        dataset_config = ray.train.DataConfig(
+            datasets_to_split=["train", "inline_val"],
+            execution_options={
+                "train": ExecutionOptions(label_selector={"ray-subcluster": "train"}),
+                "inline_val": ExecutionOptions(
+                    label_selector={"ray-subcluster": "validation"}
+                ),
+            },
+        )
+    else:
+        # Async validation: the validation dataset is consumed by a separate
+        # driver (validate_with_torch_trainer / validate_with_map_batches),
+        # which sets its own subcluster label.
+        dataset_config = ray.train.DataConfig(
+            datasets_to_split=["train"],
+            execution_options={
+                "train": ExecutionOptions(label_selector={"ray-subcluster": "train"}),
+            },
+        )
+
+    # async_save additionally requires a CPU process group alongside the GPU one
+    #   because it runs collectives in a background thread.
+    if checkpoint_save_mode == CheckpointSaveMode.TORCH_DCP_ASYNC:
+        torch_config = ray.train.torch.TorchConfig(backend="cpu:gloo,cuda:nccl")
+    else:
+        torch_config = None
+
+    trainer = ray.train.torch.TorchTrainer(
+        train_func,
+        validation_config=validation_config,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        datasets=datasets,
+        torch_config=torch_config,
+        run_config=ray.train.RunConfig(storage_path="/mnt/cluster_storage"),
+        dataset_config=dataset_config,
+    )
+    result = trainer.fit()
+    end_time = time.time()
+
+    # Return metrics
+    # TODO: consider measuring how long it takes to kick off validation,
+    #   how long checkpoint upload takes, distribution of times
+    train_func_metrics = result.return_value
+    return {
+        "e2e_time": end_time - start_time,
+        "final_validation_waiting_time": (
+            end_time - train_func_metrics["train_func_return_time"]
+        ),
+        "total_report_blocked_time": sum(train_func_metrics["report_blocked_times"]),
+        "total_validation_time": sum(
+            m["validation_time"] for c, m in result.best_checkpoints[:-1]
+        ),
+        "final_score": result.best_checkpoints[-2][1]["score"],
+    }
+
+
+def main():
+    train_dataset = ray.data.read_parquet(f"{STORAGE_PATH}/cifar10-parquet/train").map(
+        transform_cifar
+    )
+    training_rows = train_dataset.count()
+    consolidated_metrics = {}
+    num_epochs = 10
+    consolidated_metrics["sync_cp_inline_val_metrics"] = run_training_with_validation(
+        CheckpointUploadMode.SYNC,
+        ValidationType.INLINE,
+        True,
+        num_epochs,
+        train_dataset,
+        training_rows,
+        CheckpointSaveMode.TORCH_SAVE,
+    )
+    consolidated_metrics[
+        "async_cp_torch_trainer_val_metrics"
+    ] = run_training_with_validation(
+        CheckpointUploadMode.ASYNC,
+        ValidationType.TORCH_TRAINER,
+        False,
+        num_epochs,
+        train_dataset,
+        training_rows,
+        CheckpointSaveMode.TORCH_SAVE,
+    )
+    consolidated_metrics[
+        "async_cp_map_batches_val_metrics"
+    ] = run_training_with_validation(
+        CheckpointUploadMode.ASYNC,
+        ValidationType.MAP_BATCHES,
+        False,
+        num_epochs,
+        train_dataset,
+        training_rows,
+        CheckpointSaveMode.TORCH_SAVE,
+    )
+    consolidated_metrics[
+        "sync_dcp_map_batches_val_metrics"
+    ] = run_training_with_validation(
+        CheckpointUploadMode.NO_UPLOAD,
+        ValidationType.MAP_BATCHES,
+        False,
+        num_epochs,
+        train_dataset,
+        training_rows,
+        CheckpointSaveMode.TORCH_DCP_SYNC,
+    )
+    consolidated_metrics[
+        "async_dcp_map_batches_val_metrics"
+    ] = run_training_with_validation(
+        CheckpointUploadMode.ASYNC,
+        ValidationType.MAP_BATCHES,
+        False,
+        num_epochs,
+        train_dataset,
+        training_rows,
+        CheckpointSaveMode.TORCH_DCP_ASYNC,
+    )
+    safe_write_to_results_json(consolidated_metrics)
+
+    # Assert final scores aren't too far off, which would imply an inaccurate comparison
+    # Example: {'async_dcp_map_batches': 0.56, 'sync_cp_inline': 0.57, 'async_cp_map_batches': 0.57, 'async_cp_torch_trainer': 0.58, 'sync_dcp_map_batches': 0.59}
+    sync_final_score = consolidated_metrics["sync_cp_inline_val_metrics"]["final_score"]
+    async_torchtrainer_final_score = consolidated_metrics[
+        "async_cp_torch_trainer_val_metrics"
+    ]["final_score"]
+    async_map_batches_final_score = consolidated_metrics[
+        "async_cp_map_batches_val_metrics"
+    ]["final_score"]
+    sync_dcp_final_score = consolidated_metrics["sync_dcp_map_batches_val_metrics"][
+        "final_score"
+    ]
+    async_dcp_final_score = consolidated_metrics["async_dcp_map_batches_val_metrics"][
+        "final_score"
+    ]
+    logger.info(
+        "Validation metrics order=%s",
+        dict(
+            sorted(
+                (
+                    (k, round(v["final_score"], 2))
+                    for k, v in consolidated_metrics.items()
+                ),
+                key=lambda a: a[1],
+            )
+        ),
+    )
+
+    assert (
+        abs(sync_final_score - async_torchtrainer_final_score)
+        < MAXIMUM_ALLOWED_ACCURACY_DIFF
+    )
+    assert (
+        abs(sync_final_score - async_map_batches_final_score)
+        < MAXIMUM_ALLOWED_ACCURACY_DIFF
+    )
+    assert abs(sync_final_score - sync_dcp_final_score) < MAXIMUM_ALLOWED_ACCURACY_DIFF
+    assert abs(sync_final_score - async_dcp_final_score) < MAXIMUM_ALLOWED_ACCURACY_DIFF
+
+    # Assert async checkpointing/validation e2e time is faster; add multipler to account for training time variance
+    # Example: {'async_cp_map_batches': 1346.26, 'sync_dcp_map_batches': 1350.58, 'async_dcp_map_batches': 1367.41, 'async_cp_torch_trainer': 1390.7, 'sync_cp_inline': 1571.73}
+    sync_e2e_time = consolidated_metrics["sync_cp_inline_val_metrics"]["e2e_time"]
+    async_torchtrainer_e2e_time = consolidated_metrics[
+        "async_cp_torch_trainer_val_metrics"
+    ]["e2e_time"]
+    async_map_batches_e2e_time = consolidated_metrics[
+        "async_cp_map_batches_val_metrics"
+    ]["e2e_time"]
+    sync_dcp_e2e_time = consolidated_metrics["sync_dcp_map_batches_val_metrics"][
+        "e2e_time"
+    ]
+    async_dcp_e2e_time = consolidated_metrics["async_dcp_map_batches_val_metrics"][
+        "e2e_time"
+    ]
+    logger.info(
+        "Total end-to-end time order=%s",
+        dict(
+            sorted(
+                ((k, round(v["e2e_time"], 2)) for k, v in consolidated_metrics.items()),
+                key=lambda a: a[1],
+            )
+        ),
+    )
+
+    assert (
+        async_torchtrainer_e2e_time
+        < sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER
+    ), f"{async_torchtrainer_e2e_time=}, {sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER=} ({sync_e2e_time=})"
+    assert (
+        async_map_batches_e2e_time < sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER
+    ), f"{async_map_batches_e2e_time=}, {sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER=} ({sync_e2e_time=})"
+    assert (
+        sync_dcp_e2e_time < sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER
+    ), f"{sync_dcp_e2e_time=}, {sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER=} ({sync_e2e_time=})"
+    assert (
+        async_dcp_e2e_time < sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER
+    ), f"{async_dcp_e2e_time=}, {sync_e2e_time * MAXIMUM_ALLOWED_E2E_TIME_MULTIPLIER=} ({sync_e2e_time=})"
+
+    # map_batches is faster than TorchTrainer. Note that inline is the fastest but is blocking
+    # Examples: {'async_dcp_map_batches': 1.39, 'async_cp_torch_trainer': 3.19, 'async_cp_map_batches': 3.27, 'sync_dcp_map_batches': 9.02, 'sync_cp_inline': 11.75}
+    sync_validation_time = consolidated_metrics["sync_cp_inline_val_metrics"][
+        "total_validation_time"
+    ]
+
+    sync_report_blocked_time = consolidated_metrics["sync_cp_inline_val_metrics"][
+        "total_report_blocked_time"
+    ]
+    async_torchtrainer_report_blocked_time = consolidated_metrics[
+        "async_cp_torch_trainer_val_metrics"
+    ]["total_report_blocked_time"]
+    async_map_batches_report_blocked_time = consolidated_metrics[
+        "async_cp_map_batches_val_metrics"
+    ]["total_report_blocked_time"]
+    sync_dcp_report_blocked_time = consolidated_metrics[
+        "sync_dcp_map_batches_val_metrics"
+    ]["total_report_blocked_time"]
+    async_dcp_report_blocked_time = consolidated_metrics[
+        "async_dcp_map_batches_val_metrics"
+    ]["total_report_blocked_time"]
+    logger.info(
+        "Total report blocked time order=%s",
+        dict(
+            sorted(
+                (
+                    (k, round(v["total_report_blocked_time"], 2))
+                    for k, v in consolidated_metrics.items()
+                ),
+                key=lambda a: a[1],
+            )
+        ),
+    )
+
+    # Assert report blocking time is less than with async checkpointing.
+    # Example values: 3.66s vs 0.033s vs 0.028s
+    assert async_torchtrainer_report_blocked_time < sync_report_blocked_time
+    assert async_map_batches_report_blocked_time < sync_report_blocked_time
+    assert sync_dcp_report_blocked_time < sync_report_blocked_time
+    assert async_dcp_report_blocked_time < sync_dcp_report_blocked_time
+
+    # Assert sync blocking time (report + validation + final validation) is less than async blocking time (report + final validation)
+    # Example: {'async_dcp_map_batches': 25.52, 'async_cp_map_batches': 26.01, 'sync_cp_inline': 29.76, 'sync_dcp_map_batches': 31.52, 'async_cp_torch_trainer': 37.75}
+    sync_final_validation_blocking_time = consolidated_metrics[
+        "sync_cp_inline_val_metrics"
+    ]["final_validation_waiting_time"]
+    async_torchtrainer_final_validation_blocking_time = consolidated_metrics[
+        "async_cp_torch_trainer_val_metrics"
+    ]["final_validation_waiting_time"]
+    async_map_batches_final_validation_blocking_time = consolidated_metrics[
+        "async_cp_map_batches_val_metrics"
+    ]["final_validation_waiting_time"]
+    sync_dcp_final_validation_blocking_time = consolidated_metrics[
+        "sync_dcp_map_batches_val_metrics"
+    ]["final_validation_waiting_time"]
+    async_dcp_final_validation_blocking_time = consolidated_metrics[
+        "async_dcp_map_batches_val_metrics"
+    ]["final_validation_waiting_time"]
+    sync_blocking_time = (
+        sync_report_blocked_time
+        + sync_validation_time
+        + sync_final_validation_blocking_time
+    )
+    async_torchtrainer_blocking_time = (
+        async_torchtrainer_report_blocked_time
+        + async_torchtrainer_final_validation_blocking_time
+    )
+    async_map_batches_blocking_time = (
+        async_map_batches_report_blocked_time
+        + async_map_batches_final_validation_blocking_time
+    )
+    sync_dcp_blocking_time = (
+        sync_dcp_report_blocked_time + sync_dcp_final_validation_blocking_time
+    )
+    async_dcp_blocking_time = (
+        async_dcp_report_blocked_time + async_dcp_final_validation_blocking_time
+    )
+    logger.info(
+        "Total validation blocking time order=%s",
+        dict(
+            sorted(
+                (
+                    (
+                        k,
+                        round(
+                            (
+                                v["total_validation_time"]
+                                if k == "sync_cp_inline_val_metrics"
+                                else 0
+                            )
+                            + v["total_report_blocked_time"]
+                            + v["final_validation_waiting_time"],
+                            2,
+                        ),
+                    )
+                    for k, v in consolidated_metrics.items()
+                ),
+                key=lambda a: a[1],
+            )
+        ),
+    )
+
+    assert sync_blocking_time > async_torchtrainer_blocking_time
+    assert sync_blocking_time > async_map_batches_blocking_time
+    assert sync_blocking_time > sync_dcp_blocking_time
+    assert sync_blocking_time > async_dcp_blocking_time
+
+    # TODO: consider correctness checks like validating that local checkpoints get deleted
+    # TODO: track validation startup metrics: schedule validation task, autoscale nodes,
+    #   start TorchTrainer/map_batches, load checkpoint.
+
+
+if __name__ == "__main__":
+    main()

release/train_tests/benchmark/benchmark_factory.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+
+from config import BenchmarkConfig
+from dataloader_factory import BaseDataLoaderFactory
+
+
+class BenchmarkFactory(ABC):
+    def __init__(self, benchmark_config: BenchmarkConfig):
+        self.benchmark_config = benchmark_config
+        self.dataloader_factory = self.get_dataloader_factory()
+        self.dataset_creation_time = 0
+
+    @abstractmethod
+    def get_dataloader_factory(self) -> BaseDataLoaderFactory:
+        """Create the appropriate dataloader factory for this benchmark."""
+        raise NotImplementedError
+
+    # TODO: These can probably be moved to the train loop runner,
+    # since xgboost does not require instantiating the model
+    # and loss function in this way.
+    @abstractmethod
+    def get_model(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_loss_fn(self):
+        raise NotImplementedError
+
+    def get_train_dataloader(self):
+        return self.dataloader_factory.get_train_dataloader()
+
+    def get_val_dataloader(self):
+        return self.dataloader_factory.get_val_dataloader()
+
+    def get_dataloader_metrics(self):
+        return self.dataloader_factory.get_metrics()

release/train_tests/benchmark/compute_configs/compute_gpu_1x1_aws.yaml

@@ -0,0 +1,7 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+flags:
+    enable_ray_container_cgroup_isolation: false
+
+head_node:
+    instance_type: g4dn.8xlarge

release/train_tests/benchmark/compute_configs/compute_gpu_1x4_aws.yaml

@@ -0,0 +1,7 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+flags:
+    enable_ray_container_cgroup_isolation: false
+
+head_node:
+    instance_type: g4dn.12xlarge

release/train_tests/benchmark/compute_configs/compute_gpu_4x4_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/compute_gpu_4x4_cpu_4_aws.yaml

@@ -0,0 +1,22 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - name: worker_node_gpu
+      instance_type: g4dn.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND
+      # New SDK's `resources` is a full override (not a merge), so we must
+      # include `GPU: 4` explicitly to preserve the g4dn.12xlarge's natural
+      # GPU count. CPU: 0 keeps this group reserved for GPU workloads only.
+      resources:
+        CPU: 0
+        GPU: 4
+    - name: worker_node_cpu
+      instance_type: m5.4xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/compute_gpu_4x4_gce.yaml

@@ -0,0 +1,12 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-b
+
+head_node:
+    instance_type: n1-standard-16
+
+worker_nodes:
+    - instance_type: n1-standard-64-nvidia-tesla-t4-4
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/compute_gpu_8x4_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 8
+      max_nodes: 8
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/compute_gpu_8xA100_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - instance_type: p4d.24xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/compute_mock_gpu_4x4_aws.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - instance_type: m5.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND
+      # New SDK's `resources` is a full override (not a merge), so we must
+      # include `CPU: 48` explicitly to preserve the m5.12xlarge's natural
+      # vCPU count alongside the custom MOCK_GPU resource.
+      resources:
+        CPU: 48
+        MOCK_GPU: 4

release/train_tests/benchmark/compute_configs/heterogenous_autoscaling_gpu_4x4_aws.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - name: gpu-worker-node
+      instance_type: g4dn.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND
+    - name: cpu-worker-node
+      instance_type: r6i.4xlarge
+      min_nodes: 0
+      max_nodes: 4
+      market_type: ON_DEMAND

release/train_tests/benchmark/compute_configs/heterogenous_fixed_size_gpu_4x4_aws.yaml

@@ -0,0 +1,16 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - name: gpu-worker-node
+      instance_type: g4dn.12xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND
+    - name: cpu-worker-node
+      instance_type: r6i.4xlarge
+      min_nodes: 4
+      max_nodes: 4
+      market_type: ON_DEMAND

release/train_tests/benchmark/config.py

@@ -0,0 +1,181 @@
+import argparse
+import enum
+from typing import ClassVar
+
+from pydantic import BaseModel, Field
+
+
+class DataloaderType(enum.Enum):
+    RAY_DATA = "ray_data"
+    MOCK = "mock"
+    TORCH = "torch"
+
+
+class DataLoaderConfig(BaseModel):
+    train_batch_size: int = 32
+    limit_training_rows: int = 1000000  # Use -1 for unlimited
+
+    validation_batch_size: int = 256
+    limit_validation_rows: int = 50000  # Use -1 for unlimited
+
+
+class TaskConfig(BaseModel):
+    TASK_NAME: ClassVar[str] = "base"
+
+
+class ImageClassificationConfig(TaskConfig):
+    TASK_NAME: ClassVar[str] = "image_classification"
+
+    class ImageFormat(enum.Enum):
+        JPEG = "jpeg"
+        PARQUET = "parquet"
+        S3_URL = "s3_url"
+
+    image_classification_local_dataset: bool = False
+    image_classification_data_format: ImageFormat = ImageFormat.PARQUET
+    # When True and data_format=PARQUET, read from the larger
+    # IMAGENET_PARQUET_SPLIT_1T_S3_ROOT dataset instead of the default
+    # parquet_split root. Used by the slow-consumer benchmarks to sustain
+    # backpressure.
+    image_classification_use_1t_dataset: bool = False
+
+
+class RecsysConfig(TaskConfig):
+    TASK_NAME: ClassVar[str] = "recsys"
+
+
+class RayDataConfig(DataLoaderConfig):
+    # NOTE: Optional[int] doesn't play well with argparse.
+    local_buffer_shuffle_size: int = -1
+    enable_operator_progress_bars: bool = True
+    ray_data_prefetch_batches: int = 4
+    ray_data_override_num_blocks: int = -1
+    locality_with_output: bool = False
+    actor_locality_enabled: bool = True
+    enable_shard_locality: bool = True
+    preserve_order: bool = False
+    ray_data_pin_memory: bool = False
+
+
+class TorchConfig(DataLoaderConfig):
+    num_torch_workers: int = 8
+    torch_dataloader_timeout_seconds: int = 300
+    torch_pin_memory: bool = True
+    torch_non_blocking: bool = True
+    torch_prefetch_factor: int = -1
+
+
+class BenchmarkConfig(BaseModel):
+    # ScalingConfig
+    num_workers: int = 1
+    # Elastic scaling range. If both are set > 0, use (min_workers, max_workers).
+    min_workers: int = 0
+    max_workers: int = 0
+    # Run CPU training where train workers request a `MOCK_GPU` resource instead.
+    mock_gpu: bool = False
+
+    # FailureConfig
+    max_failures: int = 0
+
+    task: str = "image_classification"
+    task_config: TaskConfig = Field(
+        default_factory=lambda: TaskConfig(),
+    )
+
+    # Data
+    dataloader_type: DataloaderType = DataloaderType.RAY_DATA
+    dataloader_config: DataLoaderConfig = Field(
+        default_factory=lambda: DataLoaderConfig(),
+    )
+
+    # Training
+    num_epochs: int = 1
+    skip_train_step: bool = False
+    # Simulates a slow training consumer by sleeping for this many seconds
+    # after each training step. Used to benchmark dataloader behavior under
+    # consumer back-pressure. 0 disables the sleep.
+    train_step_sleep_s: float = 0.0
+    # Maximum number of training batches per worker per epoch. When reached,
+    # the epoch ends early regardless of dataset size. -1 disables the cap.
+    # Used with slow-consumer benchmarks to bound wall-clock without
+    # truncating the data source.
+    max_train_batches: int = -1
+
+    # Checkpointing
+    checkpoint_every_n_steps: int = -1
+
+    # Validation
+    validate_every_n_steps: int = -1
+    skip_validation_step: bool = False
+    skip_validation_at_epoch_end: bool = False
+
+    # Logging
+    log_metrics_every_n_steps: int = 512
+
+
+def _is_pydantic_model(field_type) -> bool:
+    """Check if a type is a subclass of Pydantic's BaseModel."""
+    return isinstance(field_type, type) and issubclass(field_type, BaseModel)
+
+
+def _str_to_bool(value: str) -> bool:
+    """Convert a string to a boolean value."""
+    if value.lower() == "true":
+        return True
+    elif value.lower() == "false":
+        return False
+    raise argparse.ArgumentTypeError(f"'True' or 'False' expected, got '{value}'")
+
+
+def _add_field_to_parser(parser: argparse.ArgumentParser, field: str, field_info):
+    field_type = field_info.annotation
+    if field_type is bool:
+        parser.add_argument(f"--{field}", type=_str_to_bool, default=field_info.default)
+    else:
+        parser.add_argument(f"--{field}", type=field_type, default=field_info.default)
+
+
+def cli_to_config(benchmark_config_cls=BenchmarkConfig) -> BenchmarkConfig:
+    parser = argparse.ArgumentParser()
+
+    nested_fields = []
+    for field, field_info in benchmark_config_cls.model_fields.items():
+        # Skip nested configs for now
+        if _is_pydantic_model(field_info.annotation):
+            nested_fields.append(field)
+            continue
+
+        _add_field_to_parser(parser, field, field_info)
+
+    top_level_args, _ = parser.parse_known_args()
+
+    # Handle nested configs that depend on top-level args
+    nested_configs = {}
+    for nested_field in nested_fields:
+        nested_parser = argparse.ArgumentParser()
+        nested_config_cls = benchmark_config_cls.model_fields[nested_field].annotation
+
+        if nested_config_cls == DataLoaderConfig:
+            if top_level_args.dataloader_type == DataloaderType.RAY_DATA:
+                nested_config_cls = RayDataConfig
+            elif top_level_args.dataloader_type == DataloaderType.TORCH:
+                nested_config_cls = TorchConfig
+
+        if nested_config_cls == TaskConfig:
+            if top_level_args.task == ImageClassificationConfig.TASK_NAME:
+                nested_config_cls = ImageClassificationConfig
+            elif top_level_args.task == RecsysConfig.TASK_NAME:
+                nested_config_cls = RecsysConfig
+
+        for field, field_info in nested_config_cls.model_fields.items():
+            _add_field_to_parser(nested_parser, field, field_info)
+
+        args, _ = nested_parser.parse_known_args()
+        nested_configs[nested_field] = nested_config_cls(**vars(args))
+
+    return benchmark_config_cls(**vars(top_level_args), **nested_configs)
+
+
+if __name__ == "__main__":
+    config = cli_to_config()
+    print(config)

release/train_tests/benchmark/constants.py

@@ -0,0 +1,7 @@
+"""Constants shared across the benchmarks."""
+
+
+class DatasetKey:
+    TRAIN = "train"
+    VALID = "val"
+    TEST = "test"

release/train_tests/benchmark/dataloader_factory.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Iterator, Tuple
+import logging
+
+import torch
+
+from config import BenchmarkConfig, DataLoaderConfig
+
+logger = logging.getLogger(__name__)
+
+
+class BaseDataLoaderFactory(ABC):
+    """Base class for creating and managing dataloaders."""
+
+    def __init__(self, benchmark_config: BenchmarkConfig):
+        self.benchmark_config = benchmark_config
+
+    def get_dataloader_config(self) -> DataLoaderConfig:
+        return self.benchmark_config.dataloader_config
+
+    @abstractmethod
+    def get_train_dataloader(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        pass
+
+    @abstractmethod
+    def get_val_dataloader(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        pass
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Return metrics about dataloader performance."""
+        return {}

release/train_tests/benchmark/elastic_training/__init__.py

@@ -0,0 +1 @@
+"""Ray Train benchmark elastic test helpers."""

release/train_tests/benchmark/elastic_training/elastic_e2e.py

@@ -0,0 +1,105 @@
+import json
+import os
+import pprint
+import time
+
+import ray
+import ray.train
+from ray._private.test_utils import safe_write_to_results_json
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.util import date_str
+
+from config import cli_to_config
+from image_classification.factory import ImageClassificationFactory
+from elastic_training.resource_schedule import (
+    MockResourceAvailabilityUpdater,
+    ResourceAvailabilityEvent,
+    generate_schedule,
+)
+from train_benchmark import (
+    METRICS_OUTPUT_PATH,
+    get_datasets_and_data_config,
+    train_fn_per_worker,
+)
+
+
+def main():
+    config = cli_to_config()
+    print("\nBenchmark config:\n" + pprint.pformat(config.__dict__, indent=2))
+
+    factory = ImageClassificationFactory(config)
+
+    # Resolve num_workers based on min_workers and max_workers.
+    if config.min_workers and config.max_workers:
+        num_workers = (config.min_workers, config.max_workers)
+    else:
+        num_workers = config.num_workers
+
+    updater_actor = ray.remote(num_cpus=0)(MockResourceAvailabilityUpdater).remote(
+        resource_key="GPU"
+    )
+    ray.get(updater_actor.__ray_ready__.remote())
+
+    interval_s = 60 * 5
+    schedule = generate_schedule(
+        resource_availability_options=[4, 8, 16, 32],
+        duration_s=60 * 60,
+        interval_s=interval_s,
+        seed=777777,
+    )
+    # Make sure the run can finish at the end of the schedule.
+    schedule.append(
+        ResourceAvailabilityEvent(
+            time_s=schedule[-1].time_s + interval_s, resource_units=32
+        )
+    )
+    execute_schedule_fut = updater_actor.execute_schedule.remote(schedule)
+
+    datasets, data_config = get_datasets_and_data_config(factory)
+
+    start_time = time.perf_counter()
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn_per_worker,
+        train_loop_config={"factory": factory},
+        scaling_config=ray.train.ScalingConfig(num_workers=num_workers, use_gpu=True),
+        run_config=ray.train.RunConfig(
+            storage_path=f"{os.environ['ANYSCALE_ARTIFACT_STORAGE']}/train_benchmark/",
+            name=f"{config.task}-{date_str(include_ms=True)}",
+            failure_config=ray.train.FailureConfig(max_failures=len(schedule)),
+        ),
+        datasets=datasets,
+        dataset_config=data_config,
+    )
+    trainer.fit()
+    end_time = time.perf_counter()
+    e2e_time = end_time - start_time
+
+    with open(METRICS_OUTPUT_PATH, "r") as f:
+        metrics = json.load(f)
+
+    # Includes recovery time across resource updates.
+    total_rows_processed = metrics["train/rows_processed-total"]
+    metrics["e2e_throughput"] = total_rows_processed / e2e_time
+    metrics["e2e_time"] = e2e_time
+
+    safe_write_to_results_json(metrics)
+
+    final_metrics_str = (
+        f"\nTotal training time: {e2e_time} seconds\n"
+        + "Final metrics:\n"
+        + "-" * 80
+        + "\n"
+        + pprint.pformat(metrics)
+        + "\n"
+        + "-" * 80
+    )
+    print(final_metrics_str)
+
+    ray.get(execute_schedule_fut)
+    ray.get(updater_actor.shutdown.remote())
+
+
+if __name__ == "__main__":
+    # Workers need to access the working directory module.
+    ray.init(runtime_env={"working_dir": os.path.dirname(os.path.dirname(__file__))})
+    main()

release/train_tests/benchmark/elastic_training/resource_schedule.py

@@ -0,0 +1,208 @@
+from dataclasses import dataclass
+from enum import Enum
+import logging
+import random
+import time
+from typing import List
+import uuid
+
+import psutil
+import ray
+from ray.data._internal.cluster_autoscaler import (
+    ResourceRequestPriority,
+    get_or_create_autoscaling_coordinator,
+)
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+from ray.util.state import list_actors
+
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote(num_cpus=0)
+def kill_process(pid):
+    proc = psutil.Process(pid)
+    proc.kill()
+
+
+class MockResourceRequestPriority(Enum):
+    OVERRIDE = ResourceRequestPriority.HIGH.value + 1
+
+
+@dataclass
+class ResourceAvailabilityEvent:
+    time_s: int
+    resource_units: int
+
+
+class ResourceAvailabilityUpdater:
+    def __init__(self, starting_resource_units: int = 0, resource_key: str = "GPU"):
+        self._starting_resource_units = starting_resource_units
+        self._resource_key = resource_key
+
+    def execute_schedule(self, schedule: List[ResourceAvailabilityEvent]):
+        pass
+
+    def shutdown(self):
+        pass
+
+
+class MockResourceAvailabilityUpdater(ResourceAvailabilityUpdater):
+    def __init__(self, starting_resource_units: int = 0, resource_key: str = "GPU"):
+        super().__init__(starting_resource_units, resource_key)
+
+        self._coord = get_or_create_autoscaling_coordinator()
+        self._clear_all_requests()
+
+        logging.basicConfig(level=logging.INFO)
+        logger.info(
+            "Initializing resource availability: '%s': %s",
+            resource_key,
+            starting_resource_units,
+        )
+
+        self._total_resource_units = int(ray.cluster_resources()[resource_key])
+        self._dummy_requester_ids = [
+            self._get_requester_id()
+            for _ in range(self._total_resource_units - starting_resource_units)
+        ]
+        self._request(self._dummy_requester_ids)
+
+    def _request(self, requester_ids):
+        futs = []
+        for requester_id in requester_ids:
+            fut = self._coord.request_resources.remote(
+                requester_id=requester_id,
+                resources=[{self._resource_key: 1.0}],
+                expire_after_s=10000,
+                priority=MockResourceRequestPriority.OVERRIDE,
+            )
+            futs.append(fut)
+        ray.get(futs)
+
+    def _cancel(self, requester_ids):
+        futs = []
+        for requester_id in requester_ids:
+            fut = self._coord.cancel_request.remote(requester_id=requester_id)
+            futs.append(fut)
+        ray.get(futs)
+
+    def _clear_all_requests(self):
+        def clear_all_requests(coord_self):
+            coord_self._ongoing_reqs = {}
+
+        ray.get(self._coord.__ray_call__.remote(clear_all_requests))
+
+    def _get_requester_id(self):
+        return f"dummy_{uuid.uuid4().hex[:6]}"
+
+    def _kill_random_train_worker(self):
+        actors = list_actors(
+            filters=[("class_name", "=", "RayTrainWorker"), ("state", "=", "ALIVE")]
+        )
+        if not actors:
+            return
+
+        actor_to_kill = random.choice(actors)
+        logger.info("Killing random train worker: %s", actor_to_kill)
+        strategy = NodeAffinitySchedulingStrategy(
+            node_id=actor_to_kill.node_id, soft=False
+        )
+        ray.get(
+            kill_process.options(scheduling_strategy=strategy).remote(actor_to_kill.pid)
+        )
+
+    def execute_schedule(self, schedule: List[ResourceAvailabilityEvent]):
+        schedule_str = " -> ".join(
+            f"({event.time_s:.0f}s, {self._resource_key}: {event.resource_units})"
+            for event in schedule
+        )
+        logger.info("Executing availability schedule: %s", schedule_str)
+
+        start_time = time.time()
+
+        for event in schedule:
+            curr_time_s = time.time() - start_time
+            time.sleep(max(0, event.time_s - curr_time_s))
+
+            logger.info("Executing scheduled event: %s", event)
+
+            curr_withheld = len(self._dummy_requester_ids)
+            curr_available = self._total_resource_units - curr_withheld
+            if curr_available == event.resource_units:
+                logger.info(
+                    "No change in availability: %s -> %s",
+                    curr_available,
+                    event.resource_units,
+                )
+                continue
+            if curr_available > event.resource_units:
+                num_units_to_withhold = curr_available - event.resource_units
+                new_requesters = [
+                    self._get_requester_id() for _ in range(num_units_to_withhold)
+                ]
+                logger.info(
+                    "Reducing availability from %s to %s",
+                    curr_available,
+                    event.resource_units,
+                )
+
+                # If reducing resources, kill a worker process to trigger recovery.
+                self._kill_random_train_worker()
+                self._request(new_requesters)
+                self._dummy_requester_ids += new_requesters
+            else:
+                num_to_cancel = event.resource_units - curr_available
+                self._dummy_requester_ids, ids_to_cancel = (
+                    self._dummy_requester_ids[num_to_cancel:],
+                    self._dummy_requester_ids[:num_to_cancel],
+                )
+                logger.info(
+                    "Increasing availability from %s to %s",
+                    curr_available,
+                    event.resource_units,
+                )
+                self._cancel(ids_to_cancel)
+
+    def shutdown(self):
+        self._cancel(self._dummy_requester_ids)
+
+
+def generate_schedule(
+    resource_availability_options: list,
+    duration_s: int = 60,
+    interval_s: int = 5,
+    seed: int = 42,
+) -> List[ResourceAvailabilityEvent]:
+    random.seed(seed)
+    num_updates = duration_s // interval_s
+
+    curr_idx = random.choice(range(len(resource_availability_options)))
+    schedule = [
+        ResourceAvailabilityEvent(
+            time_s=0, resource_units=resource_availability_options[curr_idx]
+        )
+    ]
+
+    for i in range(1, num_updates):
+        # Weights are chosen to bias schedules towards the max workers.
+        weights = None
+        if curr_idx == 0:
+            choices = [0, 1]
+        elif curr_idx == len(resource_availability_options) - 1:
+            choices = [-1, 0]
+            weights = [20, 80]
+        else:
+            choices = [-1, 0, 1]
+            weights = [25, 25, 50]
+
+        random_update = random.choices(choices, weights=weights)[0]
+        curr_idx += random_update
+        schedule.append(
+            ResourceAvailabilityEvent(
+                time_s=i * interval_s,
+                resource_units=resource_availability_options[curr_idx],
+            )
+        )
+
+    return schedule

release/train_tests/benchmark/image_classification/factory.py

@@ -0,0 +1,394 @@
+# Standard library imports
+import logging
+import time
+from typing import Dict, Tuple, Iterator, Generator, Optional, Union
+
+# Third-party imports
+import torch
+import torchvision
+import pyarrow
+import ray
+import ray.train
+from ray.data.collate_fn import ArrowBatchCollateFn, CollateFn
+from concurrent.futures import ThreadPoolExecutor
+from ray.data.dataset import TorchDeviceType
+
+# Local imports
+from benchmark_factory import BenchmarkFactory
+from config import BenchmarkConfig, DataloaderType, ImageClassificationConfig
+from dataloader_factory import BaseDataLoaderFactory
+from torch_dataloader_factory import TorchDataLoaderFactory
+from ray_dataloader_factory import RayDataLoaderFactory
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+def mock_dataloader(
+    num_batches: int = 64, batch_size: int = 32
+) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
+    """Generate mock image and label tensors for testing.
+
+    Args:
+        num_batches: Number of batches to generate
+        batch_size: Number of samples per batch
+
+    Yields:
+        Tuple of (image_tensor, label_tensor) for each batch
+    """
+    device = ray.train.torch.get_device()
+
+    images = torch.randn(batch_size, 3, 224, 224).to(device)
+    labels = torch.randint(0, 1000, (batch_size,)).to(device)
+
+    for _ in range(num_batches):
+        yield images, labels
+
+
+class ImageClassificationTorchDataLoaderFactory(TorchDataLoaderFactory):
+    """Factory for creating PyTorch DataLoaders for image classification tasks.
+
+    Features:
+    - Distributed file reading with round-robin worker distribution
+    - Device transfer and error handling for data batches
+    - Configurable row limits per worker for controlled processing
+    - Performance monitoring and logging
+    """
+
+    def __init__(self, benchmark_config: BenchmarkConfig):
+        super().__init__(benchmark_config)
+
+    def _calculate_rows_per_worker(
+        self, total_rows: int, num_workers: int
+    ) -> Optional[int]:
+        """Calculate rows per worker for balanced data distribution.
+
+        Args:
+            total_rows: Total rows to process across all workers (-1 for unlimited)
+            num_workers: Total workers (Ray workers × Torch workers)
+
+        Returns:
+            Rows per worker or None if no limit. Each worker gets at least 1 row.
+        """
+        if total_rows < 0:
+            return None
+
+        if num_workers == 0:
+            return total_rows
+
+        return max(1, total_rows // num_workers)
+
+    def _get_worker_row_limits(self) -> Tuple[Optional[int], Optional[int]]:
+        """Calculate row limits per worker for training and validation.
+
+        Returns:
+            Tuple of (training_rows_per_worker, validation_rows_per_worker)
+        """
+        dataloader_config = self.get_dataloader_config()
+        num_workers = max(1, dataloader_config.num_torch_workers)
+        total_workers = self.benchmark_config.num_workers * num_workers
+
+        limit_training_rows_per_worker = self._calculate_rows_per_worker(
+            self.get_dataloader_config().limit_training_rows, total_workers
+        )
+
+        limit_validation_rows_per_worker = self._calculate_rows_per_worker(
+            self.get_dataloader_config().limit_validation_rows, total_workers
+        )
+
+        return limit_training_rows_per_worker, limit_validation_rows_per_worker
+
+    def create_batch_iterator(
+        self, dataloader: torch.utils.data.DataLoader, device: torch.device
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Create iterator with device transfer and error handling.
+
+        Args:
+            dataloader: PyTorch DataLoader to iterate over
+            device: Target device for tensor transfer
+
+        Returns:
+            Iterator yielding (image_tensor, label_tensor) on target device
+        """
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(f"Worker {worker_rank}: Starting batch iteration")
+
+        try:
+            last_batch_time = time.time()
+            for batch_idx, batch in enumerate(dataloader):
+                try:
+                    # Monitor batch processing delays
+                    current_time = time.time()
+                    time_since_last_batch = current_time - last_batch_time
+                    if time_since_last_batch > 10:
+                        logger.warning(
+                            f"Worker {worker_rank}: Long delay ({time_since_last_batch:.2f}s) "
+                            f"between batches {batch_idx-1} and {batch_idx}"
+                        )
+
+                    # Process and transfer batch to device
+                    images, labels = batch
+                    logger.info(
+                        f"Worker {worker_rank}: Processing batch {batch_idx} (shape: {images.shape}, "
+                        f"time since last: {time_since_last_batch:.2f}s)"
+                    )
+
+                    # Transfer tensors to target device
+                    transfer_start = time.time()
+                    dataloader_config = self.get_dataloader_config()
+                    images = images.to(
+                        device, non_blocking=dataloader_config.torch_non_blocking
+                    )
+                    labels = labels.to(
+                        device, non_blocking=dataloader_config.torch_non_blocking
+                    )
+                    transfer_time = time.time() - transfer_start
+
+                    # Monitor device transfer performance
+                    if transfer_time > 5:
+                        logger.warning(
+                            f"Worker {worker_rank}: Slow device transfer ({transfer_time:.2f}s) "
+                            f"for batch {batch_idx}"
+                        )
+
+                    logger.info(
+                        f"Worker {worker_rank}: Completed device transfer for batch {batch_idx} in "
+                        f"{transfer_time:.2f}s"
+                    )
+
+                    last_batch_time = time.time()
+                    yield images, labels
+
+                except Exception as e:
+                    logger.error(
+                        f"Worker {worker_rank}: Error processing batch {batch_idx}: {str(e)}",
+                        exc_info=True,
+                    )
+                    raise
+
+        except Exception as e:
+            logger.error(
+                f"Worker {worker_rank}: Error in batch iterator: {str(e)}",
+                exc_info=True,
+            )
+            raise
+
+
+class CustomArrowCollateFn(ArrowBatchCollateFn):
+    """Custom collate function for converting Arrow batches to PyTorch tensors."""
+
+    _DEFAULT_NUM_WORKERS = 4
+
+    def __init__(
+        self,
+        dtypes: Optional[Union["torch.dtype", Dict[str, "torch.dtype"]]] = None,
+        device: Optional["TorchDeviceType"] = None,
+        pin_memory: bool = False,
+        num_workers: int = _DEFAULT_NUM_WORKERS,
+    ):
+        """Initialize the collate function.
+
+        Args:
+            dtypes: Optional torch dtype(s) for the tensors
+            device: Optional device to place tensors on
+            pin_memory: Whether to pin the memory of the created tensors
+            num_workers: Number of worker threads for parallel tensor conversion
+                Defaults to `_DEFAULT_NUM_WORKERS`.
+        """
+        import torch
+
+        self.dtypes = dtypes
+        if isinstance(device, (str, int)):
+            self.device = torch.device(device)
+        else:
+            self.device = device
+        self.pin_memory = pin_memory
+        self.num_workers = num_workers
+        self._threadpool: Optional[ThreadPoolExecutor] = None
+
+    def __del__(self):
+        """Clean up threadpool on destruction."""
+        if getattr(self, "_threadpool", None):
+            self._threadpool.shutdown(wait=False)
+
+    def __call__(self, batch: "pyarrow.Table") -> Tuple[torch.Tensor, torch.Tensor]:
+        """Convert an Arrow batch to PyTorch tensors.
+
+        Args:
+            batch: PyArrow Table to convert
+
+        Returns:
+            Tuple of (image_tensor, label_tensor)
+        """
+        from ray.data.util.torch_utils import (
+            arrow_batch_to_tensors,
+        )
+
+        if self.num_workers > 0 and self._threadpool is None:
+            self._threadpool = ThreadPoolExecutor(max_workers=self.num_workers)
+
+        # For GPU transfer, we can skip the combining chunked arrays. This is because
+        # we can convert the chunked arrays to corresponding numpy format and then to
+        # Tensors and transfer the corresponding list of Tensors to GPU directly.
+        # However, for CPU transfer, we need to combine the chunked arrays first
+        # before converting to numpy format and then to Tensors.
+        combine_chunks = self.device is not None and self.device.type == "cpu"
+        tensors = arrow_batch_to_tensors(
+            batch,
+            dtypes=self.dtypes,
+            combine_chunks=combine_chunks,
+            pin_memory=self.pin_memory,
+            threadpool=self._threadpool,
+        )
+        return tensors["image"], tensors["label"]
+
+
+class ImageClassificationRayDataLoaderFactory(RayDataLoaderFactory):
+    """Factory for creating Ray DataLoader for image classification tasks."""
+
+    def __init__(self, benchmark_config: BenchmarkConfig):
+        super().__init__(benchmark_config)
+
+    def _get_collate_fn(self) -> Optional[CollateFn]:
+        return CustomArrowCollateFn(
+            device=ray.train.torch.get_device(),
+            pin_memory=self.get_dataloader_config().ray_data_pin_memory,
+        )
+
+
+class ImageClassificationMockDataLoaderFactory(BaseDataLoaderFactory):
+    """Factory for creating mock dataloaders for testing.
+
+    Provides mock implementations of training and validation dataloaders
+    that generate random image and label tensors.
+    """
+
+    def get_train_dataloader(
+        self,
+    ) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
+        """Get mock training dataloader.
+
+        Returns:
+            Generator yielding (image_tensor, label_tensor) batches
+        """
+        dataloader_config = self.get_dataloader_config()
+        return mock_dataloader(
+            num_batches=1024, batch_size=dataloader_config.train_batch_size
+        )
+
+    def get_val_dataloader(
+        self,
+    ) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
+        """Get mock validation dataloader.
+
+        Returns:
+            Generator yielding (image_tensor, label_tensor) batches
+        """
+        dataloader_config = self.get_dataloader_config()
+        return mock_dataloader(
+            num_batches=512, batch_size=dataloader_config.validation_batch_size
+        )
+
+
+def get_imagenet_data_dirs(task_config: ImageClassificationConfig) -> Dict[str, str]:
+    """Returns a dict with the root imagenet dataset directories for train/val/test,
+    corresponding to the data format and local/s3 dataset location."""
+    from image_classification.imagenet import IMAGENET_LOCALFS_SPLIT_DIRS
+    from image_classification.jpeg.imagenet import (
+        IMAGENET_JPEG_SPLIT_S3_DIRS,
+    )
+    from image_classification.parquet.imagenet import (
+        IMAGENET_PARQUET_SPLIT_S3_DIRS,
+        IMAGENET_PARQUET_SPLIT_1T_S3_DIRS,
+    )
+    from image_classification.s3_url.imagenet import (
+        IMAGENET_S3_URL_SPLIT_DIRS,
+    )
+
+    data_format = task_config.image_classification_data_format
+
+    if task_config.image_classification_local_dataset:
+        return IMAGENET_LOCALFS_SPLIT_DIRS
+
+    if data_format == ImageClassificationConfig.ImageFormat.JPEG:
+        return IMAGENET_JPEG_SPLIT_S3_DIRS
+    elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
+        if task_config.image_classification_use_1t_dataset:
+            return IMAGENET_PARQUET_SPLIT_1T_S3_DIRS
+        return IMAGENET_PARQUET_SPLIT_S3_DIRS
+    elif data_format == ImageClassificationConfig.ImageFormat.S3_URL:
+        return IMAGENET_S3_URL_SPLIT_DIRS
+    else:
+        raise ValueError(f"Unknown data format: {data_format}")
+
+
+class ImageClassificationFactory(BenchmarkFactory):
+    def get_dataloader_factory(self) -> BaseDataLoaderFactory:
+        dataloader_type = self.benchmark_config.dataloader_type
+        task_config = self.benchmark_config.task_config
+        assert isinstance(task_config, ImageClassificationConfig)
+
+        data_dirs = get_imagenet_data_dirs(task_config)
+
+        data_format = task_config.image_classification_data_format
+
+        if dataloader_type == DataloaderType.MOCK:
+            return ImageClassificationMockDataLoaderFactory(self.benchmark_config)
+
+        elif dataloader_type == DataloaderType.RAY_DATA:
+            if data_format == ImageClassificationConfig.ImageFormat.JPEG:
+                from image_classification.jpeg.factory import (
+                    ImageClassificationJpegRayDataLoaderFactory,
+                )
+
+                return ImageClassificationJpegRayDataLoaderFactory(
+                    self.benchmark_config, data_dirs
+                )
+            elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
+                from image_classification.parquet.factory import (
+                    ImageClassificationParquetRayDataLoaderFactory,
+                )
+
+                return ImageClassificationParquetRayDataLoaderFactory(
+                    self.benchmark_config, data_dirs
+                )
+            elif data_format == ImageClassificationConfig.ImageFormat.S3_URL:
+                # NOTE: This format downloads images via ray data expressions,
+                # which is less efficient than native Ray Data S3 reading (JPEG format or Parquet format).
+                # Use this primarily for testing the S3 URL download pattern.
+                from image_classification.s3_url.factory import (
+                    ImageClassificationS3UrlRayDataLoaderFactory,
+                )
+
+                return ImageClassificationS3UrlRayDataLoaderFactory(
+                    self.benchmark_config, data_dirs
+                )
+
+        elif dataloader_type == DataloaderType.TORCH:
+            if data_format == ImageClassificationConfig.ImageFormat.JPEG:
+                from image_classification.jpeg.factory import (
+                    ImageClassificationJpegTorchDataLoaderFactory,
+                )
+
+                return ImageClassificationJpegTorchDataLoaderFactory(
+                    self.benchmark_config, data_dirs
+                )
+            elif data_format == ImageClassificationConfig.ImageFormat.PARQUET:
+                from image_classification.parquet.factory import (
+                    ImageClassificationParquetTorchDataLoaderFactory,
+                )
+
+                return ImageClassificationParquetTorchDataLoaderFactory(
+                    self.benchmark_config, data_dirs
+                )
+
+        raise ValueError(
+            f"Invalid dataloader configuration: {dataloader_type}\n"
+            f"{task_config}\n{self.benchmark_config.dataloader_config}"
+        )
+
+    def get_model(self) -> torch.nn.Module:
+        return torchvision.models.resnet50(weights=None)
+
+    def get_loss_fn(self) -> torch.nn.Module:
+        return torch.nn.CrossEntropyLoss()

release/train_tests/benchmark/image_classification/jpeg/download_input_data_from_s3.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+set -e  # Exit on any error
+
+DATA_DIR="/mnt/local_storage/imagenet"
+ZIP_NAME="imagenet-64k.zip"
+ZIP_URL="s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC/$ZIP_NAME"
+ZIP_PATH="$DATA_DIR/$ZIP_NAME"
+TRAIN_DIR="$DATA_DIR/train"
+
+if [ ! -d "$TRAIN_DIR" ]; then
+    echo "Downloading and extracting ImageNet subset to $DATA_DIR..."
+    mkdir -p "$DATA_DIR"
+    pushd "$DATA_DIR" || exit
+
+    echo "Fetching $ZIP_URL..."
+    aws s3 cp "$ZIP_URL" "$ZIP_PATH"
+
+    echo "Unzipping..."
+    unzip -q "$ZIP_NAME"
+    rm "$ZIP_NAME"
+
+    popd || exit
+else
+    echo "Dataset already exists at $TRAIN_DIR. Skipping download and unzip."
+fi
+
+echo "Duplicating images in-place..."
+
+python3 <<EOF
+import shutil
+from pathlib import Path
+from tqdm import tqdm
+
+dataset_dir = Path("$TRAIN_DIR")
+
+for class_dir in tqdm(sorted(dataset_dir.iterdir()), desc="Processing classes"):
+    if not class_dir.is_dir():
+        continue
+
+    # Skip if already duplicated
+    if any(class_dir.glob("*_copy1.JPEG")):
+        print(f"Skipping {class_dir.name} (already duplicated)")
+        continue
+
+    for img_path in class_dir.glob("*.JPEG"):
+        for i in range(1, 8):
+            copy_name = img_path.stem + f"_copy{i}" + img_path.suffix
+            copy_path = class_dir / copy_name
+            shutil.copy2(img_path, copy_path)
+EOF
+
+echo "Image duplication complete."

release/train_tests/benchmark/image_classification/jpeg/factory.py

@@ -0,0 +1,214 @@
+# Standard library imports
+import logging
+from typing import Dict
+
+# Third-party imports
+import torchvision
+from torch.utils.data import IterableDataset
+import pyarrow.fs
+
+# Ray imports
+import ray.train
+from ray.data.datasource.partitioning import Partitioning
+
+# Local imports
+from constants import DatasetKey
+from config import BenchmarkConfig
+from image_classification.factory import (
+    ImageClassificationRayDataLoaderFactory,
+    ImageClassificationTorchDataLoaderFactory,
+)
+from image_classification.imagenet import get_transform
+from s3_reader import AWS_REGION
+from .imagenet import get_preprocess_map_fn
+from .jpeg_iterable_dataset import S3JpegImageIterableDataset
+from s3_jpeg_reader import S3JpegReader
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class ImageClassificationJpegRayDataLoaderFactory(
+    ImageClassificationRayDataLoaderFactory
+):
+    """Factory for creating Ray DataLoader for JPEG image classification.
+
+    Extends ImageClassificationRayDataLoaderFactory to provide:
+    1. S3 filesystem configuration with boto credentials
+    2. Ray dataset creation with partitioning by class
+    3. Resource allocation for concurrent validation
+    4. Image preprocessing with optional random transforms
+    """
+
+    def __init__(self, benchmark_config: BenchmarkConfig, dataset_dirs: Dict[str, str]):
+        super().__init__(benchmark_config)
+        self._dataset_dirs = dataset_dirs
+
+    def get_s3fs_with_boto_creds(
+        self, connection_timeout: int = 60, request_timeout: int = 60
+    ) -> pyarrow.fs.S3FileSystem:
+        """Create S3 filesystem with boto credentials.
+
+        Args:
+            connection_timeout: Timeout for establishing connection in seconds
+            request_timeout: Timeout for requests in seconds
+
+        Returns:
+            Configured S3FileSystem instance with boto credentials
+        """
+        import boto3
+
+        credentials = boto3.Session().get_credentials()
+
+        s3fs = pyarrow.fs.S3FileSystem(
+            access_key=credentials.access_key,
+            secret_key=credentials.secret_key,
+            session_token=credentials.token,
+            region=AWS_REGION,
+            connect_timeout=connection_timeout,
+            request_timeout=request_timeout,
+        )
+        return s3fs
+
+    def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
+        """Get Ray datasets for training and validation.
+
+        Creates training and validation datasets with:
+        1. Partitioning by class for efficient data loading
+        2. Image preprocessing with optional random transforms
+        3. Resource allocation for concurrent validation
+        4. Row limits based on benchmark configuration
+
+        Returns:
+            Dictionary containing:
+                - "train": Training dataset with random transforms
+                - "val": Validation dataset without transforms
+        """
+        train_dir = self._dataset_dirs[DatasetKey.TRAIN]
+        # TODO: The validation dataset directory is not partitioned by class.
+        val_dir = train_dir
+
+        filesystem = (
+            self.get_s3fs_with_boto_creds() if train_dir.startswith("s3://") else None
+        )
+
+        # Create training dataset with class-based partitioning
+        train_partitioning = Partitioning(
+            "dir", base_dir=train_dir, field_names=["class"]
+        )
+        train_ds = ray.data.read_images(
+            train_dir,
+            mode="RGB",
+            include_paths=False,
+            partitioning=train_partitioning,
+            filesystem=filesystem,
+        ).map(get_preprocess_map_fn(random_transforms=True))
+
+        if self.get_dataloader_config().limit_training_rows > 0:
+            train_ds = train_ds.limit(self.get_dataloader_config().limit_training_rows)
+
+        # Create validation dataset with same partitioning
+        val_partitioning = Partitioning("dir", base_dir=val_dir, field_names=["class"])
+        val_ds = ray.data.read_images(
+            val_dir,
+            mode="RGB",
+            include_paths=False,
+            partitioning=val_partitioning,
+            filesystem=filesystem,
+        ).map(get_preprocess_map_fn(random_transforms=False))
+
+        if self.get_dataloader_config().limit_validation_rows > 0:
+            val_ds = val_ds.limit(self.get_dataloader_config().limit_validation_rows)
+
+        return {
+            DatasetKey.TRAIN: train_ds,
+            DatasetKey.VALID: val_ds,
+        }
+
+
+class ImageClassificationJpegTorchDataLoaderFactory(
+    ImageClassificationTorchDataLoaderFactory, S3JpegReader
+):
+    """Factory for creating PyTorch DataLoaders for JPEG image classification.
+
+    Features:
+    - S3-based JPEG file reading with round-robin worker distribution
+    - Device transfer and error handling for data batches
+    - Row limits per worker for controlled processing
+    - Dataset caching for efficiency
+    """
+
+    def __init__(self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]):
+        super().__init__(benchmark_config)
+        S3JpegReader.__init__(self)  # Initialize S3JpegReader to set up _s3_client
+        self._data_dirs = data_dirs
+        self._cached_datasets = None
+
+    def get_iterable_datasets(self) -> Dict[str, IterableDataset]:
+        """Get train and validation datasets with worker-specific configurations.
+
+        Returns:
+            Dictionary containing:
+                - "train": Training dataset with random transforms
+                - "val": Validation dataset without transforms
+        """
+        if self._cached_datasets is not None:
+            return self._cached_datasets
+
+        if self._data_dirs[DatasetKey.TRAIN].startswith("s3://"):
+            return self._get_iterable_datasets_s3()
+        else:
+            return self._get_iterable_datasets_local()
+
+    def _get_iterable_datasets_local(self) -> Dict[str, IterableDataset]:
+        """Get train and validation datasets from local filesystem."""
+        train_dir = self._data_dirs[DatasetKey.TRAIN]
+        val_dir = self._data_dirs[DatasetKey.VALID]
+
+        train_dataset = torchvision.datasets.ImageFolder(
+            root=train_dir,
+            transform=get_transform(to_torch_tensor=True, random_transforms=True),
+        )
+
+        val_dataset = torchvision.datasets.ImageFolder(
+            root=val_dir,
+            transform=get_transform(to_torch_tensor=True, random_transforms=False),
+        )
+
+        return {
+            DatasetKey.TRAIN: train_dataset,
+            DatasetKey.VALID: val_dataset,
+        }
+
+    def _get_iterable_datasets_s3(self) -> Dict[str, IterableDataset]:
+        """Get train and validation datasets from S3."""
+
+        train_dir = self._data_dirs[DatasetKey.TRAIN]
+
+        # Get row limits for workers and total processing
+        (
+            limit_training_rows_per_worker,
+            limit_validation_rows_per_worker,
+        ) = self._get_worker_row_limits()
+
+        # Get file URLs for training and validation
+        train_file_urls = val_file_urls = self._get_file_urls(train_dir)
+        train_ds = S3JpegImageIterableDataset(
+            file_urls=train_file_urls,
+            random_transforms=True,
+            limit_rows_per_worker=limit_training_rows_per_worker,
+        )
+
+        # TODO: IMAGENET_JPEG_SPLIT_S3_DIRS["val"] does not have the label
+        # partitioning like "train" does. So we use "train" for validation.
+        val_ds = S3JpegImageIterableDataset(
+            file_urls=val_file_urls,
+            random_transforms=False,
+            limit_rows_per_worker=limit_validation_rows_per_worker,
+        )
+
+        self._cached_datasets = {
+            DatasetKey.TRAIN: train_ds,
+            DatasetKey.VALID: val_ds,
+        }
+        return self._cached_datasets

release/train_tests/benchmark/image_classification/jpeg/imagenet.py

@@ -0,0 +1,57 @@
+import torch
+import numpy as np
+from typing import Dict, Union, Callable
+from PIL import Image
+
+from constants import DatasetKey
+from image_classification.imagenet import (
+    get_transform,
+    IMAGENET_WNID_TO_ID,
+)
+
+
+IMAGENET_JPEG_SPLIT_S3_ROOT = "s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC"
+IMAGENET_JPEG_SPLIT_S3_DIRS = {
+    DatasetKey.TRAIN: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/train",
+    DatasetKey.VALID: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/val",
+    DatasetKey.TEST: f"{IMAGENET_JPEG_SPLIT_S3_ROOT}/test",
+}
+
+
+def get_preprocess_map_fn(
+    random_transforms: bool = True,
+) -> Callable[[Dict[str, Union[np.ndarray, str]]], Dict[str, torch.Tensor]]:
+    """Get a map function that transforms a row to the format expected by the training loop.
+
+    Args:
+        random_transforms: Whether to use random transforms for training
+
+    Returns:
+        A function that takes a row dict with:
+        - "image": numpy array in HWC format
+        - "class": WNID string
+
+    The output is a dict with "image" and "label" keys.
+    """
+    crop_resize_transform = get_transform(
+        to_torch_tensor=True, random_transforms=random_transforms
+    )
+
+    def map_fn(row: Dict[str, Union[np.ndarray, str]]) -> Dict[str, torch.Tensor]:
+        """Process a single row into the expected format.
+
+        Args:
+            row: Dict containing "image" and "class" keys
+
+        Returns:
+            Dict with "image" and "label" keys
+        """
+        # Convert NumPy HWC image to PIL
+        image_pil = Image.fromarray(row["image"])
+        # Apply transform (includes ToTensor + Normalize)
+        image = crop_resize_transform(image_pil)
+        # Convert label
+        label = IMAGENET_WNID_TO_ID[row["class"]]
+        return {"image": image, "label": label}
+
+    return map_fn

release/train_tests/benchmark/image_classification/jpeg/jpeg_iterable_dataset.py

@@ -0,0 +1,266 @@
+# Standard library imports
+import io
+import logging
+import time
+from typing import Iterator, List, Optional, Tuple, Callable
+
+# Third-party imports
+import numpy as np
+from PIL import Image as PILImage
+import torch
+from torch.utils.data import IterableDataset
+
+# Ray imports
+import ray
+import ray.train
+
+# Local imports
+from s3_jpeg_reader import S3JpegReader
+from .imagenet import get_preprocess_map_fn
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class S3JpegImageIterableDataset(S3JpegReader, IterableDataset):
+    """An iterable dataset that loads images from S3-stored JPEG files.
+
+    Features:
+    - Direct image fetching from S3
+    - Optional random transforms for training
+    - Row limits per worker for controlled processing
+    - Progress logging and performance metrics
+    """
+
+    # Constants
+    LOG_FREQUENCY = 1000  # Log progress every 1000 rows
+
+    def __init__(
+        self,
+        file_urls: List[str],
+        random_transforms: bool = True,
+        limit_rows_per_worker: Optional[int] = None,
+    ):
+        """Initialize the dataset.
+
+        Args:
+            file_urls: List of S3 URLs to load
+            random_transforms: Whether to use random transforms for training
+            limit_rows_per_worker: Maximum number of rows to process per worker
+        """
+        super().__init__()
+        self.file_urls = file_urls
+        self.limit_rows_per_worker = limit_rows_per_worker
+        self.random_transforms = random_transforms
+
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(
+            f"Worker {worker_rank}: Initialized with {len(file_urls)} files"
+            f"{f' (limit: {limit_rows_per_worker} rows)' if limit_rows_per_worker else ''}"
+        )
+
+    def _get_worker_info(self) -> Tuple[int, int]:
+        """Get current worker information.
+
+        Returns:
+            Tuple of (worker_id, num_workers)
+        """
+        worker_info = torch.utils.data.get_worker_info()
+        worker_id = worker_info.id if worker_info else 0
+        num_workers = worker_info.num_workers if worker_info else 1
+        return worker_id, num_workers
+
+    def _has_reached_row_limit(self, rows_processed: int) -> bool:
+        """Check if we've reached the row limit per worker.
+
+        Args:
+            rows_processed: Number of rows processed so far
+
+        Returns:
+            True if we've reached the limit, False otherwise
+        """
+        return (
+            self.limit_rows_per_worker is not None
+            and rows_processed >= self.limit_rows_per_worker
+        )
+
+    def _log_progress(
+        self, worker_id: int, rows_processed: int, last_log_time: float
+    ) -> float:
+        """Log processing progress and return updated last_log_time.
+
+        Args:
+            worker_id: ID of the current worker
+            rows_processed: Number of rows processed so far
+            last_log_time: Time of last progress log
+
+        Returns:
+            Updated last_log_time
+        """
+        if rows_processed % self.LOG_FREQUENCY == 0:
+            current_time = time.time()
+            elapsed_time = current_time - last_log_time
+            rows_per_second = (
+                self.LOG_FREQUENCY / elapsed_time if elapsed_time > 0 else 0
+            )
+            logger.info(
+                f"Worker {worker_id}: Processed {rows_processed} rows "
+                f"({rows_per_second:.2f} rows/sec)"
+            )
+            return current_time
+        return last_log_time
+
+    def _fetch_image(self, file_url: str) -> Tuple[str, Optional[PILImage.Image]]:
+        """Fetch a single image from S3.
+
+        Args:
+            file_url: S3 URL to fetch (e.g., "s3://bucket/path/to/image.jpg")
+
+        Returns:
+            Tuple of (file_url, PIL Image) where PIL Image is None if fetch failed
+        """
+        worker_id, _ = self._get_worker_info()
+        try:
+            bucket = file_url.replace("s3://", "").split("/")[0]
+            key = "/".join(file_url.replace("s3://", "").split("/")[1:])
+
+            response = self.s3_client.get_object(Bucket=bucket, Key=key)
+            image_data = response["Body"].read()
+            image = PILImage.open(io.BytesIO(image_data))
+            return file_url, image
+
+        except Exception as e:
+            logger.error(
+                f"Worker {worker_id}: Error fetching image from {file_url}: {str(e)}",
+                exc_info=True,
+            )
+            return file_url, None
+
+    def _process_image(
+        self,
+        image: PILImage.Image,
+        file_url: str,
+        preprocess_fn: Callable,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Process a single image and convert to tensors.
+
+        Args:
+            image: PIL Image to process
+            file_url: URL of the image file
+            preprocess_fn: Preprocessing function to apply
+
+        Returns:
+            Tuple of (image_tensor, label_tensor)
+        """
+        try:
+            # Convert to RGB and numpy array
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            image_array = np.array(image, dtype=np.uint8)
+
+            # Ensure HWC format (Height x Width x Channels)
+            if len(image_array.shape) == 2:  # Grayscale
+                image_array = np.stack([image_array] * 3, axis=-1)
+            elif len(image_array.shape) == 3 and image_array.shape[0] == 3:  # CHW
+                image_array = np.transpose(image_array, (1, 2, 0))
+
+            wnid = file_url.split("/")[-2]  # Extract WNID from path
+
+            processed = preprocess_fn({"image": image_array, "class": wnid})
+
+            image = torch.as_tensor(processed["image"], dtype=torch.float32)
+            label = torch.as_tensor(processed["label"], dtype=torch.int64)
+
+            return image, label
+
+        except Exception as e:
+            logger.error(
+                f"Error processing {file_url}: {str(e)}",
+                exc_info=True,
+            )
+            raise
+
+    def _process_files(
+        self, files_to_read: List[str], preprocess_fn: Callable, worker_id: int
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Process multiple files and yield processed rows.
+
+        Args:
+            files_to_read: List of file URLs to process
+            preprocess_fn: Preprocessing function to apply
+            worker_id: ID of the current worker
+
+        Yields:
+            Tuple of (image_tensor, label_tensor)
+        """
+        rows_processed = 0
+        last_log_time = time.time()
+        total_start_time = time.time()
+
+        for file_url in files_to_read:
+            if self._has_reached_row_limit(rows_processed):
+                logger.info(f"Worker {worker_id}: Reached row limit: {rows_processed}")
+                break
+
+            file_url, image = self._fetch_image(file_url)
+            if image is None:
+                continue
+
+            try:
+                image, label = self._process_image(
+                    image,
+                    file_url,
+                    preprocess_fn,
+                )
+                rows_processed += 1
+                last_log_time = self._log_progress(
+                    worker_id, rows_processed, last_log_time
+                )
+                yield image, label
+            except Exception:
+                continue
+
+        # Log final statistics
+        total_time = time.time() - total_start_time
+        logger.info(
+            f"Worker {worker_id}: Finished: {rows_processed} rows in {total_time:.2f}s "
+            f"({rows_processed/total_time:.2f} rows/sec)"
+        )
+
+    def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Iterate through the dataset and yield (image, label) tensors.
+
+        Yields:
+            Tuple[torch.Tensor, torch.Tensor]: (image, label) tensors
+
+        Raises:
+            Exception: If there's a fatal error during iteration
+        """
+        try:
+            # Get worker info for file distribution
+            worker_id, num_workers = self._get_worker_info()
+            logger.info(f"Worker {worker_id}/{num_workers}: Starting")
+
+            # Distribute files among workers
+            files_to_read = (
+                self.file_urls
+                if num_workers == 1
+                else self.file_urls[worker_id::num_workers]
+            )
+
+            logger.info(f"Worker {worker_id}: Processing {len(files_to_read)} files")
+
+            # Initialize preprocessing function
+            preprocess_fn = get_preprocess_map_fn(
+                random_transforms=self.random_transforms
+            )
+
+            # Process files and yield results
+            yield from self._process_files(files_to_read, preprocess_fn, worker_id)
+
+        except Exception as e:
+            logger.error(
+                f"Worker {worker_id}: Fatal error: {str(e)}",
+                exc_info=True,
+            )
+            raise

release/train_tests/benchmark/image_classification/parquet/factory.py

@@ -0,0 +1,139 @@
+# Standard library imports
+import logging
+from typing import Dict, Optional
+
+# Third-party imports
+from torch.utils.data import IterableDataset
+import ray
+import ray.data
+import ray.train
+
+# Local imports
+from constants import DatasetKey
+from config import BenchmarkConfig
+from image_classification.factory import (
+    ImageClassificationRayDataLoaderFactory,
+    ImageClassificationTorchDataLoaderFactory,
+)
+from .imagenet import get_preprocess_map_fn
+from .parquet_iterable_dataset import S3ParquetImageIterableDataset
+from s3_parquet_reader import S3ParquetReader
+
+logger = logging.getLogger(__name__)
+
+
+class ImageClassificationParquetRayDataLoaderFactory(
+    ImageClassificationRayDataLoaderFactory
+):
+    """Factory for creating Ray DataLoader for Parquet image classification.
+
+    Features:
+    - Parquet file reading with column selection
+    - Image decoding and preprocessing
+    - Resource allocation for concurrent validation
+    - Row limits based on benchmark configuration
+    """
+
+    def __init__(
+        self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
+    ) -> None:
+        super().__init__(benchmark_config)
+        self._data_dirs = data_dirs
+
+    def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
+        """Get Ray datasets for training and validation.
+
+        Returns:
+            Dictionary containing:
+                - "train": Training dataset with random transforms
+                - "val": Validation dataset without transforms
+        """
+        # Create training dataset with image decoding and transforms
+        train_ds = ray.data.read_parquet(
+            self._data_dirs[DatasetKey.TRAIN],
+            columns=["image", "label"],
+        ).map(get_preprocess_map_fn(decode_image=True, random_transforms=True))
+
+        if self.get_dataloader_config().limit_training_rows > 0:
+            train_ds = train_ds.limit(self.get_dataloader_config().limit_training_rows)
+
+        # Create validation dataset without random transforms
+        val_ds = ray.data.read_parquet(
+            self._data_dirs[DatasetKey.TRAIN],
+            columns=["image", "label"],
+        ).map(get_preprocess_map_fn(decode_image=True, random_transforms=False))
+
+        if self.get_dataloader_config().limit_validation_rows > 0:
+            val_ds = val_ds.limit(self.get_dataloader_config().limit_validation_rows)
+
+        return {
+            DatasetKey.TRAIN: train_ds,
+            DatasetKey.VALID: val_ds,
+        }
+
+
+class ImageClassificationParquetTorchDataLoaderFactory(
+    ImageClassificationTorchDataLoaderFactory, S3ParquetReader
+):
+    """Factory for creating PyTorch DataLoaders for Parquet image classification.
+
+    Features:
+    - Parquet file reading with row count-based distribution
+    - Worker-based file distribution for balanced workloads
+    - Row limits per worker for controlled processing
+    - Dataset instance caching for efficiency
+    """
+
+    def __init__(
+        self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
+    ) -> None:
+        """Initialize factory with benchmark configuration.
+
+        Args:
+            benchmark_config: Configuration for benchmark parameters
+        """
+        super().__init__(benchmark_config)
+        S3ParquetReader.__init__(
+            self
+        )  # Initialize S3ParquetReader to set up _s3_client
+        self.train_url = data_dirs[DatasetKey.TRAIN]
+        self._cached_datasets: Optional[Dict[str, IterableDataset]] = None
+
+    def get_iterable_datasets(self) -> Dict[str, IterableDataset]:
+        """Get train and validation datasets with worker-specific configurations.
+
+        Returns:
+            Dictionary containing:
+                - "train": Training dataset with random transforms
+                - "val": Validation dataset without transforms
+        """
+        if self._cached_datasets is not None:
+            return self._cached_datasets
+
+        # Get row limits for workers and total processing
+        (
+            limit_training_rows_per_worker,
+            limit_validation_rows_per_worker,
+        ) = self._get_worker_row_limits()
+
+        # Create training dataset
+        train_file_urls = self._get_file_urls(self.train_url)
+        train_ds = S3ParquetImageIterableDataset(
+            file_urls=train_file_urls,
+            random_transforms=True,
+            limit_rows_per_worker=limit_training_rows_per_worker,
+        )
+
+        # Create validation dataset
+        val_file_urls = train_file_urls
+        val_ds = S3ParquetImageIterableDataset(
+            file_urls=val_file_urls,
+            random_transforms=False,
+            limit_rows_per_worker=limit_validation_rows_per_worker,
+        )
+
+        self._cached_datasets = {
+            DatasetKey.TRAIN: train_ds,
+            DatasetKey.VALID: val_ds,
+        }
+        return self._cached_datasets

release/train_tests/benchmark/image_classification/parquet/imagenet.py

@@ -0,0 +1,69 @@
+import io
+import numpy as np
+from typing import Dict, Union, Callable
+from PIL import Image
+from torchvision.transforms.functional import pil_to_tensor
+
+from constants import DatasetKey
+from image_classification.imagenet import (
+    get_transform,
+    IMAGENET_WNID_TO_ID,
+)
+
+IMAGENET_PARQUET_SPLIT_S3_ROOT = (
+    "s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet_split"
+)
+IMAGENET_PARQUET_SPLIT_S3_DIRS = {
+    DatasetKey.TRAIN: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/train",
+    DatasetKey.VALID: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/val",
+    DatasetKey.TEST: f"{IMAGENET_PARQUET_SPLIT_S3_ROOT}/test",
+}
+
+# Much larger parquet dataset used to sustain Ray Data backpressure in the
+# slow-consumer ingest benchmarks.
+IMAGENET_PARQUET_SPLIT_1T_S3_ROOT = (
+    "s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet_split_1t"
+)
+IMAGENET_PARQUET_SPLIT_1T_S3_DIRS = {
+    DatasetKey.TRAIN: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/train",
+    DatasetKey.VALID: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/val",
+    DatasetKey.TEST: f"{IMAGENET_PARQUET_SPLIT_1T_S3_ROOT}/test",
+}
+
+
+def get_preprocess_map_fn(
+    decode_image: bool = True, random_transforms: bool = True
+) -> Callable[[Dict[str, Union[bytes, str]]], Dict[str, Union[np.ndarray, int]]]:
+    """Get a map function that transforms a row of the dataset to the format
+    expected by the training loop.
+
+    Args:
+        decode_image: Whether to decode the image bytes into a tensor
+        random_transforms: Whether to use random transforms for training
+
+    Returns:
+        A function that takes a row dict and returns a processed dict.
+        Input row dict should have:
+        - "image": bytes or tensor in CHW format
+        - "label": WNID string
+
+        Output dict has:
+        - "image": np.array of the transformed, normalized image
+        - "label": An integer index of the WNID
+    """
+    crop_resize_transform = get_transform(
+        to_torch_tensor=False, random_transforms=random_transforms
+    )
+
+    def map_fn(row: Dict[str, Union[bytes, str]]) -> Dict[str, Union[np.ndarray, int]]:
+        assert "image" in row and "label" in row, row.keys()
+
+        if decode_image:
+            row["image"] = pil_to_tensor(Image.open(io.BytesIO(row["image"]))) / 255.0
+
+        row["image"] = np.array(crop_resize_transform(row["image"]))
+        row["label"] = IMAGENET_WNID_TO_ID[row["label"]]
+
+        return {"image": row["image"], "label": row["label"]}
+
+    return map_fn

release/train_tests/benchmark/image_classification/parquet/parquet_iterable_dataset.py

@@ -0,0 +1,273 @@
+# Standard library imports
+from typing import List, Tuple, Optional, Iterator, Callable
+import logging
+import io
+import time
+
+# Third-party imports
+import pandas as pd
+import pyarrow.parquet as pq
+import torch
+from torch.utils.data import IterableDataset
+
+# Ray imports
+import ray
+import ray.train
+
+# Local imports
+from s3_parquet_reader import S3ParquetReader
+from .imagenet import get_preprocess_map_fn
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+# TODO Look into https://github.com/webdataset/webdataset for more canonical way to do data
+# distribution between Ray Train and Torch Dataloader workers.
+
+
+class S3ParquetImageIterableDataset(S3ParquetReader, IterableDataset):
+    """An iterable dataset that loads images from S3-stored Parquet files.
+
+    This dataset:
+    1. Reads Parquet files from S3 one row group at a time
+    2. Processes images with optional random transforms
+    3. Yields (image, label) tensors
+    4. Supports row limits per worker for controlled data processing
+    """
+
+    LOG_FREQUENCY = 1000  # Log progress every 1000 rows
+
+    def __init__(
+        self,
+        file_urls: List[str],
+        random_transforms: bool = True,
+        limit_rows_per_worker: Optional[int] = None,
+    ):
+        """Initialize the dataset.
+
+        Args:
+            file_urls: List of S3 URLs to load
+            random_transforms: Whether to use random transforms for training
+            limit_rows_per_worker: Maximum number of rows to process per worker (None for all rows)
+        """
+        super().__init__()
+        self.file_urls = file_urls
+        self.limit_rows_per_worker = limit_rows_per_worker
+        self.random_transforms = random_transforms
+
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(
+            f"Worker {worker_rank}: Initialized with {len(file_urls)} files"
+            f"{f' (limit: {limit_rows_per_worker} rows)' if limit_rows_per_worker else ''}"
+        )
+
+    def _get_worker_info(self) -> Tuple[int, int]:
+        """Get current worker information.
+
+        Returns:
+            Tuple of (worker_id, num_workers)
+        """
+        worker_info = torch.utils.data.get_worker_info()
+        worker_id = worker_info.id if worker_info else 0
+        num_workers = worker_info.num_workers if worker_info else 1
+        return worker_id, num_workers
+
+    def _has_reached_row_limit(self, rows_processed: int) -> bool:
+        """Check if we've reached the row limit per worker.
+
+        Args:
+            rows_processed: Number of rows processed so far
+
+        Returns:
+            True if we've reached the limit, False otherwise
+        """
+        return (
+            self.limit_rows_per_worker is not None
+            and rows_processed >= self.limit_rows_per_worker
+        )
+
+    def _log_progress(
+        self, worker_id: int, rows_processed: int, last_log_time: float
+    ) -> float:
+        """Log processing progress and return updated last_log_time.
+
+        Args:
+            worker_id: ID of the current worker
+            rows_processed: Number of rows processed so far
+            last_log_time: Time of last progress log
+
+        Returns:
+            Updated last_log_time
+        """
+        if rows_processed % self.LOG_FREQUENCY == 0:
+            current_time = time.time()
+            elapsed_time = current_time - last_log_time
+            rows_per_second = (
+                self.LOG_FREQUENCY / elapsed_time if elapsed_time > 0 else 0
+            )
+            logger.info(
+                f"Worker {worker_id}: Processed {rows_processed} rows "
+                f"({rows_per_second:.2f} rows/sec)"
+            )
+            return current_time
+        return last_log_time
+
+    def _read_parquet_file(self, file_url: str) -> Iterator[pd.DataFrame]:
+        """Read a Parquet file from S3 one row group at a time.
+
+        This method:
+        1. Fetches the Parquet file from S3
+        2. Reads it row group by row group
+        3. Converts each row group to a pandas DataFrame
+
+        Args:
+            file_url: S3 URL of the Parquet file
+
+        Yields:
+            DataFrame containing one row group at a time
+
+        Raises:
+            Exception: If there's an error reading the file
+        """
+        try:
+            start_time = time.time()
+            worker_id, _ = self._get_worker_info()
+            logger.info(f"Worker {worker_id}: Reading Parquet file: {file_url}")
+
+            # Get parquet file metadata
+            bucket, key = self._parse_s3_url(file_url)
+            response = self.s3_client.get_object(Bucket=bucket, Key=key)
+            parquet_file = pq.ParquetFile(io.BytesIO(response["Body"].read()))
+            num_row_groups = parquet_file.num_row_groups
+
+            logger.info(
+                f"Worker {worker_id}: Found {num_row_groups} row groups in {file_url}"
+            )
+
+            for row_group in range(num_row_groups):
+                # Read row group and convert to pandas
+                table = parquet_file.read_row_group(row_group)
+                df = table.to_pandas()
+                yield df
+
+            total_time = time.time() - start_time
+            logger.info(
+                f"Worker {worker_id}: Completed reading {file_url} in {total_time:.2f}s"
+            )
+
+        except Exception as e:
+            worker_id, _ = self._get_worker_info()
+            logger.error(
+                f"Worker {worker_id}: Error reading file {file_url}: {str(e)}",
+                exc_info=True,
+            )
+            raise
+
+    def _process_file(
+        self,
+        file_url: str,
+        preprocess_fn: Callable,
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Process a single file and yield processed rows.
+
+        Args:
+            file_url: URL of the file to process
+            preprocess_fn: Preprocessing function to apply
+
+        Yields:
+            Tuple of (image_tensor, label_tensor)
+        """
+        for df in self._read_parquet_file(file_url):
+            for _, row in df.iterrows():
+                try:
+                    # Process row and convert to tensors
+                    processed = preprocess_fn(row)
+                    image = torch.as_tensor(processed["image"], dtype=torch.float32)
+                    label = torch.as_tensor(processed["label"], dtype=torch.int64)
+                    yield image, label
+                except Exception:
+                    continue
+
+    def _process_files(
+        self, files_to_read: List[str], preprocess_fn: Callable, worker_id: int
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Process multiple files and yield processed rows.
+
+        Args:
+            files_to_read: List of file URLs to process
+            preprocess_fn: Preprocessing function to apply
+            worker_id: ID of the current worker
+
+        Yields:
+            Tuple of (image_tensor, label_tensor)
+        """
+        rows_processed = 0
+        last_log_time = time.time()
+        total_start_time = time.time()
+
+        for file_url in files_to_read:
+            if self._has_reached_row_limit(rows_processed):
+                logger.info(f"Worker {worker_id}: Reached row limit: {rows_processed}")
+                break
+
+            for image, label in self._process_file(file_url, preprocess_fn):
+                if self._has_reached_row_limit(rows_processed):
+                    break
+
+                rows_processed += 1
+                last_log_time = self._log_progress(
+                    worker_id, rows_processed, last_log_time
+                )
+                yield image, label
+
+        # Log final statistics
+        total_time = time.time() - total_start_time
+        logger.info(
+            f"Worker {worker_id}: Finished: {rows_processed} rows in {total_time:.2f}s "
+            f"({rows_processed/total_time:.2f} rows/sec)"
+        )
+
+    def __iter__(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Main iteration method that processes files and yields (image, label) tensors.
+
+        This method:
+        1. Distributes files among workers
+        2. Processes rows with image transforms
+        3. Converts to tensors
+        4. Respects row limits per worker
+
+        Yields:
+            Tuple of (image_tensor, label_tensor)
+
+        Raises:
+            Exception: If there's a fatal error during processing
+        """
+        try:
+            # Get worker info for file distribution
+            worker_id, num_workers = self._get_worker_info()
+            logger.info(f"Worker {worker_id}/{num_workers}: Starting")
+
+            # Initialize preprocessing function
+            preprocess_fn = get_preprocess_map_fn(
+                decode_image=True, random_transforms=self.random_transforms
+            )
+
+            # Distribute files among workers
+            files_to_read = (
+                self.file_urls
+                if num_workers == 1
+                else self.file_urls[worker_id::num_workers]
+            )
+
+            logger.info(f"Worker {worker_id}: Processing {len(files_to_read)} files")
+
+            # Process files and yield results
+            yield from self._process_files(files_to_read, preprocess_fn, worker_id)
+
+        except Exception as e:
+            logger.error(
+                f"Worker {worker_id}: Fatal error: {str(e)}",
+                exc_info=True,
+            )
+            raise

release/train_tests/benchmark/image_classification/s3_url/factory.py

@@ -0,0 +1,77 @@
+# Standard library imports
+import logging
+from typing import Dict
+
+# Third-party imports
+import ray.data
+
+# Local imports
+from constants import DatasetKey
+from config import BenchmarkConfig
+from image_classification.factory import ImageClassificationRayDataLoaderFactory
+from .imagenet import (
+    create_s3_url_dataset,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ImageClassificationS3UrlRayDataLoaderFactory(
+    ImageClassificationRayDataLoaderFactory
+):
+    """Factory for creating Ray DataLoader that downloads images from S3 URLs.
+
+    This factory:
+    1. Lists JPEG files from S3 using boto3
+    2. Creates a Ray dataset from the file records
+    3. Uses map_batches to download and process images from S3
+
+    This approach separates file listing from image downloading, which can be
+    more efficient for certain workloads as it allows parallel downloads during
+    the map_batches execution on CPU workers.
+    """
+
+    def __init__(
+        self, benchmark_config: BenchmarkConfig, data_dirs: Dict[str, str]
+    ) -> None:
+        super().__init__(benchmark_config)
+        self._data_dirs = data_dirs
+
+    def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
+        """Get Ray datasets for training and validation.
+
+        Returns:
+            Dictionary containing:
+                - "train": Training dataset with random transforms
+                - "val": Validation dataset without transforms
+        """
+        dataloader_config = self.get_dataloader_config()
+
+        # Create training dataset
+        train_limit = (
+            dataloader_config.limit_training_rows
+            if dataloader_config.limit_training_rows > 0
+            else None
+        )
+        train_ds = create_s3_url_dataset(
+            data_dir=self._data_dirs[DatasetKey.TRAIN],
+            random_transforms=True,
+            limit_rows=train_limit,
+        )
+
+        # Create validation dataset
+        val_limit = (
+            dataloader_config.limit_validation_rows
+            if dataloader_config.limit_validation_rows > 0
+            else None
+        )
+        val_ds = create_s3_url_dataset(
+            data_dir=self._data_dirs[DatasetKey.TRAIN],
+            random_transforms=False,
+            limit_rows=val_limit,
+        )
+
+        return {
+            DatasetKey.TRAIN: train_ds,
+            DatasetKey.VALID: val_ds,
+        }

release/train_tests/benchmark/image_classification/s3_url/imagenet.py

@@ -0,0 +1,251 @@
+"""ImageNet dataset loading via S3 URL download with Ray Data expressions.
+
+This module provides dataset loading that:
+1. Lists JPEG files from S3 using boto3 (parallelized via Ray tasks)
+2. Creates a Ray dataset from the file records
+3. Uses Ray Data expressions (alpha) to download image bytes efficiently
+4. Uses map_batches to decode and process images
+
+This approach leverages Ray Data's expressions API for optimized parallel I/O,
+separating the download step from image processing for better throughput.
+"""
+
+import io
+import logging
+from functools import lru_cache
+from typing import Callable, Dict, List, Optional, Tuple
+
+import boto3
+import numpy as np
+from PIL import Image
+from torchvision.transforms.functional import pil_to_tensor
+
+import ray.data
+from ray.data.expressions import download
+
+from constants import DatasetKey
+from image_classification.imagenet import (
+    get_transform,
+    IMAGENET_WNID_TO_ID,
+)
+
+logger = logging.getLogger(__name__)
+
+# S3 configuration for ImageNet JPEG data
+AWS_REGION = "us-west-2"
+S3_ROOT = "s3://anyscale-imagenet/ILSVRC/Data/CLS-LOC"
+IMAGENET_S3_URL_SPLIT_DIRS = {
+    DatasetKey.TRAIN: f"{S3_ROOT}/train",
+    DatasetKey.VALID: f"{S3_ROOT}/val",
+    DatasetKey.TEST: f"{S3_ROOT}/test",
+}
+
+
+def _get_class_labels(bucket: str, prefix: str) -> List[str]:
+    """Get all class label directories from S3.
+
+    Args:
+        bucket: S3 bucket name
+        prefix: S3 prefix path
+
+    Returns:
+        List of class label directory names
+    """
+    from typing import Set
+
+    # Ensure prefix ends with /
+    if prefix and not prefix.endswith("/"):
+        prefix += "/"
+
+    # List directories using delimiter
+    s3_client = boto3.client("s3", region_name=AWS_REGION)
+    paginator = s3_client.get_paginator("list_objects_v2")
+
+    # Use delimiter to get "directory" level
+    labels: Set[str] = set()
+    for page in paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/"):
+        # CommonPrefixes contains the "directories"
+        for common_prefix in page.get("CommonPrefixes", []):
+            prefix_path = common_prefix["Prefix"]
+            # Extract the directory name
+            label = prefix_path.rstrip("/").split("/")[-1]
+            labels.add(label)
+
+    return sorted(labels)
+
+
+@ray.remote
+def _list_files_for_label(
+    bucket: str, prefix: str, label: str
+) -> List[Tuple[str, str]]:
+    """Ray task to list all image files for a specific label.
+
+    Args:
+        bucket: S3 bucket name
+        prefix: S3 prefix (parent directory)
+        label: Class label (subdirectory name)
+
+    Returns:
+        List of tuples with (file_path, class_name)
+    """
+    s3_client = boto3.client("s3", region_name=AWS_REGION)
+    paginator = s3_client.get_paginator("list_objects_v2")
+
+    # Construct the full prefix for this label
+    label_prefix = f"{prefix}/{label}/" if prefix else f"{label}/"
+
+    file_records = []
+    for page in paginator.paginate(Bucket=bucket, Prefix=label_prefix):
+        for obj in page.get("Contents", []):
+            key = obj["Key"]
+            if key.lower().endswith((".jpg", ".jpeg")):
+                file_path = f"s3://{bucket}/{key}"
+                file_records.append((file_path, label))
+
+    return file_records
+
+
+@lru_cache(maxsize=8)
+def _list_s3_image_files_cached(data_dir: str) -> Tuple[Tuple[str, str], ...]:
+    """Cached implementation of S3 file listing using Ray tasks for parallelism.
+
+    Returns a tuple of tuples for hashability (required by lru_cache).
+    """
+    logger.info(f"Listing JPEG files from {data_dir}...")
+
+    # Parse S3 URL: s3://bucket/prefix
+    s3_path = data_dir
+    if s3_path.startswith("s3://"):
+        s3_path = s3_path[5:]
+    parts = s3_path.split("/", 1)
+    bucket = parts[0]
+    prefix = parts[1].rstrip("/") if len(parts) > 1 else ""
+
+    # Get all class labels
+    labels = _get_class_labels(bucket, prefix)
+    logger.info(
+        f"Found {len(labels)} class labels, launching Ray tasks for parallel listing..."
+    )
+
+    # Launch Ray tasks for each label
+    futures = [_list_files_for_label.remote(bucket, prefix, label) for label in labels]
+
+    # Wait for all tasks to complete and aggregate results
+    results = ray.get(futures)
+
+    # Flatten the list of lists
+    file_records = []
+    for records in results:
+        file_records.extend(records)
+
+    logger.info(f"Listed and cached {len(file_records)} JPEG files")
+    return tuple(file_records)
+
+
+def list_s3_image_files(data_dir: str) -> List[Dict[str, str]]:
+    """List JPEG files from S3 with class labels extracted from path.
+
+    Results are cached to avoid repeated S3 listings.
+
+    Args:
+        data_dir: S3 path to list files from (e.g., "s3://bucket/prefix")
+
+    Returns:
+        List of dicts with "path" (S3 URL) and "class" (WNID) keys
+    """
+    cached_records = _list_s3_image_files_cached(data_dir)
+    return [{"path": path, "class": cls} for path, cls in cached_records]
+
+
+def get_process_batch_fn(
+    random_transforms: bool = True,
+    label_to_id_map: Optional[Dict[str, int]] = None,
+) -> Callable[[Dict[str, np.ndarray]], Dict[str, np.ndarray]]:
+    """Get a map_batches function that processes pre-downloaded image bytes.
+
+    This function expects image bytes to already be downloaded (via Ray Data
+    expressions) and handles decoding and transformations.
+
+    Args:
+        random_transforms: Whether to use random transforms for training
+        label_to_id_map: Mapping from WNID strings to integer IDs
+
+    Returns:
+        A function suitable for use with dataset.map_batches()
+    """
+    if label_to_id_map is None:
+        label_to_id_map = IMAGENET_WNID_TO_ID
+
+    transform = get_transform(
+        to_torch_tensor=False, random_transforms=random_transforms
+    )
+
+    def process_batch(
+        batch: Dict[str, np.ndarray],
+    ) -> Dict[str, np.ndarray]:
+        """Process pre-downloaded image bytes.
+
+        Args:
+            batch: Dict with "bytes" (image data) and "class" arrays
+
+        Returns:
+            Dict with "image" (numpy array) and "label" (int) arrays
+        """
+        processed_images = []
+        labels = []
+
+        image_bytes_list = list(batch["bytes"])
+        classes = list(batch["class"])
+
+        for data, wnid in zip(image_bytes_list, classes):
+            # Decode and transform image
+            image_pil = Image.open(io.BytesIO(data)).convert("RGB")
+            image_tensor = pil_to_tensor(image_pil) / 255.0
+            processed_image = np.array(transform(image_tensor))
+            processed_images.append(processed_image)
+
+            # Convert label
+            labels.append(label_to_id_map[wnid])
+
+        return {
+            "image": np.stack(processed_images),
+            "label": np.array(labels),
+        }
+
+    return process_batch
+
+
+def create_s3_url_dataset(
+    data_dir: str,
+    random_transforms: bool = True,
+    limit_rows: Optional[int] = None,
+) -> ray.data.Dataset:
+    """Create a Ray dataset that downloads images from S3 URLs.
+
+    Uses Ray Data expressions (alpha) for efficient parallel downloads,
+    then map_batches for image decoding and transformations.
+
+    Args:
+        data_dir: S3 path to the image directory
+        random_transforms: Whether to use random transforms
+        limit_rows: Optional row limit
+
+    Returns:
+        Ray dataset with "image" and "label" columns
+    """
+    file_records = list_s3_image_files(data_dir)
+
+    ds = ray.data.from_items(file_records)
+
+    if limit_rows is not None and limit_rows > 0:
+        ds = ds.limit(limit_rows)
+
+    # Download image bytes using Ray Data expressions (alpha)
+    # This enables optimized parallel I/O managed by Ray Data
+    ds = ds.with_column("bytes", download("path"))
+
+    # Process downloaded bytes (decode and transform)
+    process_fn = get_process_batch_fn(random_transforms=random_transforms)
+    ds = ds.map_batches(process_fn)
+
+    return ds

release/train_tests/benchmark/logger_utils.py

@@ -0,0 +1,55 @@
+import logging
+import inspect
+from typing import Any, Dict, Tuple, Union
+
+
+class ContextLoggerAdapter(logging.LoggerAdapter):
+    def __init__(
+        self, logger: logging.Logger, extra: Union[Dict[str, Any], None] = None
+    ) -> None:
+        """Initialize the logger adapter.
+
+        Args:
+            logger: The logger to wrap
+            extra: Extra data to include in log records
+        """
+        super().__init__(logger, extra or {})
+
+    def process(self, msg: str, kwargs: Dict[str, Any]) -> Tuple[str, Dict[str, Any]]:
+        """Process a log message and add context information.
+
+        Args:
+            msg: The log message to process
+            kwargs: Additional keyword arguments for logging
+
+        Returns:
+            Tuple containing:
+                - The processed message with context prefix
+                - The original kwargs
+        """
+        # Get the frame that called the logging method
+        # Go up 3 frames: process -> log -> info/error/etc -> actual caller
+        frame = inspect.currentframe()
+        if (
+            frame
+            and frame.f_back
+            and frame.f_back.f_back
+            and frame.f_back.f_back.f_back
+        ):
+            frame = frame.f_back.f_back.f_back
+            class_name = getattr(frame.f_locals.get("self"), "__class__", None)
+            func_name = frame.f_code.co_name
+
+            # Create the prefix with class and function context
+            prefix = (
+                f"[{class_name.__name__}.{func_name}]"
+                if class_name
+                else f"[{func_name}]"
+            )
+        else:
+            prefix = "[unknown]"
+
+        # Add any extra context from the adapter
+        # Don't modify kwargs directly as it causes issues with level handling
+
+        return f"{prefix} {msg}", kwargs

release/train_tests/benchmark/ray_dataloader_factory.py

@@ -0,0 +1,281 @@
+import logging
+import threading
+import time
+from abc import abstractmethod
+from typing import Any, Dict, Optional
+
+import ray
+import ray.train
+from ray._private.internal_api import get_memory_info_reply, get_state_from_address
+from ray.data import Dataset
+from ray.data.collate_fn import CollateFn
+
+from constants import DatasetKey
+from config import BenchmarkConfig, RayDataConfig
+from dataloader_factory import BaseDataLoaderFactory
+
+logger = logging.getLogger(__name__)
+
+SPILL_MONITOR_ACTOR_NAME = "spill_metrics_monitor"
+SPILL_MONITOR_ACTOR_NAMESPACE = "_spill_metrics_monitor"
+
+
+@ray.remote(num_cpus=0)
+class SpillMetricsMonitor:
+    """Actor that periodically polls object store spill metrics
+    to compute peak and average spilling rates (GB/min).
+
+    GB/min is used instead of GB/s because object store spilling rates are
+    typically small fractions of a GB per second, making GB/s values hard to
+    read and compare (e.g. 0.0021 GB/s vs 0.13 GB/min). GB/min produces
+    more human-friendly numbers while still matching the 60-second poll
+    interval naturally.
+
+    A single instance is shared across all workers via a named actor.
+    """
+
+    def __init__(self, poll_interval_s: float = 60.0):
+        self._poll_interval_s = poll_interval_s
+        self._count = 0
+        self._sum_gb_min = 0.0
+        self._max_gb_min = 0.0
+        self._lock = threading.Lock()
+
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def _get_spilled_bytes(self) -> int:
+        memory_info = get_memory_info_reply(
+            get_state_from_address(ray.get_runtime_context().gcs_address)
+        )
+        return memory_info.store_stats.spilled_bytes_total
+
+    def _poll_loop(self) -> None:
+        prev_spilled_bytes: Optional[int] = None
+        prev_time: Optional[float] = None
+
+        while True:
+            time.sleep(self._poll_interval_s)
+            try:
+                current_bytes = self._get_spilled_bytes()
+                current_time = time.monotonic()
+
+                if prev_spilled_bytes is not None and prev_time is not None:
+                    delta_bytes = current_bytes - prev_spilled_bytes
+                    delta_time = current_time - prev_time
+
+                    if delta_time > 0 and delta_bytes >= 0:
+                        rate_gb_min = (delta_bytes / (1024**3)) / delta_time * 60
+                        with self._lock:
+                            self._count += 1
+                            self._sum_gb_min += rate_gb_min
+                            self._max_gb_min = max(self._max_gb_min, rate_gb_min)
+
+                prev_spilled_bytes = current_bytes
+                prev_time = current_time
+            except Exception as e:
+                logger.warning(f"SpillMetricsMonitor: poll failed: {e}")
+
+    def get_metrics(self) -> Dict[str, float]:
+        with self._lock:
+            count = self._count
+            sum_gb_min = self._sum_gb_min
+            max_gb_min = self._max_gb_min
+
+        if count == 0:
+            return {}
+
+        return {
+            "object_store_spilling_peak_gb_min": round(max_gb_min, 4),
+            "object_store_spilling_avg_gb_min": round(sum_gb_min / count, 4),
+        }
+
+
+def get_or_create_spill_metrics_monitor(
+    poll_interval_s: float = 60.0,
+) -> ray.actor.ActorHandle:
+    return SpillMetricsMonitor.options(
+        name=SPILL_MONITOR_ACTOR_NAME,
+        namespace=SPILL_MONITOR_ACTOR_NAMESPACE,
+        get_if_exists=True,
+        lifetime="detached",
+    ).remote(poll_interval_s=poll_interval_s)
+
+
+class RayDataLoaderFactory(BaseDataLoaderFactory):
+    def __init__(self, benchmark_config: BenchmarkConfig) -> None:
+        super().__init__(benchmark_config)
+        self._ray_ds_iterators = {}
+        self._spill_monitor: Optional[ray.actor.ActorHandle] = None
+
+        dataloader_config = self.get_dataloader_config()
+        assert isinstance(dataloader_config, RayDataConfig), type(dataloader_config)
+
+        # Configure Ray Data settings.
+        data_context = ray.data.DataContext.get_current()
+        data_context.enable_operator_progress_bars = (
+            dataloader_config.enable_operator_progress_bars
+        )
+        # Retry transient S3 errors that sometimes show up due to
+        # throttling during read operations.
+        data_context.retried_io_errors.append("AWS Error ACCESS_DENIED")
+        data_context.retried_io_errors.append("AWS Error UNKNOWN (HTTP status 500)")
+
+        data_context.execution_options.locality_with_output = (
+            dataloader_config.locality_with_output
+        )
+        data_context.execution_options.actor_locality_enabled = (
+            dataloader_config.actor_locality_enabled
+        )
+        data_context.execution_options.preserve_order = dataloader_config.preserve_order
+
+    @abstractmethod
+    def get_ray_datasets(self) -> Dict[str, Dataset]:
+        """Get Ray datasets."""
+        raise NotImplementedError
+
+    def _get_collate_fn(self) -> Optional[CollateFn]:
+        """Return the collate function for the dataloader."""
+        return None
+
+    def get_ray_data_config(self) -> ray.train.DataConfig:
+        return ray.train.DataConfig(
+            enable_shard_locality=self.get_dataloader_config().enable_shard_locality,
+        )
+
+    def get_train_dataloader(self):
+        """Get the training dataloader.
+
+        Returns:
+            Iterator of training batches
+        """
+        # Get or create the shared spill monitor actor on first call.
+        if self._spill_monitor is None:
+            self._spill_monitor = get_or_create_spill_metrics_monitor()
+
+        ds_iterator = ray.train.get_dataset_shard(DatasetKey.TRAIN)
+        self._ray_ds_iterators[DatasetKey.TRAIN] = ds_iterator
+
+        dataloader_config = self.get_dataloader_config()
+        return iter(
+            ds_iterator.iter_torch_batches(
+                batch_size=dataloader_config.train_batch_size,
+                local_shuffle_buffer_size=(
+                    dataloader_config.local_buffer_shuffle_size
+                    if dataloader_config.local_buffer_shuffle_size > 0
+                    else None
+                ),
+                collate_fn=self._get_collate_fn(),
+                prefetch_batches=dataloader_config.ray_data_prefetch_batches,
+                drop_last=True,
+                pin_memory=dataloader_config.ray_data_pin_memory,
+            )
+        )
+
+    def get_val_dataloader(self):
+        """Get the validation dataloader.
+
+        Returns:
+            Iterator of validation batches
+        """
+        ds_iterator = ray.train.get_dataset_shard(DatasetKey.VALID)
+        self._ray_ds_iterators[DatasetKey.VALID] = ds_iterator
+
+        dataloader_config = self.get_dataloader_config()
+        return iter(
+            ds_iterator.iter_torch_batches(
+                batch_size=dataloader_config.validation_batch_size,
+                collate_fn=self._get_collate_fn(),
+                prefetch_batches=dataloader_config.ray_data_prefetch_batches,
+                drop_last=True,
+            )
+        )
+
+    def get_metrics(self) -> Dict[str, Any]:
+        metrics = {}
+        for ds_key, ds_iterator in self._ray_ds_iterators.items():
+            stats = ray.get(ds_iterator._coord_actor.stats.remote())
+            summary = stats.to_summary()
+            summary.iter_stats = ds_iterator._iter_stats.to_summary().iter_stats
+            summary.iter_stats.streaming_split_coord_time.add(
+                stats.streaming_split_coordinator_s.get()
+            )
+
+            if not summary.parents:
+                continue
+
+            # The split() operator has no metrics, so pull the stats
+            # from the final dataset stage.
+            ds_output_summary = summary.parents[0]
+            ds_throughput = (
+                ds_output_summary.operators_stats[-1].output_num_rows.sum
+                / ds_output_summary.get_total_wall_time()
+            )
+
+            iter_stats = summary.iter_stats
+
+            metrics[f"dataloader/{ds_key}"] = {
+                "producer_throughput": ds_throughput,
+                "iter_stats": {
+                    "prefetch_block-avg": iter_stats.wait_time.avg(),
+                    "prefetch_block-min": iter_stats.wait_time.min(),
+                    "prefetch_block-max": iter_stats.wait_time.max(),
+                    "prefetch_block-total": iter_stats.wait_time.get(),
+                    "get_ref_bundles-avg": iter_stats.get_ref_bundles_time.avg(),
+                    "get_ref_bundles-min": iter_stats.get_ref_bundles_time.min(),
+                    "get_ref_bundles-max": iter_stats.get_ref_bundles_time.max(),
+                    "get_ref_bundles-total": iter_stats.get_ref_bundles_time.get(),
+                    "fetch_block-avg": iter_stats.get_time.avg(),
+                    "fetch_block-min": iter_stats.get_time.min(),
+                    "fetch_block-max": iter_stats.get_time.max(),
+                    "fetch_block-total": iter_stats.get_time.get(),
+                    "block_to_batch-avg": iter_stats.next_time.avg(),
+                    "block_to_batch-min": iter_stats.next_time.min(),
+                    "block_to_batch-max": iter_stats.next_time.max(),
+                    "block_to_batch-total": iter_stats.next_time.get(),
+                    "format_batch-avg": iter_stats.format_time.avg(),
+                    "format_batch-min": iter_stats.format_time.min(),
+                    "format_batch-max": iter_stats.format_time.max(),
+                    "format_batch-total": iter_stats.format_time.get(),
+                    "collate-avg": iter_stats.collate_time.avg(),
+                    "collate-min": iter_stats.collate_time.min(),
+                    "collate-max": iter_stats.collate_time.max(),
+                    "collate-total": iter_stats.collate_time.get(),
+                    "finalize-avg": iter_stats.finalize_batch_time.avg(),
+                    "finalize-min": iter_stats.finalize_batch_time.min(),
+                    "finalize-max": iter_stats.finalize_batch_time.max(),
+                    "finalize-total": iter_stats.finalize_batch_time.get(),
+                    "time_spent_blocked-avg": iter_stats.block_time.avg(),
+                    "time_spent_blocked-min": iter_stats.block_time.min(),
+                    "time_spent_blocked-max": iter_stats.block_time.max(),
+                    "time_spent_blocked-total": iter_stats.block_time.get(),
+                    "time_spent_training-avg": iter_stats.user_time.avg(),
+                    "time_spent_training-min": iter_stats.user_time.min(),
+                    "time_spent_training-max": iter_stats.user_time.max(),
+                    "time_spent_training-total": iter_stats.user_time.get(),
+                },
+            }
+
+        # Collect object store spill metrics.
+        try:
+            memory_info = get_memory_info_reply(
+                get_state_from_address(ray.get_runtime_context().gcs_address)
+            )
+            spilled_bytes_total = memory_info.store_stats.spilled_bytes_total
+            metrics["object_store_spilled_total_gb"] = round(
+                spilled_bytes_total / (1024**3), 4
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to collect object_store_spilled_total_gb metric: {e}"
+            )
+
+        # Collect peak and average spilling rate from the background monitor.
+        if self._spill_monitor is not None:
+            try:
+                spill_metrics = ray.get(self._spill_monitor.get_metrics.remote())
+                metrics.update(spill_metrics)
+            except Exception as e:
+                logger.warning(f"Failed to collect spill rate metrics: {e}")
+
+        return metrics

release/train_tests/benchmark/recsys/criteo.py

@@ -0,0 +1,292 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Dict, List, Tuple
+
+import boto3
+import json
+import numpy as np
+import pyarrow.csv
+
+import ray.data
+
+from constants import DatasetKey
+
+if TYPE_CHECKING:
+    from torchrec.datasets.utils import Batch
+
+logger = logging.getLogger(__name__)
+
+
+S3_BUCKET = "ray-benchmark-data-internal-us-west-2"
+CRITEO_S3_URI = f"s3://{S3_BUCKET}/criteo/tsv.gz"
+CAT_FEATURE_VALUE_COUNT_JSON_PATH_PATTERN = (
+    "criteo/tsv.gz/categorical_feature_value_counts/{}-value_counts.json"
+)
+
+
+INT_FEATURE_COUNT = 13
+CAT_FEATURE_COUNT = 26
+DAYS = 24
+DEFAULT_LABEL_NAME = "label"
+DEFAULT_INT_NAMES: List[str] = [f"int_{idx}" for idx in range(INT_FEATURE_COUNT)]
+DEFAULT_CAT_NAMES: List[str] = [f"cat_{idx}" for idx in range(CAT_FEATURE_COUNT)]
+DEFAULT_COLUMN_NAMES: List[str] = [
+    DEFAULT_LABEL_NAME,
+    *DEFAULT_INT_NAMES,
+    *DEFAULT_CAT_NAMES,
+]
+CRITEO_NUM_EMBEDDINGS_PER_FEATURE: List[int] = [
+    45833188,
+    36746,
+    17245,
+    7413,
+    20243,
+    3,
+    7114,
+    1441,
+    62,
+    29275261,
+    1572176,
+    345138,
+    10,
+    2209,
+    11267,
+    128,
+    4,
+    974,
+    14,
+    48937457,
+    11316796,
+    40094537,
+    452104,
+    12606,
+    104,
+    35,
+]
+
+DATASET_PATHS = {
+    DatasetKey.TRAIN: f"{CRITEO_S3_URI}/train",
+    DatasetKey.VALID: f"{CRITEO_S3_URI}/valid",
+    DatasetKey.TEST: f"{CRITEO_S3_URI}/test",
+}
+
+
+def fill_missing(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+    """Fill in missing feature values with defaults.
+    Default to 0 for dense features, empty string "" for categorical features.
+    """
+    for feature_name in DEFAULT_INT_NAMES:
+        batch[feature_name] = np.nan_to_num(batch[feature_name], nan=0)
+    for feature_name in DEFAULT_CAT_NAMES:
+        features = batch[feature_name]
+        features[np.equal(features, None)] = ""
+    return batch
+
+
+def concat_and_normalize_dense_features(
+    batch: Dict[str, np.ndarray],
+) -> Dict[str, np.ndarray]:
+    """Concatenate dense and sparse features together.
+    Apply log transformation to dense features."""
+
+    out = {}
+
+    out["dense"] = np.column_stack(
+        [batch[feature_name] for feature_name in DEFAULT_INT_NAMES]
+    )
+    out["sparse"] = np.column_stack(
+        [batch[feature_name] for feature_name in DEFAULT_CAT_NAMES]
+    )
+
+    out["dense"] += 3  # Prevent log(0)
+    out["dense"] = np.log(out["dense"], dtype=np.float32)
+    out["label"] = batch["label"]
+
+    return out
+
+
+def mock_dataloader(num_batches: int, batch_size: int):
+    """Creates a dummy batch of size `batch_size` and yields it `num_batches` times."""
+    dense = np.random.randn(batch_size, INT_FEATURE_COUNT).astype(np.float32)
+    sparse = np.random.randint(
+        1,
+        np.array(CRITEO_NUM_EMBEDDINGS_PER_FEATURE),
+        (batch_size, CAT_FEATURE_COUNT),
+    ).astype(np.int32)
+    labels = np.random.randint(0, 1, (batch_size,)).astype(np.int32)
+    batch = convert_to_torchrec_batch_format(
+        {"dense": dense, "sparse": sparse, "label": labels}
+    )
+    batch = batch.pin_memory()
+
+    for _ in range(num_batches):
+        yield batch
+
+
+def convert_to_torchrec_batch_format(batch: Dict[str, np.ndarray]) -> "Batch":
+    """Convert to a Batch, packaging sparse features as a KJT."""
+    import torch
+
+    from torchrec.datasets.utils import Batch
+    from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
+
+    dense = batch["dense"]
+    sparse = batch["sparse"]
+    labels = batch["label"]
+
+    batch_size = len(dense)
+    lengths = torch.ones((batch_size * CAT_FEATURE_COUNT,), dtype=torch.int32)
+    offsets = torch.arange(0, batch_size * CAT_FEATURE_COUNT + 1, dtype=torch.int32)
+    length_per_key: List[int] = [batch_size] * CAT_FEATURE_COUNT
+    offset_per_key = [batch_size * i for i in range(CAT_FEATURE_COUNT + 1)]
+    index_per_key = {key: i for i, key in enumerate(DEFAULT_CAT_NAMES)}
+
+    # Handle partial batches (last batch).
+    # if batch_size == self.batch_size:
+    #     length_per_key = self.length_per_key
+    #     offset_per_key = self.offset_per_key
+    # else:
+    #     # handle last batch in dataset when it's an incomplete batch.
+    #     length_per_key = CAT_FEATURE_COUNT * [batch_size]
+    #     offset_per_key = [batch_size * i for i in range(CAT_FEATURE_COUNT + 1)]
+
+    return Batch(
+        dense_features=torch.from_numpy(dense),
+        sparse_features=KeyedJaggedTensor(
+            keys=DEFAULT_CAT_NAMES,
+            # transpose().reshape(-1) introduces a copy
+            values=torch.from_numpy(sparse.transpose(1, 0).reshape(-1)),
+            lengths=lengths,
+            offsets=offsets,
+            stride=batch_size,
+            length_per_key=length_per_key,
+            offset_per_key=offset_per_key,
+            index_per_key=index_per_key,
+        ),
+        labels=torch.from_numpy(labels.reshape(-1)),
+    )
+
+
+def read_json_from_s3(bucket_name, key):
+    s3 = boto3.client("s3")
+
+    # Download object content
+    response = s3.get_object(Bucket=bucket_name, Key=key)
+    content = response["Body"].read().decode("utf-8")
+
+    # Parse JSON
+    data = json.loads(content)
+    return data
+
+
+def _get_base_dataset(stage: DatasetKey = DatasetKey.TRAIN):
+    ds_path = DATASET_PATHS[stage]
+
+    ds = ray.data.read_csv(
+        ds_path,
+        read_options=pyarrow.csv.ReadOptions(column_names=DEFAULT_COLUMN_NAMES),
+        parse_options=pyarrow.csv.ParseOptions(delimiter="\t"),
+        shuffle=(
+            "files" if stage == DatasetKey.TRAIN else None
+        ),  # coarse file-level shuffle
+    )
+    return ds
+
+
+def get_ray_dataset(stage: DatasetKey = DatasetKey.TRAIN):
+    ds = _get_base_dataset(stage)
+
+    # Convert categorical features to integers.
+
+    # Fetch cached value counts instead of "fitting" the preprocessor from scratch.
+    COMPUTE_VALUE_COUNTS_FROM_SCRATCH: bool = False
+
+    FREQUENCY_THRESHOLD = 3
+    LOW_FREQUENCY_INDEX = 1  # map low frequency values -> 1
+    categorical_value_to_index = {}
+    for cat_feature in DEFAULT_CAT_NAMES:
+        if COMPUTE_VALUE_COUNTS_FROM_SCRATCH:
+            value_counts = _compute_value_counts(ds, cat_feature)
+        else:
+            json_filepath = CAT_FEATURE_VALUE_COUNT_JSON_PATH_PATTERN.format(
+                cat_feature
+            )
+            logger.info(f"Downloading value counts file: {json_filepath}")
+            value_counts = read_json_from_s3(bucket_name=S3_BUCKET, key=json_filepath)
+
+        value_counts = filter(lambda x: x[1] >= FREQUENCY_THRESHOLD, value_counts)
+        categorical_value_to_index[cat_feature] = {
+            val: i for i, (val, _) in enumerate(value_counts, start=2)
+        }
+
+    # TODO: This will not scale well for the full dataset, since this dict might be 10s of GBs,
+    # which is expensive to copy to each map task.
+    # This mapping is large, so put a shared copy in the object store for all the map tasks to use.
+    categorical_value_to_index_ref = ray.put(categorical_value_to_index)
+
+    # Clean data.
+    ds = ds.map_batches(fill_missing)
+
+    def categorical_values_to_indices(
+        batch: Dict[str, np.ndarray], mapping_ref: ray.ObjectRef
+    ):
+        mapping: Dict[str, int] = ray.get(mapping_ref)
+        for cat_feature in DEFAULT_CAT_NAMES:
+            batch[cat_feature] = np.vectorize(
+                lambda k: mapping.get(cat_feature, {}).get(k, LOW_FREQUENCY_INDEX)
+            )(batch[cat_feature])
+        return batch
+
+    ds = ds.map_batches(
+        categorical_values_to_indices, fn_args=(categorical_value_to_index_ref,)
+    )
+
+    # HACK: Dummy encoding for quicker testing.
+    # def dummy_categorical_encoder(batch):
+    #     for feature_name in DEFAULT_CAT_NAMES:
+    #         batch[feature_name] = np.random.randint(0, 3, size=(len(batch[feature_name]),))
+    #     return batch
+    # ds = ds.map_batches(dummy_categorical_encoder)
+
+    ds = ds.map_batches(concat_and_normalize_dense_features)
+
+    # TODO: Need to shuffle the data.
+
+    return ds
+
+
+def _compute_value_counts(ds, feature_name) -> List[Tuple]:
+    logger.info(f"Computing value counts for: {feature_name}")
+
+    # TODO: This needs to be optimized in order to run on the full dataset.
+    # Need to fill missing values with empty string.
+    value_counts = [
+        (
+            group[feature_name] if group[feature_name] is not None else "",
+            group["count()"],
+        )
+        for group in (
+            ds.select_columns(feature_name).groupby(key=feature_name).count().take_all()
+        )
+    ]
+
+    return value_counts
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+
+    ds = _get_base_dataset(stage="train")
+
+    # Create a directory for the value counts files.
+    save_dir = "/mnt/cluster_storage/criteo"
+    os.makedirs(save_dir, exist_ok=True)
+
+    for cat_feature in DEFAULT_CAT_NAMES:
+        value_counts = _compute_value_counts(ds, cat_feature)
+
+        json_filepath = os.path.join(save_dir, f"{cat_feature}-value_counts.json")
+        logger.info(f"Writing value counts to: {json_filepath}")
+
+        with open(json_filepath, "w") as f:
+            json.dump(value_counts, f)

release/train_tests/benchmark/recsys/recsys_factory.py

@@ -0,0 +1,214 @@
+from typing import Dict, List, Optional
+import logging
+
+import numpy as np
+from pydantic import BaseModel
+import torch
+import torch.distributed as torch_dist
+
+from ray.data.collate_fn import CollateFn, NumpyBatchCollateFn
+import ray.train
+import ray.train.torch
+
+from constants import DatasetKey
+from config import DataloaderType, BenchmarkConfig
+from benchmark_factory import BenchmarkFactory
+from dataloader_factory import (
+    BaseDataLoaderFactory,
+)
+from logger_utils import ContextLoggerAdapter
+from ray_dataloader_factory import RayDataLoaderFactory
+from recsys.criteo import (
+    CRITEO_NUM_EMBEDDINGS_PER_FEATURE,
+    convert_to_torchrec_batch_format,
+    get_ray_dataset,
+    mock_dataloader,
+)
+
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class RecsysMockDataLoaderFactory(BaseDataLoaderFactory):
+    def get_train_dataloader(self):
+        return mock_dataloader(
+            2048, self.benchmark_config.dataloader_config.train_batch_size
+        )
+
+    def get_val_dataloader(self):
+        return mock_dataloader(
+            256, self.benchmark_config.dataloader_config.validation_batch_size
+        )
+
+
+class RecsysRayDataLoaderFactory(RayDataLoaderFactory):
+    def get_ray_datasets(self) -> Dict[str, ray.data.Dataset]:
+        # TODO: Use the train dataset for validation as well.
+        ds = get_ray_dataset(DatasetKey.VALID)
+        return {
+            DatasetKey.TRAIN: ds,
+            DatasetKey.VALID: ds,
+        }
+
+    def _get_collate_fn(self) -> Optional[CollateFn]:
+        from torchrec.datasets.utils import Batch
+
+        class TorchRecCollateFn(NumpyBatchCollateFn):
+            def __call__(self, batch: Dict[str, np.ndarray]) -> Batch:
+                return convert_to_torchrec_batch_format(batch)
+
+        return TorchRecCollateFn()
+
+
+class TorchRecConfig(BaseModel):
+    embedding_dim: int = 128
+    num_embeddings_per_feature: List[int] = CRITEO_NUM_EMBEDDINGS_PER_FEATURE
+    over_arch_layer_sizes: List[int] = [1024, 1024, 512, 256, 1]
+    dense_arch_layer_sizes: List[int] = [512, 256, 128]
+    interaction_type: str = "dcn"
+    dcn_num_layers: int = 3
+    dcn_low_rank_dim: int = 512
+
+
+class RecsysFactory(BenchmarkFactory):
+    def __init__(self, benchmark_config: BenchmarkConfig):
+        super().__init__(benchmark_config)
+
+        self.torchrec_config = TorchRecConfig()
+
+    def get_dataloader_factory(self) -> BaseDataLoaderFactory:
+        data_factory_cls = {
+            DataloaderType.MOCK: RecsysMockDataLoaderFactory,
+            DataloaderType.RAY_DATA: RecsysRayDataLoaderFactory,
+        }[self.benchmark_config.dataloader_type]
+
+        return data_factory_cls(self.benchmark_config)
+
+    def get_model(self) -> torch.nn.Module:
+        # NOTE: These imports error on a CPU-only driver node.
+        # Delay the import to happen on the GPU train workers instead.
+        from torchrec import EmbeddingBagCollection
+        from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES
+        from torchrec.distributed.model_parallel import (
+            DistributedModelParallel,
+            get_default_sharders,
+        )
+        from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
+        from torchrec.distributed.planner.storage_reservations import (
+            HeuristicalStorageReservation,
+        )
+        from torchrec.models.dlrm import DLRM, DLRM_DCN, DLRM_Projection, DLRMTrain
+        from torchrec.modules.embedding_configs import EmbeddingBagConfig
+        from torchrec.optim.apply_optimizer_in_backward import (
+            apply_optimizer_in_backward,
+        )
+
+        args = self.torchrec_config
+        device = ray.train.torch.get_device()
+        local_world_size = ray.train.get_context().get_local_world_size()
+        global_world_size = ray.train.get_context().get_world_size()
+
+        eb_configs = [
+            EmbeddingBagConfig(
+                name=f"t_{feature_name}",
+                embedding_dim=args.embedding_dim,
+                num_embeddings=args.num_embeddings_per_feature[feature_idx],
+                feature_names=[feature_name],
+            )
+            for feature_idx, feature_name in enumerate(DEFAULT_CAT_NAMES)
+        ]
+        sharded_module_kwargs = {}
+        if args.over_arch_layer_sizes is not None:
+            sharded_module_kwargs["over_arch_layer_sizes"] = args.over_arch_layer_sizes
+
+        if args.interaction_type == "original":
+            dlrm_model = DLRM(
+                embedding_bag_collection=EmbeddingBagCollection(
+                    tables=eb_configs, device=torch.device("meta")
+                ),
+                dense_in_features=len(DEFAULT_INT_NAMES),
+                dense_arch_layer_sizes=args.dense_arch_layer_sizes,
+                over_arch_layer_sizes=args.over_arch_layer_sizes,
+                dense_device=device,
+            )
+        elif args.interaction_type == "dcn":
+            dlrm_model = DLRM_DCN(
+                embedding_bag_collection=EmbeddingBagCollection(
+                    tables=eb_configs, device=torch.device("meta")
+                ),
+                dense_in_features=len(DEFAULT_INT_NAMES),
+                dense_arch_layer_sizes=args.dense_arch_layer_sizes,
+                over_arch_layer_sizes=args.over_arch_layer_sizes,
+                dcn_num_layers=args.dcn_num_layers,
+                dcn_low_rank_dim=args.dcn_low_rank_dim,
+                dense_device=device,
+            )
+        elif args.interaction_type == "projection":
+            raise NotImplementedError
+
+            dlrm_model = DLRM_Projection(
+                embedding_bag_collection=EmbeddingBagCollection(
+                    tables=eb_configs, device=torch.device("meta")
+                ),
+                dense_in_features=len(DEFAULT_INT_NAMES),
+                dense_arch_layer_sizes=args.dense_arch_layer_sizes,
+                over_arch_layer_sizes=args.over_arch_layer_sizes,
+                interaction_branch1_layer_sizes=args.interaction_branch1_layer_sizes,
+                interaction_branch2_layer_sizes=args.interaction_branch2_layer_sizes,
+                dense_device=device,
+            )
+        else:
+            raise ValueError(
+                "Unknown interaction option set. Should be original, dcn, or projection."
+            )
+
+        train_model = DLRMTrain(dlrm_model)
+        embedding_optimizer = torch.optim.Adagrad
+        # This will apply the Adagrad optimizer in the backward pass for the embeddings (sparse_arch). This means that
+        # the optimizer update will be applied in the backward pass, in this case through a fused op.
+        # TorchRec will use the FBGEMM implementation of EXACT_ADAGRAD. For GPU devices, a fused CUDA kernel is invoked. For CPU, FBGEMM_GPU invokes CPU kernels
+        # https://github.com/pytorch/FBGEMM/blob/2cb8b0dff3e67f9a009c4299defbd6b99cc12b8f/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py#L676-L678
+
+        # Note that lr_decay, weight_decay and initial_accumulator_value for Adagrad optimizer in FBGEMM v0.3.2
+        # cannot be specified below. This equivalently means that all these parameters are hardcoded to zero.
+        optimizer_kwargs = {"lr": 15.0, "eps": 1e-8}
+        apply_optimizer_in_backward(
+            embedding_optimizer,
+            train_model.model.sparse_arch.parameters(),
+            optimizer_kwargs,
+        )
+        planner = EmbeddingShardingPlanner(
+            topology=Topology(
+                local_world_size=local_world_size,
+                world_size=global_world_size,
+                compute_device=device.type,
+            ),
+            batch_size=self.benchmark_config.dataloader_config.train_batch_size,
+            # If experience OOM, increase the percentage. see
+            # https://pytorch.org/torchrec/torchrec.distributed.planner.html#torchrec.distributed.planner.storage_reservations.HeuristicalStorageReservation
+            storage_reservation=HeuristicalStorageReservation(percentage=0.05),
+        )
+        plan = planner.collective_plan(
+            train_model, get_default_sharders(), torch_dist.GroupMember.WORLD
+        )
+
+        model = DistributedModelParallel(
+            module=train_model,
+            device=device,
+            plan=plan,
+        )
+
+        if ray.train.get_context().get_world_rank() == 0:
+            for collectionkey, plans in model._plan.plan.items():
+                logger.info(collectionkey)
+                for table_name, plan in plans.items():
+                    logger.info(table_name)
+                    logger.info(plan)
+
+        return model
+
+    def get_loss_fn(self) -> torch.nn.Module:
+        raise NotImplementedError(
+            "torchrec model should return the loss directly in forward. "
+            "See the `DLRMTrain` wrapper class."
+        )

release/train_tests/benchmark/recsys/torchrec_runner.py

@@ -0,0 +1,141 @@
+import logging
+import gc
+import os
+
+import torch
+import torch.nn
+
+from torchrec.distributed.train_pipeline import StagedTrainPipeline, SparseDataDistUtil
+from torchrec.distributed.train_pipeline.utils import PipelineStage
+from torchrec.optim.keyed import CombinedOptimizer, KeyedOptimizerWrapper
+from torchrec.optim.optimizers import in_backward_optimizer_filter
+
+import ray.train
+import ray.train.torch
+
+from runner import TrainLoopRunner
+
+
+logger = logging.getLogger(__name__)
+
+
+class TorchRecRunner(TrainLoopRunner):
+    def _setup(self):
+        if self.factory.benchmark_config.mock_gpu:
+            raise ValueError("Mock GPU is not supported for running TorchRec.")
+
+        self.model = self.factory.get_model()
+
+        # TODO: This code depends on the model having a fused_optimizer,
+        # which is hidden in the `get_model` method of the factory.
+        dense_optimizer = KeyedOptimizerWrapper(
+            dict(in_backward_optimizer_filter(self.model.named_parameters())),
+            lambda params: torch.optim.Adagrad(params, lr=15.0, eps=1e-8),
+        )
+        self.optimizer = CombinedOptimizer(
+            [self.model.fused_optimizer, dense_optimizer]
+        )
+
+        self._data_dist_stream = torch.cuda.Stream()
+        self._h2d_stream = torch.cuda.Stream()
+
+    def _wrap_dataloader(self, dataloader, train: bool = True):
+        dataloader_iter = iter(dataloader)
+
+        device = ray.train.torch.get_device()
+        sdd = SparseDataDistUtil(
+            model=self.model,
+            data_dist_stream=self._data_dist_stream,
+            # prefetch_stream=torch.cuda.Stream(),
+        )
+        pipeline = [
+            PipelineStage(
+                name="data_copy",
+                runnable=lambda batch: batch.to(device, non_blocking=True),
+                stream=self._h2d_stream,
+            ),
+            PipelineStage(
+                name="start_sparse_data_dist",
+                runnable=sdd.start_sparse_data_dist,
+                stream=sdd.data_dist_stream,
+                fill_callback=sdd.wait_sparse_data_dist,
+            ),
+            # PipelineStage(
+            #     name="prefetch",
+            #     runnable=sdd.prefetch,
+            #     stream=sdd.prefetch_stream,
+            #     fill_callback=sdd.load_prefetch,
+            # ),
+        ]
+
+        pipeline = StagedTrainPipeline(pipeline_stages=pipeline)
+
+        def dataloader_with_torchrec_pipeline():
+            while batch := pipeline.progress(dataloader_iter):
+                yield batch
+
+            pipeline.flush_end()
+
+        return super()._wrap_dataloader(
+            dataloader_with_torchrec_pipeline(), train=train
+        )
+
+    def _train_step(self, batch):
+        self.model.train()
+
+        self.optimizer.zero_grad()
+        loss, out = self.model(batch)
+        loss.backward()
+        self.optimizer.step()
+
+    def _validate_step(self, batch):
+        self.model.eval()
+
+        with torch.no_grad():
+            loss, out = self.model(batch)
+
+        return loss
+
+    def _get_model_and_optim_filenames(self):
+        rank = ray.train.get_context().get_world_rank()
+        return f"model_shard_{rank=}.pt", f"optimizer_shard_{rank=}.pt"
+
+    def _save_training_state(self, local_dir: str):
+        # NOTE: Embedding table shards are on different GPUs,
+        # so we need to do distributed checkpointing.
+        # This checkpoint format must be loaded on the same number
+        # of workers and GPU types, since it was sharded with a compute-specific plan.
+        model_filename, optimizer_filename = self._get_model_and_optim_filenames()
+
+        torch.save(self.model.state_dict(), os.path.join(local_dir, model_filename))
+        torch.save(
+            self.optimizer.state_dict(), os.path.join(local_dir, optimizer_filename)
+        )
+
+    def _load_training_state(self, local_dir: str):
+        model_filename, optimizer_filename = self._get_model_and_optim_filenames()
+
+        self.model.load_state_dict(
+            torch.load(
+                os.path.join(local_dir, model_filename),
+                map_location=self.model.device,
+            )
+        )
+        self.optimizer.load_state_dict(
+            torch.load(
+                os.path.join(local_dir, optimizer_filename),
+                map_location=self.model.device,
+            )
+        )
+
+    def _cleanup(self):
+        # NOTE: This cleanup is needed to avoid zombie Train worker processes
+        # that hang on gc collect on python teardown.
+        del self.model
+        del self.optimizer
+        del self._data_dist_stream
+        del self._h2d_stream
+
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+        gc.collect()

release/train_tests/benchmark/runner.py

@@ -0,0 +1,443 @@
+import collections
+import json
+import logging
+import os
+import pprint
+import time
+import tempfile
+from typing import Dict, Optional
+
+import ray.train
+from ray.data._internal.stats import Timer
+import torch
+from logger_utils import ContextLoggerAdapter
+
+from benchmark_factory import BenchmarkFactory
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class TrainLoopRunner:
+    """Generic runner that sets up the training loop scaffolding.
+
+    Collects perf metrics and handles periodic checkpointing and validation.
+    """
+
+    def __init__(self, factory: BenchmarkFactory):
+        self.factory = factory
+        self.benchmark_config = factory.benchmark_config
+
+        self._setup()
+
+        # Training progress state.
+        self._train_batch_idx: int = 0
+        self._train_epoch_idx: int = 0
+        self._global_rows_processed_this_epoch: int = 0
+
+        # Performance metrics
+        self._metrics = collections.defaultdict(lambda: Timer())
+
+        checkpoint = ray.train.get_checkpoint()
+        if checkpoint:
+            self._restore_from_checkpoint(checkpoint)
+
+    # Methods for subclasses to implement.
+    def _setup(self):
+        """Subclasses should override this to setup the model, optimizer, etc.
+        The attributes initialized in this method should only be used in the
+        other overridden methods."""
+        pass
+
+    def _cleanup(self):
+        """Subclasses can override this to cleanup any resources."""
+        pass
+
+    def _train_step(self, train_dataloader):
+        """Subclasses should override this to implement the training step.
+        A training step represents a single forward and backward pass on a batch of data.
+        """
+        raise NotImplementedError
+
+    def _validate_step(self, val_dataloader):
+        """Subclasses should override this to implement the validation step.
+        A validation step represents a single forward pass on a batch of data."""
+        raise NotImplementedError
+
+    def _save_training_state(self, local_dir: str):
+        """Subclasses should override this to save the training state.
+        This should reference the model and optimizer state initialized
+        in the `_setup` method."""
+        pass
+
+    def _load_training_state(self, local_dir: str):
+        """Subclasses should override this to load the training state.
+        This should reference the model and optimizer state initialized
+        in the `_setup` method."""
+        pass
+
+    def _restore_from_checkpoint(self, checkpoint: ray.train.Checkpoint):
+        logger.info(
+            f"Restoring from checkpoint: {checkpoint} for worker "
+            f"{ray.train.get_context().get_world_rank()}"
+        )
+        with tempfile.TemporaryDirectory(
+            dir="/mnt/local_storage"
+        ) as temp_checkpoint_dir:
+            download_start = time.perf_counter()
+            checkpoint.to_directory(temp_checkpoint_dir)
+            download_time = time.perf_counter() - download_start
+
+            load_start = time.perf_counter()
+            self._load_checkpoint(temp_checkpoint_dir)
+            load_time = time.perf_counter() - load_start
+
+            self._metrics["checkpoint/download"].add(download_time)
+            self._metrics["checkpoint/load"].add(load_time)
+
+    def _wrap_dataloader(self, dataloader, train: bool = True):
+        dataloader_iter = iter(dataloader)
+
+        prefix = "train" if train else "validation"
+
+        def dataloader_with_timers():
+            try:
+                with self._metrics[f"{prefix}/iter_first_batch"].timer():
+                    batch = next(dataloader_iter)
+                    if train:
+                        self._train_batch_idx += 1
+            except StopIteration:
+                return
+
+            while True:
+                yield batch
+
+                try:
+                    with self._metrics[f"{prefix}/iter_batch"].timer():
+                        batch = next(dataloader_iter)
+                        if train:
+                            self._train_batch_idx += 1
+                except StopIteration:
+                    return
+
+        return dataloader_with_timers()
+
+    @property
+    def _num_batches_to_skip(self) -> int:
+        """Calculate the number of batches to skip based on the number of rows already processed in this epoch."""
+
+        global_batch_size = (
+            self.benchmark_config.dataloader_config.train_batch_size
+            * ray.train.get_context().get_world_size()
+        )
+
+        return self._global_rows_processed_this_epoch // global_batch_size
+
+    def _train_epoch(self):
+        """Subclasses can override the entrire `_train_epoch` method for more training
+        logic customization."""
+        if ray.train.get_context().get_world_rank() == 0:
+            logger.info(f"Training starting @ epoch={self._train_epoch_idx}")
+
+        train_dataloader = self.factory.get_train_dataloader()
+        train_dataloader = self._wrap_dataloader(train_dataloader, train=True)
+
+        # Skip through batches if we restored to a middle of the epoch.
+        # TODO: Compare this baseline to the data checkpointing approach once we have it.
+        if self._num_batches_to_skip:
+            if ray.train.get_context().get_world_rank() == 0:
+                logger.info(f"Skipping {self._num_batches_to_skip} batches...")
+
+            # Zero before the skip loop drives the wrapper, which would
+            # otherwise double-count against the value restored from the
+            # checkpoint. After the skip, _train_batch_idx is rebuilt to
+            # _num_batches_to_skip — matching the restored value.
+            self._train_batch_idx = 0
+
+            for _ in range(self._num_batches_to_skip):
+                with self._metrics["train/iter_skip_batch"].timer():
+                    next(train_dataloader)
+
+        for batch in train_dataloader:
+            with self._metrics["train/step"].timer():
+                if not self.benchmark_config.skip_train_step:
+                    self._train_step(batch)
+                if self.benchmark_config.train_step_sleep_s > 0:
+                    time.sleep(self.benchmark_config.train_step_sleep_s)
+
+            # TODO: This is slightly off if the last batch is a partial batch (if drop_last=False)
+            global_batch_size = (
+                self.benchmark_config.dataloader_config.train_batch_size
+                * ray.train.get_context().get_world_size()
+            )
+            self._metrics["train/rows_processed"].add(global_batch_size)
+
+            self._global_rows_processed_this_epoch += global_batch_size
+
+            if self._should_checkpoint_during_epoch():
+                self._checkpoint()
+
+            if self._should_validate_during_epoch():
+                validation_metrics = self._validate()
+                self._checkpoint(validation_metrics)
+
+            if self._should_log_metrics():
+                logger.info(pprint.pformat(self.get_metrics(), indent=2))
+
+            if (
+                self.benchmark_config.max_train_batches > 0
+                and self._train_batch_idx >= self.benchmark_config.max_train_batches
+            ):
+                break
+
+        self._train_epoch_idx += 1
+        self._train_batch_idx = 0
+        self._global_rows_processed_this_epoch = 0
+
+    def _validate_epoch(self) -> Dict[str, float]:
+        if ray.train.get_context().get_world_rank() == 0:
+            logger.info(
+                f"Validation starting @ epoch={self._train_epoch_idx}, "
+                f"batch={self._train_batch_idx}"
+            )
+
+        val_dataloader = self.factory.get_val_dataloader()
+        val_dataloader = self._wrap_dataloader(val_dataloader, train=False)
+
+        total_loss = torch.tensor(0.0).to(ray.train.torch.get_device())
+        num_rows = 0
+
+        for batch in val_dataloader:
+            with self._metrics["validation/step"].timer():
+                if not self.benchmark_config.skip_validation_step:
+                    total_loss += self._validate_step(batch)
+
+            num_rows += self.benchmark_config.dataloader_config.validation_batch_size
+            self._metrics["validation/rows_processed"].add(
+                self.benchmark_config.dataloader_config.validation_batch_size
+            )
+        assert num_rows > 0, "Validation dataset yielded no batches."
+
+        return {"validation/loss": total_loss.item() / num_rows}
+
+    def _should_checkpoint_during_epoch(self) -> bool:
+        """Handles the checkpoint_every_n_steps logic."""
+        return (
+            self.benchmark_config.checkpoint_every_n_steps > 0
+            and self._train_batch_idx % self.benchmark_config.checkpoint_every_n_steps
+            == 0
+        )
+
+    def _should_validate_during_epoch(self) -> bool:
+        """Handles the validate_every_n_steps logic."""
+        return (
+            self.benchmark_config.validate_every_n_steps > 0
+            and self._train_batch_idx % self.benchmark_config.validate_every_n_steps
+            == 0
+        )
+
+    def _should_log_metrics(self) -> bool:
+        """Handles the log_metrics_every_n_steps logic."""
+        return (
+            self.benchmark_config.log_metrics_every_n_steps > 0
+            and self._train_batch_idx % self.benchmark_config.log_metrics_every_n_steps
+            == 0
+        )
+
+    def _validate(self) -> Dict[str, float]:
+        with self._metrics["validation/epoch"].timer():
+            validation_metrics = self._validate_epoch()
+            return validation_metrics
+
+    def _checkpoint(self, metrics: Optional[Dict[str, float]] = None):
+        with tempfile.TemporaryDirectory(
+            dir="/mnt/local_storage"
+        ) as temp_checkpoint_dir:
+            with self._metrics["checkpoint/save"].timer():
+                self._save_checkpoint(temp_checkpoint_dir)
+
+            with self._metrics["checkpoint/report"].timer():
+                self._report_checkpoint(
+                    metrics=metrics or {},
+                    checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir),
+                )
+
+    def _load_checkpoint(self, local_dir: str):
+        self._load_training_state(local_dir)
+
+        run_state = torch.load(os.path.join(local_dir, "run_state.pt"))
+        self._train_epoch_idx = run_state["epoch"]
+        self._train_batch_idx = run_state["batch_idx"]
+        self._global_rows_processed_this_epoch = run_state[
+            "global_rows_processed_this_epoch"
+        ]
+
+        with open(os.path.join(local_dir, "metrics.json"), "r") as f:
+            metrics_json = json.load(f)
+
+        for k, v in metrics_json.items():
+            self._metrics[k].from_dict(v)
+
+        if ray.train.get_context().get_world_rank() == 0:
+            logger.info(
+                f"Restored to epoch={self._train_epoch_idx}, "
+                f"train_batch_idx={self._train_batch_idx} from checkpoint: "
+                f"{ray.train.get_checkpoint()}"
+            )
+
+    def _save_checkpoint(self, local_dir: str):
+        logger.info(
+            f"Saving checkpoint @ epoch={self._train_epoch_idx}, "
+            f"train_batch_idx={self._train_batch_idx}"
+        )
+
+        self._save_training_state(local_dir)
+
+        if ray.train.get_context().get_world_rank() == 0:
+            run_state = {
+                "epoch": self._train_epoch_idx,
+                "batch_idx": self._train_batch_idx,
+                "global_rows_processed_this_epoch": self._global_rows_processed_this_epoch,
+            }
+            torch.save(run_state, os.path.join(local_dir, "run_state.pt"))
+
+            metrics_json = {k: v.as_dict() for k, v in self._metrics.items()}
+            with open(os.path.join(local_dir, "metrics.json"), "w") as f:
+                json.dump(metrics_json, f)
+
+    def _report_checkpoint(self, metrics, checkpoint):
+        logger.info(
+            f"Uploading checkpoint @ epoch={self._train_epoch_idx}, "
+            f"train_batch_idx={self._train_batch_idx}"
+        )
+
+        checkpoint_dir_name = (
+            f"checkpoint_epoch={self._train_epoch_idx}_batch={self._train_batch_idx}"
+        )
+
+        ray.train.report(
+            metrics,
+            checkpoint=checkpoint,
+            checkpoint_dir_name=checkpoint_dir_name,
+        )
+
+    def run(self):
+        starting_epoch = self._train_epoch_idx
+
+        for _ in range(starting_epoch, self.benchmark_config.num_epochs):
+            with self._metrics["train/epoch"].timer():
+                self._train_epoch()
+
+            if not self.benchmark_config.skip_validation_at_epoch_end:
+                validation_metrics = self._validate()
+                self._checkpoint(validation_metrics)
+
+            if ray.train.get_context().get_world_rank() == 0:
+                logger.info(pprint.pformat(self.get_metrics(), indent=2))
+
+        self._cleanup()
+
+    def get_metrics(self, dataset_creation_time: float = 0.0) -> Dict[str, float]:
+        # TODO: These metrics should be aggregated across training workers.
+        metrics = {}
+        for key, metric in self._metrics.items():
+            metrics.update(
+                {
+                    f"{key}-avg": metric.avg(),
+                    f"{key}-min": metric.min(),
+                    f"{key}-max": metric.max(),
+                    f"{key}-total": metric.get(),
+                }
+            )
+
+        metrics["train/dataset_creation_time"] = dataset_creation_time
+        metrics["validation/dataset_creation_time"] = dataset_creation_time
+
+        # Throughput
+        # TODO: Ray Data can provide these throughput metrics automatically.
+        train_time = (
+            metrics["train/dataset_creation_time"]
+            + self._metrics["train/step"].get()
+            # Include the time it takes to get the first batch.
+            + self._metrics["train/iter_first_batch"].get()
+            + self._metrics["train/iter_batch"].get()
+        )
+        if train_time > 0:
+            metrics["train/global_throughput"] = (
+                self._metrics["train/rows_processed"].get() / train_time
+            )
+
+        validation_time = (
+            metrics["validation/dataset_creation_time"]
+            + self._metrics["validation/step"].get()
+            # Include the time it takes to get the first batch.
+            + self._metrics["validation/iter_first_batch"].get()
+            + self._metrics["validation/iter_batch"].get()
+        )
+        if validation_time > 0:
+            metrics["validation/global_throughput"] = (
+                self._metrics["validation/rows_processed"].get() / validation_time
+            )
+
+        # Extra time that each worker spends to restore from checkpoint,
+        # which includes downloading the checkpoint, loading the checkpoint,
+        # and skipping through batches that were already processed.
+        restoration_time = (
+            self._metrics["checkpoint/download"].get()
+            + self._metrics["checkpoint/load"].get()
+            + self._metrics["train/iter_skip_batch"].get()
+        )
+        if restoration_time > 0:
+            metrics["checkpoint/restoration_time"] = restoration_time
+
+        # Dataloader metrics (ex: Ray Data stats)
+        metrics.update(self.factory.get_dataloader_metrics())
+
+        return metrics
+
+
+class VanillaTorchRunner(TrainLoopRunner):
+    """A simple runner that uses a PyTorch model, optimizer, and loss function."""
+
+    def _setup(self):
+        model = self.factory.get_model()
+        self.model = ray.train.torch.prepare_model(model)
+        self.loss_fn = self.factory.get_loss_fn()
+        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3)
+
+    def _train_step(self, batch):
+        self.model.train()
+
+        input_batch, labels = batch
+
+        self.model.train()
+        self.optimizer.zero_grad()
+        out = self.model(input_batch)
+        loss = self.loss_fn(out, labels)
+        loss.backward()
+        self.optimizer.step()
+
+    def _validate_step(self, batch):
+        self.model.eval()
+
+        input_batch, labels = batch
+
+        with torch.no_grad():
+            out = self.model(input_batch)
+            loss = self.loss_fn(out, labels)
+        return loss
+
+    def _save_training_state(self, local_dir: str):
+        # Standard DDP checkpointing.
+        if ray.train.get_context().get_world_rank() == 0:
+            torch.save(self.model.state_dict(), os.path.join(local_dir, "model.pt"))
+            torch.save(
+                self.optimizer.state_dict(), os.path.join(local_dir, "optimizer.pt")
+            )
+
+    def _load_training_state(self, local_dir: str):
+        self.model.load_state_dict(
+            torch.load(os.path.join(local_dir, "model.pt"), map_location="cpu")
+        )
+        self.optimizer.load_state_dict(
+            torch.load(os.path.join(local_dir, "optimizer.pt"), map_location="cpu")
+        )

release/train_tests/benchmark/s3_jpeg_reader.py

@@ -0,0 +1,71 @@
+# Standard library imports
+from typing import List
+import logging
+
+# Third-party imports
+from botocore.exceptions import NoCredentialsError
+import ray
+import ray.train
+
+# Local imports
+from s3_reader import S3Reader
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class S3JpegReader(S3Reader):
+    """Extended S3Reader class for JPEG-specific functionality.
+
+    Provides specialized methods for:
+    1. Collecting JPEG file metadata (sizes) from S3
+    2. Distributing files among workers based on file sizes
+    3. Managing parallel S3 operations with Ray tasks
+    """
+
+    def _get_file_urls(self, url: str) -> List[str]:
+        """Get file URLs from S3 and distribute them among Ray workers.
+
+        Collects file metadata from S3 and distributes files among workers based on
+        file sizes to ensure balanced workload distribution.
+
+        Args:
+            url: S3 URL to list files from (e.g., "s3://bucket/path/to/directory")
+
+        Returns:
+            List of S3 URLs assigned to the current Ray worker
+
+        Raises:
+            S3CredentialsError: If AWS credentials are not found or invalid
+            S3FileError: If there's an error listing files from S3
+        """
+        try:
+            # Get Ray worker configuration
+            worker_rank = ray.train.get_context().get_world_rank()
+            num_workers = ray.train.get_context().get_world_size()
+
+            # Parse S3 URL components
+            bucket, prefix = self._parse_s3_url(url)
+
+            # Collect file metadata for balanced distribution
+            logger.info(
+                f"Worker {worker_rank}: Collecting file metadata for balanced distribution"
+            )
+            file_urls, file_size_bytes = self._list_s3_files(bucket, prefix)
+            logger.info(f"Found {len(file_urls)} files in {url}")
+
+            # Distribute files based on size
+            return self._distribute_files(
+                file_urls=file_urls,
+                file_weights=file_size_bytes,
+                worker_rank=worker_rank,
+                num_workers=num_workers,
+                weight_unit="bytes",
+            )
+
+        except NoCredentialsError:
+            raise self.S3CredentialsError(
+                "AWS credentials not found. Ensure you have configured them."
+            )
+        except Exception as e:
+            raise self.S3FileError(f"Error listing files from {url}: {str(e)}")

release/train_tests/benchmark/s3_parquet_reader.py

@@ -0,0 +1,155 @@
+# Standard library imports
+from typing import List, Tuple
+import logging
+
+# Third-party imports
+import boto3
+from botocore.exceptions import NoCredentialsError
+import ray
+import ray.train
+
+# Local imports
+from s3_reader import S3Reader, AWS_REGION
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+@ray.remote(num_cpus=0.25)
+def _fetch_parquet_metadata(bucket: str, key: str, file_url: str) -> Tuple[str, int]:
+    """Fetch Parquet row count using S3 Select in parallel.
+
+    Uses S3 Select to efficiently count rows in a Parquet file without reading
+    the entire file contents, enabling balanced distribution based on row counts.
+
+    Args:
+        bucket: S3 bucket name containing the Parquet file
+        key: S3 object key for the Parquet file
+        file_url: Full S3 URL for logging purposes
+
+    Returns:
+        Tuple containing:
+            - file_url: Full S3 URL of the processed file
+            - row_count: Number of rows in the Parquet file
+    """
+    s3_client = boto3.client("s3", region_name=AWS_REGION)
+
+    # Execute S3 Select query to count rows
+    query = "SELECT COUNT(*) FROM s3object"
+    response = s3_client.select_object_content(
+        Bucket=bucket,
+        Key=key,
+        Expression=query,
+        ExpressionType="SQL",
+        InputSerialization={"Parquet": {}},
+        OutputSerialization={"CSV": {}},
+    )
+
+    # Extract row count from response
+    row_count = 0
+    for event in response["Payload"]:
+        if "Records" in event:
+            row_count = int(event["Records"]["Payload"].decode("utf-8").strip())
+
+    logger.info(f"File {file_url} has {row_count} rows")
+    return file_url, row_count
+
+
+class S3ParquetReader(S3Reader):
+    """Extended S3Reader class for Parquet-specific functionality.
+
+    Provides specialized methods for:
+    1. Collecting Parquet file metadata (row counts) from S3
+    2. Distributing files among workers based on row counts
+    3. Managing parallel S3 operations with Ray tasks
+    """
+
+    def _collect_file_info(
+        self, bucket: str, prefix: str
+    ) -> Tuple[List[str], List[int]]:
+        """Collect file URLs and their row counts in parallel using Ray tasks.
+
+        Lists all Parquet files in the specified S3 prefix and launches parallel
+        tasks to count rows in each file using S3 Select for efficient metadata
+        collection.
+
+        Args:
+            bucket: S3 bucket name to list files from
+            prefix: S3 prefix to filter files
+
+        Returns:
+            Tuple containing:
+                - List of file URLs (e.g., "s3://bucket/path/to/file")
+                - List of row counts for each file
+        """
+        file_urls, _ = self._list_s3_files(bucket, prefix)
+
+        # Launch parallel metadata collection tasks
+        tasks = []
+        for file_url in file_urls:
+            # Extract key from file_url
+            key = file_url.replace(f"s3://{bucket}/", "")
+            task = _fetch_parquet_metadata.remote(bucket, key, file_url)
+            tasks.append(task)
+
+        # Wait for all tasks to complete
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(
+            f"Worker {worker_rank}: Waiting for metadata from {len(tasks)} files..."
+        )
+        results = ray.get(tasks)
+
+        # Process results
+        file_urls, file_rows = zip(*results) if results else ([], [])
+
+        logger.info(
+            f"Worker {worker_rank}: Collected metadata for {len(file_urls)} files"
+        )
+        return list(file_urls), list(file_rows)
+
+    def _get_file_urls(self, url: str) -> List[str]:
+        """Get file URLs from S3 and distribute them among Ray workers.
+
+        Collects file metadata from S3 and distributes files among workers based on
+        row counts to ensure balanced workload distribution.
+
+        Args:
+            url: S3 URL to list files from (e.g., "s3://bucket/path/to/directory")
+
+        Returns:
+            List of S3 URLs assigned to the current Ray worker
+
+        Raises:
+            S3CredentialsError: If AWS credentials are not found or invalid
+            S3FileError: If there's an error listing files from S3
+        """
+        try:
+            # Get Ray worker configuration
+            worker_rank = ray.train.get_context().get_world_rank()
+            num_workers = ray.train.get_context().get_world_size()
+
+            # Parse S3 URL components
+            bucket, prefix = self._parse_s3_url(url)
+
+            # Collect file metadata for balanced distribution
+            logger.info(
+                f"Worker {worker_rank}: Collecting file metadata for balanced distribution"
+            )
+            file_urls, file_rows = self._collect_file_info(bucket, prefix)
+            logger.info(f"Found {len(file_urls)} files in {url}")
+
+            # Distribute files based on row counts
+            return self._distribute_files(
+                file_urls=file_urls,
+                file_weights=file_rows,
+                worker_rank=worker_rank,
+                num_workers=num_workers,
+                weight_unit="rows",
+            )
+
+        except NoCredentialsError:
+            raise self.S3CredentialsError(
+                "AWS credentials not found. Ensure you have configured them."
+            )
+        except Exception as e:
+            raise self.S3FileError(f"Error listing files from {url}: {str(e)}")

release/train_tests/benchmark/s3_reader.py

@@ -0,0 +1,255 @@
+# Standard library imports
+from typing import Tuple, List, Optional
+import logging
+
+# Third-party imports
+import boto3
+import ray
+import ray.train
+
+# Local imports
+from logger_utils import ContextLoggerAdapter
+
+# AWS configuration
+AWS_REGION = "us-west-2"
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+@ray.remote(num_cpus=0.25)
+def _list_s3_batch(
+    bucket: str,
+    prefix: str,
+    continuation_token: Optional[str] = None,
+    batch_size: int = 1000,
+) -> Tuple[List[Tuple[str, int]], Optional[str]]:
+    """List a batch of files from S3 in parallel.
+
+    Makes a paginated request to S3's list_objects_v2 API to efficiently list files
+    in batches. Each file is returned with its size in bytes.
+
+    Args:
+        bucket: S3 bucket name to list files from
+        prefix: S3 prefix to filter files (e.g., "path/to/directory/")
+        continuation_token: Token from previous request for pagination
+        batch_size: Maximum number of files to return in one request (default: 1000)
+
+    Returns:
+        Tuple containing:
+            - List of (file_url, size) tuples, where:
+                - file_url: Full S3 URL (e.g., "s3://bucket/path/to/file")
+                - size: File size in bytes
+            - Optional[str]: Token for the next batch (None if no more files)
+    """
+    s3_client = boto3.client("s3")
+
+    # Prepare request parameters
+    list_params = {
+        "Bucket": bucket,
+        "Prefix": prefix,
+        "MaxKeys": batch_size,
+    }
+    if continuation_token:
+        list_params["ContinuationToken"] = continuation_token
+
+    # List objects from S3
+    response = s3_client.list_objects_v2(**list_params)
+
+    # Return empty results if no files found
+    if "Contents" not in response:
+        return [], None
+
+    # Extract file URLs and sizes
+    batch_files = response["Contents"]
+    results = [(f"s3://{bucket}/{f['Key']}", f["Size"]) for f in batch_files]
+
+    # Get token for next batch
+    next_token = (
+        response.get("NextContinuationToken") if response.get("IsTruncated") else None
+    )
+
+    return results, next_token
+
+
+class S3Reader:
+    """Base class for reading files from S3.
+
+    Provides common functionality for:
+    1. S3 client initialization and management
+    2. URL parsing and validation
+    3. File listing with pagination
+    4. Worker-based file distribution
+    5. Error handling for S3 operations
+    """
+
+    class S3Error(Exception):
+        """Base exception for S3-related errors."""
+
+        pass
+
+    class S3CredentialsError(S3Error):
+        """Raised when AWS credentials are not found or invalid."""
+
+        pass
+
+    class S3FileError(S3Error):
+        """Raised when there's an error accessing S3 files."""
+
+        pass
+
+    def __init__(self) -> None:
+        """Initialize the S3Reader with lazy client initialization."""
+        self._s3_client = None
+
+    @property
+    def s3_client(self) -> "boto3.client":
+        """Get or create the S3 client with AWS region configuration.
+
+        Uses lazy initialization to avoid serialization issues with Ray.
+
+        Returns:
+            boto3.client: Configured S3 client
+        """
+        if self._s3_client is None:
+            self._s3_client = boto3.client("s3", region_name=AWS_REGION)
+        return self._s3_client
+
+    def _parse_s3_url(self, s3_url: str) -> Tuple[str, str]:
+        """Parse an S3 URL into bucket and key components.
+
+        Args:
+            s3_url: S3 URL in format "s3://bucket/key"
+
+        Returns:
+            Tuple[str, str]: (bucket, key) components
+
+        Raises:
+            S3FileError: If URL is not a valid S3 URL
+        """
+        if not s3_url.startswith("s3://"):
+            raise self.S3FileError(f"Invalid S3 URL format: {s3_url}")
+        s3_parts = s3_url.replace("s3://", "").split("/", 1)
+        return s3_parts[0], s3_parts[1]
+
+    def _list_s3_files(self, bucket: str, prefix: str) -> Tuple[List[str], List[int]]:
+        """List files in an S3 bucket with the given prefix.
+
+        Uses Ray tasks to make parallel requests to S3's list_objects_v2 API,
+        handling pagination automatically. Returns file URLs and their sizes.
+
+        Args:
+            bucket: S3 bucket name
+            prefix: S3 prefix to filter files
+
+        Returns:
+            Tuple containing:
+                - List of file URLs (e.g., "s3://bucket/path/to/file")
+                - List of file sizes in bytes
+        """
+        file_urls = []
+        file_sizes = []
+        continuation_token = None
+        batch_size = 1000  # Maximum allowed by S3 API
+
+        while True:
+            # Get next batch of files
+            batch_results, next_token = ray.get(
+                _list_s3_batch.remote(
+                    bucket=bucket,
+                    prefix=prefix,
+                    continuation_token=continuation_token,
+                    batch_size=batch_size,
+                )
+            )
+
+            # Handle empty results
+            if not batch_results:
+                if not file_urls:  # Only warn on first request
+                    logger.info(
+                        f"No files found in s3://{bucket}/{prefix}", level="warning"
+                    )
+                break
+
+            # Process batch results
+            batch_urls, batch_sizes = zip(*batch_results)
+            file_urls.extend(batch_urls)
+            file_sizes.extend(batch_sizes)
+
+            # Log progress
+            logger.info(f"Listed {len(file_urls)} files from s3://{bucket}/{prefix}")
+
+            # Continue if there are more files
+            if not next_token:
+                break
+            continuation_token = next_token
+
+        return file_urls, file_sizes
+
+    def _distribute_files(
+        self,
+        file_urls: List[str],
+        file_weights: List[int],
+        worker_rank: int,
+        num_workers: int,
+        weight_unit: str = "units",
+    ) -> List[str]:
+        """Distribute files among workers based on weights.
+
+        Uses a greedy algorithm to distribute files among workers while trying to
+        minimize the difference in total weight between workers. Files are sorted
+        by weight (descending) before distribution for better balance.
+
+        Args:
+            file_urls: List of file URLs to distribute
+            file_weights: List of weights for each file (e.g., size, row count)
+            worker_rank: Current worker's rank
+            num_workers: Total number of workers
+            weight_unit: Unit of measurement for weights (e.g., "bytes", "rows")
+
+        Returns:
+            List of file URLs assigned to this worker
+        """
+        # Sort files by weight
+        files_with_weights = sorted(
+            zip(file_urls, file_weights), key=lambda x: x[1], reverse=True
+        )
+        file_urls = [f[0] for f in files_with_weights]
+        file_weights = [f[1] for f in files_with_weights]
+
+        # Handle single worker case
+        if num_workers <= 1 or not file_urls:
+            logger.info(
+                f"Worker {worker_rank}: Single worker or no files, "
+                f"returning all {len(file_urls)} files with total {sum(file_weights)} "
+                f"{weight_unit}"
+            )
+            return file_urls
+
+        # Calculate target weight per worker
+        total_weight = sum(file_weights)
+        target_weight_per_worker = total_weight / num_workers
+        logger.info(
+            f"Worker {worker_rank}: Total {weight_unit}: {total_weight}, "
+            f"Target per worker: {target_weight_per_worker:.0f} {weight_unit}"
+        )
+
+        # Initialize worker assignments
+        worker_files = [[] for _ in range(num_workers)]
+        worker_weights = [0] * num_workers
+
+        # Distribute files using greedy algorithm
+        for file_url, weight in zip(file_urls, file_weights):
+            min_weight_worker = min(range(num_workers), key=lambda w: worker_weights[w])
+            worker_files[min_weight_worker].append(file_url)
+            worker_weights[min_weight_worker] += weight
+
+        # Get this worker's assignment
+        my_files = worker_files[worker_rank]
+        my_weight = worker_weights[worker_rank]
+
+        logger.info(
+            f"Worker {worker_rank}: Assigned {len(my_files)}/{len(file_urls)} "
+            f"files with {my_weight}/{total_weight} {weight_unit} "
+            f"({my_weight/total_weight*100:.1f}%)"
+        )
+        return my_files

release/train_tests/benchmark/setup_chaos.py

@@ -0,0 +1 @@
+../../nightly_tests/setup_chaos.py

release/train_tests/benchmark/torch_dataloader_factory.py

@@ -0,0 +1,186 @@
+from typing import Dict, Iterator, Tuple
+import logging
+from abc import ABC, abstractmethod
+import multiprocessing
+
+import torch
+from torch.utils.data import IterableDataset
+
+import ray
+import ray.train
+
+from constants import DatasetKey
+from config import BenchmarkConfig, TorchConfig
+from dataloader_factory import BaseDataLoaderFactory
+from logger_utils import ContextLoggerAdapter
+
+logger = ContextLoggerAdapter(logging.getLogger(__name__))
+
+
+class TorchDataLoaderFactory(BaseDataLoaderFactory, ABC):
+    """Factory for creating PyTorch DataLoaders."""
+
+    @staticmethod
+    def worker_init_fn(worker_id: int):
+        """Initialize each worker with proper CUDA settings and seed.
+
+        Args:
+            worker_id: The ID of the worker being initialized
+        """
+        # Set worker-specific seed for reproducibility
+        worker_seed = torch.initial_seed() % 2**32
+        torch.manual_seed(worker_seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(worker_seed)
+            torch.cuda.manual_seed_all(worker_seed)
+
+        logger.info(f"Initialized worker {worker_id} with seed {worker_seed}")
+
+    def __init__(
+        self,
+        benchmark_config: BenchmarkConfig,
+    ):
+        """Initialize the factory.
+
+        Args:
+            benchmark_config: Configuration for the benchmark
+        """
+        super().__init__(benchmark_config)
+
+        dataloader_config = self.get_dataloader_config()
+        assert isinstance(dataloader_config, TorchConfig), type(dataloader_config)
+
+        # Get worker configuration
+        num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
+        self.num_torch_workers = dataloader_config.num_torch_workers
+        self.num_ray_workers = benchmark_config.num_workers
+
+        # Log configuration without worker rank since context may not be initialized
+        logger.info(
+            f"Configuration: {self.num_ray_workers * self.num_torch_workers} total workers "
+            f"({self.num_ray_workers} Ray × {self.num_torch_workers} Torch) "
+            f"across {num_gpus} GPUs"
+        )
+
+    def _get_device(self) -> torch.device:
+        """Get the device for the current worker using Ray Train's device management."""
+        try:
+            device = ray.train.torch.get_device()
+        except RuntimeError:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(f"Worker {worker_rank}: Using device: {device}")
+        return device
+
+    @abstractmethod
+    def create_batch_iterator(
+        self, dataloader: torch.utils.data.DataLoader, device: torch.device
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Create a safe iterator that handles device transfer and error handling.
+
+        Args:
+            dataloader: The PyTorch DataLoader to iterate over
+            device: The device to move tensors to
+
+        Returns:
+            An iterator that yields batches moved to the specified device
+        """
+        pass
+
+    @abstractmethod
+    def get_iterable_datasets(self) -> Dict[str, IterableDataset]:
+        """Get the train and validation datasets.
+
+        Returns:
+            A dictionary containing the train and validation datasets.
+        """
+        pass
+
+    def _create_multiprocessing_context(self):
+        # Importing libs in torch dataloader worker subprocesses is very slow.
+        # Preload some modules to speed up subprocess forking.
+        ctx = multiprocessing.get_context("forkserver")
+        modules = ["torch", "torchvision", "pandas", "numpy", "boto3", "fsspec"]
+        ctx.set_forkserver_preload(modules)
+        return ctx
+
+    def _create_dataloader(self, dataset_key: DatasetKey, batch_size: int):
+        worker_rank = ray.train.get_context().get_world_rank()
+        dataloader_config = self.get_dataloader_config()
+
+        # Create dataset and dataloader
+        ds = self.get_iterable_datasets()[dataset_key]
+
+        device = self._get_device()
+
+        # Adjust worker settings for 0 workers case
+        num_workers = max(0, self.num_torch_workers)
+        persistent_workers = num_workers > 0
+        pin_memory = dataloader_config.torch_pin_memory
+
+        if dataloader_config.torch_prefetch_factor >= 0:
+            prefetch_factor = dataloader_config.torch_prefetch_factor
+        else:
+            prefetch_factor = None
+
+        timeout = (
+            dataloader_config.torch_dataloader_timeout_seconds if num_workers > 0 else 0
+        )
+
+        logger.info(
+            f"Worker {worker_rank}: Creating train DataLoader with "
+            f"num_workers={num_workers}, pin_memory={pin_memory}, "
+            f"persistent_workers={persistent_workers}, prefetch_factor={prefetch_factor}, "
+            f"timeout={timeout}, batch_size={batch_size}"
+        )
+
+        multiprocessing_args = {}
+        if num_workers > 0:
+            multiprocessing_args = dict(
+                multiprocessing_context=self._create_multiprocessing_context(),
+                worker_init_fn=self.worker_init_fn,
+                persistent_workers=persistent_workers,
+            )
+        dataloader = torch.utils.data.DataLoader(
+            dataset=ds,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+            prefetch_factor=prefetch_factor,
+            timeout=timeout,
+            drop_last=False,
+            **multiprocessing_args,
+        )
+        # Add a DistributedSampler to the dataloader if possible (map-style datasets)
+        dataloader = ray.train.torch.prepare_data_loader(
+            dataloader, move_to_device=False
+        )
+
+        return self.create_batch_iterator(dataloader, device)
+
+    def get_train_dataloader(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Create a DataLoader for training data.
+
+        Returns:
+            An iterator that yields (image, label) tensors for training
+        """
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(f"Worker {worker_rank}: Creating train dataloader")
+
+        return self._create_dataloader(
+            DatasetKey.TRAIN, self.get_dataloader_config().train_batch_size
+        )
+
+    def get_val_dataloader(self) -> Iterator[Tuple[torch.Tensor, torch.Tensor]]:
+        """Create a DataLoader for validation data.
+
+        Returns:
+            An iterator that yields (image, label) tensors for validation
+        """
+        worker_rank = ray.train.get_context().get_world_rank()
+        logger.info(f"Worker {worker_rank}: Creating validation dataloader")
+
+        return self._create_dataloader(
+            DatasetKey.VALID, self.get_dataloader_config().validation_batch_size
+        )

release/train_tests/benchmark/train_benchmark.py

@@ -0,0 +1,123 @@
+import json
+import logging
+import os
+import pprint
+import time
+
+import ray.train
+from ray._private.test_utils import safe_write_to_results_json
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.util import date_str
+
+from config import BenchmarkConfig, cli_to_config
+from benchmark_factory import BenchmarkFactory
+from ray_dataloader_factory import RayDataLoaderFactory
+
+
+logger = logging.getLogger(__name__)
+
+
+METRICS_OUTPUT_PATH = "/mnt/cluster_storage/train_benchmark_metrics.json"
+
+
+def train_fn_per_worker(config):
+    factory: BenchmarkFactory = config["factory"]
+
+    if factory.benchmark_config.task == "recsys":
+        from recsys.torchrec_runner import TorchRecRunner
+
+        runner = TorchRecRunner(factory)
+    else:
+        from runner import VanillaTorchRunner
+
+        runner = VanillaTorchRunner(factory)
+
+    runner.run()
+
+    metrics = runner.get_metrics(
+        dataset_creation_time=config.get("dataset_creation_time", 0)
+    )
+    if ray.train.get_context().get_world_rank() == 0:
+        with open(METRICS_OUTPUT_PATH, "w") as f:
+            json.dump(metrics, f)
+
+
+def get_datasets_and_data_config(factory: BenchmarkFactory):
+    dataloader_factory = factory.get_dataloader_factory()
+    if isinstance(dataloader_factory, RayDataLoaderFactory):
+        datasets = dataloader_factory.get_ray_datasets()
+        data_config = dataloader_factory.get_ray_data_config()
+    else:
+        datasets = {}
+        data_config = None
+
+    return datasets, data_config
+
+
+def main():
+    start_time = time.perf_counter()
+    logging.basicConfig(level=logging.INFO)
+
+    benchmark_config: BenchmarkConfig = cli_to_config()
+    logger.info(
+        "\nBenchmark config:\n" + pprint.pformat(benchmark_config.__dict__, indent=2)
+    )
+
+    if benchmark_config.task == "image_classification":
+        from image_classification.factory import ImageClassificationFactory
+
+        factory = ImageClassificationFactory(benchmark_config)
+    elif benchmark_config.task == "recsys":
+        from recsys.recsys_factory import RecsysFactory
+
+        factory = RecsysFactory(benchmark_config)
+    else:
+        raise ValueError(f"Unknown task: {benchmark_config.task}")
+
+    datasets, data_config = get_datasets_and_data_config(factory)
+
+    dataset_creation_time = time.perf_counter() - start_time
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn_per_worker,
+        train_loop_config={
+            "factory": factory,
+            "dataset_creation_time": dataset_creation_time,
+        },
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=benchmark_config.num_workers,
+            use_gpu=not benchmark_config.mock_gpu,
+            resources_per_worker={"MOCK_GPU": 1} if benchmark_config.mock_gpu else None,
+        ),
+        run_config=ray.train.RunConfig(
+            storage_path=f"{os.environ['ANYSCALE_ARTIFACT_STORAGE']}/train_benchmark/",
+            name=f"{benchmark_config.task}-{date_str(include_ms=True)}",
+            failure_config=ray.train.FailureConfig(
+                max_failures=benchmark_config.max_failures
+            ),
+        ),
+        datasets=datasets,
+        dataset_config=data_config,
+    )
+
+    trainer.fit()
+
+    end_time = time.perf_counter()
+
+    with open(METRICS_OUTPUT_PATH, "r") as f:
+        metrics = json.load(f)
+
+    final_metrics_str = (
+        f"\nTotal training time: {end_time - start_time} seconds\n"
+        "Final metrics:\n" + "-" * 80 + "\n" + pprint.pformat(metrics) + "\n" + "-" * 80
+    )
+    logger.info(final_metrics_str)
+
+    # Write metrics as a release test result.
+    safe_write_to_results_json(metrics)
+
+
+if __name__ == "__main__":
+    # Workers need to access the working directory module.
+    ray.init(runtime_env={"working_dir": os.path.dirname(__file__)})
+    main()

release/train_tests/colocate_trainer/compute_aws.yaml

@@ -0,0 +1,22 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'
+
+head_node:
+    instance_type: g4dn.4xlarge
+    # GPU head — expose natural resources (16 vCPU + 1 T4 GPU) so the test
+    # (which colocates the trainer driver on the head) can use GPU.
+    resources:
+      CPU: 16
+      GPU: 1
+
+worker_nodes:
+    - instance_type: g4dn.xlarge
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND

release/train_tests/colocate_trainer/test_colocate_trainer.py

@@ -0,0 +1,75 @@
+"""Ray Train release test: Colocate Trainer and Rank 0 worker
+
+Setup:
+- 1 x g4dn.4xlarge (16 CPU, 1 GPU, 64 GB Memory)
+- 3 x g4dn.xlarge (4 CPU, 1 GPU, 16 GB memory)
+
+Test owner: woshiyyya
+"""
+
+import ray
+import ray.train
+import pytest
+
+from ray.train.data_parallel_trainer import DataParallelTrainer
+from ray.train.backend import Backend, BackendConfig
+from ray.train import ScalingConfig
+
+
+@pytest.mark.parametrize(
+    "trainer_resources", [None, {"memory": 40 * 1024**3}, {"CPU": 10}]
+)
+@pytest.mark.parametrize(
+    "resources_per_worker_and_use_gpu",
+    [
+        (None, True),
+        ({"CPU": 1}, False),
+        ({"GPU": 1}, True),
+    ],
+)
+def test_colocate_trainer_and_rank0_worker(
+    trainer_resources,
+    resources_per_worker_and_use_gpu,
+):
+    ray.init(ignore_reinit_error=True)
+
+    resources_per_worker, use_gpu = resources_per_worker_and_use_gpu
+
+    def train_func():
+        pass
+
+    class CustomBackend(Backend):
+        def on_training_start(self, worker_group, backend_config):
+            trainer_node_ip = ray.util.get_node_ip_address()
+
+            def check_node_ip():
+                if ray.train.get_context().get_world_rank() == 0:
+                    assert trainer_node_ip == ray.util.get_node_ip_address()
+
+            worker_group.execute(check_node_ip)
+
+    class CustomBackendConfig(BackendConfig):
+        @property
+        def backend_cls(self):
+            return CustomBackend
+
+    for num_workers in [1, 2, 4]:
+        scale_config = ScalingConfig(
+            num_workers=num_workers,
+            use_gpu=use_gpu,
+            trainer_resources=trainer_resources,
+            resources_per_worker=resources_per_worker,
+        )
+
+        trainer = DataParallelTrainer(
+            train_func,
+            scaling_config=scale_config,
+            backend_config=CustomBackendConfig(),
+        )
+        trainer.fit()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

release/train_tests/elastic_training/compute_gpu_aws.yaml

@@ -0,0 +1,30 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            VolumeSize: 800
+            Iops: 5000
+            Throughput: 1000
+            VolumeType: gp3
+    TagSpecifications:
+        - ResourceType: "instance"
+          Tags:
+            - Key: chaos-test-name
+              Value: "train-chaos-test"
+
+head_node:
+    instance_type: g4dn.12xlarge
+    # GPU head — expose natural resources (48 vCPU + 4 T4 GPU) so elastic
+    # training can use the head as one of its GPU nodes.
+    resources:
+      CPU: 48
+      GPU: 4
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 0
+      max_nodes: 3
+      market_type: ON_DEMAND

release/train_tests/elastic_training/elastic_util.py

@@ -0,0 +1,68 @@
+import subprocess
+
+import requests
+from torch import nn
+
+import ray
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28 * 28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+
+def terminate_current_instance():
+    """Use AWS CLI to terminate current instance."""
+
+    token = requests.put(
+        "http://169.254.169.254/latest/api/token",
+        headers={"X-aws-ec2-metadata-token-ttl-seconds": "300"},
+        timeout=10,
+    ).text
+    instance_id = requests.get(
+        "http://169.254.169.254/latest/meta-data/instance-id",
+        headers={"X-aws-ec2-metadata-token": token},
+        timeout=10,
+    ).text
+    region = requests.get(
+        "http://169.254.169.254/latest/meta-data/placement/region",
+        headers={"X-aws-ec2-metadata-token": token},
+        timeout=10,
+    ).text
+    return subprocess.run(
+        [
+            "aws",
+            "ec2",
+            "terminate-instances",
+            "--instance-ids",
+            instance_id,
+            "--region",
+            region,
+        ],
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+
+
+def terminate_node(node_id: str):
+    killer_task = ray.remote(terminate_current_instance).options(
+        num_cpus=0,
+        scheduling_strategy=NodeAffinitySchedulingStrategy(node_id, soft=False),
+    )
+    ray.get(killer_task.remote())

release/train_tests/elastic_training/torch_elastic_e2e.py

@@ -0,0 +1,373 @@
+import logging
+import os
+import tempfile
+import time
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import click
+from elastic_util import NeuralNetwork, terminate_node
+from filelock import FileLock
+import ray
+import ray.train as train
+from ray.tune.utils import date_str
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import ToTensor
+
+
+logger = logging.getLogger(__name__)
+
+
+CONFIG = {"lr": 1e-3, "batch_size": 64}
+LOG_FILE = "/tmp/driver.log"
+DATA_DIR = "/tmp/fashion_mnist"
+
+
+def get_default_storage_path():
+    remote_default_artifact_storage_prefix = os.environ.get(
+        "ANYSCALE_ARTIFACT_STORAGE", "artifact_storage"
+    )
+    return f"{remote_default_artifact_storage_prefix}/train_release_tests/elastic_e2e"
+
+
+STORAGE_PATH = get_default_storage_path()
+
+
+def load_data(data_dir):
+    with FileLock(f"{DATA_DIR}.data.lock"):
+        trainset = datasets.FashionMNIST(
+            root=data_dir, train=True, download=True, transform=ToTensor()
+        )
+        testset = datasets.FashionMNIST(
+            root=data_dir, train=False, download=True, transform=ToTensor()
+        )
+    return trainset, testset
+
+
+def train_epoch(
+    dataloader, model, loss_fn, optimizer, world_size: int, world_rank: int
+):
+    size = len(dataloader.dataset) // world_size
+    model.train()
+    for batch_index, (inputs, labels) in enumerate(dataloader):
+        predictions = model(inputs)
+        loss = loss_fn(predictions, labels)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        if batch_index % 100 == 0:
+            current = batch_index * len(inputs)
+            print(
+                f"[rank={world_rank}] loss: {loss.item():>7f}  [{current:>5d}/{size:>5d}]"
+            )
+
+
+def validate_epoch(dataloader, model, loss_fn, world_size: int, world_rank: int):
+    size = len(dataloader.dataset) // world_size
+    num_batches = len(dataloader)
+    model.eval()
+    test_loss, correct = 0, 0
+    with torch.no_grad():
+        for inputs, labels in dataloader:
+            predictions = model(inputs)
+            test_loss += loss_fn(predictions, labels).item()
+            correct += (predictions.argmax(1) == labels).type(torch.float).sum().item()
+    test_loss /= num_batches
+    correct /= size
+    print(
+        f"[rank={world_rank}] Test Error: \n "
+        f"Accuracy: {(100 * correct):>0.1f}%, "
+        f"Avg loss: {test_loss:>8f} \n"
+    )
+    return test_loss
+
+
+def save_checkpoint(local_dir, model, optimizer, epoch):
+    checkpoint = {
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "epoch": epoch,
+    }
+    torch.save(checkpoint, os.path.join(local_dir, "checkpoint.pt"))
+
+
+def load_checkpoint(local_ckpt_path, model, optimizer) -> int:
+    checkpoint = torch.load(os.path.join(local_ckpt_path, "checkpoint.pt"))
+    model.load_state_dict(checkpoint["model"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    return checkpoint["epoch"] + 1
+
+
+def train_func(config: Dict):
+    local_start_time = time.monotonic()
+
+    batch_size = config["batch_size"]
+    lr = config["lr"]
+    epochs = config["epochs"]
+    shuffle = config.get("shuffle", False)
+
+    world_size = train.get_context().get_world_size()
+    world_rank = train.get_context().get_world_rank()
+
+    worker_batch_size = batch_size // world_size
+    if world_rank == 0:
+        print(f"global batch size is {worker_batch_size * world_size}")
+
+    training_data, test_data = load_data(DATA_DIR)
+
+    train_dataloader = DataLoader(
+        training_data, shuffle=shuffle, batch_size=worker_batch_size
+    )
+    test_dataloader = DataLoader(
+        test_data, shuffle=shuffle, batch_size=worker_batch_size
+    )
+
+    train_dataloader = train.torch.prepare_data_loader(train_dataloader)
+    test_dataloader = train.torch.prepare_data_loader(test_dataloader)
+
+    model = train.torch.prepare_model(NeuralNetwork())
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+
+    start_epoch = 1
+    checkpoint = ray.train.get_checkpoint()
+    if checkpoint:
+        with checkpoint.as_directory() as temp_ckpt_dir:
+            print("Found checkpoint: ", checkpoint)
+            start_epoch = load_checkpoint(temp_ckpt_dir, model, optimizer)
+            print(f"Restoration done! Resuming training from {start_epoch=}")
+
+    for epoch in range(start_epoch, epochs + 1):
+        if world_size > 1:
+            train_dataloader.sampler.set_epoch(epoch)
+
+        train_epoch(
+            train_dataloader,
+            model,
+            loss_fn,
+            optimizer,
+            world_size=world_size,
+            world_rank=world_rank,
+        )
+        loss = validate_epoch(
+            test_dataloader,
+            model,
+            loss_fn,
+            world_size=world_size,
+            world_rank=world_rank,
+        )
+
+        local_time_taken = time.monotonic() - local_start_time
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            checkpoint = None
+            if world_rank == 0:
+                print("Saving checkpoint...")
+                save_checkpoint(temp_checkpoint_dir, model, optimizer, epoch)
+                checkpoint = train.Checkpoint.from_directory(temp_checkpoint_dir)
+
+            train.report(
+                metrics={"loss": loss, "local_time_taken": local_time_taken},
+                checkpoint=checkpoint,
+                checkpoint_dir_name=f"checkpoint-epoch={epoch}",
+            )
+
+
+def train_torch_ray_train(
+    config: dict,
+    num_workers: Tuple[int, int] = (4, 12),
+    use_gpu: bool = True,
+) -> train.Result:
+    from ray.train.torch import TorchTrainer
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=lambda c: train_func(config=c),
+        train_loop_config=config,
+        scaling_config=ray.train.ScalingConfig(
+            num_workers=num_workers, use_gpu=use_gpu
+        ),
+        run_config=ray.train.RunConfig(
+            name=f"elastic_train_experiment-{date_str()}",
+            storage_path=STORAGE_PATH,
+            checkpoint_config=ray.train.CheckpointConfig(num_to_keep=2),
+            failure_config=ray.train.FailureConfig(max_failures=3),
+        ),
+    )
+    return trainer.fit()
+
+
+@ray.remote(num_cpus=0)
+def run_cluster_node_killing_events(target_gpu_count: int):
+    logging.basicConfig(level=logging.INFO)
+    terminator_logger = logging.getLogger(__name__)
+    terminator_logger.addHandler(get_file_handler())
+
+    start = time.time()
+    head_node_id = ray.get_runtime_context().get_node_id()
+
+    def get_cluster_resources() -> Dict[str, float]:
+        return {
+            resource: value
+            for resource, value in ray.cluster_resources().items()
+            if resource in ("CPU", "GPU")
+        }
+
+    def get_worker_nodes() -> List[Dict]:
+        return [
+            node
+            for node in ray.nodes()
+            if node["Alive"] and node["NodeID"] != head_node_id
+        ]
+
+    def kill_nodes(nodes_to_kill):
+        terminator_logger.info(
+            "Nodes to kill: %s", [n["NodeID"] for n in nodes_to_kill]
+        )
+        for node in nodes_to_kill:
+            terminator_logger.info(
+                "Killing node: %s (alive=%s)", node["NodeID"], node["Alive"]
+            )
+            terminate_node(node["NodeID"])
+
+    def all_nodes_dead(dying_nodes) -> bool:
+        dying_node_ids = [n["NodeID"] for n in dying_nodes]
+        return all(
+            not node["Alive"]
+            for node in ray.nodes()
+            if node["NodeID"] in dying_node_ids
+        )
+
+    def log_status(message):
+        elapsed = time.time() - start
+        status_str = "\n"
+        status_str += "-" * 80 + "\n"
+        status_str += (
+            f"[elapsed={elapsed:.1f}s] cluster_resources={get_cluster_resources()}\n"
+        )
+        status_str += message + "\n"
+        status_str += "-" * 80 + "\n\n"
+        terminator_logger.info(status_str)
+
+    log_status(f"Waiting to upscale back to {target_gpu_count} GPUs...")
+    while get_cluster_resources().get("GPU", 0) < target_gpu_count:
+        time.sleep(1)
+
+    log_status("Waiting for 30s before modifying cluster resources...")
+    time.sleep(30)
+
+    log_status("Killing all nodes in the current cluster...")
+    nodes_to_kill = get_worker_nodes()
+    kill_nodes(nodes_to_kill)
+    while not all_nodes_dead(nodes_to_kill):
+        time.sleep(1)
+
+    log_status(f"Waiting to upscale back to {target_gpu_count} GPUs...")
+    while get_cluster_resources().get("GPU", 0) < target_gpu_count:
+        time.sleep(1)
+
+    log_status("Waiting for 30s before modifying cluster resources...")
+    time.sleep(30)
+
+    log_status("Killing two worker nodes...")
+    nodes_to_kill = get_worker_nodes()[-2:]
+    kill_nodes(nodes_to_kill)
+    while not all_nodes_dead(nodes_to_kill):
+        time.sleep(1)
+
+    log_status(f"Waiting to upscale back to {target_gpu_count} GPUs...")
+    while get_cluster_resources().get("GPU", 0) < target_gpu_count:
+        time.sleep(1)
+
+    log_status("Waiting for 30s before modifying cluster resources...")
+    time.sleep(30)
+
+    log_status("Killing 1 worker node...")
+    nodes_to_kill = [get_worker_nodes()[-1]]
+    kill_nodes(nodes_to_kill)
+    while not all_nodes_dead(nodes_to_kill):
+        time.sleep(1)
+
+    log_status("All node killing events generated, waiting for training finish...")
+
+
+@click.group(help="Run Torch benchmarks")
+def cli():
+    pass
+
+
+@cli.command(help="Kick off Ray Train elastic benchmark")
+@click.option("--num-epochs", type=int, default=50)
+@click.option("--num-workers", type=tuple, default=(4, 12))
+@click.option("--use-gpu", is_flag=True, default=True)
+@click.option("--batch-size", type=int, default=64)
+def run(
+    num_epochs: int = 50,
+    num_workers: Tuple[int, int] = (4, 12),
+    use_gpu: bool = True,
+    batch_size: int = 64,
+):
+    config = CONFIG.copy()
+    config["epochs"] = num_epochs
+    config["batch_size"] = batch_size
+
+    ray.init(log_to_driver=True, runtime_env={"working_dir": os.path.dirname(__file__)})
+
+    head_node_id = ray.get_runtime_context().get_node_id()
+    event_future = run_cluster_node_killing_events.options(
+        scheduling_strategy=NodeAffinitySchedulingStrategy(
+            node_id=head_node_id, soft=False
+        ),
+        runtime_env={"env_vars": {"RAY_TRAIN_V2_ENABLED": "1"}},
+    ).remote(target_gpu_count=num_workers[1])
+
+    result = train_torch_ray_train(
+        config=config,
+        num_workers=num_workers,
+        use_gpu=use_gpu,
+    )
+    ray.get(event_future)
+
+    logger.info(
+        "`trainer.fit` finished with (error, checkpoint):\nerror = %s\ncheckpoint = %s",
+        result.error,
+        result.checkpoint,
+    )
+    assert not result.error, result.error
+    assert result.checkpoint
+
+    checkpoint_dir_name = Path(result.checkpoint.path).name
+    expected_checkpoint_dir_name = f"checkpoint-epoch={num_epochs}"
+    assert (
+        checkpoint_dir_name == expected_checkpoint_dir_name
+    ), f"{checkpoint_dir_name=} != {expected_checkpoint_dir_name=}"
+
+    with open(LOG_FILE, "r") as f:
+        print(f.read())
+
+
+def get_file_handler() -> logging.FileHandler:
+    handler = logging.FileHandler(LOG_FILE)
+    handler.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(asctime)s [%(levelname)s] :: %(message)s")
+    handler.setFormatter(formatter)
+    return handler
+
+
+def setup_logging():
+    file_handler = get_file_handler()
+    logger.addHandler(file_handler)
+    logging.getLogger("ray.train").addHandler(file_handler)
+
+
+def main():
+    setup_logging()
+    cli()
+
+
+if __name__ == "__main__":
+    main()

release/train_tests/horovod/compute_tpl_aws.yaml

@@ -0,0 +1,23 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west-2
+
+max_workers: 1
+
+head_node_type:
+    name: head_node
+    # 4 cpus, 16G mem, $0.224/hr on demand
+    instance_type: m5.xlarge
+
+worker_node_types:
+    - name: worker_node
+      instance_type: m5.xlarge
+      max_workers: 1
+      min_workers: 1
+      use_spot: false
+
+advanced_configurations_json:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'

release/train_tests/horovod/compute_tpl_gce.yaml

@@ -0,0 +1,25 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-west1
+allowed_azs:
+    - us-west1-b
+
+max_workers: 1
+
+head_node_type:
+    name: head_node
+    # 4 cpus, 16G mem, $0.224/hr on demand
+    instance_type: n1-standard-4
+
+worker_node_types:
+    - name: worker_node
+      instance_type: n1-standard-4
+      max_workers: 1
+      min_workers: 1
+      use_spot: false
+
+#advanced_configurations_json:
+#  TagSpecifications:
+#    - ResourceType: "instance"
+#      Tags:
+#        - Key: ttl-hours
+#          Value: '24'

release/train_tests/horovod/train_horovod_multi_node_test.py

@@ -0,0 +1,37 @@
+import json
+import os
+import time
+
+import ray
+from ray.train import ScalingConfig
+from ray.air.constants import TRAINING_ITERATION
+from ray.train.examples.horovod.horovod_example import (
+    train_func as horovod_torch_train_func,
+)
+from ray.train.horovod.horovod_trainer import HorovodTrainer
+
+if __name__ == "__main__":
+    ray.init(address=os.environ.get("RAY_ADDRESS", "auto"))
+    start_time = time.time()
+
+    num_workers = 8
+    num_epochs = 10
+    trainer = HorovodTrainer(
+        horovod_torch_train_func,
+        train_loop_config={"num_epochs": num_epochs, "lr": 1e-3},
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            trainer_resources={"CPU": 0},
+        ),
+    )
+    results = trainer.fit()
+    result = results.metrics
+    assert result[TRAINING_ITERATION] == num_epochs
+
+    loss = list(results.metrics_dataframe["loss"])
+    assert len(loss) == num_epochs
+    assert loss[-1] < loss[0]
+
+    delta = time.time() - start_time
+    with open(os.environ["TEST_OUTPUT_JSON"], "w") as f:
+        f.write(json.dumps({"train_time": delta, "success": True}))

release/train_tests/huggingface_accelerate/compute_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/train_tests/huggingface_accelerate/test_huggingface_accelerate.py

@@ -0,0 +1,170 @@
+import tempfile
+
+import torch
+import evaluate
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AdamW,
+    get_linear_schedule_with_warmup,
+)
+from accelerate import Accelerator
+
+import ray
+import ray.train
+from ray.train import Checkpoint, ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+def train_func():
+    # Instantiate the accelerator
+    accelerator = Accelerator()
+
+    # Datasets
+    dataset = load_dataset("yelp_review_full")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def tokenize_function(examples):
+        outputs = tokenizer(examples["text"], padding="max_length", truncation=True)
+        outputs["labels"] = examples["label"]
+        return outputs
+
+    small_train_dataset = (
+        dataset["train"].select(range(100)).map(tokenize_function, batched=True)
+    )
+    small_eval_dataset = (
+        dataset["test"].select(range(100)).map(tokenize_function, batched=True)
+    )
+
+    # Remove unwanted columns and convert datasets to PyTorch format
+    columns_to_remove = [
+        "text",
+        "label",
+    ]  # Remove original columns, keep tokenized ones
+    small_train_dataset = small_train_dataset.remove_columns(columns_to_remove)
+    small_eval_dataset = small_eval_dataset.remove_columns(columns_to_remove)
+
+    small_train_dataset.set_format("torch")
+    small_eval_dataset.set_format("torch")
+
+    # Create data loaders
+    train_dataloader = torch.utils.data.DataLoader(
+        small_train_dataset, batch_size=16, shuffle=True
+    )
+    eval_dataloader = torch.utils.data.DataLoader(
+        small_eval_dataset, batch_size=16, shuffle=False
+    )
+
+    # Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", num_labels=5
+    )
+
+    # Optimizer and scheduler
+    optimizer = AdamW(model.parameters(), lr=2e-5)
+
+    num_training_steps = len(train_dataloader) * 3  # 3 epochs
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=num_training_steps,
+    )
+
+    # Prepare everything for distributed training
+    (
+        model,
+        optimizer,
+        train_dataloader,
+        eval_dataloader,
+        lr_scheduler,
+    ) = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # Evaluation metric
+    metric = evaluate.load("accuracy")
+
+    # Start training
+    num_epochs = 3
+
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        total_loss = 0
+
+        for batch in train_dataloader:
+            outputs = model(**batch)
+            loss = outputs.loss
+            accelerator.backward(loss)
+
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+            total_loss += loss.item()
+
+        # Evaluation
+        model.eval()
+        for batch in eval_dataloader:
+            with torch.no_grad():
+                outputs = model(**batch)
+
+            predictions = outputs.logits.argmax(dim=-1)
+            predictions, references = accelerator.gather_for_metrics(
+                (predictions, batch["labels"])
+            )
+            metric.add_batch(predictions=predictions, references=references)
+
+        eval_results = metric.compute()
+        accelerator.print(f"Epoch {epoch + 1}: {eval_results}")
+
+        # Report metrics and checkpoint to Ray Train
+        metrics = {
+            "epoch": epoch + 1,
+            "train_loss": total_loss / len(train_dataloader),
+            "eval_accuracy": eval_results["accuracy"],
+        }
+
+        # Create checkpoint
+        with tempfile.TemporaryDirectory() as tmpdir:
+            if accelerator.is_main_process:
+                unwrapped_model = accelerator.unwrap_model(model)
+                unwrapped_model.save_pretrained(tmpdir)
+                tokenizer.save_pretrained(tmpdir)
+                checkpoint = Checkpoint.from_directory(tmpdir)
+            else:
+                checkpoint = None
+
+            ray.train.report(metrics=metrics, checkpoint=checkpoint)
+
+
+def test_huggingface_accelerate():
+    # Define a Ray TorchTrainer to launch `train_func` on all workers
+    trainer = TorchTrainer(
+        train_func,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        run_config=ray.train.RunConfig(
+            storage_path="/mnt/cluster_storage/huggingface_accelerate_run"
+        ),
+    )
+    result: ray.train.Result = trainer.fit()
+
+    # Verify training completed successfully
+    assert result.metrics is not None
+    assert "eval_accuracy" in result.metrics
+    assert result.checkpoint is not None
+
+    # Load the trained model from checkpoint
+    with result.checkpoint.as_directory() as checkpoint_dir:
+        model = AutoModelForSequenceClassification.from_pretrained(  # noqa: F841
+            checkpoint_dir
+        )
+        tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir)  # noqa: F841
+
+
+if __name__ == "__main__":
+    test_huggingface_accelerate()

release/train_tests/huggingface_transformers/compute_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/train_tests/huggingface_transformers/test_huggingface_transformers.py

@@ -0,0 +1,104 @@
+import os
+
+import numpy as np
+import evaluate
+from datasets import load_dataset
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+)
+
+import ray.train.huggingface.transformers
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+# [1] Encapsulate data preprocessing, training, and evaluation
+# logic in a training function
+# ============================================================
+def train_func():
+    # Datasets
+    dataset = load_dataset("yelp_review_full")
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    small_train_dataset = (
+        dataset["train"].select(range(100)).map(tokenize_function, batched=True)
+    )
+    small_eval_dataset = (
+        dataset["test"].select(range(100)).map(tokenize_function, batched=True)
+    )
+
+    # Model
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "bert-base-cased", num_labels=5
+    )
+
+    # Evaluation Metrics
+    metric = evaluate.load("accuracy")
+
+    def compute_metrics(eval_pred):
+        logits, labels = eval_pred
+        predictions = np.argmax(logits, axis=-1)
+        return metric.compute(predictions=predictions, references=labels)
+
+    # Hugging Face Trainer
+    training_args = TrainingArguments(
+        output_dir="test_trainer",
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=small_train_dataset,
+        eval_dataset=small_eval_dataset,
+        compute_metrics=compute_metrics,
+    )
+
+    # [2] Report Metrics and Checkpoints to Ray Train
+    # ===============================================
+    callback = ray.train.huggingface.transformers.RayTrainReportCallback()
+    trainer.add_callback(callback)
+
+    # [3] Prepare Transformers Trainer
+    # ================================
+    trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)
+
+    # Start Training
+    trainer.train()
+
+
+def test_huggingface_transformers():
+    # [4] Define a Ray TorchTrainer to launch `train_func` on all workers
+    # ===================================================================
+    ray_trainer = TorchTrainer(
+        train_func,
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
+        # [4a] For multi-node clusters, configure persistent storage that is
+        # accessible across all worker nodes
+        run_config=ray.train.RunConfig(
+            storage_path="/mnt/cluster_storage/huggingface_run"
+        ),
+    )
+    result: ray.train.Result = ray_trainer.fit()
+
+    # [5] Load the trained model
+    with result.checkpoint.as_directory() as checkpoint_dir:
+        checkpoint_path = os.path.join(  # noqa: F841
+            checkpoint_dir,
+            ray.train.huggingface.transformers.RayTrainReportCallback.CHECKPOINT_NAME,
+        )
+        model = AutoModelForSequenceClassification.from_pretrained(  # noqa: F841
+            checkpoint_path
+        )
+
+
+if __name__ == "__main__":
+    test_huggingface_transformers()

release/train_tests/local_mode/compute_gpu_2x4_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.4xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 2
+      max_nodes: 2
+      market_type: ON_DEMAND

release/train_tests/local_mode/torch_local_mode_launcher.py

@@ -0,0 +1,99 @@
+"""Ray Train release test: local mode launched by torchrun.
+
+Setup:
+- 2 x g4dn.12xlarge (4 GPU)
+
+Test owner: xinyuangui2
+
+The test launches a ray cluster with 2 nodes, and launches a torchrun job on each node.
+"""
+import os
+import ray
+import subprocess
+import logging
+from ray.air.util.node import _force_on_node
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+def _write(stream: bytes, path: str):
+    Path(path).parent.mkdir(parents=True, exist_ok=True)
+
+    with open(path, "wb") as f:
+        f.write(stream)
+
+
+@ray.remote
+def _torch_run_launch(
+    master_address: str,
+    node_rank: int,
+    absolute_path: str,
+    n_nodes: int,
+    n_processes_per_node: int,
+    master_port: int,
+):
+    cmd = [
+        "torchrun",
+        f"--nnodes={n_nodes}",
+        f"--nproc-per-node={n_processes_per_node}",
+        f"--node_rank={node_rank}",
+        "--rdzv_backend=c10d",
+        f"--rdzv_endpoint={master_address}:{master_port}",
+        "--rdzv_id=local_mode_job",
+        absolute_path,
+    ]
+
+    env = os.environ.copy()
+    env["RAY_TRAIN_V2_ENABLED"] = "1"
+
+    subprocess.check_call(cmd, env=env)
+
+
+def torch_run_launch_on_nodes():
+    head_ip = ray.util.get_node_ip_address()
+    node_id_ips = []
+    for node in ray.nodes():
+        if not node["Alive"]:
+            continue
+
+        node_ip = node["NodeManagerAddress"]
+
+        if node_ip == head_ip:
+            continue
+
+        node_id = node["NodeID"]
+        node_id_ips.append((node_id, node_ip))
+
+    assert len(node_id_ips) == 2, f"Expected 2 nodes, got {len(node_id_ips)}"
+    master_address = node_id_ips[0][1]
+    futures = []
+    absolute_path = os.path.abspath("torch_local_mode_test.py")
+    with open(absolute_path, "rb") as f:
+        stream = f.read()
+    logger.info(f"Uploading file to all nodes: {absolute_path}")
+    for i in range(len(node_id_ips)):
+        futures.append(
+            _force_on_node(node_id_ips[i][0], _write).remote(stream, absolute_path)
+        )
+    ray.get(futures)
+    logger.info("Uploaded file to all nodes, starting torch run launch")
+    futures = []
+    for i in range(len(node_id_ips)):
+        futures.append(
+            _force_on_node(node_id_ips[i][0], _torch_run_launch).remote(
+                master_address, i, absolute_path, len(node_id_ips), 4, 29500
+            )
+        )
+    ray.get(futures)
+
+
+if __name__ == "__main__":
+    # https://docs.ray.io/en/latest/ray-core/scheduling/accelerators.html#using-accelerators-in-tasks-and-actors
+    # we don't want actors to override CUDA_VISIBLE_DEVICES
+    ray.init(
+        "auto",
+        runtime_env={"env_vars": {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}},
+    )
+    torch_run_launch_on_nodes()

release/train_tests/local_mode/torch_local_mode_test.py

@@ -0,0 +1,162 @@
+import os
+import tempfile
+
+import logging
+
+import torch
+from torch.nn import CrossEntropyLoss
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from torchvision.models import resnet18
+from torchvision.datasets import FashionMNIST
+from torchvision.transforms import ToTensor, Normalize, Compose
+from filelock import FileLock
+import torch.distributed as dist
+
+import ray
+from ray.train import (
+    Checkpoint,
+    CheckpointConfig,
+    RunConfig,
+    ScalingConfig,
+    get_context,
+)
+from ray.train.torch import TorchTrainer
+
+logger = logging.getLogger(__name__)
+DATA_ROOT = "/tmp/test_data"
+
+
+def train_func(config):
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    lock = FileLock(os.path.join(DATA_ROOT, "fashionmnist.lock"))
+    # [1] Prepare model.
+    model = ray.train.torch.prepare_model(model)
+
+    # model.to("cuda")  # This is done by `prepare_model`
+    criterion = CrossEntropyLoss()
+    optimizer = Adam(model.parameters(), lr=config["lr"])
+
+    # Data
+    transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
+    local_rank = get_context().get_local_rank()
+    if local_rank == 0:
+        logger.info(f"Downloading FashionMNIST data to {DATA_ROOT}")
+        with lock:
+            _ = FashionMNIST(
+                root=DATA_ROOT, train=True, download=True, transform=transform
+            )
+    dist.barrier()
+    logger.info(f"Loading FashionMNIST data from {DATA_ROOT}")
+    train_data = FashionMNIST(
+        root=DATA_ROOT, train=True, download=False, transform=transform
+    )
+
+    train_loader = DataLoader(train_data, batch_size=config["batch_size"], shuffle=True)
+    # [2] Prepare dataloader.
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    epoch_losses = []
+    for epoch in range(config["num_epochs"]):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        epoch_loss = 0.0
+        num_batches = 0
+        for images, labels in train_loader:
+            # This is done by `prepare_data_loader`!
+            # images, labels = images.to("cuda"), labels.to("cuda")
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            num_batches += 1
+
+        # Calculate average loss for the epoch
+        avg_epoch_loss = epoch_loss / num_batches if num_batches > 0 else float("inf")
+        epoch_losses.append(avg_epoch_loss)
+
+        # [3] Report metrics and checkpoint.
+        metrics = {
+            "loss": avg_epoch_loss,
+            "epoch": epoch,
+            "epoch_losses": epoch_losses.copy(),  # Track all losses for validation
+        }
+        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+            torch.save(
+                model.state_dict(), os.path.join(temp_checkpoint_dir, "model.pt")
+            )
+            ray.train.report(
+                metrics,
+                checkpoint=Checkpoint.from_directory(temp_checkpoint_dir),
+            )
+        if ray.train.get_context().get_world_rank() == 0:
+            logger.info(f"metrics: {metrics}")
+
+
+def fit_func():
+    # Define configurations.
+    train_loop_config = {"num_epochs": 20, "lr": 0.01, "batch_size": 32}
+    scaling_config = ScalingConfig(num_workers=0, use_gpu=True)
+    run_config = RunConfig(checkpoint_config=CheckpointConfig(num_to_keep=1))
+
+    # Initialize the Trainer.
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_func,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        run_config=run_config,
+    )
+
+    # Train the model.
+    result = trainer.fit()
+
+    # Inspect the results and validate loss makes sense
+    final_loss = result.metrics["loss"]
+    epoch_losses = result.metrics.get("epoch_losses", [])
+
+    logger.info(f"final_loss: {final_loss}")
+    logger.info(f"all epoch losses: {epoch_losses}")
+
+    # Validation 1: Check loss is finite and not NaN
+    assert not torch.isnan(torch.tensor(final_loss)), f"Final loss is NaN: {final_loss}"
+    assert torch.isfinite(
+        torch.tensor(final_loss)
+    ), f"Final loss is not finite: {final_loss}"
+
+    # Validation 2: Check loss convergence - final loss should be lower than initial loss
+    if len(epoch_losses) >= 2:
+        initial_loss = epoch_losses[0]
+        assert (
+            final_loss < initial_loss
+        ), f"Loss didn't decrease: initial={initial_loss}, final={final_loss}"
+        logger.info(
+            f"Loss successfully decreased from {initial_loss:.4f} to {final_loss:.4f}"
+        )
+
+        # Additional check: loss should show general decreasing trend
+        # Allow for some fluctuation but overall trend should be downward
+        mid_point = len(epoch_losses) // 2
+        early_avg = sum(epoch_losses[:mid_point]) / mid_point
+        late_avg = sum(epoch_losses[mid_point:]) / (len(epoch_losses) - mid_point)
+        assert (
+            late_avg < early_avg
+        ), f"Loss trend not decreasing: early_avg={early_avg:.4f}, late_avg={late_avg:.4f}"
+        logger.info(
+            f"Loss trend validation passed: early_avg={early_avg:.4f}, late_avg={late_avg:.4f}"
+        )
+
+    logger.info("All loss validation checks passed!")
+    return result
+
+
+if __name__ == "__main__":
+    fit_func()

release/train_tests/multinode_persistence/compute_aws.yaml

@@ -0,0 +1,20 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+  TagSpecifications:
+    - ResourceType: "instance"
+      Tags:
+        - Key: ttl-hours
+          Value: '24'
+
+head_node:
+    instance_type: m5.2xlarge
+    # Test's wait_for_nodes: 4 > 3 workers — head must count as a usable node.
+    resources:
+      CPU: 8
+
+worker_nodes:
+    - instance_type: m5.2xlarge
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND

release/train_tests/multinode_persistence/compute_gce.yaml

@@ -0,0 +1,14 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+zones:
+    - us-west1-b
+
+head_node:
+    instance_type: n2-standard-8
+    resources:
+      CPU: 8
+
+worker_nodes:
+    - instance_type: n2-standard-8
+      min_nodes: 3
+      max_nodes: 3
+      market_type: ON_DEMAND

release/train_tests/multinode_persistence/test_persistence.py

@@ -0,0 +1,429 @@
+"""Train multi-node persistence/checkpoint release test.
+
+This test is a multi-node version of `test_new_persistence.py`/`test_persistence.py`
+and is meant to be run on a cluster with NFS or S3 storage configured.
+
+This test also records timing metrics on checkpoint save (to disk), save (to storage),
+and load (from storage) operations and outputs them as release test metrics.
+
+Setup:
+- 4x 8 CPU instances
+- 8 workers, each allocated 4 CPUs
+
+Test owner: justinvyu
+"""
+
+import collections
+from contextlib import contextmanager
+from datetime import datetime
+import json
+import os
+from pathlib import Path
+import pickle
+import shutil
+import subprocess
+import time
+from typing import Any, Dict
+import uuid
+
+import fsspec
+import numpy as np
+import pyarrow.fs
+import pytest
+import torch
+import torch.distributed as dist
+
+import ray
+from ray import train
+from ray._private.dict import flatten_dict
+from ray.air.constants import TRAINING_ITERATION
+from ray.air._internal.uri_utils import URI
+from ray.train import Checkpoint
+from ray.train.torch import TorchTrainer
+from ray.train.v2._internal.constants import is_v2_enabled
+
+if is_v2_enabled():
+    from test_v2_persistence import (
+        train_fn,
+        _assert_storage_contents,
+    )
+    from ray.train.v2.api.exceptions import WorkerGroupError
+else:
+    from test_v1_persistence import (
+        train_fn,
+        _assert_storage_contents,
+        _resume_from_checkpoint,
+    )
+    from ray.train.base_trainer import TrainingFailedError
+
+
+# Add a unique ID to the storage path to avoid collisions between release test runs.
+TEST_ID = uuid.uuid4().hex[:4] + "_" + datetime.today().strftime("%Y-%m-%d_%H-%M-%S")
+CLOUD_TEST_DIR = (
+    os.environ["ANYSCALE_ARTIFACT_STORAGE"] + f"/test-persistence-{TEST_ID}/"
+)
+NFS_TEST_DIR = f"/mnt/cluster_storage/test-persistence-{TEST_ID}/"
+
+
+class TestConstants:
+    NUM_ITERATIONS = 10  # == num_checkpoints == num_artifacts
+    NUM_TRIALS = 2
+
+    # 4 * 8 = 32 CPUs total
+    NUM_WORKERS = 8
+    NUM_CPUS_PER_WORKER = 4
+
+    SCORE_KEY = "score"
+
+    NUM_GB = 2
+    NUM_MB = 10
+    NUM_KB = 10
+
+
+def update_output_json(metrics: Dict[str, Any]):
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/release_test_out.json")
+    data = {}
+    if os.path.exists(test_output_json):
+        with open(test_output_json, "r") as f:
+            data = json.load(f)
+
+    data.update(metrics)
+    with open(test_output_json, "w") as f:
+        json.dump(data, f)
+
+
+def create_checkpoint(checkpoint_dir: str) -> float:
+    """Create a somewhat realistic checkpoint of a given size.
+
+    Returns the time it takes to dump this checkpoint to disk."""
+    start = time.perf_counter()
+    # Small (1kb) files
+    for i in range(TestConstants.NUM_KB):
+        with open(os.path.join(checkpoint_dir, f"1kb-{i}.txt"), "w") as f:
+            f.write("a" * 1024)
+
+    # Medium files (1 mb)
+    for i in range(TestConstants.NUM_MB):
+        with open(os.path.join(checkpoint_dir, f"1mb-{i}.txt"), "w") as f:
+            f.write("a" * 1024 * 1024)
+
+    # Large files (1 gb)
+    for i in range(TestConstants.NUM_GB):
+        with open(os.path.join(checkpoint_dir, f"1gb-{i}.txt"), "w") as f:
+            f.write("a" * 1024 * 1024 * 1024)
+    return time.perf_counter() - start
+
+
+def custom_restore_fn(checkpoint: Checkpoint):
+    start = time.perf_counter()
+    with checkpoint.as_directory() as checkpoint_dir:
+        time_to_load = time.perf_counter() - start
+
+        dist.barrier()
+        time_tensor = torch.tensor([time_to_load, 1.0])
+        dist.reduce(time_tensor, dst=0, op=dist.ReduceOp.SUM)
+
+        if train.get_context().get_world_rank() == 0:
+            aggregated_metrics = {"load": time_tensor[0].item() / time_tensor[1].item()}
+            checkpoint.update_metadata(aggregated_metrics)
+            print("[checkpoint] Restore metrics:\n", aggregated_metrics)
+
+        # This is a file populated by the default saving logic in `train_fn`.
+        with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
+            state = pickle.load(f)
+            return state
+
+
+@contextmanager
+def custom_save_fn(temp_checkpoint_dir: str):
+    time_to_save = create_checkpoint(temp_checkpoint_dir)
+
+    start = time.perf_counter()
+    yield  # train.report happens here
+    time_to_report = time.perf_counter() - start
+
+    # Do an all-gather and have rank 0 write the aggregated timing metrics
+    dist.barrier()
+    timing_metrics = torch.tensor([time_to_save, time_to_report, 1.0])
+    dist.reduce(timing_metrics, dst=0, op=dist.ReduceOp.SUM)
+
+    if train.get_context().get_world_rank() == 0:
+        persisted_checkpoint = train.get_checkpoint()
+        aggregated_metrics = {
+            "save_to_disk": timing_metrics[0].item() / timing_metrics[2].item(),
+            "report": timing_metrics[1].item() / timing_metrics[2].item(),
+        }
+        persisted_checkpoint.update_metadata(aggregated_metrics)
+        print("[checkpoint] Save metrics:\n", aggregated_metrics)
+
+
+def get_custom_cloud_fs() -> pyarrow.fs.FileSystem:
+    fsspec_fs, _ = fsspec.core.url_to_fs(os.environ["ANYSCALE_ARTIFACT_STORAGE"])
+    return pyarrow.fs.PyFileSystem(pyarrow.fs.FSSpecHandler(fsspec_fs))
+
+
+def strip_prefix(path: str) -> str:
+    return path.replace("s3://", "").replace("gs://", "")
+
+
+def delete_at_uri(uri: str):
+    if uri.startswith("s3://"):
+        subprocess.check_output(["aws", "s3", "rm", "--recursive", uri])
+    elif uri.startswith("gs://"):
+        subprocess.check_output(["gsutil", "-m", "rm", "-r", uri])
+    else:
+        raise NotImplementedError(f"Invalid URI: {uri}")
+
+
+def download_from_uri(uri: str, local_path: str):
+    if uri.startswith("s3://"):
+        subprocess.check_output(["aws", "s3", "cp", "--recursive", uri, local_path])
+    elif uri.startswith("gs://"):
+        subprocess.check_output(
+            ["gsutil", "-m", "cp", "-r", uri.rstrip("/") + "/*", local_path]
+        )
+    else:
+        raise NotImplementedError(f"Invalid URI: {uri}")
+
+
+@pytest.mark.parametrize(
+    "root_path_storage_filesystem_label",
+    [
+        (CLOUD_TEST_DIR, None, "cloud"),
+        (NFS_TEST_DIR, None, "nfs"),
+        (strip_prefix(CLOUD_TEST_DIR), get_custom_cloud_fs(), "cloud+custom_fs"),
+    ],
+)
+def test_trainer(root_path_storage_filesystem_label, tmp_path, monkeypatch):
+    """Tests that a data parallel trainer can save and restore checkpoints to
+    various storage types properly. Also records checkpoint save/restore timing.
+
+    Here's the rundown of what this test does:
+    1. Passes in a `custom_save_fn` and `custom_restore_fn` to the trainer to
+       record how long the operations take, as well as save a large checkpoint.
+       See `create_checkpoint` for details on the checkpoint contents.
+    2. Configures the training loop to fail 3 times.
+    3. Runs the trainer, which will fail 2 times and recover via FailureConfig.
+       This first run will exit on the 3rd failure.
+    4. Manually restores the trainer, which will restore from the 3rd failure and
+       run to completion.
+    5. Downloads the results from the storage path and asserts that the contents
+       are all correct. See `ray.train.test_new_persistence` for the expected filetree.
+    6. Tests a new run with `resume_from_checkpoint`.
+    """
+    ray.init(runtime_env={"working_dir": "."}, ignore_reinit_error=True)
+
+    root_path, storage_filesystem, label = root_path_storage_filesystem_label
+    storage_path = root_path + label
+    num_to_keep = TestConstants.NUM_ITERATIONS // 2
+    checkpoint_config = train.CheckpointConfig(num_to_keep=num_to_keep)
+    exp_name = "test_trainer"
+
+    print(
+        "\nSaving results under (storage_path, exp_name) = "
+        f"({storage_path}, {exp_name})\n"
+    )
+    train_loop_config = {
+        "fail_iters": [3, 6, 8],
+        "time_per_iter": 1.0,
+        "num_iterations": TestConstants.NUM_ITERATIONS,
+        "custom_save_fn": custom_save_fn,
+        "custom_restore_fn": custom_restore_fn,
+        "num_to_keep": num_to_keep,
+    }
+    scaling_config = train.ScalingConfig(
+        num_workers=TestConstants.NUM_WORKERS,
+        resources_per_worker={"CPU": TestConstants.NUM_CPUS_PER_WORKER},
+    )
+    run_config = train.RunConfig(
+        failure_config=train.FailureConfig(max_failures=2),
+        name=exp_name,
+        storage_path=storage_path,
+        storage_filesystem=storage_filesystem,
+        checkpoint_config=checkpoint_config,
+    )
+    if not is_v2_enabled():
+        train_loop_config["in_trainer"] = True
+        scaling_config.trainer_resources = {"CPU": 0}
+        run_config.sync_config = train.SyncConfig(sync_artifacts=True)
+    trainer = TorchTrainer(
+        train_fn,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        run_config=run_config,
+    )
+    print("\nStarting initial run.\n")
+    if is_v2_enabled():
+        with pytest.raises(WorkerGroupError):
+            trainer.fit()
+    else:
+        with pytest.raises(TrainingFailedError):
+            result = trainer.fit()
+
+    print("\nStarting manually restored run.\n")
+    if is_v2_enabled():
+        restored_trainer = TorchTrainer(
+            train_fn,
+            train_loop_config=train_loop_config,
+            scaling_config=scaling_config,
+            run_config=run_config,
+        )
+    else:
+        restored_trainer = TorchTrainer.restore(
+            path=str(URI(storage_path) / exp_name),
+            storage_filesystem=storage_filesystem,
+        )
+    result = restored_trainer.fit()
+    print(result)
+
+    print("\nAsserting contents of uploaded results.\n")
+    local_inspect_dir = tmp_path / "inspect_dir"
+    local_inspect_dir.mkdir()
+    # Download the results from storage
+    if "cloud" in label:
+        # NOTE: Use the CLI to download, since the python libraries
+        # (pyarrow, fsspec) aren't consistent across cloud platforms (s3, gs).
+        cloud_uri = CLOUD_TEST_DIR + label
+        print("\nDownloading from cloud URI:", cloud_uri, "\n")
+        download_from_uri(cloud_uri, str(local_inspect_dir))
+    elif label == "nfs":
+        local_inspect_dir = Path(storage_path)
+    else:
+        raise NotImplementedError(f"Invalid storage type: {label}")
+
+    if is_v2_enabled():
+        _assert_storage_contents(
+            local_inspect_dir,
+            exp_name,
+            checkpoint_config,
+            constants=TestConstants,
+        )
+    else:
+        _assert_storage_contents(
+            local_inspect_dir,
+            exp_name,
+            checkpoint_config,
+            "TorchTrainer",
+            test_trainer=True,
+            constants=TestConstants,
+        )
+
+    # Test `resume_from_checkpoint`
+    if not is_v2_enabled():
+        _resume_from_checkpoint(
+            result.checkpoint,
+            expected_state={"iter": TestConstants.NUM_ITERATIONS - 1},
+            storage_path=storage_path,
+            storage_filesystem=storage_filesystem,
+        )
+
+    # Upload checkpoint save and restore timing release test metrics
+    all_checkpoint_timing_metrics = collections.defaultdict(list)
+    for checkpoint, _ in result.best_checkpoints:
+        metadata = checkpoint.get_metadata()
+        for metric, value in metadata.items():
+            all_checkpoint_timing_metrics[metric].append(value)
+
+    aggregated_metrics = {
+        key: np.mean(values) for key, values in all_checkpoint_timing_metrics.items()
+    }
+    checkpoint_size_mb = (
+        TestConstants.NUM_GB * 1000 + TestConstants.NUM_MB + TestConstants.NUM_KB / 1000
+    )
+    speeds = {
+        key + "_speed_mbps": checkpoint_size_mb / time_s
+        for key, time_s in aggregated_metrics.items()
+    }
+    # Add units as the suffix
+    aggregated_metrics = {
+        key + "_avg_s": time_s for key, time_s in aggregated_metrics.items()
+    }
+    aggregated_metrics.update(speeds)
+    aggregated_metrics["checkpoint_size_mb"] = checkpoint_size_mb
+
+    print(aggregated_metrics)
+    update_output_json(flatten_dict({label: aggregated_metrics}))
+
+    print("Deleting files from the run...")
+    if "cloud" in label:
+        # NOTE: Use the CLI to delete files on cloud, since the python libraries
+        # (pyarrow, fsspec) aren't consistent across cloud platforms (s3, gs).
+        delete_at_uri(CLOUD_TEST_DIR)
+    elif label == "nfs":
+        shutil.rmtree(NFS_TEST_DIR, ignore_errors=True)
+    else:
+        raise NotImplementedError(f"Invalid storage type: {label}")
+
+
+def test_no_storage_error(tmp_path, monkeypatch):
+    """Tests that an error is raised if you do multi-node checkpointing
+    w/ no persistent storage configured."""
+    ray.init(runtime_env={"working_dir": "."}, ignore_reinit_error=True)
+
+    train_loop_config = {
+        "time_per_iter": 1.0,
+        "num_iterations": TestConstants.NUM_ITERATIONS,
+    }
+    scaling_config = train.ScalingConfig(
+        num_workers=TestConstants.NUM_WORKERS,
+        resources_per_worker={"CPU": TestConstants.NUM_CPUS_PER_WORKER},
+    )
+    if not is_v2_enabled():
+        train_loop_config["in_trainer"] = True
+        scaling_config.trainer_resources = {"CPU": 0}
+    trainer = TorchTrainer(
+        train_fn,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        run_config=train.RunConfig(name="test_trainer", storage_path=None),
+    )
+    if is_v2_enabled():
+        with pytest.raises(WorkerGroupError):
+            trainer.fit()
+    else:
+        with pytest.raises(TrainingFailedError):
+            trainer.fit()
+
+
+def test_no_storage_no_checkpoints(tmp_path, monkeypatch):
+    """Tests that it's ok to run multi-node with no persistent storage
+    if you never report checkpoints."""
+    ray.init(runtime_env={"working_dir": "."}, ignore_reinit_error=True)
+
+    train_loop_config = {
+        "time_per_iter": 1.0,
+        "num_iterations": TestConstants.NUM_ITERATIONS,
+        # Don't report any checkpoints
+        "no_checkpoint_ranks": list(range(TestConstants.NUM_WORKERS)),
+    }
+    scaling_config = train.ScalingConfig(
+        num_workers=TestConstants.NUM_WORKERS,
+        resources_per_worker={"CPU": TestConstants.NUM_CPUS_PER_WORKER},
+    )
+    run_config = train.RunConfig(
+        failure_config=train.FailureConfig(max_failures=2),
+        name="test_trainer",
+        storage_path=None,
+    )
+    if not is_v2_enabled():
+        train_loop_config["in_trainer"] = True
+        scaling_config.trainer_resources = {"CPU": 0}
+        run_config.sync_config = train.SyncConfig(sync_artifacts=True)
+    trainer = TorchTrainer(
+        train_fn,
+        train_loop_config=train_loop_config,
+        scaling_config=scaling_config,
+        run_config=run_config,
+    )
+    result = trainer.fit()
+
+    # v2 does not support free floating metrics
+    if not is_v2_enabled():
+        assert result.metrics[TRAINING_ITERATION] == TestConstants.NUM_ITERATIONS
+        assert len(result.metrics_dataframe) == TestConstants.NUM_ITERATIONS
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

release/train_tests/multinode_persistence/test_v1_persistence.py

@@ -0,0 +1 @@
+../../../python/ray/train/tests/test_new_persistence.py

release/train_tests/multinode_persistence/test_v2_persistence.py

@@ -0,0 +1 @@
+../../../python/ray/train/v2/tests/test_persistence.py

release/train_tests/pytorch_lightning/compute_aws.yaml

@@ -0,0 +1,10 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: g4dn.12xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/train_tests/pytorch_lightning/test_lightning.py

@@ -0,0 +1,93 @@
+import os
+import tempfile
+
+import torch
+from torch.utils.data import DataLoader
+from torchvision.models import resnet18
+from torchvision.datasets import FashionMNIST
+from torchvision.transforms import ToTensor, Normalize, Compose
+import lightning.pytorch as pl
+
+import ray.train.lightning
+from ray.train.torch import TorchTrainer
+
+# Model, Loss, Optimizer
+class ImageClassifier(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = resnet18(num_classes=10)
+        self.model.conv1 = torch.nn.Conv2d(
+            1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+        )
+        self.criterion = torch.nn.CrossEntropyLoss()
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        outputs = self.forward(x)
+        loss = self.criterion(outputs, y)
+        self.log("loss", loss, on_step=True, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.model.parameters(), lr=0.001)
+
+
+def train_func():
+    # Data
+    transform = Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
+    data_dir = os.path.join(tempfile.gettempdir(), "data")
+    train_data = FashionMNIST(
+        root=data_dir, train=True, download=True, transform=transform
+    )
+    train_dataloader = DataLoader(train_data, batch_size=128, shuffle=True)
+
+    # Training
+    model = ImageClassifier()
+    # [1] Configure PyTorch Lightning Trainer.
+    trainer = pl.Trainer(
+        max_epochs=10,
+        devices="auto",
+        accelerator="auto",
+        strategy=ray.train.lightning.RayDDPStrategy(),
+        plugins=[ray.train.lightning.RayLightningEnvironment()],
+        callbacks=[ray.train.lightning.RayTrainReportCallback()],
+        # [1a] Optionally, disable the default checkpointing behavior
+        # in favor of the `RayTrainReportCallback` above.
+        enable_checkpointing=False,
+    )
+    trainer = ray.train.lightning.prepare_trainer(trainer)
+    trainer.fit(model, train_dataloaders=train_dataloader)
+
+
+def test_lightning_train_run():
+    # [2] Configure scaling and resource requirements.
+    scaling_config = ray.train.ScalingConfig(num_workers=4, use_gpu=True)
+
+    # [3] Launch distributed training job.
+    trainer = TorchTrainer(
+        train_func,
+        scaling_config=scaling_config,
+        # [3a] If running in a multi-node cluster, this is where you
+        # should configure the run's persistent storage that is accessible
+        # across all worker nodes.
+        run_config=ray.train.RunConfig(
+            storage_path="/mnt/cluster_storage/lightning_run"
+        ),
+    )
+    result: ray.train.Result = trainer.fit()
+
+    # [4] Load the trained model.
+    with result.checkpoint.as_directory() as checkpoint_dir:
+        model = ImageClassifier.load_from_checkpoint(  # noqa: F841
+            os.path.join(
+                checkpoint_dir,
+                ray.train.lightning.RayTrainReportCallback.CHECKPOINT_NAME,
+            ),
+        )
+
+
+if __name__ == "__main__":
+    test_lightning_train_run()

release/train_tests/torchft/compute_aws.yaml

@@ -0,0 +1,8 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+# Single GPU node. The torchft linear example trains on CPU (gloo) but the
+# image ships GPU torch (torch==2.9.0+cu128), so a CUDA node is representative.
+head_node:
+    instance_type: g4dn.4xlarge
+
+worker_nodes: []

release/train_tests/torchft/torchft_hello_world.py

@@ -0,0 +1,39 @@
+"""Hello-world release test for the torchft Ray ML image variant.
+
+This is a reference test showing how to run a release test on the custom
+torchft BYOD image. The image is the core Ray CUDA image (py3.13) with the
+torchft ML dependency lock installed on top:
+
+  - lock: release/ray_release/byod/ml_torchft_py3.13.lock
+          (generated by ci/raydepsets/configs/release_ml_torchft_tests.depsets.yaml)
+
+The lock is referenced from the release test via `byod.python_depset`; the BYOD
+build installs it automatically (`uv pip install --system --no-deps -r
+python_depset.lock`), so no post_build_script is required.
+
+See the matching `torchft_hello_world` entry in release/release_tests.yaml for
+the cluster wiring (byod.type / byod.python_depset).
+
+The test simply proves the image works end to end: torch 2.9.0 + torchft import
+cleanly and a tiny Ray Train v2 + torchft training loop runs to completion.
+"""
+
+import ray
+import torch
+import torchft  # noqa: F401  -- provided by torchft-nightly in the image
+
+from ray.train.v2.examples.pytorch.torchft_linear_example import train_torchft
+
+
+def main() -> None:
+    print(f"ray=={ray.__version__} torch=={torch.__version__} torchft OK")
+
+    # torchft_linear_example trains a tiny linear model on CPU (gloo backend)
+    # under Ray Train v2 with a torchft TorchftConfig. A short run is enough to
+    # validate that the image's torch + torchft + Ray Train stack works.
+    metrics = train_torchft(num_workers=2, num_steps=20)
+    print(f"torchft hello world succeeded: {metrics}")
+
+
+if __name__ == "__main__":
+    main()

release/train_tests/xgboost_lightgbm/compute_aws_100G.yaml

@@ -0,0 +1,20 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            Iops: 5000
+            Throughput: 1000
+            VolumeSize: 200
+            VolumeType: gp3
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: m5.4xlarge
+      min_nodes: 10
+      max_nodes: 10
+      market_type: ON_DEMAND

release/train_tests/xgboost_lightgbm/compute_aws_10G.yaml

@@ -0,0 +1,20 @@
+cloud: {{env["ANYSCALE_CLOUD_NAME"]}}
+
+advanced_instance_config:
+    BlockDeviceMappings:
+        - DeviceName: /dev/sda1
+          Ebs:
+            DeleteOnTermination: true
+            Iops: 5000
+            Throughput: 1000
+            VolumeSize: 200
+            VolumeType: gp3
+
+head_node:
+    instance_type: m5.2xlarge
+
+worker_nodes:
+    - instance_type: m5.4xlarge
+      min_nodes: 1
+      max_nodes: 1
+      market_type: ON_DEMAND

release/train_tests/xgboost_lightgbm/train_batch_inference_benchmark.py

@@ -0,0 +1,246 @@
+import json
+import numpy as np
+import os
+import pandas as pd
+import time
+from typing import Dict
+
+import xgboost as xgb
+import lightgbm as lgb
+
+import ray
+from ray import data
+from ray.train.lightgbm import (
+    LightGBMTrainer,
+    RayTrainReportCallback as LightGBMReportCallback,
+    normalize_pandas_for_lightgbm,
+)
+from ray.train.xgboost import (
+    RayTrainReportCallback as XGBoostReportCallback,
+    XGBoostTrainer,
+)
+from ray.train import RunConfig, ScalingConfig
+
+_TRAINING_TIME_THRESHOLD = 600
+_PREDICTION_TIME_THRESHOLD = 450
+
+_EXPERIMENT_PARAMS = {
+    "smoke_test": {
+        "data": (
+            "https://air-example-data-2.s3.us-west-2.amazonaws.com/"
+            "10G-xgboost-data.parquet/8034b2644a1d426d9be3bbfa78673dfa_000000.parquet"
+        ),
+        "num_workers": 1,
+        "cpus_per_worker": 1,
+    },
+    "10G": {
+        "data": "s3://air-example-data-2/10G-xgboost-data.parquet/",
+        "num_workers": 1,
+        "cpus_per_worker": 12,
+    },
+    "100G": {
+        "data": "s3://air-example-data-2/100G-xgboost-data.parquet/",
+        "num_workers": 10,
+        "cpus_per_worker": 12,
+    },
+}
+
+
+class BasePredictor:
+    def __init__(self, report_callback_cls, result: ray.train.Result):
+        self.model = report_callback_cls.get_model(result.checkpoint)
+
+    def __call__(self, data):
+        raise NotImplementedError
+
+
+class XGBoostPredictor(BasePredictor):
+    def __call__(self, data: pd.DataFrame) -> Dict[str, np.ndarray]:
+        dmatrix = xgb.DMatrix(data)
+        return {"predictions": self.model.predict(dmatrix)}
+
+
+class LightGBMPredictor(BasePredictor):
+    def __call__(self, data: pd.DataFrame) -> Dict[str, np.ndarray]:
+        return {"predictions": self.model.predict(normalize_pandas_for_lightgbm(data))}
+
+
+def xgboost_train_loop_function(config: Dict):
+    train_ds_iter = ray.train.get_dataset_shard("train")
+    train_df = train_ds_iter.materialize().to_pandas()
+
+    label_column, params = config["label_column"], config["params"]
+    train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
+
+    dtrain = xgb.DMatrix(train_X, label=train_y)
+
+    report_callback = config["report_callback_cls"]
+    xgb.train(
+        params,
+        dtrain=dtrain,
+        num_boost_round=10,
+        callbacks=[report_callback()],
+    )
+
+
+def lightgbm_train_loop_function(config: Dict):
+    train_ds_iter = ray.train.get_dataset_shard("train")
+    train_df = normalize_pandas_for_lightgbm(train_ds_iter.materialize().to_pandas())
+
+    label_column, params = config["label_column"], config["params"]
+    train_X, train_y = train_df.drop(label_column, axis=1), train_df[label_column]
+    train_set = lgb.Dataset(train_X, label=train_y)
+
+    report_callback = config["report_callback_cls"]
+    network_params = ray.train.lightgbm.get_network_params()
+    params.update(network_params)
+
+    lgb.train(
+        params,
+        train_set=train_set,
+        num_boost_round=10,
+        callbacks=[report_callback()],
+    )
+
+
+_FRAMEWORK_PARAMS = {
+    "xgboost": {
+        "trainer_cls": XGBoostTrainer,
+        "predictor_cls": XGBoostPredictor,
+        "train_loop_function": xgboost_train_loop_function,
+        "train_loop_config": {
+            "params": {
+                "objective": "binary:logistic",
+                "eval_metric": ["logloss", "error"],
+            },
+            "label_column": "labels",
+            "report_callback_cls": XGBoostReportCallback,
+        },
+    },
+    "lightgbm": {
+        "trainer_cls": LightGBMTrainer,
+        "predictor_cls": LightGBMPredictor,
+        "train_loop_function": lightgbm_train_loop_function,
+        "train_loop_config": {
+            "params": {
+                "objective": "binary",
+                "metric": ["binary_logloss", "binary_error"],
+            },
+            "label_column": "labels",
+            "report_callback_cls": LightGBMReportCallback,
+        },
+    },
+}
+
+
+def train(
+    framework: str, data_path: str, num_workers: int, cpus_per_worker: int
+) -> ray.train.Result:
+    ds = data.read_parquet(data_path)
+    framework_params = _FRAMEWORK_PARAMS[framework]
+
+    trainer_cls = framework_params["trainer_cls"]
+    framework_train_loop_fn = framework_params["train_loop_function"]
+
+    trainer = trainer_cls(
+        train_loop_per_worker=framework_train_loop_fn,
+        train_loop_config=framework_params["train_loop_config"],
+        scaling_config=ScalingConfig(
+            num_workers=num_workers,
+            resources_per_worker={"CPU": cpus_per_worker},
+        ),
+        datasets={"train": ds},
+        run_config=RunConfig(
+            storage_path="/mnt/cluster_storage", name=f"{framework}_benchmark"
+        ),
+    )
+    result = trainer.fit()
+    return result
+
+
+def predict(framework: str, result: ray.train.Result, data_path: str):
+    framework_params = _FRAMEWORK_PARAMS[framework]
+
+    predictor_cls = framework_params["predictor_cls"]
+
+    ds = data.read_parquet(data_path)
+    ds = ds.drop_columns(["labels"])
+
+    concurrency = int(ray.cluster_resources()["CPU"] // 2)
+    ds.map_batches(
+        predictor_cls,
+        # Improve prediction throughput with larger batch size than default 4096
+        batch_size=8192,
+        concurrency=concurrency,
+        fn_constructor_kwargs={
+            "report_callback_cls": framework_params["train_loop_config"][
+                "report_callback_cls"
+            ],
+            "result": result,
+        },
+        batch_format="pandas",
+    ).write_parquet("/mnt/cluster_storage/predictions")
+
+
+def main(args):
+    framework = args.framework
+
+    experiment = args.size if not args.smoke_test else "smoke_test"
+    experiment_params = _EXPERIMENT_PARAMS[experiment]
+
+    data_path, num_workers, cpus_per_worker = (
+        experiment_params["data"],
+        experiment_params["num_workers"],
+        experiment_params["cpus_per_worker"],
+    )
+
+    print(f"Running {framework} training benchmark...")
+    training_start = time.perf_counter()
+    result = train(framework, data_path, num_workers, cpus_per_worker)
+    training_time = time.perf_counter() - training_start
+
+    print(f"Running {framework} prediction benchmark...")
+    prediction_start = time.perf_counter()
+    predict(framework, result, data_path)
+    prediction_time = time.perf_counter() - prediction_start
+
+    times = {"training_time": training_time, "prediction_time": prediction_time}
+    print("Training result:\n", result)
+    print("Training/prediction times:", times)
+    test_output_json = os.environ.get("TEST_OUTPUT_JSON", "/tmp/result.json")
+    with open(test_output_json, "wt") as f:
+        json.dump(times, f)
+
+    if not args.disable_check:
+        if training_time > _TRAINING_TIME_THRESHOLD:
+            raise RuntimeError(
+                f"Training is taking {training_time} seconds, "
+                f"which is longer than expected ({_TRAINING_TIME_THRESHOLD} seconds)."
+            )
+
+        if prediction_time > _PREDICTION_TIME_THRESHOLD:
+            raise RuntimeError(
+                f"Batch prediction is taking {prediction_time} seconds, "
+                f"which is longer than expected ({_PREDICTION_TIME_THRESHOLD} seconds)."
+            )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "framework", type=str, choices=["xgboost", "lightgbm"], default="xgboost"
+    )
+    parser.add_argument("--size", type=str, choices=["10G", "100G"], default="100G")
+    # Add a flag for disabling the timeout error.
+    # Use case: running the benchmark as a documented example, in infra settings
+    # different from the formal benchmark's EC2 setup.
+    parser.add_argument(
+        "--disable-check",
+        action="store_true",
+        help="disable runtime error on benchmark timeout",
+    )
+    parser.add_argument("--smoke-test", action="store_true")
+    args = parser.parse_args()
+    main(args)