Warn and resolve conflict when both blocksize and files_per_partition are specified

nemo_curator/stages/file_partitioning.py

@@ -64,6 +64,12 @@ class FilePartitioningStage(ProcessingStage[_EmptyTask, FileGroupTask]):
 
     def __post_init__(self):
         """Initialize default values."""
+        if self.files_per_partition is not None and self.blocksize is not None:
+            logger.warning(
+                "Both 'files_per_partition' and 'blocksize' were specified. "
+                "'files_per_partition' will take precedence and 'blocksize' will be ignored."
+            )
+            self.blocksize = None
         if self.file_extensions is None:
             self.file_extensions = [".jsonl", ".json", ".parquet"]
         if self.storage_options is None:

tests/stages/common/test_file_partitioning.py

@@ -69,19 +69,35 @@ def test_initialization_default_values(self):
         assert stage.limit is None
         assert stage.name == "file_partitioning"
 
-    def test_initialization_custom_values(self):
-        """Test initialization with custom parameter values."""
+    def test_initialization_custom_values_with_files_per_partition(self):
+        """Test initialization with custom parameter values using files_per_partition."""
         stage = FilePartitioningStage(
             file_paths="/custom/path",
             files_per_partition=5,
-            blocksize="128MB",
             file_extensions=[".txt", ".json"],
             storage_options={"key": "value"},
             limit=3,
         )
 
         assert stage.file_paths == "/custom/path"
         assert stage.files_per_partition == 5
+        assert stage.blocksize is None
+        assert stage.file_extensions == [".txt", ".json"]
+        assert stage.storage_options == {"key": "value"}
+        assert stage.limit == 3
+
+    def test_initialization_custom_values_with_blocksize(self):
+        """Test initialization with custom parameter values using blocksize."""
+        stage = FilePartitioningStage(
+            file_paths="/custom/path",
+            blocksize="128MB",
+            file_extensions=[".txt", ".json"],
+            storage_options={"key": "value"},
+            limit=3,
+        )
+
+        assert stage.file_paths == "/custom/path"
+        assert stage.files_per_partition is None
         assert stage.blocksize == "128MB"
         assert stage.file_extensions == [".txt", ".json"]
         assert stage.storage_options == {"key": "value"}
@@ -194,6 +210,19 @@ def test_process_with_blocksize(self, empty_task: _EmptyTask, tmp_path: Path):
             assert len(task.data) == 1
             assert task.data[0] == test_files[i]
 
+    def test_both_blocksize_and_files_per_partition_warns(self, caplog: pytest.LogCaptureFixture):
+        """Test that specifying both blocksize and files_per_partition logs a warning and ignores blocksize."""
+        with caplog.at_level("WARNING"):
+            stage = FilePartitioningStage(
+                file_paths="/test/path",
+                files_per_partition=2,
+                blocksize="128MB",
+            )
+        assert stage.files_per_partition == 2
+        assert stage.blocksize is None
+        assert "files_per_partition" in caplog.text
+        assert "blocksize" in caplog.text
+
     def test_process_empty_file_list(self, empty_task: _EmptyTask):
         """Test processing with empty file list."""
         stage = FilePartitioningStage(file_paths=[])