added check on rank 0 for input files on readers

huggingface · Dec 6, 2023 · c834995 · c834995
1 parent 3e3f0c8
commit c834995
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/datatrove/pipeline/readers/base.py b/src/datatrove/pipeline/readers/base.py
@@ -82,6 +82,12 @@ def read_files_shard(self, shard):
     def run(self, data: DocumentsPipeline = None, rank: int = 0, world_size: int = 1) -> DocumentsPipeline:
         if data:
             yield from data
-        for doc in self.read_files_shard(self.data_folder.get_files_shard(rank, world_size)):
+        files_shard = self.data_folder.get_files_shard(rank, world_size)
+        if len(files_shard) == 0:
+            if rank == 0:
+                raise RuntimeError(f"No files found on {self.data_folder.path}!")
+            # otherwise just a warning
+            logger.warning(f"No files found on {self.data_folder.path} for {rank=}")
+        for doc in self.read_files_shard(files_shard):
             self.update_doc_stats(doc)
             yield doc