Support inserting based on batchsize into shuffler - #1369
Conversation
Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 41ba9c0 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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) |
There was a problem hiding this comment.
The fallback to read_and_insert will always fail since this method now raises NotImplementedError. This creates a breaking change for any code that might call read_and_insert directly outside the adapter context.
| 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 read_and_insert(self, task: FileGroupTask) -> FileGroupTask: | |
| """Process a single task by delegating to batch method. | |
| Parameters | |
| ---------- | |
| task | |
| Single FileGroupTask to process. | |
| Returns | |
| ------- | |
| FileGroupTask | |
| The input task unchanged. | |
| """ | |
| return self.read_and_insert_batch([task])[0] |
| 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 |
There was a problem hiding this comment.
Missing validation for empty tasks list. If an empty list is passed, line 217 will cause an IndexError when accessing tasks[0].
| 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_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]: | |
| """Batch process multiple file group tasks for exact deduplication. | |
| 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. | |
| ValueError | |
| If tasks list is empty. | |
| """ | |
| if not tasks: | |
| msg = "tasks list must contain at least one task" | |
| raise ValueError(msg) | |
| 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 |
| dfs = [self._read_files(task.data) for task in tasks] | ||
| df = cudf.concat(dfs, ignore_index=True) | ||
| self.dataset_name = tasks[0].dataset_name |
There was a problem hiding this comment.
When batching multiple tasks with different dataset_name values, only the first task's name is used. This could cause issues if tasks from different datasets are batched together. Consider validating that all tasks have the same dataset name:
| dfs = [self._read_files(task.data) for task in tasks] | |
| df = cudf.concat(dfs, ignore_index=True) | |
| self.dataset_name = tasks[0].dataset_name | |
| else: | |
| dfs = [self._read_files(task.data) for task in tasks] | |
| # Validate all tasks have same dataset name | |
| dataset_names = {task.dataset_name for task in tasks} | |
| if len(dataset_names) > 1: | |
| msg = f"All tasks in a batch must have the same dataset_name, found: {dataset_names}" | |
| raise ValueError(msg) | |
| df = cudf.concat(dfs, ignore_index=True) | |
| self.dataset_name = tasks[0].dataset_name |
| 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 |
There was a problem hiding this comment.
Empty tasks list will cause IndexError on line 217/222. Add validation:
| 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 | |
| if not tasks: | |
| msg = "tasks list must contain at least one task" | |
| raise ValueError(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 |
| 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)]) |
There was a problem hiding this comment.
Test coverage reduced from 4 combinations (2 assign_id × 2 total_nparts) to just 2 specific combinations. Consider restoring full combinatorial testing or add more edge cases like (True, 2, 1) and (False, 4, 5)
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!
| if hasattr(self.stage, "read_and_insert_batch"): | ||
| return self.stage.read_and_insert_batch(tasks, **insert_kwargs) |
There was a problem hiding this comment.
band_range kwarg forwarded to read_and_insert_batch which doesn't accept it
When band_range is not None, the executor builds insert_kwargs = {"band_range": band_range} and passes it via **insert_kwargs to self.stage.read_and_insert_batch(tasks, **insert_kwargs). However, ExactDuplicateIdentification.read_and_insert_batch only accepts tasks and will raise a TypeError if band_range is ever non-None.
This is currently safe because band_range is only non-None for LSH stages, and LSH stages don't implement read_and_insert_batch yet. However, if read_and_insert_batch is added to any stage that also uses band_range (e.g. an LSH stage), it will silently break. Consider adding band_range as an optional parameter (defaulting to None) to read_and_insert_batch, or using **kwargs:
def read_and_insert_batch(self, tasks: list[FileGroupTask], band_range: tuple[int, int] | None = None) -> list[FileGroupTask]:Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
| If an integer is provided, it will be interpreted as bytes. | ||
| If a string is provided, it will be interpreted as a size with a unit. | ||
| If not provided, the default blocksize of 1GiB will be used. | ||
| identification_batchsize: str | int = "auto" |
There was a problem hiding this comment.
Docstring default value contradicts actual default
The docstring states identification_batchsize: str | int = "auto" but the actual parameter default on line 56 is 1. This will mislead callers who read the docs expecting "auto" to be the default behaviour, but in practice the pipeline runs with a batch size of 1 unless they explicitly set it.
| identification_batchsize: str | int = "auto" | |
| identification_batchsize: str | int = 1 |
| if self.identification_batchsize == "auto": | ||
| msg = "Auto batch size is not implemented yet" | ||
| raise NotImplementedError(msg) |
There was a problem hiding this comment.
"auto" batchsize guard runs after expensive I/O work
When initial_tasks=None (the common workflow path), the code first runs the entire input filegroups pipeline (file discovery + partitioning, lines 243-249), and only then raises NotImplementedError for "auto". This wastes potentially significant computation time.
The check should be moved to _validate_inputs() so it fails immediately on construction, before any work is done:
def _validate_inputs(self) -> None:
if self.perform_removal:
msg = "Removal is not implemented yet"
raise NotImplementedError(msg)
if self.identification_batchsize == "auto":
msg = "Auto batch size is not implemented yet"
raise NotImplementedError(msg)|
|
||
| from typing import TYPE_CHECKING, Any, Literal | ||
|
|
||
| import cudf |
There was a problem hiding this comment.
cudf promoted to hard runtime import
Previously cudf was guarded under TYPE_CHECKING so it was only imported by type checkers. Now it is a top-level runtime import (needed for cudf.concat in read_and_insert_batch). This is a breaking change for any code path that imports this module in a non-GPU environment — the import will raise ModuleNotFoundError even if the class is never instantiated. The test file already handles this with suppress(ImportError), but workflow.py imports ExactDuplicateIdentification unconditionally.
If non-GPU imports of the package need to remain possible, consider lazy-importing cudf inside the method that needs it:
def read_and_insert_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]:
import cudf # only needed at runtime for concat
...
| 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) | ||
|
|
There was a problem hiding this comment.
The call to _check_actor_obj() is deferred until _hash_and_insert, which runs after file reads (lines 190–197) and cudf.concat() (line 198). If the actor is not initialized, this check fails only after expensive GPU I/O and memory operations have completed.
Move the check to the top of read_and_insert_batch to catch misconfiguration immediately:
| 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) | |
| def read_and_insert_batch(self, tasks: list[FileGroupTask]) -> list[FileGroupTask]: | |
| """Batch process multiple file group tasks for exact deduplication. | |
| 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. | |
| """ | |
| self._check_actor_obj() | |
| if self.assign_id_field and self.id_generator is None: |
Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
… into batched-exact-dedup Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
| identification_batchsize: int = 1 | ||
| Number of batches to process in a single call for identification. | ||
| For example: A input_blocksize of 256MiB and identification_batchsize of 4 will result in ~1GB of data processed in a single call. |
There was a problem hiding this comment.
Misleading parameter description
The docstring says "Number of batches to process in a single call" but the parameter is actually a batch size — the number of tasks processed per call. "Number of batches to process" implies multiple batches per call, which is the opposite of what happens (multiple tasks are collapsed into one batch call). The example clarifies the intent, but the description sentence will confuse callers.
| identification_batchsize: int = 1 | |
| Number of batches to process in a single call for identification. | |
| For example: A input_blocksize of 256MiB and identification_batchsize of 4 will result in ~1GB of data processed in a single call. | |
| identification_batchsize: int = 1 | |
| Number of tasks to process in a single batch call for identification. | |
| For example: A input_blocksize of 256MiB and identification_batchsize of 4 will result in ~1GB of data processed in a single call. |
The same misleading wording appears in benchmarking/scripts/exact_dedup_identification_benchmark.py at the --identification-batchsize help string.
| 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 |
There was a problem hiding this comment.
All tasks' files loaded into GPU memory simultaneously before concat
When len(tasks) > 1, the code reads every task's data into a separate cuDF DataFrame via a synchronous list comprehension, holds all of them in GPU memory at once, and only then calls cudf.concat. For a batch_size of N with input_blocksize of B, the peak GPU memory required is N × B before any of it is released.
This is the correct design trade-off for throughput, but it is not documented and can silently cause GPU OOM errors for callers who set a large identification_batchsize without accounting for the multiplicative memory impact. At minimum, the docstring for read_and_insert_batch and identification_batchsize should warn about peak memory usage.
| hashed_df = df[[self.id_field]] | ||
| hashed_df[EXACT_DUPLICATE_GROUP_FIELD] = df[self.text_field].hash_values(method="md5") |
There was a problem hiding this comment.
SettingWithCopyWarning / silent no-op risk on cuDF slice
hashed_df = df[[self.id_field]]
hashed_df[EXACT_DUPLICATE_GROUP_FIELD] = df[self.text_field].hash_values(method="md5")df[[col]] returns a copy in cuDF (the double-bracket selection), so the column assignment does land on hashed_df and not on df. However, depending on the cuDF version this can raise a SettingWithCopyWarning. The safer and clearer idiom is to build the result in one shot:
| hashed_df = df[[self.id_field]] | |
| hashed_df[EXACT_DUPLICATE_GROUP_FIELD] = df[self.text_field].hash_values(method="md5") | |
| hashed_df = cudf.DataFrame({ | |
| self.id_field: df[self.id_field], | |
| EXACT_DUPLICATE_GROUP_FIELD: df[self.text_field].hash_values(method="md5"), | |
| }) |
This avoids any copy-on-write ambiguity and is explicit about intent.
Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>
Additional Comments (1)
Initialize workflow_result = None
try:
workflow = ExactDeduplicationWorkflow(...)
workflow_result = workflow.run(initial_tasks=None)
...And then the return line is already fine, since |
3 participants
Description
Unrelated to the core of this PR:
Usage
# Add snippet demonstrating usageChecklist