diff --git a/deepspec/data/jsonl_dataset.py b/deepspec/data/jsonl_dataset.py index a3c2219..fd1f8a8 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) ): @@ -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, diff --git a/scripts/data/generate_train_data.py b/scripts/data/generate_train_data.py index 2ab156d..9c60f6e 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 @@ -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.")