Prevent training data corruption during generation resumptions

deepspec/data/jsonl_dataset.py

@@ -9,6 +9,7 @@
 from tqdm import tqdm
 
 CACHE_DIR = os.path.expanduser("~/.cache/deepspec")
+LINE_INDEX_CACHE_VERSION = 2
 
 
 class JsonLineDataset(torch.utils.data.Dataset):
@@ -96,6 +97,7 @@ def _build_all_line_starts(self):
                     cached = pickle.load(handle)
                 if (
                     isinstance(cached, dict)
+                    and cached.get("version") == LINE_INDEX_CACHE_VERSION
                     and cached.get("file_key") == file_key
                     and isinstance(cached.get("line_starts"), list)
                 ):
@@ -105,25 +107,41 @@ def _build_all_line_starts(self):
                     continue
 
             handle = open(path, "rb")
+            if os.path.getsize(path) == 0:
+                starts = []
+                handle.close()
+                self.line_starts_per_file[idx] = starts
+                self.num_data_per_file.append(0)
+                if cache_path is not None:
+                    self._atomic_pickle_dump(
+                        {
+                            "version": LINE_INDEX_CACHE_VERSION,
+                            "file_key": file_key,
+                            "file_path": os.path.abspath(path),
+                            "line_starts": starts,
+                        },
+                        cache_path,
+                    )
+                continue
+
             mm = mmap.mmap(handle.fileno(), 0, access=mmap.ACCESS_READ)
             self.files[idx] = handle
             self.mmaps[idx] = mm
             starts = []
             mm.seek(0)
-            pos = 0
             while True:
-                starts.append(pos)
+                pos = mm.tell()
                 line = mm.readline()
                 if not line:
                     break
-                pos = mm.tell()
-            if starts and mm.size() == pos:
-                starts.pop()
+                if line.strip():
+                    starts.append(pos)
             self.line_starts_per_file[idx] = starts
             self.num_data_per_file.append(len(starts))
             if cache_path is not None:
                 self._atomic_pickle_dump(
                     {
+                        "version": LINE_INDEX_CACHE_VERSION,
                         "file_key": file_key,
                         "file_path": os.path.abspath(path),
                         "line_starts": starts,

scripts/data/generate_train_data.py

@@ -36,6 +36,8 @@ def parse_args():
 
 
 def validate_args(args):
+    if not args.output_file_path.endswith(".jsonl"):
+        raise ValueError("output-file-path must end with .jsonl")
     if not 0.0 <= args.temperature <= 1.0:
         raise ValueError("temperature must be between 0.0 and 1.0")
     if args.top_p is not None and not 0.0 <= args.top_p <= 1.0:
@@ -149,6 +151,13 @@ def count_lines(path):
         return sum(1 for _ in handle)
 
 
+def build_error_path(output_path):
+    root, ext = os.path.splitext(output_path)
+    if ext != ".jsonl":
+        raise ValueError("output-file-path must end with .jsonl")
+    return f"{root}_error.jsonl"
+
+
 def find_resume_offset(output_path, error_path):
     if not os.path.exists(output_path):
         return 0, 0, 0
@@ -231,14 +240,14 @@ def validate_servers(args):
 
 
 def write_finished_result(
-    future,
+    sample,
     output_handle,
     error_handle,
     stats,
 ):
-    sample = future.result()
     if sample["status"] == "error":
         error_handle.write(json.dumps(sample, ensure_ascii=False) + "\n")
+        error_handle.flush()
         stats["errors"] += 1
         return
 
@@ -252,6 +261,40 @@ def write_finished_result(
     stats["context_max"] = max(stats["context_max"], context_length)
     stats["success"] += 1
     output_handle.write(json.dumps(sample, ensure_ascii=False) + "\n")
+    output_handle.flush()
+
+
+def collect_finished_results(queues, pending_results):
+    for queue in queues.values():
+        for future in list(queue):
+            if future.done():
+                sample_index = getattr(future, "sample_index")
+                pending_results[sample_index] = future.result()
+                queue.remove(future)
+
+
+def count_outstanding_results(queues, pending_results):
+    return len(pending_results) + sum(len(queue) for queue in queues.values())
+
+
+def write_ready_results(
+    *,
+    next_write_index,
+    pending_results,
+    output_handle,
+    error_handle,
+    stats,
+):
+    while next_write_index in pending_results:
+        sample = pending_results.pop(next_write_index)
+        write_finished_result(
+            sample,
+            output_handle,
+            error_handle,
+            stats,
+        )
+        next_write_index += 1
+    return next_write_index
 
 
 def print_config(args):
@@ -275,7 +318,7 @@ def main():
     print_config(args)
 
     total_lines = count_lines(args.input_file_path)
-    error_path = args.output_file_path.replace(".jsonl", "_error.jsonl")
+    error_path = build_error_path(args.output_file_path)
     skip_lines, existing_success, existing_errors = (
         find_resume_offset(args.output_file_path, error_path)
         if args.resume
@@ -303,8 +346,11 @@ def main():
         "context_max": 0,
     }
     queues = {server_address: [] for server_address in valid_servers}
+    pending_results = {}
     next_server_index = 0
+    next_write_index = skip_lines
     submitted_count = 0
+    max_outstanding_results = args.concurrency * len(valid_servers)
 
     with (
         open(args.input_file_path, "r", encoding="utf-8") as input_handle,
@@ -329,27 +375,47 @@ def main():
             server_address = valid_servers[next_server_index]
             next_server_index = (next_server_index + 1) % len(valid_servers)
 
-            while len(queues[server_address]) >= args.concurrency:
-                wrote_result = False
-                for future in list(queues[server_address]):
-                    if future.done():
-                        write_finished_result(
-                            future, output_handle, error_handle, stats
-                        )
-                        queues[server_address].remove(future)
-                        wrote_result = True
-                        break
-                if not wrote_result:
+            while (
+                len(queues[server_address]) >= args.concurrency
+                or count_outstanding_results(queues, pending_results)
+                >= max_outstanding_results
+            ):
+                outstanding_count = count_outstanding_results(queues, pending_results)
+                collect_finished_results(queues, pending_results)
+                next_write_index = write_ready_results(
+                    next_write_index=next_write_index,
+                    pending_results=pending_results,
+                    output_handle=output_handle,
+                    error_handle=error_handle,
+                    stats=stats,
+                )
+                if count_outstanding_results(queues, pending_results) >= outstanding_count:
                     time.sleep(0.05)
 
             future = executor.submit(call_sglang, args, server_address, sample)
+            future.sample_index = skip_lines + submitted_count
             queues[server_address].append(future)
             submitted_count += 1
             progress.update(1)
 
-        for server_address in valid_servers:
-            for future in queues[server_address]:
-                write_finished_result(future, output_handle, error_handle, stats)
+        while any(queues.values()):
+            collect_finished_results(queues, pending_results)
+            next_write_index = write_ready_results(
+                next_write_index=next_write_index,
+                pending_results=pending_results,
+                output_handle=output_handle,
+                error_handle=error_handle,
+                stats=stats,
+            )
+            if any(queues.values()):
+                time.sleep(0.05)
+        next_write_index = write_ready_results(
+            next_write_index=next_write_index,
+            pending_results=pending_results,
+            output_handle=output_handle,
+            error_handle=error_handle,
+            stats=stats,
+        )
         progress.close()
 
     print("Processing completed.")