diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 1d9bf1a3fd..7c10151695 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/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: >- diff --git a/nemo_curator/stages/image/io/image_reader.py b/nemo_curator/stages/image/io/image_reader.py index f099794952..b65576e6c4 100644 --- a/nemo_curator/stages/image/io/image_reader.py +++ b/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()