From bf3fd8533c506f52834542b994fe354e3c29eb36 Mon Sep 17 00:00:00 2001 From: Huy Vu2 Date: Wed, 4 Feb 2026 07:53:49 -0800 Subject: [PATCH] cherry pick commit, no benchmarking needed --- 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 +- .../image/getting-started/image_dedup_example.py | 4 ++-- 4 files changed, 12 insertions(+), 12 deletions(-) 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,