Flush/upload incomplete buffer only on run end when enabled

crawler/shards.py

@@ -52,32 +52,8 @@ class ParquetShardWriter:
         await self.uploader.initialize()
 
     async def consume(self, record_queue: asyncio.Queue[dict[str, Any] | None]) -> None:
-        allow_incomplete_upload = (
-            self.config.enable_hf_upload and self.config.upload_incomplete_shards
-        )
-        loop = asyncio.get_running_loop()
-        last_flush_started = loop.time()
         while True:
-            try:
-                if allow_incomplete_upload:
-                    timeout = self.config.incomplete_shard_flush_seconds
-                    if self.buffer:
-                        elapsed = loop.time() - last_flush_started
-                        timeout = max(
-                            0.0,
-                            self.config.incomplete_shard_flush_seconds - elapsed,
-                        )
-                    item = await asyncio.wait_for(
-                        record_queue.get(),
-                        timeout=timeout,
-                    )
-                else:
-                    item = await record_queue.get()
-            except asyncio.TimeoutError:
-                if self.buffer:
-                    await self.flush()
-                    last_flush_started = loop.time()
-                continue
+            item = await record_queue.get()
 
             if item is None:
                 record_queue.task_done()
@@ -87,17 +63,18 @@ class ParquetShardWriter:
                 self.buffer.append(item)
                 if len(self.buffer) >= self.config.shard_size_rows:
                     await self.flush()
-                    last_flush_started = loop.time()
-                elif allow_incomplete_upload and self.buffer:
-                    elapsed = loop.time() - last_flush_started
-                    if elapsed >= self.config.incomplete_shard_flush_seconds:
-                        await self.flush()
-                        last_flush_started = loop.time()
             finally:
                 record_queue.task_done()
 
         if self.buffer:
-            await self.flush()
+            should_flush_incomplete = (
+                (not self.config.enable_hf_upload)
+                or self.config.upload_incomplete_shards
+            )
+            if should_flush_incomplete:
+                await self.flush()
+            else:
+                self.buffer = []
 
     async def flush(self) -> None:
         if not self.buffer: