diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 9356ce2580..c195bfe17e 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -437,21 +437,31 @@ entries: - name: image_curation enabled: true - script: "{curator_repo_dir}/tutorials/image/getting-started/image_curation_example.py" + script: image_pipeline_benchmark.py args: >- - --input-wds-dataset-dir {dataset:mscoco,wds} - --output-dataset-dir {session_entry_dir}/scratch/output - --model-dir {dataset:mscoco_model_weights,files} - --batch-size 100 - --embedding-batch-size 100 - --aesthetic-batch-size 100 - --nsfw-batch-size 100 - --tar-files-per-partition 10 - --aesthetic-threshold 0.9 - --nsfw-threshold 0.9 - --skip-download + --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=xenna + --tar-files-per-partition=1 + --batch-size=1000 + --embedding-batch-size=500 + --aesthetic-batch-size=500 + --aesthetic-threshold=0.9 --verbose - + timeout_s: 1500 + 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: audio_fleurs enabled: true script: audio_fleurs_benchmark.py diff --git a/benchmarking/scripts/image_pipeline_benchmark.py b/benchmarking/scripts/image_pipeline_benchmark.py new file mode 100644 index 0000000000..4cc8744918 --- /dev/null +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -0,0 +1,314 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Image pipeline benchmarking script. + +This script runs an image curation pipeline benchmark with comprehensive +metrics collection and various executor support. +""" + +import argparse +import time +import traceback +from pathlib import Path +from typing import Any + +from loguru import logger +from utils import setup_executor, write_benchmark_results + +from nemo_curator.core.client import RayClient +from nemo_curator.pipeline import Pipeline +from nemo_curator.stages.file_partitioning import FilePartitioningStage +from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage +from nemo_curator.stages.image.filters.aesthetic_filter import ImageAestheticFilterStage +from nemo_curator.stages.image.io.image_reader import ImageReaderStage +from nemo_curator.stages.image.io.image_writer import ImageWriterStage + + +def create_image_curation_pipeline(args: argparse.Namespace) -> Pipeline: + """Create image curation pipeline with file partitioning, image reading, embedding, and aesthetic scoring stages.""" + + # Define pipeline + pipeline = Pipeline(name="image_curation", description="Curate images with embeddings and quality scoring") + + # Stage 0: Partition tar files for parallel processing + pipeline.add_stage(FilePartitioningStage( + file_paths=args.input_wds_dataset_dir, + files_per_partition=args.tar_files_per_partition, + file_extensions=[".tar"], + )) + + # Stage 1: Read images from webdataset tar files (now runs in parallel) + pipeline.add_stage(ImageReaderStage( + dali_batch_size=args.batch_size, + verbose=args.verbose, # Force verbose to see debug info + num_threads=args.reader_num_threads, # More threads for I/O + num_gpus_per_worker=args.reader_gpus_per_worker, + )) + + # Stage 2: Generate CLIP embeddings for images + pipeline.add_stage(ImageEmbeddingStage( + model_dir=args.model_dir, + num_gpus_per_worker=args.embedding_gpus_per_worker, + model_inference_batch_size=args.embedding_batch_size, + remove_image_data=False, + verbose=args.verbose, + )) + + # Stage 3: Generate aesthetic quality scores and filter + pipeline.add_stage(ImageAestheticFilterStage( + model_dir=args.model_dir, + num_gpus_per_worker=args.aesthetic_gpus_per_worker, + model_inference_batch_size=args.aesthetic_batch_size, + score_threshold=args.aesthetic_threshold, + verbose=args.verbose, + )) + + # Stage 4: Write down to disk + pipeline.add_stage(ImageWriterStage( + output_dir=args.output_dataset_dir, + images_per_tar=args.images_per_tar, + remove_image_data=True, + verbose=args.verbose, + )) + + return pipeline + + +def run_image_pipeline_benchmark(args: argparse.Namespace) -> dict[str, Any]: + """Run the image pipeline benchmark and collect comprehensive metrics.""" + executor = setup_executor(args.executor) + + input_wds_dir = Path(args.input_wds_dataset_dir).absolute() + output_dir = Path(args.output_dataset_dir).absolute() + output_dir.mkdir(parents=True, exist_ok=True) + + logger.info(f"Input webdataset directory: {input_wds_dir}") + logger.info(f"Output dataset directory: {output_dir}") + logger.info(f"Model directory: {args.model_dir}") + logger.info(f"Tar files per partition: {args.tar_files_per_partition}") + logger.info(f"Task batch size: {args.batch_size}") + logger.info(f"Embedding batch size: {args.embedding_batch_size}") + logger.info(f"Aesthetic threshold: {args.aesthetic_threshold}") + logger.debug(f"Executor: {executor}") + + # Create pipeline + pipeline = create_image_curation_pipeline(args) + + run_start_time = time.perf_counter() + + try: + logger.info("Running image curation pipeline...") + logger.info(f"Pipeline description:\n{pipeline.describe()}") + + output_tasks = pipeline.run(executor) + run_time_taken = time.perf_counter() - run_start_time + + # Calculate metrics from output tasks + # Count total images processed (sum of images in each ImageBatch) + num_images_processed = sum( + len(task.data) for task in output_tasks if task.data is not None + ) + + logger.success(f"Benchmark completed in {run_time_taken:.2f}s") + logger.success(f"Processed {num_images_processed} images") + logger.success(f"Output tasks: {len(output_tasks)}") + success = True + + except Exception as e: # noqa: BLE001 + error_traceback = traceback.format_exc() + logger.error(f"Benchmark failed: {e}") + logger.debug(f"Full traceback:\n{error_traceback}") + output_tasks = [] + run_time_taken = time.perf_counter() - run_start_time + num_images_processed = 0 + success = False + + return { + "params": { + "executor": args.executor, + "input_wds_dataset_dir": str(input_wds_dir), + "output_dataset_dir": str(output_dir), + "benchmark_results_path": str(args.benchmark_results_path), + "model_dir": args.model_dir, + "tar_files_per_partition": args.tar_files_per_partition, + "batch_size": args.batch_size, + "embedding_batch_size": args.embedding_batch_size, + "embedding_gpus_per_worker": args.embedding_gpus_per_worker, + "aesthetic_batch_size": args.aesthetic_batch_size, + "aesthetic_gpus_per_worker": args.aesthetic_gpus_per_worker, + "aesthetic_threshold": args.aesthetic_threshold, + "images_per_tar": args.images_per_tar, + }, + "metrics": { + "is_success": success, + "time_taken_s": run_time_taken, + "num_images_processed": num_images_processed, + "num_output_tasks": len(output_tasks), + "throughput_images_per_sec": num_images_processed / run_time_taken if run_time_taken > 0 else 0, + }, + "tasks": output_tasks, + } + + +def main() -> int: + """Main entry point for image pipeline benchmark.""" + ray_client = RayClient() + ray_client.start() + + parser = argparse.ArgumentParser( + description="Image curation pipeline benchmark with embedding generation and quality scoring" + ) + + # Benchmark-specific arguments + parser.add_argument( + "--benchmark-results-path", + type=Path, + required=True, + help="Path to write benchmark results", + ) + parser.add_argument( + "--executor", + default="xenna", + choices=["xenna", "ray_data"], + help="Executor to use for pipeline execution", + ) + + # Dataset arguments + parser.add_argument( + "--input-wds-dataset-dir", + type=str, + required=True, + help="Directory containing the input webdataset" + ) + parser.add_argument( + "--output-dataset-dir", + type=str, + required=True, + help="Directory to save the resulting webdataset" + ) + + # Image reader arguments + parser.add_argument( + "--tar-files-per-partition", + type=int, + default=1, + help="Number of tar files to process per partition (controls parallelism) for FilePartitioningStage" + ) + parser.add_argument( + "--batch-size", + type=int, + default=100, + help="Number of images per ImageBatch for the reader stage" + ) + + # General arguments + parser.add_argument( + "--model-dir", + type=str, + required=True, + help="Path to model directory containing all model weights" + ) + parser.add_argument( + "--verbose", + action="store_true", + default=False, + help="Enable verbose logging for all stages" + ) + + # Image reader arguments + parser.add_argument( + "--reader-num-threads", + type=int, + default=16, + help="Number of threads for image reading" + ) + parser.add_argument( + "--reader-gpus-per-worker", + type=float, + default=0.25, + help="GPU allocation per worker for image reading" + ) + + # Embedding stage arguments + parser.add_argument( + "--embedding-batch-size", + type=int, + default=32, + help="Batch size for embedding generation" + ) + parser.add_argument( + "--embedding-gpus-per-worker", + type=float, + default=0.25, + help="GPU allocation per worker for embedding generation" + ) + + # Aesthetic scoring arguments + parser.add_argument( + "--aesthetic-batch-size", + type=int, + default=32, + help="Batch size for aesthetic scoring" + ) + parser.add_argument( + "--aesthetic-gpus-per-worker", + type=float, + default=0.25, + help="GPU allocation per worker for aesthetic scoring" + ) + parser.add_argument( + "--aesthetic-threshold", + type=float, + default=0.5, + help="Aesthetic score threshold for filtering (images below this score will be filtered out)" + ) + + # Output dataset arguments + parser.add_argument( + "--images-per-tar", + type=int, + default=100, + help="Number of images per tar file in output dataset" + ) + + args = parser.parse_args() + + logger.info("=== Image Pipeline Benchmark Starting ===") + logger.info(f"Arguments: {vars(args)}") + + try: + results = run_image_pipeline_benchmark(args) + + except Exception as e: # noqa: BLE001 + error_traceback = traceback.format_exc() + print(f"Benchmark failed: {e}") + logger.debug(f"Full traceback:\n{error_traceback}") + results = { + "params": vars(args), + "metrics": { + "is_success": False, + }, + "tasks": [], + } + finally: + write_benchmark_results(results, args.benchmark_results_path) + ray_client.stop() + + # Return proper exit code based on success + return 0 if results["metrics"]["is_success"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/nemo_curator/stages/image/io/image_reader.py b/nemo_curator/stages/image/io/image_reader.py index 88f90ece56..f099794952 100644 --- a/nemo_curator/stages/image/io/image_reader.py +++ b/nemo_curator/stages/image/io/image_reader.py @@ -33,7 +33,7 @@ class ImageReaderStage(ProcessingStage[FileGroupTask, ImageBatch]): otherwise falls back to CPU decoding. """ - batch_size: int = 100 + dali_batch_size: int = 100 verbose: bool = True num_threads: int = 8 num_gpus_per_worker: float = 0.25 @@ -68,7 +68,7 @@ def _create_dali_pipeline(self, tar_paths: list[str]) -> object: raise RuntimeError(msg) from exc @pipeline_def( - batch_size=self.batch_size, + batch_size=self.dali_batch_size, num_threads=self.num_threads, device_id=0, # First device; unused for CPU-only DALI builds ) diff --git a/tests/stages/image/io/test_image_reader.py b/tests/stages/image/io/test_image_reader.py index bbae49312f..251f24ccf9 100644 --- a/tests/stages/image/io/test_image_reader.py +++ b/tests/stages/image/io/test_image_reader.py @@ -124,7 +124,7 @@ class _Types: def test_inputs_outputs_and_name() -> None: from nemo_curator.stages.image.io.image_reader import ImageReaderStage with patch("torch.cuda.is_available", return_value=True): - stage = ImageReaderStage(batch_size=3, verbose=False) + stage = ImageReaderStage(dali_batch_size=3, verbose=False) assert stage.inputs() == ([], []) assert stage.outputs() == (["data"], ["image_data", "image_path", "image_id"]) assert stage.name == "image_reader" @@ -134,7 +134,7 @@ def test_init_allows_cpu_when_no_cuda() -> None: from nemo_curator.stages.image.io.image_reader import ImageReaderStage # When CUDA is unavailable, the stage should initialize and use CPU DALI with patch("torch.cuda.is_available", return_value=False): - stage = ImageReaderStage(batch_size=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, verbose=False) assert stage is not None @@ -148,7 +148,7 @@ def test_process_streams_batches_from_dali() -> None: ) with patch("torch.cuda.is_available", return_value=True): - stage = ImageReaderStage(batch_size=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, verbose=False) with patch.object( ImageReaderStage, @@ -171,7 +171,7 @@ def test_process_raises_on_empty_task() -> None: empty = FileGroupTask(task_id="e1", dataset_name="ds", data=[]) with patch("torch.cuda.is_available", return_value=True): - stage = ImageReaderStage(batch_size=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, verbose=False) with pytest.raises(ValueError, match="No tar file paths"): stage.process(empty) @@ -182,7 +182,7 @@ def test_resources_with_cuda_available() -> None: from nemo_curator.stages.image.io.image_reader import ImageReaderStage # Instantiate with CUDA available so __post_init__ passes with patch("torch.cuda.is_available", return_value=True): - stage = ImageReaderStage(batch_size=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, verbose=False) res = stage.resources assert res.gpus == stage.num_gpus_per_worker @@ -193,7 +193,7 @@ def test_resources_without_cuda() -> None: from nemo_curator.stages.image.io.image_reader import ImageReaderStage # Create the stage without CUDA available with patch("torch.cuda.is_available", return_value=False): - stage = ImageReaderStage(batch_size=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, verbose=False) res = stage.resources assert res.gpus == 0 @@ -215,7 +215,7 @@ def test_dali_image_reader_on_gpu() -> None: from nemo_curator.stages.image.io.image_reader import ImageReaderStage from nemo_curator.tasks import FileGroupTask - stage = ImageReaderStage(batch_size=2, num_threads=2, verbose=False) + stage = ImageReaderStage(dali_batch_size=2, num_threads=2, verbose=False) task = FileGroupTask(task_id="t0", dataset_name="ds", data=[str(tar_path)]) batches = stage.process(task) diff --git a/tutorials/image/getting-started/image_curation_example.py b/tutorials/image/getting-started/image_curation_example.py index db03876ef5..6d58740cac 100644 --- a/tutorials/image/getting-started/image_curation_example.py +++ b/tutorials/image/getting-started/image_curation_example.py @@ -44,7 +44,7 @@ def create_image_curation_pipeline(args: argparse.Namespace) -> Pipeline: # Stage 1: Read images from webdataset tar files (now runs in parallel) pipeline.add_stage(ImageReaderStage( - batch_size=args.batch_size, + dali_batch_size=args.batch_size, verbose=args.verbose, # Force verbose to see debug info num_threads=16, # More threads for I/O num_gpus_per_worker=0.25, diff --git a/tutorials/image/getting-started/image_dedup_example.py b/tutorials/image/getting-started/image_dedup_example.py index b2cd4d0076..d8199cd7dd 100644 --- a/tutorials/image/getting-started/image_dedup_example.py +++ b/tutorials/image/getting-started/image_dedup_example.py @@ -45,7 +45,7 @@ def create_image_embedding_pipeline(args: argparse.Namespace) -> Pipeline: # Stage 1: Read images from webdataset tar files (now runs in parallel) pipeline.add_stage(ImageReaderStage( - batch_size=args.batch_size, + dali_batch_size=args.batch_size, verbose=args.verbose, num_threads=16, # More threads for I/O num_gpus_per_worker=0.25, @@ -97,7 +97,7 @@ def create_image_deduplication_pipeline(args: argparse.Namespace) -> Pipeline: # Stage 1: Read images from webdataset tar files (now runs in parallel) pipeline.add_stage(ImageReaderStage( - batch_size=args.batch_size, + dali_batch_size=args.batch_size, verbose=args.verbose, num_threads=16, # More threads for I/O num_gpus_per_worker=0.25,