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