Add ray data for image

benchmarking/nightly-benchmark.yaml

@@ -553,7 +553,40 @@ entries:
       num_gpus: 0
       enable_object_spilling: false
 
-  - name: image_curation
+  - name: image_curation_raydata
+    enabled: true
+    script: image_pipeline_benchmark.py
+    args: >-
+      --benchmark-results-path={session_entry_dir}
+      --input-wds-dataset-dir={dataset:mscoco,wds}
+      --output-dataset-dir={session_entry_dir}/scratch/output
+      --model-dir={dataset:mscoco_model_weights,files}
+      --executor=ray_data
+      --tar-files-per-partition=1
+      --batch-size=1000
+      --embedding-batch-size=500
+      --aesthetic-batch-size=500
+      --aesthetic-threshold=0.9
+      --verbose
+    timeout_s: 1500
+    sink_data:
+      - name: slack
+        additional_metrics:
+          - throughput_images_per_sec
+        ping_on_failure:
+          - U020C9VC05N  # Ao Tang
+    ray:
+      num_cpus: 64
+      num_gpus: 4
+      enable_object_spilling: false
+    requirements:
+      # ensure the total number of documents processed is correct
+      - metric: num_images_processed
+        exact_value: 3800
+      - metric: throughput_images_per_sec
+        min_value: 3.0
+
+  - name: image_curation_xenna
     enabled: true
     script: image_pipeline_benchmark.py
     args: >-

nemo_curator/stages/image/io/image_reader.py

@@ -15,11 +15,13 @@
 import pathlib
 from collections.abc import Generator
 from dataclasses import dataclass
+from typing import Any
 
 import numpy as np
 import torch
 from loguru import logger
 
+from nemo_curator.backends.experimental.utils import RayStageSpecKeys
 from nemo_curator.stages.base import ProcessingStage
 from nemo_curator.stages.resources import Resources
 from nemo_curator.tasks import FileGroupTask, ImageBatch, ImageObject
@@ -51,6 +53,12 @@ def __post_init__(self) -> None:
         else:
             self.resources = Resources()
 
+    def ray_stage_spec(self) -> dict[str, Any]:
+        """Ray stage specification for this stage."""
+        return {
+            RayStageSpecKeys.IS_FANOUT_STAGE: True,
+        }
+
     def inputs(self) -> tuple[list[str], list[str]]:
         return [], []
 
@@ -98,11 +106,7 @@ def _read_tars_with_dali(self, tar_paths: list[pathlib.Path]) -> Generator[list[
         # Use the tar filename stem as the id prefix for single shards; for grouped shards,
         # synthesize a group prefix and place generated image paths under the tars' parent dir.
         base_path = tar_paths[0] if len(tar_paths) == 1 else tar_paths[0].parent
-        id_prefix = (
-            tar_paths[0].stem
-            if len(tar_paths) == 1
-            else f"group_{tar_paths[0].stem}_x{len(tar_paths)}"
-        )
+        id_prefix = tar_paths[0].stem if len(tar_paths) == 1 else f"group_{tar_paths[0].stem}_x{len(tar_paths)}"
 
         while samples_completed < total_samples:
             img_batch = pipe.run()