import argparse import json import os import random import time from concurrent.futures import ThreadPoolExecutor, as_completed from openai import OpenAI from tqdm import tqdm def parse_args(): parser = argparse.ArgumentParser( description="Regenerate JSONL conversations through OpenAI-compatible sglang servers." ) parser.add_argument("--model", required=True) parser.add_argument("--server-address", nargs="+", required=True) parser.add_argument("--input-file-path", required=True) parser.add_argument("--output-file-path", required=True) parser.add_argument("--concurrency", type=int, default=64) parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--top-p", type=float, default=None) parser.add_argument("--top-k", type=int, default=None) parser.add_argument("--min-p", type=float, default=None) parser.add_argument("--repetition-penalty", type=float, default=None) parser.add_argument("--max-tokens", type=int, default=4096) parser.add_argument("--num-samples", type=int, default=None) parser.add_argument("--resume", action="store_true") parser.add_argument("--is-reasoning-model", action="store_true") parser.add_argument("--is-gpt-oss", action="store_true") thinking_group = parser.add_mutually_exclusive_group() thinking_group.add_argument("--enable-thinking", action="store_true") thinking_group.add_argument("--disable-thinking", action="store_true") return parser.parse_args() def validate_args(args): 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: raise ValueError("top-p must be between 0.0 and 1.0") if args.top_k is not None and args.top_k <= 0: raise ValueError("top-k must be greater than 0") if args.min_p is not None and not 0.0 <= args.min_p <= 1.0: raise ValueError("min-p must be between 0.0 and 1.0") if args.max_tokens <= 0: raise ValueError("max-tokens must be greater than 0") if args.concurrency <= 0: raise ValueError("concurrency must be greater than 0") def get_random_reasoning_effort(): return random.choices(["low", "medium", "high"], weights=[4, 4, 2], k=1)[0] def compute_context_length(conversations): length = 0 for message in conversations: content = message.get("content") if isinstance(content, str): length += len(content.split()) elif isinstance(content, list): for part in content: if isinstance(part, dict) and isinstance(part.get("text"), str): length += len(part["text"].split()) return length def build_query_kwargs(args, messages, max_tokens=None): query_kwargs = { "model": args.model, "messages": messages, "max_tokens": args.max_tokens if max_tokens is None else max_tokens, "temperature": args.temperature, "stream": False, } if args.top_p is not None: query_kwargs["top_p"] = args.top_p if args.repetition_penalty is not None: query_kwargs["presence_penalty"] = args.repetition_penalty extra_body = {} if args.top_k is not None: extra_body["top_k"] = args.top_k if args.min_p is not None: extra_body["min_p"] = args.min_p if args.enable_thinking: extra_body.setdefault("chat_template_kwargs", {})["enable_thinking"] = True if args.disable_thinking: extra_body.setdefault("chat_template_kwargs", {})["enable_thinking"] = False if extra_body: query_kwargs["extra_body"] = extra_body if args.is_gpt_oss: query_kwargs["reasoning_effort"] = get_random_reasoning_effort() return query_kwargs def error_sample(sample, message): sample["status"] = "error" sample["error"] = message return sample def call_sglang(args, server_address, sample, max_tokens=None): conversations = sample.get("conversations") if not conversations: return error_sample(sample, "Missing conversations") if conversations[0].get("role") == "assistant": return error_sample(sample, "Data starts with an assistant message") client = OpenAI(base_url=f"http://{server_address}/v1", api_key="None") regenerated = [] for message in conversations: role = message.get("role") if role == "system": regenerated.append(message) continue if role == "assistant": continue if role != "user": return error_sample(sample, f"Invalid message role: {role}") regenerated.append(message) try: response = client.chat.completions.create( **build_query_kwargs(args, regenerated, max_tokens=max_tokens) ) except Exception as exc: return error_sample(sample, str(exc)) response_message = { "role": "assistant", "content": response.choices[0].message.content, } if args.is_reasoning_model: response_message["thinking"] = response.choices[0].message.reasoning_content regenerated.append(response_message) sample["conversations"] = regenerated sample["status"] = "success" return sample def count_lines(path): with open(path, "r", encoding="utf-8") as handle: return sum(1 for _ in handle) def find_resume_offset(output_path, error_path): if not os.path.exists(output_path): return 0, 0, 0 success_count = count_lines(output_path) error_count = count_lines(error_path) if os.path.exists(error_path) else 0 return success_count + error_count, success_count, error_count def validate_server(args, server_address, probe): start_time = time.perf_counter() try: result = call_sglang(args, server_address, dict(probe), max_tokens=1) except Exception as exc: result = {"status": "error", "error": str(exc)} elapsed = time.perf_counter() - start_time return server_address, result, elapsed def validate_servers(args): probe = {"conversations": [{"role": "user", "content": "Hello"}]} server_results = {} server_count = len(args.server_address) print(f"Validating {server_count} sglang servers in parallel...", flush=True) with ThreadPoolExecutor(max_workers=server_count) as executor: future_to_server = { executor.submit( validate_server, args, server_address, probe, ): server_address for server_address in args.server_address } for completed_count, future in enumerate(as_completed(future_to_server), 1): server_address = future_to_server[future] try: _, result, elapsed = future.result() except Exception as exc: result = {"status": "error", "error": str(exc)} elapsed = 0.0 if result.get("status") == "success": server_results[server_address] = True print( f"[validate {completed_count}/{server_count}] " f"ok server {server_address} elapsed={elapsed:.2f}s", flush=True, ) else: server_results[server_address] = False print( f"[validate {completed_count}/{server_count}] " f"skip server {server_address} elapsed={elapsed:.2f}s: " f"{result.get('error')}", flush=True, ) valid_servers = [ server_address for server_address in args.server_address if server_results[server_address] ] invalid_servers = [ server_address for server_address in args.server_address if not server_results[server_address] ] print(f"Available servers ({len(valid_servers)}/{server_count}): {valid_servers}") print( f"Unavailable servers ({len(invalid_servers)}/{server_count}): " f"{invalid_servers}" ) if not valid_servers: raise RuntimeError("No available sglang server") return valid_servers def write_finished_result( future, output_handle, error_handle, stats, ): sample = future.result() if sample["status"] == "error": error_handle.write(json.dumps(sample, ensure_ascii=False) + "\n") stats["errors"] += 1 return context_length = compute_context_length(sample.get("conversations", [])) stats["context_sum"] += context_length stats["context_min"] = ( context_length if stats["context_min"] is None else min(stats["context_min"], context_length) ) stats["context_max"] = max(stats["context_max"], context_length) stats["success"] += 1 output_handle.write(json.dumps(sample, ensure_ascii=False) + "\n") def print_config(args): print("Configuration:") print(f" model: {args.model}") print(f" servers: {args.server_address}") print(f" input: {args.input_file_path}") print(f" output: {args.output_file_path}") print(f" concurrency: {args.concurrency}") print(f" max_tokens: {args.max_tokens}") print(f" temperature: {args.temperature}") print(f" top_p: {args.top_p}") print(f" top_k: {args.top_k}") print(f" min_p: {args.min_p}") print(f" resume: {args.resume}") def main(): args = parse_args() validate_args(args) print_config(args) total_lines = count_lines(args.input_file_path) error_path = args.output_file_path.replace(".jsonl", "_error.jsonl") skip_lines, existing_success, existing_errors = ( find_resume_offset(args.output_file_path, error_path) if args.resume else (0, 0, 0) ) if skip_lines >= total_lines: print(f"All {total_lines} samples are already processed.") return if args.resume and skip_lines > 0: print( "Resume mode: " f"{existing_success} success, {existing_errors} errors, skip {skip_lines}" ) valid_servers = validate_servers(args) print(f"Using servers: {valid_servers}") file_mode = "a" if args.resume and skip_lines > 0 else "w" stats = { "success": 0, "errors": 0, "context_sum": 0, "context_min": None, "context_max": 0, } queues = {server_address: [] for server_address in valid_servers} next_server_index = 0 submitted_count = 0 with ( open(args.input_file_path, "r", encoding="utf-8") as input_handle, open(args.output_file_path, file_mode, encoding="utf-8") as output_handle, open(error_path, file_mode, encoding="utf-8") as error_handle, ThreadPoolExecutor(max_workers=args.concurrency * len(valid_servers)) as executor, ): for _ in range(skip_lines): next(input_handle, None) progress_total = ( total_lines if args.num_samples is None else min(total_lines, skip_lines + args.num_samples) ) progress = tqdm(total=progress_total, initial=skip_lines, desc="Processing") for line in input_handle: if args.num_samples is not None and submitted_count >= args.num_samples: break sample = json.loads(line) 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: time.sleep(0.05) future = executor.submit(call_sglang, args, server_address, sample) 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) progress.close() print("Processing completed.") print(f" success: {stats['success']}") print(f" errors: {stats['errors']}") if stats["success"] > 0: avg_context = stats["context_sum"] / stats["success"] print(f" context_min: {stats['context_min']}") print(f" context_max: {stats['context_max']}") print(f" context_avg: {avg_context:.2f}") if __name__ == "__main__": main()