From 9c9f2b96a9ce55b0e6ac6083268648e55f9840b8 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 20:39:15 +0000 Subject: [PATCH 1/5] Write generated data in input order for safe resume --- scripts/data/generate_train_data.py | 78 +++++++++++++++++++++++------ 1 file changed, 63 insertions(+), 15 deletions(-) diff --git a/scripts/data/generate_train_data.py b/scripts/data/generate_train_data.py index 2ab156d0..dbd679ec 100644 --- a/scripts/data/generate_train_data.py +++ b/scripts/data/generate_train_data.py @@ -231,14 +231,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 +252,36 @@ 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 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): @@ -303,7 +333,9 @@ 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 with ( @@ -330,26 +362,42 @@ def main(): 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: + queue_len = len(queues[server_address]) + 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 len(queues[server_address]) >= queue_len: 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.") From 24863443ecbf9d5cb829916bc540feeee444ad3e Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 20:39:38 +0000 Subject: [PATCH 2/5] Validate train data output JSONL paths --- scripts/data/generate_train_data.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/scripts/data/generate_train_data.py b/scripts/data/generate_train_data.py index dbd679ec..0fad5561 100644 --- a/scripts/data/generate_train_data.py +++ b/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 @@ -305,7 +314,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 From 7962e8d7174bb80f0cae0b139b11359069b026c3 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 20:41:43 +0000 Subject: [PATCH 3/5] Skip blank lines when indexing JSONL files --- deepspec/data/jsonl_dataset.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/deepspec/data/jsonl_dataset.py b/deepspec/data/jsonl_dataset.py index a3c22191..71febfba 100644 --- a/deepspec/data/jsonl_dataset.py +++ b/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) ): @@ -110,20 +112,19 @@ def _build_all_line_starts(self): 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, From 3c61533f415b84cd2f09a5cda1473f9dc8c2af06 Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 20:42:00 +0000 Subject: [PATCH 4/5] Handle empty JSONL files as zero records --- deepspec/data/jsonl_dataset.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/deepspec/data/jsonl_dataset.py b/deepspec/data/jsonl_dataset.py index 71febfba..fd1f8a83 100644 --- a/deepspec/data/jsonl_dataset.py +++ b/deepspec/data/jsonl_dataset.py @@ -107,6 +107,23 @@ 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 From a3971c60c22f65f4d5479cad9028253478c5337c Mon Sep 17 00:00:00 2001 From: morluto Date: Tue, 30 Jun 2026 21:30:45 +0000 Subject: [PATCH 5/5] Bound ordered JSONL generation buffering --- scripts/data/generate_train_data.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/scripts/data/generate_train_data.py b/scripts/data/generate_train_data.py index 0fad5561..9c60f6e9 100644 --- a/scripts/data/generate_train_data.py +++ b/scripts/data/generate_train_data.py @@ -273,6 +273,10 @@ def collect_finished_results(queues, pending_results): 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, @@ -346,6 +350,7 @@ def main(): 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, @@ -370,8 +375,12 @@ 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: - queue_len = len(queues[server_address]) + 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, @@ -380,7 +389,7 @@ def main(): error_handle=error_handle, stats=stats, ) - if len(queues[server_address]) >= queue_len: + if count_outstanding_results(queues, pending_results) >= outstanding_count: time.sleep(0.05) future = executor.submit(call_sglang, args, server_address, sample)