From 09207f78f99c44f003609521f376e2da5de6a3ff Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Wed, 28 Jan 2026 19:35:32 -0800 Subject: [PATCH 1/9] workable benchmarking code --- benchmarking/nightly-benchmark.yaml | 12 +- .../scripts/image_pipeline_benchmark.py | 277 ++++++++++++++++++ 2 files changed, 283 insertions(+), 6 deletions(-) create mode 100644 benchmarking/scripts/image_pipeline_benchmark.py diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 9356ce2580..fe90c4e802 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -33,7 +33,7 @@ datasets: - name: "mscoco" formats: - type: "wds" - path: "{datasets_path}/mscoco/wds/truncated_100K_mscoco_benchmarking" + path: "{datasets_path}/mscoco/wds/full_mscoco_benchmarking" - name: "mscoco_model_weights" formats: - type: "files" @@ -437,15 +437,15 @@ 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 + --batch-size 1000 + --embedding-batch-size 500 + --aesthetic-batch-size 500 + --nsfw-batch-size 500 --tar-files-per-partition 10 --aesthetic-threshold 0.9 --nsfw-threshold 0.9 diff --git a/benchmarking/scripts/image_pipeline_benchmark.py b/benchmarking/scripts/image_pipeline_benchmark.py new file mode 100644 index 0000000000..248329feee --- /dev/null +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -0,0 +1,277 @@ +# 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. + +import argparse +import os +import time + +from nemo_curator.backends.xenna import XennaExecutor +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.filters.nsfw_filter import ImageNSFWFilterStage +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, aesthetic scoring, and NSFW detection 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( + 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, + )) + + # 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: Generate NSFW probability scores and filter + pipeline.add_stage(ImageNSFWFilterStage( + model_dir=args.model_dir, + num_gpus_per_worker=args.nsfw_gpus_per_worker, + model_inference_batch_size=args.nsfw_batch_size, + score_threshold=args.nsfw_threshold, + verbose=args.verbose, + )) + + # Stage 5: 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 main(args: argparse.Namespace) -> None: + """Main execution function for image curation pipeline.""" + + ray_client = RayClient() + ray_client.start() + + print("Starting image curation pipeline...") + print(f"Input parquet file: {args.input_parquet}") + print(f"Input webdataset directory: {args.input_wds_dataset_dir}") + print(f"Output webdataset directory: {args.output_dataset_dir}") + print(f"Model directory: {args.model_dir}") + print(f"Tar files per partition: {args.tar_files_per_partition}") + print(f"Task batch size: {args.batch_size}") + print("\n" + "=" * 50 + "\n") + + # Step 1: Download and prepare webdataset from parquet file + if not args.skip_download: + assert False, "Downloading is not supported in pipeline benchmark." + else: + print("Step 1: Skipping download (using existing dataset)") + print(f"Using existing dataset at: {args.input_wds_dataset_dir}") + print("\n" + "=" * 50 + "\n") + + # Step 2: Create and run curation pipeline + print("Step 2: Running image curation pipeline...") + start_time = time.time() + pipeline = create_image_curation_pipeline(args) + + # Print pipeline description + print(pipeline.describe()) + print("\n" + "=" * 50 + "\n") + + # Create executor + executor = XennaExecutor() + + # Execute pipeline + pipeline.run(executor) + + end_time = time.time() + + # Calculate and print execution time + execution_time = end_time - start_time + hours, remainder = divmod(execution_time, 3600) + minutes, seconds = divmod(remainder, 60) + + print("\nImage curation pipeline completed!") + print(f"Total execution time: {int(hours):02d}:{int(minutes):02d}:{seconds:.2f}") + print(f"Total execution time: {execution_time:.2f} seconds") + print(f"\nProcessed dataset available at: {args.output_dataset_dir}") + + ray_client.stop() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Image curation pipeline with embedding generation and quality scoring" + ) + + # Dataset arguments + parser.add_argument( + "--input-parquet", + type=str, + required=False, + default=None, + help="Path to input parquet file containing image URLs and metadata" + ) + parser.add_argument( + "--input-wds-dataset-dir", + type=str, + required=True, + help="Directory to save the downloaded webdataset" + ) + parser.add_argument( + "--output-dataset-dir", + type=str, + required=True, + help="Directory to save the resulting webdataset" + ) + parser.add_argument( + "--download-processes", + type=int, + default=8, + help="Number of parallel processes for downloading images" + ) + parser.add_argument( + "--entries-per-tar", + type=int, + default=1000, + help="Number of entries per tar shard during download" + ) + parser.add_argument( + "--skip-download", + action="store_true", + default=False, + help="Skip dataset download and use existing 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" + ) + + # 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)" + ) + + # NSFW scoring arguments + parser.add_argument( + "--nsfw-batch-size", + type=int, + default=32, + help="Batch size for NSFW scoring" + ) + parser.add_argument( + "--nsfw-gpus-per-worker", + type=float, + default=0.25, + help="GPU allocation per worker for NSFW scoring" + ) + parser.add_argument( + "--nsfw-threshold", + type=float, + default=0.5, + help="NSFW score threshold for filtering (images above this score will be filtered out as NSFW)" + ) + + # 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() + main(args) From d7248bcdb8576abb2a3cba344a9942a166f76372 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Wed, 28 Jan 2026 19:42:25 -0800 Subject: [PATCH 2/9] fix format error --- benchmarking/scripts/image_pipeline_benchmark.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/benchmarking/scripts/image_pipeline_benchmark.py b/benchmarking/scripts/image_pipeline_benchmark.py index 248329feee..438baebc44 100644 --- a/benchmarking/scripts/image_pipeline_benchmark.py +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -13,7 +13,6 @@ # limitations under the License. import argparse -import os import time from nemo_curator.backends.xenna import XennaExecutor @@ -103,7 +102,7 @@ def main(args: argparse.Namespace) -> None: # Step 1: Download and prepare webdataset from parquet file if not args.skip_download: - assert False, "Downloading is not supported in pipeline benchmark." + raise AssertionError("Downloading is not supported in pipeline benchmark.") else: print("Step 1: Skipping download (using existing dataset)") print(f"Using existing dataset at: {args.input_wds_dataset_dir}") From 629b125346abb1181d11b27480addce3f084a565 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Wed, 28 Jan 2026 19:45:26 -0800 Subject: [PATCH 3/9] fix format error --- benchmarking/scripts/image_pipeline_benchmark.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarking/scripts/image_pipeline_benchmark.py b/benchmarking/scripts/image_pipeline_benchmark.py index 438baebc44..eb7c635c01 100644 --- a/benchmarking/scripts/image_pipeline_benchmark.py +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -102,7 +102,8 @@ def main(args: argparse.Namespace) -> None: # Step 1: Download and prepare webdataset from parquet file if not args.skip_download: - raise AssertionError("Downloading is not supported in pipeline benchmark.") + msg = "Downloading is not supported in pipeline benchmark." + raise AssertionError(msg) else: print("Step 1: Skipping download (using existing dataset)") print(f"Using existing dataset at: {args.input_wds_dataset_dir}") From 8707bfefeed06615031583f67bcc6d79dd27c429 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Wed, 28 Jan 2026 21:19:17 -0800 Subject: [PATCH 4/9] actual benchmarking script --- benchmarking/nightly-benchmark.yaml | 34 ++- .../scripts/image_pipeline_benchmark.py | 243 ++++++++++-------- 2 files changed, 155 insertions(+), 122 deletions(-) diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index fe90c4e802..0e63011e99 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -439,19 +439,29 @@ entries: enabled: true 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 1000 - --embedding-batch-size 500 - --aesthetic-batch-size 500 - --nsfw-batch-size 500 - --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=10 + --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_generated + exact_value: 11836 + - metric: throughput_images_per_sec + min_value: 8.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 index eb7c635c01..d7b6ea3695 100644 --- a/benchmarking/scripts/image_pipeline_benchmark.py +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -12,22 +12,32 @@ # 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.backends.xenna import XennaExecutor 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.filters.nsfw_filter import ImageNSFWFilterStage 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, aesthetic scoring, and NSFW detection stages.""" + """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") @@ -65,16 +75,7 @@ def create_image_curation_pipeline(args: argparse.Namespace) -> Pipeline: verbose=args.verbose, )) - # Stage 4: Generate NSFW probability scores and filter - pipeline.add_stage(ImageNSFWFilterStage( - model_dir=args.model_dir, - num_gpus_per_worker=args.nsfw_gpus_per_worker, - model_inference_batch_size=args.nsfw_batch_size, - score_threshold=args.nsfw_threshold, - verbose=args.verbose, - )) - - # Stage 5: Write down to disk + # Stage 4: Write down to disk pipeline.add_stage(ImageWriterStage( output_dir=args.output_dataset_dir, images_per_tar=args.images_per_tar, @@ -85,78 +86,111 @@ def create_image_curation_pipeline(args: argparse.Namespace) -> Pipeline: return pipeline -def main(args: argparse.Namespace) -> None: - """Main execution function for image curation pipeline.""" - - ray_client = RayClient() - ray_client.start() - - print("Starting image curation pipeline...") - print(f"Input parquet file: {args.input_parquet}") - print(f"Input webdataset directory: {args.input_wds_dataset_dir}") - print(f"Output webdataset directory: {args.output_dataset_dir}") - print(f"Model directory: {args.model_dir}") - print(f"Tar files per partition: {args.tar_files_per_partition}") - print(f"Task batch size: {args.batch_size}") - print("\n" + "=" * 50 + "\n") - - # Step 1: Download and prepare webdataset from parquet file - if not args.skip_download: - msg = "Downloading is not supported in pipeline benchmark." - raise AssertionError(msg) - else: - print("Step 1: Skipping download (using existing dataset)") - print(f"Using existing dataset at: {args.input_wds_dataset_dir}") - print("\n" + "=" * 50 + "\n") - - # Step 2: Create and run curation pipeline - print("Step 2: Running image curation pipeline...") - start_time = time.time() - pipeline = create_image_curation_pipeline(args) - - # Print pipeline description - print(pipeline.describe()) - print("\n" + "=" * 50 + "\n") - - # Create executor - executor = XennaExecutor() - - # Execute pipeline - pipeline.run(executor) +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) - end_time = time.time() + 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) - # Calculate and print execution time - execution_time = end_time - start_time - hours, remainder = divmod(execution_time, 3600) - minutes, seconds = divmod(remainder, 60) + 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}") - print("\nImage curation pipeline completed!") - print(f"Total execution time: {int(hours):02d}:{int(minutes):02d}:{seconds:.2f}") - print(f"Total execution time: {execution_time:.2f} seconds") - print(f"\nProcessed dataset available at: {args.output_dataset_dir}") - - ray_client.stop() + # 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() -if __name__ == "__main__": parser = argparse.ArgumentParser( - description="Image curation pipeline with embedding generation and quality scoring" + description="Image curation pipeline benchmark with embedding generation and quality scoring" ) - # Dataset arguments + # Benchmark-specific arguments parser.add_argument( - "--input-parquet", - type=str, - required=False, - default=None, - help="Path to input parquet file containing image URLs and metadata" + "--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 to save the downloaded webdataset" + help="Directory containing the input webdataset" ) parser.add_argument( "--output-dataset-dir", @@ -164,24 +198,6 @@ def main(args: argparse.Namespace) -> None: required=True, help="Directory to save the resulting webdataset" ) - parser.add_argument( - "--download-processes", - type=int, - default=8, - help="Number of parallel processes for downloading images" - ) - parser.add_argument( - "--entries-per-tar", - type=int, - default=1000, - help="Number of entries per tar shard during download" - ) - parser.add_argument( - "--skip-download", - action="store_true", - default=False, - help="Skip dataset download and use existing webdataset" - ) # Image reader arguments parser.add_argument( @@ -245,26 +261,6 @@ def main(args: argparse.Namespace) -> None: help="Aesthetic score threshold for filtering (images below this score will be filtered out)" ) - # NSFW scoring arguments - parser.add_argument( - "--nsfw-batch-size", - type=int, - default=32, - help="Batch size for NSFW scoring" - ) - parser.add_argument( - "--nsfw-gpus-per-worker", - type=float, - default=0.25, - help="GPU allocation per worker for NSFW scoring" - ) - parser.add_argument( - "--nsfw-threshold", - type=float, - default=0.5, - help="NSFW score threshold for filtering (images above this score will be filtered out as NSFW)" - ) - # Output dataset arguments parser.add_argument( "--images-per-tar", @@ -274,4 +270,31 @@ def main(args: argparse.Namespace) -> None: ) args = parser.parse_args() - main(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()) From 4f24401c0ee2690adad77f5251b3b46324dc91c9 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Thu, 29 Jan 2026 06:49:06 -0800 Subject: [PATCH 5/9] fix error --- benchmarking/nightly-benchmark.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 0e63011e99..3633d17287 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -457,7 +457,7 @@ entries: enable_object_spilling: false requirements: # ensure the total number of documents processed is correct - - metric: num_images_generated + - metric: num_images_processed exact_value: 11836 - metric: throughput_images_per_sec min_value: 8.0 From d31625910421f45cd3dca5793fa56fc31cb9ec79 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Thu, 29 Jan 2026 20:52:18 -0800 Subject: [PATCH 6/9] tested benchmarking run on test node --- benchmarking/nightly-benchmark.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 3633d17287..31e126c343 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -450,7 +450,7 @@ entries: --aesthetic-batch-size=500 --aesthetic-threshold=0.9 --verbose - timeout_s: 1500 + timeout_s: 5000 ray: num_cpus: 64 num_gpus: 4 @@ -460,7 +460,7 @@ entries: - metric: num_images_processed exact_value: 11836 - metric: throughput_images_per_sec - min_value: 8.0 + min_value: 2.0 - name: audio_fleurs enabled: true From 301118649c1b577c852568effce6225c54ea1b1c Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Mon, 2 Feb 2026 13:21:51 -0800 Subject: [PATCH 7/9] addressing PR comments --- .../scripts/image_pipeline_benchmark.py | 20 ++++++++++++++++--- nemo_curator/stages/image/io/image_reader.py | 4 ++-- tests/stages/image/io/test_image_reader.py | 14 ++++++------- .../getting-started/image_curation_example.py | 2 +- .../getting-started/image_dedup_example.py | 4 ++-- 5 files changed, 29 insertions(+), 15 deletions(-) diff --git a/benchmarking/scripts/image_pipeline_benchmark.py b/benchmarking/scripts/image_pipeline_benchmark.py index d7b6ea3695..4cc8744918 100644 --- a/benchmarking/scripts/image_pipeline_benchmark.py +++ b/benchmarking/scripts/image_pipeline_benchmark.py @@ -51,10 +51,10 @@ 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, + 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 @@ -227,6 +227,20 @@ def main() -> int: 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", 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, From 6c2ce9addc3deada039d6818338c487e91c4df53 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Mon, 2 Feb 2026 16:08:03 -0800 Subject: [PATCH 8/9] update nightly-benchmark.yaml after testing on test node --- benchmarking/nightly-benchmark.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 31e126c343..13db6d18dc 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -444,13 +444,13 @@ entries: --output-dataset-dir={session_entry_dir}/scratch/output --model-dir={dataset:mscoco_model_weights,files} --executor=xenna - --tar-files-per-partition=10 + --tar-files-per-partition=1 --batch-size=1000 --embedding-batch-size=500 --aesthetic-batch-size=500 --aesthetic-threshold=0.9 --verbose - timeout_s: 5000 + timeout_s: 1500 ray: num_cpus: 64 num_gpus: 4 @@ -458,9 +458,9 @@ entries: requirements: # ensure the total number of documents processed is correct - metric: num_images_processed - exact_value: 11836 + exact_value: 3800 - metric: throughput_images_per_sec - min_value: 2.0 + min_value: 3.0 - name: audio_fleurs enabled: true From 79f1b930466fd41068706a9e5186ba1b054b57d2 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Mon, 2 Feb 2026 18:52:26 -0800 Subject: [PATCH 9/9] update nightly-benchmark.yaml --- benchmarking/nightly-benchmark.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarking/nightly-benchmark.yaml b/benchmarking/nightly-benchmark.yaml index 13db6d18dc..c195bfe17e 100644 --- a/benchmarking/nightly-benchmark.yaml +++ b/benchmarking/nightly-benchmark.yaml @@ -33,7 +33,7 @@ datasets: - name: "mscoco" formats: - type: "wds" - path: "{datasets_path}/mscoco/wds/full_mscoco_benchmarking" + path: "{datasets_path}/mscoco/wds/truncated_100K_mscoco_benchmarking" - name: "mscoco_model_weights" formats: - type: "files"