chore: import upstream snapshot with attribution

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