chore: import upstream snapshot with attribution

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()