Repository: ray-project/ray
File: release/train_tests/benchmark/image_classification/factory.py (Python)

Code before:
        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 boto3 in map_batches,
                # 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(

Review comment (suggestion):
```suggestion
                # NOTE: This format downloads images via ray data expressions,
```

Code after:
        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(