[Batched Storage] Overwrite file when batch_step=0 (#432)

Author: fatty-belly

dataflow/utils/storage.py

@@ -1060,7 +1060,8 @@ def clean_surrogates(obj):
         os.makedirs(os.path.dirname(file_path), exist_ok=True)
         self.logger.success(f"Writing data to {file_path} with type {self.cache_type}")
         if self.cache_type == "jsonl":
-            with open(file_path, 'a', encoding='utf-8') as f:
+            open_mode = 'w' if self.batch_step == 0 else 'a'
+            with open(file_path, open_mode, encoding='utf-8') as f:
                 dataframe.to_json(f, orient="records", lines=True, force_ascii=False)
         elif self.cache_type == "csv":
             if self.batch_step == 0:

test/test_batched_pipeline.py

@@ -56,4 +56,5 @@ def forward(self):
 if __name__ == "__main__":
     pipeline = AutoOPPipeline()
     pipeline.compile()
-    pipeline.forward(batch_size=2, resume_from_last=True)
+    pipeline.forward(batch_size=2, resume_from_last=True)
+    pipeline.forward(batch_size=2, resume_from_last=False)  # should overwrite