Support inserting based on batchsize into shuffler

nemo_curator/backends/experimental/ray_actor_pool/shuffle_adapter.py

@@ -129,10 +129,10 @@ def read_and_insert(
     ) -> list[FileGroupTask]:
         """Read and insert tasks into the shuffler."""
         insert_kwargs = {"band_range": band_range} if band_range is not None else {}
-        results = []
-        for task in tasks:
-            results.append(self.stage.read_and_insert(task, **insert_kwargs))
-        return results
+        if hasattr(self.stage, "read_and_insert_batch"):
+            return self.stage.read_and_insert_batch(tasks, **insert_kwargs)
+        else:
+            return [self.stage.read_and_insert(task, **insert_kwargs) for task in tasks]
 
     def insert_finished(self) -> None:
         """Finish the insertion phase and trigger shuffle."""

nemo_curator/stages/deduplication/exact/identification.py

@@ -14,6 +14,8 @@
 
 from typing import TYPE_CHECKING, Any, Literal
 
+import cudf
+
 from nemo_curator.stages.deduplication.id_generator import CURATOR_DEDUP_ID_STR, get_id_generator_actor
 from nemo_curator.stages.deduplication.io_utils import DeduplicationIO
 from nemo_curator.stages.deduplication.shuffle_utils.rapidsmpf_shuffler import pylibcudf_to_cudf_dataframe
@@ -22,8 +24,6 @@
 from nemo_curator.utils.file_utils import get_fs
 
 if TYPE_CHECKING:
-    import cudf
-
     from nemo_curator.backends.base import WorkerMetadata
 
 
@@ -145,32 +145,99 @@ def inputs(self) -> tuple[list[str], list[str]]:
     def outputs(self) -> tuple[list[str], list[str]]:
         return (["data"], [])
 
-    def read_and_insert(self, task: FileGroupTask) -> FileGroupTask:
+    def _hash_and_insert(self, df: "cudf.DataFrame") -> None:
+        """Hash the text field and insert into the shuffle actor.
+
+        Parameters
+        ----------
+        df
+            DataFrame containing the id_field and text_field columns.
+        """
         self._check_actor_obj()
+        hashed_df = df[[self.id_field]]
+        hashed_df[EXACT_DUPLICATE_GROUP_FIELD] = df[self.text_field].hash_values(method="md5")
+        self.output_columns = list(hashed_df.columns)
+        self._actor_obj.insert_chunk(hashed_df, self.output_columns)
 
-        if self.assign_id_field and self.id_generator is None:
-            msg = "ID generator not initialized. Call setup() first."
-            raise RuntimeError(msg)
+    def _read_files(self, filepaths: list[str]) -> "cudf.DataFrame":
+        """Read files and return a DataFrame.
+
+        Parameters
+        ----------
+        filepaths
+            List of file paths to read.
+
+        Returns
+        -------
+        cudf.DataFrame
+            DataFrame containing the id_field and text_field columns.
+        """
         input_columns = [self.text_field] if self.assign_id_field else [self.text_field, self.id_field]
         if self.input_filetype == "jsonl":
-            df = self.read_jsonl(
-                filepath=task.data, columns=input_columns, assign_id=self.assign_id_field, **self.read_kwargs
-            )
+            read_func = self.read_jsonl
         elif self.input_filetype == "parquet":
-            df = self.read_parquet(
-                filepath=task.data, columns=input_columns, assign_id=self.assign_id_field, **self.read_kwargs
-            )
+            read_func = self.read_parquet
         else:
             msg = f"Unsupported input filetype: {self.input_filetype}"
             raise ValueError(msg)
+        return read_func(filepaths, columns=input_columns, assign_id=self.assign_id_field, **self.read_kwargs)
 
-        hashed_df = df[[self.id_field]]
-        hashed_df[EXACT_DUPLICATE_GROUP_FIELD] = df[self.text_field].hash_values(method="md5")
-        self.output_columns = list(hashed_df.columns)
-        self.dataset_name = task.dataset_name
+    def read_and_insert_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]:
+        """Batch process multiple file group tasks for exact deduplication.
 
-        self._actor_obj.insert_chunk(hashed_df, self.output_columns)
-        return task
+        This method reads all files from all tasks, concatenates them (if needed),
+        hashes the text field using MD5, and inserts into the shuffle actor for
+        deduplication. Processing tasks in batches significantly improves
+        throughput by reducing actor call overhead and enabling more efficient
+        GPU operations.
+
+        Parameters
+        ----------
+        tasks
+            List of FileGroupTask objects containing files to process.
+            Must contain at least one task.
+
+        Returns
+        -------
+        list[FileGroupTask]
+            The input tasks unchanged. The actual deduplication results are
+            written through the shuffle actor as a side effect.
+
+        Raises
+        ------
+        RuntimeError
+            If ID generator is not initialized when assign_id is True.
+        """
+        if self.assign_id_field and self.id_generator is None:
+            msg = "ID generator not initialized. Call setup() first."
+            raise RuntimeError(msg)
+
+        # Optimize for single-task batch to avoid unnecessary concat overhead
+        if len(tasks) == 1:
+            df = self._read_files(tasks[0].data)
+            self.dataset_name = tasks[0].dataset_name
+        else:
+            dfs = [self._read_files(task.data) for task in tasks]
+            df = cudf.concat(dfs, ignore_index=True)
+            self.dataset_name = tasks[0].dataset_name
+
+        self._hash_and_insert(df)
+        return tasks
+
+    def read_and_insert(self, task: FileGroupTask) -> FileGroupTask:
+        """Single task processing is not supported.
+
+        This stage requires batch processing via read_and_insert_batch for
+        optimal performance. The shuffle adapter will automatically route
+        tasks to the batch method.
+
+        Raises
+        ------
+        NotImplementedError
+            Always raised as this stage only supports batch processing.
+        """
+        msg = "ExactDuplicateIdentification only supports batch processing via read_and_insert_batch"
+        raise NotImplementedError(msg)
 
     def insert_finished(self) -> None:
         super().insert_finished()

tests/stages/deduplication/exact/test_identification.py

@@ -118,10 +118,14 @@ def exact_no_dedup_data_jsonl(tmp_path: Path) -> list[FileGroupTask]:
 @pytest.mark.gpu
 @pytest.mark.usefixtures("ray_client_with_id_generator")
 class TestExactDuplicates:
-    @pytest.mark.parametrize("assign_id", [True, False])
-    @pytest.mark.parametrize("total_nparts", [2, 4])
+    @pytest.mark.parametrize(("assign_id", "total_nparts", "batch_size"), [(False, 2, 1), (True, 4, 5)])
     def test_dup(
-        self, exact_dedup_data_parquet: list[FileGroupTask], tmpdir: Path, total_nparts: int, assign_id: bool
+        self,
+        exact_dedup_data_parquet: list[FileGroupTask],
+        tmpdir: Path,
+        total_nparts: int,
+        assign_id: bool,
+        batch_size: int,
     ) -> None:
         stage = ExactDuplicateIdentification(
             text_field="text",
@@ -133,7 +137,7 @@ def test_dup(
             enable_statistics=False,
             assign_id=assign_id,
             id_field="id" if not assign_id else None,
-        )
+        ).with_(batch_size=batch_size)
         pipeline = Pipeline(name="test_exact_dedup", stages=[stage])
         executor = RayActorPoolExecutor()
         pipeline.run(executor, initial_tasks=exact_dedup_data_parquet)