Add TorchDataLoader to Train Benchmark (#51456)

## Why are these changes needed?
Add TorchDataLoader to Train Benchmark.

---------

Signed-off-by: Srinath Krishnamachari <srinath.krishnamachari@anyscale.com>

Author: srinathk10

release/release_tests.yaml

@@ -2118,11 +2118,21 @@
         timeout: 2000
         script: RAY_TRAIN_V2_ENABLED=1 python train_benchmark.py --task=image_classification --dataloader_type=ray_data --num_workers=16
 
+    - __suffix__: full_training_torch_dataloader
+      run:
+        timeout: 2000
+        script: RAY_TRAIN_V2_ENABLED=1 python train_benchmark.py --task=image_classification --dataloader_type=torch --num_workers=16
+
     - __suffix__: skip_training
       run:
         timeout: 1200
         script: RAY_TRAIN_V2_ENABLED=1 python train_benchmark.py --task=image_classification --dataloader_type=ray_data --num_workers=16 --skip_train_step --skip_validation_at_epoch_end
 
+    - __suffix__: skip_training_torch_dataloader
+      run:
+        timeout: 1200
+        script: RAY_TRAIN_V2_ENABLED=1 python train_benchmark.py --task=image_classification --dataloader_type=torch --num_workers=16 --skip_train_step --skip_validation_at_epoch_end
+
     - __suffix__: skip_training.fault_tolerance
       run:
         timeout: 2700

release/train_tests/benchmark/config.py

@@ -7,21 +7,31 @@
 class DataloaderType(enum.Enum):
     RAY_DATA = "ray_data"
     MOCK = "mock"
+    TORCH = "torch"
 
 
 class DataLoaderConfig(BaseModel):
     train_batch_size: int = 32
     validation_batch_size: int = 256
+    prefetch_batches: int = 1
 
 
 class RayDataConfig(DataLoaderConfig):
     # NOTE: Optional[int] doesn't play well with argparse.
     local_buffer_shuffle_size: int = -1
 
 
+class TorchConfig(DataLoaderConfig):
+    num_torch_workers: int = 8
+    torch_dataloader_timeout_seconds: int = 300
+    torch_pin_memory: bool = True
+    torch_non_blocking: bool = True
+
+
 class BenchmarkConfig(BaseModel):
     # ScalingConfig
     num_workers: int = 1
+
     # Run CPU training where train workers request a `MOCK_GPU` resource instead.
     mock_gpu: bool = False
 
@@ -39,11 +49,14 @@ class BenchmarkConfig(BaseModel):
     # Training
     num_epochs: int = 1
     skip_train_step: bool = False
+    train_step_anomaly_detection: bool = False
+    limit_training_rows: int = 500000
 
     # Validation
     validate_every_n_steps: int = -1
     skip_validation_step: bool = False
     skip_validation_at_epoch_end: bool = False
+    limit_validation_rows: int = 50000
 
     # Logging
     log_metrics_every_n_steps: int = 512
@@ -57,11 +70,10 @@ def _is_pydantic_model(field_type) -> bool:
 def _add_field_to_parser(parser: argparse.ArgumentParser, field: str, field_info):
     field_type = field_info.annotation
     if field_type is bool:
-        assert (
-            not field_info.default
-        ), "Only supports bool flags that are False by default."
         parser.add_argument(
-            f"--{field}", action="store_true", default=field_info.default
+            f"--{field}",
+            type=lambda x: x.lower() == "true",
+            default=field_info.default,
         )
     else:
         parser.add_argument(f"--{field}", type=field_type, default=field_info.default)
@@ -87,11 +99,11 @@ def cli_to_config() -> BenchmarkConfig:
         nested_parser = argparse.ArgumentParser()
         config_cls = BenchmarkConfig.model_fields[nested_field].annotation
 
-        if (
-            config_cls == DataLoaderConfig
-            and top_level_args.dataloader_type == DataloaderType.RAY_DATA
-        ):
-            config_cls = RayDataConfig
+        if config_cls == DataLoaderConfig:
+            if top_level_args.dataloader_type == DataloaderType.RAY_DATA:
+                config_cls = RayDataConfig
+            elif top_level_args.dataloader_type == DataloaderType.TORCH:
+                config_cls = TorchConfig
 
         for field, field_info in config_cls.model_fields.items():
             _add_field_to_parser(nested_parser, field, field_info)

release/train_tests/benchmark/dataloader_factory.py

@@ -1,13 +1,13 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterator, Tuple
+import logging
 
 import torch
-
-import ray.data
-import ray.train
 from ray.data import Dataset
 
-from config import BenchmarkConfig, DataLoaderConfig, RayDataConfig
+from config import BenchmarkConfig, DataLoaderConfig
+
+logger = logging.getLogger(__name__)
 
 
 class BaseDataLoaderFactory(ABC):
@@ -34,124 +34,3 @@ def get_metrics(self) -> Dict[str, Any]:
     def get_ray_datasets(self) -> Dict[str, Dataset]:
         """Get Ray datasets if this loader type uses Ray Data."""
         return {}
-
-
-class RayDataLoaderFactory(BaseDataLoaderFactory):
-    def __init__(self, benchmark_config: BenchmarkConfig):
-        super().__init__(benchmark_config)
-        self._ray_ds_iterators = {}
-
-        assert isinstance(self.get_dataloader_config(), RayDataConfig), type(
-            self.get_dataloader_config()
-        )
-
-        # Configure Ray Data settings.
-        data_context = ray.data.DataContext.get_current()
-        data_context.enable_operator_progress_bars = False
-
-    @abstractmethod
-    def get_ray_datasets(self) -> Dict[str, Dataset]:
-        """Get the Ray datasets for training and validation.
-
-        Returns:
-            Dict with "train" and "val" Dataset objects
-        """
-        pass
-
-    @abstractmethod
-    def collate_fn(self) -> Dict[str, Dataset]:
-        """Get the collate function for the dataloader.
-
-        Returns:
-            A function that takes a batch and returns a tuple of tensors.
-        """
-        pass
-
-    def get_train_dataloader(self):
-        ds_iterator = self._ray_ds_iterators["train"] = ray.train.get_dataset_shard(
-            "train"
-        )
-        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.collate_fn,
-            )
-        )
-
-    def get_val_dataloader(self):
-        ds_iterator = self._ray_ds_iterators["val"] = ray.train.get_dataset_shard("val")
-        dataloader_config = self.get_dataloader_config()
-        return iter(
-            ds_iterator.iter_torch_batches(
-                batch_size=dataloader_config.validation_batch_size,
-                collate_fn=self.collate_fn,
-            )
-        )
-
-    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(),
-                    "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(),
-                },
-            }
-        return metrics