Benchmarking script for image pipeline by huvunvidia · Pull Request #1441 · NVIDIA-NeMo/Curator

huvunvidia · 2026-01-29T03:38:43Z

Description

Usage

# Add snippet demonstrating usage

Checklist

I am familiar with the Contributing Guide.
New or Existing tests cover these changes.
The documentation is up to date with these changes.

copy-pr-bot · 2026-01-29T03:38:46Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

greptile-apps · 2026-01-29T03:41:19Z

Greptile Overview

Greptile Summary

This PR adds a comprehensive benchmarking script for the image curation pipeline and renames the batch_size parameter to dali_batch_size in ImageReaderStage for better clarity.

Key Changes

New benchmark script: image_pipeline_benchmark.py implements a complete image pipeline benchmark with proper metrics collection (num_images_processed, throughput_images_per_sec), benchmark framework integration via write_benchmark_results(), and support for multiple executors (Xenna, RayData)
Benchmark configuration: Updated nightly-benchmark.yaml to use the new script with proper arguments, timeout, resource allocation, and metric requirements (3800 images, 3.0+ images/sec)
API improvement: Renamed batch_size to dali_batch_size in ImageReaderStage to clarify it's specifically for DALI batching (not general batch size). This is a breaking change for external users
Consistent updates: All tests and tutorials updated to use dali_batch_size

The benchmark script properly addresses all previous feedback by including --benchmark-results-path, collecting metrics from output tasks, and writing results in the finally block.

Confidence Score: 4/5

This PR is safe to merge with minimal risk
Score reflects well-structured implementation that addresses all previous reviewer feedback. The benchmark script follows established patterns from other benchmark scripts, includes proper error handling and metrics collection. The parameter rename is a breaking change but is consistently applied across all internal code. Minor deduction for the breaking API change that may affect external users.
No files require special attention. The parameter rename in nemo_curator/stages/image/io/image_reader.py is a breaking change but intentional and well-executed.

Important Files Changed

Filename	Overview
benchmarking/scripts/image_pipeline_benchmark.py	New benchmark script with proper metrics collection and framework integration, addresses all previous feedback
benchmarking/nightly-benchmark.yaml	Updated config to use new benchmark script with proper args and requirements
nemo_curator/stages/image/io/image_reader.py	Renamed parameter from `batch_size` to `dali_batch_size` for clarity, breaking change for external users

Sequence Diagram

sequenceDiagram
    participant BF as Benchmark Framework
    participant Script as image_pipeline_benchmark.py
    participant Ray as RayClient
    participant Exec as Executor (Xenna/RayData)
    participant Pipeline as Image Pipeline
    participant FP as FilePartitioningStage
    participant IR as ImageReaderStage
    participant IE as ImageEmbeddingStage
    participant AF as AestheticFilterStage
    participant IW as ImageWriterStage
    participant Utils as utils.write_benchmark_results

    BF->>Script: Execute with --benchmark-results-path
    Script->>Ray: start()
    Script->>Script: parse arguments
    Script->>Script: run_image_pipeline_benchmark(args)
    Script->>Exec: setup_executor(args.executor)
    Exec-->>Script: executor instance
    Script->>Pipeline: create_image_curation_pipeline(args)
    Pipeline->>FP: add FilePartitioningStage
    Pipeline->>IR: add ImageReaderStage (dali_batch_size)
    Pipeline->>IE: add ImageEmbeddingStage
    Pipeline->>AF: add ImageAestheticFilterStage
    Pipeline->>IW: add ImageWriterStage
    Pipeline-->>Script: pipeline instance
    Script->>Pipeline: pipeline.run(executor)
    Pipeline->>FP: process tar files
    FP->>IR: ImageBatch with tar paths
    IR->>IE: ImageBatch with image_data
    IE->>AF: ImageBatch with embeddings
    AF->>IW: Filtered ImageBatch
    IW-->>Pipeline: output_tasks
    Pipeline-->>Script: output_tasks
    Script->>Script: Calculate metrics (num_images_processed, throughput)
    Script->>Utils: write_benchmark_results(results, path)
    Utils->>BF: Write params.json, metrics.json, tasks.pkl
    Script->>Ray: stop()
    Script-->>BF: Exit code (0=success, 1=failure)

greptile-apps

_{2 files reviewed, 5 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-01-29T03:41:31Z

Additional Comments (1)

benchmarking/nightly-benchmark.yaml
Missing required benchmark configuration. This entry lacks:

ray: section (num_cpus, num_gpus, enable_object_spilling)
timeout_s: value
sink_data: configuration
requirements: section for validating metrics

Other benchmark entries define these to properly configure resources and validate results.

greptile-apps

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-01-29T03:45:40Z

Additional Comments (1)

benchmarking/nightly-benchmark.yaml
Missing --benchmark-results-path={session_entry_dir} argument. All benchmark scripts require this argument to output metrics for the framework to parse. See other entries like domain_classification_raydata:84 for the expected pattern.

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

huvunvidia · 2026-01-29T03:56:06Z

/ok to test 96098b2

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

huvunvidia · 2026-01-30T04:53:42Z

/ok to test d316259

greptile-apps

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-01-30T04:56:59Z

+        num_threads=16,  # More threads for I/O
+        num_gpus_per_worker=0.25,


hardcoded num_threads=16 and num_gpus_per_worker=0.25 - consider making these configurable via command-line arguments for flexibility

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

I agree with this @huvunvidia

abhinavg4

Overall, Looks good to me.

greptile-apps

_{6 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-02T21:25:47Z

+
+from loguru import logger
+from utils import setup_executor, write_benchmark_results


[P0] utils import likely fails when invoked by the benchmarking runner

This script uses from utils import setup_executor, write_benchmark_results, but in this repo utils.py lives at benchmarking/scripts/utils.py and isn't a package. If the benchmark runner executes image_pipeline_benchmark.py without setting the working directory / PYTHONPATH to include benchmarking/scripts, this will raise ModuleNotFoundError: No module named 'utils'.

This is environment-dependent: it will work if cwd is benchmarking/scripts/ (or that dir is on PYTHONPATH), and fail otherwise. Consider importing via an explicit module path (or otherwise aligning with how the benchmarking harness invokes scripts).

greptile-apps · 2026-02-02T21:25:48Z

+    """Main entry point for image pipeline benchmark."""
+    ray_client = RayClient()
+    ray_client.start()


[P2] Ray is started unconditionally even for non-Ray executors

main() always does ray_client.start()/stop(), even when --executor=xenna. If the benchmark framework runs multiple entries back-to-back or Ray is already started externally, this can add overhead and may error depending on Ray configuration. Consider starting/stopping Ray only when the selected executor needs it (or making it idempotent).

greptile-apps · 2026-02-02T21:25:50Z

+        # 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
+        )


[P2] num_images_processed is derived from output tasks, which may undercount total processed

The benchmark reports num_images_processed as sum(len(task.data) for task in output_tasks ...). If later stages filter/drop items (e.g., aesthetic filter), the final output tasks may represent kept/written images rather than read/processed images. This makes the metric ambiguous and could regress silently if filter thresholds change.

If the intention is “images written”, it may be worth naming it accordingly or also emitting a separate “images_read/seen” metric from an earlier stage.

greptile-apps · 2026-02-02T21:25:51Z

    otherwise falls back to CPU decoding.
    """

-    batch_size: int = 100
+    dali_batch_size: int = 100
    verbose: bool = True


[P2] Renaming batch_size to dali_batch_size may break external call sites

Within this repo you updated tests/tutorials, but any downstream users instantiating ImageReaderStage(batch_size=...) will now get an unexpected keyword argument error. If this is intended, ignore; otherwise consider keeping batch_size as a deprecated alias (or documenting the breaking change in release notes).

greptile-apps

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

huvunvidia · 2026-02-03T00:19:08Z

/ok to test 6c2ce9a

huvunvidia · 2026-02-03T02:55:44Z

/ok to test ae0b502

greptile-apps

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

huvunvidia · 2026-02-03T19:38:54Z

/ok to test c859e54

greptile-apps

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

VibhuJawa · 2026-02-03T20:28:37Z

+    ray_client = RayClient()
+    ray_client.start()


Wont this start the ray cluster, dont we assume cluster is started somewhere else in the framework ? THis seems incorrect to me.

Right.
This is redundant.
As I understand from here, link, ray_client.start() will be skipped when there is an already running Ray, so it wouldn't affect the run.
But I will remove this in the amendment PR.

VibhuJawa · 2026-02-03T20:30:56Z

+            "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,


We dont track reader_num_threads on the reader here . Benchmark is dependent on reader_num_threads. Please track that too.

Can you add "args": vars(args) too , to catch problems

I can add this into an amendment PR.

workable benchmarking code

09207f7

greptile-apps Bot reviewed Jan 29, 2026

View reviewed changes

Comment thread benchmarking/scripts/image_pipeline_benchmark.py Outdated

Comment thread benchmarking/scripts/image_pipeline_benchmark.py Outdated

Comment thread benchmarking/scripts/image_pipeline_benchmark.py Outdated

Comment thread benchmarking/nightly-benchmark.yaml

Huy Vu2 added 2 commits January 28, 2026 19:42

fix format error

d7248bc

fix format error

629b125

greptile-apps Bot reviewed Jan 29, 2026

View reviewed changes

Comment thread benchmarking/scripts/image_pipeline_benchmark.py Outdated

greptile-apps Bot reviewed Jan 29, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into huvu/image_benchmarking

96098b2

copy-pr-bot Bot temporarily deployed to test January 29, 2026 03:56 Inactive

copy-pr-bot Bot had a problem deploying to nemo-ci January 29, 2026 03:56 Failure

copy-pr-bot Bot temporarily deployed to nemo-ci January 29, 2026 03:56 Inactive

greptile-apps Bot reviewed Jan 29, 2026

View reviewed changes

tested benchmarking run on test node

d316259

copy-pr-bot Bot temporarily deployed to test January 30, 2026 04:54 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci January 30, 2026 04:54 Inactive

greptile-apps Bot reviewed Jan 30, 2026

View reviewed changes

abhinavg4 approved these changes Jan 31, 2026

View reviewed changes

addressing PR comments

3011186

greptile-apps Bot reviewed Feb 2, 2026

View reviewed changes

update nightly-benchmark.yaml after testing on test node

6c2ce9a

greptile-apps Bot reviewed Feb 3, 2026

View reviewed changes

Huy Vu2 added 2 commits February 2, 2026 18:52

update nightly-benchmark.yaml

79f1b93

Merge remote-tracking branch 'origin/main' into huvu/image_benchmarking

ae0b502

greptile-apps Bot reviewed Feb 3, 2026

View reviewed changes

Merge branch 'main' into huvu/image_benchmarking

c859e54

greptile-apps Bot reviewed Feb 3, 2026

View reviewed changes

VibhuJawa reviewed Feb 3, 2026

View reviewed changes

huvunvidia mentioned this pull request Feb 4, 2026

Cherry pick of Image Reader efficiency fix #1461

Merged

3 tasks

sarahyurick mentioned this pull request Feb 11, 2026

Add relevant 26.02 docs to r1.1.0 #1493

Merged

44 tasks

		num_threads=16, # More threads for I/O
		num_gpus_per_worker=0.25,


		from loguru import logger
		from utils import setup_executor, write_benchmark_results

Uh oh!

Conversation

huvunvidia commented Jan 29, 2026

Description

Usage

Checklist

Uh oh!

copy-pr-bot Bot commented Jan 29, 2026

Uh oh!

greptile-apps Bot commented Jan 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Key Changes

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot commented Jan 29, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

greptile-apps Bot commented Jan 29, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

huvunvidia commented Jan 29, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

huvunvidia commented Jan 30, 2026

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Jan 30, 2026

Choose a reason for hiding this comment

Uh oh!

abhinavg4 Jan 31, 2026

Choose a reason for hiding this comment

Uh oh!

huvunvidia Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

abhinavg4 left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

huvunvidia commented Feb 3, 2026

Uh oh!

huvunvidia commented Feb 3, 2026

Uh oh!

greptile-apps Bot commented Jan 29, 2026 •

edited

Loading