import argparse import hashlib import json import platform import time from importlib.metadata import version from pathlib import Path from common import ROOT, inside, rows, save_json, append, stable_seed, permit_compute from grading import extract, correct, repetition, grade_ifeval from resources import snapshot from inference_queue import completed_outputs def recover_generations(path): if not path.exists(): return [] result, valid_bytes = [], 0 with path.open('rb') as f: size = path.stat().st_size for line in f: try: result.append(json.loads(line)) except (ValueError, UnicodeDecodeError): if f.tell() != size or line.endswith(b'\n'): raise backup = path.with_name(path.name + '.interrupted_tail') backup.write_bytes(line) break valid_bytes = f.tell() if valid_bytes < path.stat().st_size: with path.open('r+b') as f: f.truncate(valid_bytes) return result def model_fingerprint(model): files = sorted(list(model.glob('*.safetensors')) + list(model.glob('*.json')) + list(model.glob('*.jinja'))) entries = [] for file in files: with file.open('rb') as f: entries.append([file.name, hashlib.file_digest(f, 'sha256').hexdigest()]) return entries def main(): parser = argparse.ArgumentParser() parser.add_argument('--model', default=str(ROOT / 'models/original')) parser.add_argument('--data', required=True) parser.add_argument('--out', required=True) parser.add_argument('--set', default='practice', choices=['practice', 'math500', 'gsm8k', 'ifeval']) parser.add_argument('--samples', type=int) parser.add_argument('--limit', type=int) parser.add_argument('--max-tokens', type=int) parser.add_argument('--config', default=str(ROOT / 'configs/protocol.json')) parser.add_argument('--queue-limit', type=int) args = parser.parse_args() permit_compute() from transformers import AutoTokenizer from vllm import LLM, SamplingParams model, data, out = inside(args.model), inside(args.data), inside(args.out) protocol = json.loads(inside(args.config).read_text()) if args.set != 'practice' and (args.limit or args.max_tokens): raise ValueError('Сокращённый запуск финального бенчмарка запрещён; используй practice') records = rows(data) if args.limit: records = records[:args.limit] if len({str(r.get('id', r.get('key'))) for r in records}) != len(records): raise ValueError('Повторяющиеся ID') samples = args.samples or protocol['samples'].get(args.set, 1) cap = args.max_tokens or protocol['max_tokens'] if not 1 <= cap <= 32768 or samples < 1: raise ValueError('Некорректный протокол') out.mkdir(parents=True, exist_ok=True) job = {'model': str(model), 'fingerprint': model_fingerprint(model), 'data': str(data), 'data_sha256': hashlib.sha256(data.read_bytes()).hexdigest(), 'protocol': protocol, 'samples': samples, 'max_tokens': cap, 'set': args.set, 'limit': args.limit, 'runtime_versions': {x: version(x) for x in ('vllm', 'torch', 'transformers', 'math-verify')}, 'platform': platform.platform(), 'seed_scheme': 'sha256(base_seed:item_id:sample)31'} path = out / 'job.json' if path.exists() and json.loads(path.read_text()) != job: raise ValueError('В каталоге сохранён другой запуск') save_json(path, job) completed = recover_generations(out / 'generations.jsonl') done = {(str(r['id']), r['sample']) for r in completed} if len(done) != len(completed): raise ValueError('Повторные записи генераций') sample_order = protocol.get('sample_order', 'problem_major') if sample_order not in ('problem_major', 'sample_major'): raise ValueError('Неизвестный порядок образцов') ordered = ([(r, s) for s in range(samples) for r in records] if sample_order == 'sample_major' else [(r, s) for r in records for s in range(samples)]) size = protocol.get('request_batch_size', 1) batches = [ordered[i:i+size] for i in range(0, len(ordered), size)] pending = [batch for batch in batches if any((str(r.get('id', r.get('key'))), s) not in done for r, s in batch)] if args.queue_limit is not None: if args.queue_limit < 1: raise ValueError('queue-limit must be positive') pending = pending[:args.queue_limit] if pending: tokenizer = AutoTokenizer.from_pretrained(str(model), local_files_only=True) memory_options = {} if protocol.get('kv_cache_memory_bytes') is not None: memory_options['kv_cache_memory_bytes'] = protocol['kv_cache_memory_bytes'] engine = LLM(model=str(model), dtype=protocol['dtype'], seed=protocol['seed'], max_model_len=protocol['max_model_len'], gpu_memory_utilization=protocol['gpu_memory_utilization'], max_num_seqs=protocol['max_num_seqs'], max_num_batched_tokens=protocol['max_num_batched_tokens'], enable_prefix_caching=protocol['enable_prefix_caching'], enforce_eager=protocol['enforce_eager'], limit_mm_per_prompt={'image': 0, 'video': 0}, **memory_options) append(out/'resources.jsonl',snapshot('engine_initialized')) for batch in pending: prompts, samplings = [], [] for row, sample in batch: identity = str(row.get('id', row.get('key'))) seed = stable_seed(identity, sample, protocol['seed']) user = row['prompt'] if args.set == 'ifeval' else row['problem'] + '\n' + protocol['instruction'] prompts.append(tokenizer.apply_chat_template([{'role': 'user', 'content': user}], tokenize=False, add_generation_prompt=True, enable_thinking=True)) samplings.append(SamplingParams(n=1, temperature=protocol['temperature'], top_p=protocol['top_p'], top_k=protocol['top_k'], presence_penalty=protocol['presence_penalty'], max_tokens=cap, seed=seed)) start = time.time() batch_tokens = 0 for index,output in completed_outputs(engine,prompts,samplings,protocol.get('stream_completed',False)): elapsed = time.time()-start row,sample = batch[index] prompt = prompts[index] identity = str(row.get('id', row.get('key'))) generation = output.outputs[0] batch_tokens += len(generation.token_ids) if generation.finish_reason not in ('stop','length'): raise RuntimeError('Незавершённая генерация не засчитывается') if (identity, sample) in done: continue generation = output.outputs[0] prediction = extract(generation.text) good = grade_ifeval(row, generation.text) if args.set == 'ifeval' else correct(prediction, row['answer']) result = {'id': identity, 'sample': sample, 'seed': stable_seed(identity, sample, protocol['seed']), 'correct': bool(good), 'prediction': prediction, 'prompt': prompt, 'text': generation.text, 'tokens': len(generation.token_ids), 'prompt_token_ids': tokenizer.encode(prompt, add_special_tokens=False), 'token_ids': list(generation.token_ids), 'finish_reason': generation.finish_reason, 'batch_seconds': elapsed, **repetition(generation.text)} append(out / 'generations.jsonl', result) completed.append(result) print(json.dumps({k: result[k] for k in ('id', 'sample', 'correct', 'tokens', 'finish_reason')}), flush=True) append(out/'resources.jsonl',snapshot('batch_completed')) append(out/'timings.jsonl',{'requests':len(batch),'seconds':time.time()-start,'tokens':batch_tokens}) n = len(completed) if not n: raise ValueError('Нет результатов') save_json(out / 'summary.json', {'generations': n,'planned_generations':len(records)*samples, 'complete':n==len(records)*samples,'problems':len({r['id'] for r in completed}), 'planned_problems':len(records),'accuracy': sum(r['correct'] for r in completed)/n, 'cut_share': sum(r['finish_reason'] == 'length' for r in completed)/n, 'mean_tokens': sum(r['tokens'] for r in completed)/n, 'loop_share': sum(r['repetitive_end'] for r in completed)/n, 'duplicate_8gram_tail_mean': sum(r['duplicate_8gram_tail'] for r in completed)/n, 'whitespace_tail_share': sum(r['whitespace_tail'] for r in completed)/n, 'thinking_closed_share': sum(r['thinking_closed'] for r in completed)/n, 'missing_box_share': sum(r['prediction'] is None for r in completed)/n if args.set != 'ifeval' else None, 'note': 'Фиксированные группы запросов и стабильный seed на задачу и образец; численное совпадение с исходным vLLM не гарантируется.'}) if __name__ == '__main__': main()