271 lines
15 KiB
Python
271 lines
15 KiB
Python
"""Bounded generation, transactional checkpoints, streaming sharded export."""
|
|
import asyncio
|
|
import csv
|
|
from collections import Counter
|
|
from contextlib import contextmanager
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
from itertools import zip_longest
|
|
from pathlib import Path
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
|
|
from .config import Config
|
|
from .language import validate_texts
|
|
from .csv_export import COLUMNS, csv_record
|
|
from .progress import ProgressDisplay
|
|
from .world import canonical, digest, make_episode, training_rows, validate_plan, validate_branch
|
|
|
|
VERSION = 3
|
|
STORAGE_VERSION = 2
|
|
|
|
|
|
@contextmanager
|
|
def storage(config):
|
|
root = config.output_dir
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
with (root / '.lock').open('a') as lock:
|
|
try:
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
except BlockingIOError:
|
|
raise ValueError('Другой процесс уже работает с output_dir') from None
|
|
db = sqlite3.connect(root / 'episodes.sqlite3')
|
|
try:
|
|
db.execute('PRAGMA journal_mode=WAL')
|
|
db.execute('CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)')
|
|
db.execute('CREATE TABLE IF NOT EXISTS episodes (seq INTEGER PRIMARY KEY, signature TEXT UNIQUE NOT NULL, payload TEXT NOT NULL, rows INTEGER NOT NULL)')
|
|
settings = config.model_dump(mode='json', exclude={'target_examples', 'output_dir', 'shard_size', 'prompts_file'})
|
|
# Structured prompts only affect export, so existing v2 episodes remain reusable.
|
|
settings['prompts'].pop('structured_system', None)
|
|
fingerprint = digest({'version': STORAGE_VERSION, 'settings': settings})
|
|
saved = db.execute("SELECT value FROM meta WHERE key='fingerprint'").fetchone()
|
|
if saved and saved[0] != fingerprint:
|
|
raise ValueError('Изменён config: для других настроек используйте новый output_dir')
|
|
if not saved:
|
|
with db:
|
|
db.execute('INSERT INTO meta VALUES (?, ?)', ('fingerprint', fingerprint))
|
|
db.execute('INSERT INTO meta VALUES (?, ?)', ('settings', canonical(settings)))
|
|
db.execute('INSERT INTO meta VALUES (?, ?)', ('next_index', '0'))
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
async def generate(config: Config, agent=None, *, show_progress=False, teacher_only=False):
|
|
with storage(config) as db:
|
|
count = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
|
|
index = int(db.execute("SELECT value FROM meta WHERE key='next_index'").fetchone()[0])
|
|
resume_until = db.execute('SELECT COALESCE(MAX(seq), 0) FROM episodes').fetchone()[0]
|
|
with ProgressDisplay('Генерация', config.target_examples, initial=count, enabled=show_progress) as progress:
|
|
if count >= config.target_examples and index >= resume_until:
|
|
return count
|
|
if config.teacher.enabled and config.teacher.fraction > 0 and agent is None:
|
|
from .teacher import create_agent
|
|
progress.update(status='подключение LLM')
|
|
agent = await create_agent(config.teacher, config.prompts)
|
|
batch_size = config.teacher.concurrency if agent is not None else 32
|
|
duplicates = 0
|
|
while count < config.target_examples or index < resume_until:
|
|
# Never construct 500k coroutines or keep the entire dataset in RAM.
|
|
size = min(batch_size, max(1, config.target_examples - count))
|
|
episodes = [make_episode(i, config) for i in range(index, index + size)]
|
|
async def save_episode(episode):
|
|
nonlocal count, duplicates
|
|
# A restarted batch can contain responses already committed before cancellation.
|
|
if db.execute('SELECT 1 FROM episodes WHERE signature=?',
|
|
(episode['signature'],)).fetchone():
|
|
duplicates += 1
|
|
return
|
|
if agent is not None:
|
|
from .teacher import enrich
|
|
episode, _ = await enrich(episode, config.teacher, agent, config.prompts)
|
|
if teacher_only and (config.teacher.llm is None or
|
|
episode['source']['provider'] != config.teacher.llm.provider or
|
|
'teacher_error' in episode):
|
|
raise ValueError('Режим --teacher-only: ответ модели не прошёл проверку; символическая замена не сохранена')
|
|
rows = sum(1 for _ in training_rows(episode, config))
|
|
# No await between receiving a validated response and committing it.
|
|
with db:
|
|
cursor = db.execute('INSERT OR IGNORE INTO episodes(seq,signature,payload,rows) VALUES (?,?,?,?)',
|
|
(episode['candidate_index'] + 1, episode['signature'], canonical(episode), rows))
|
|
if cursor.rowcount:
|
|
count += rows
|
|
duplicates = 0
|
|
else:
|
|
duplicates += 1
|
|
progress.update(count, status='ожидание LLM' if agent is not None else '')
|
|
|
|
if agent is not None:
|
|
progress.update(status='ожидание LLM')
|
|
tasks = [asyncio.create_task(save_episode(episode)) for episode in episodes]
|
|
try:
|
|
await asyncio.gather(*tasks)
|
|
finally:
|
|
for task in tasks:
|
|
if not task.done():
|
|
task.cancel()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
# Advance only after the whole batch succeeds; holes are retried on resume.
|
|
index += size
|
|
with db:
|
|
db.execute("UPDATE meta SET value=? WHERE key='next_index'", (str(index),))
|
|
if duplicates > 10000:
|
|
raise ValueError('Пространство сценариев исчерпано: расширьте диапазон компонентов или domains')
|
|
progress.update(count, status='')
|
|
return count
|
|
|
|
|
|
class Shards:
|
|
def __init__(self, root, size):
|
|
self.root, self.size = root, size
|
|
self.counts = Counter()
|
|
self.handles = {}
|
|
self.csv_handles = {}
|
|
self.csv_writers = {}
|
|
self.files = []
|
|
|
|
def write(self, row, structured_row):
|
|
key = row['split'], row['mode']
|
|
count = self.counts[key]
|
|
for variant, record in (('text', row), ('structured', structured_row)):
|
|
handle_key = (*key, variant)
|
|
if count % self.size == 0:
|
|
if handle_key in self.handles:
|
|
self.handles[handle_key].close()
|
|
self.csv_handles[handle_key].close()
|
|
suffix = '.jsonl' if variant == 'text' else '.structured.jsonl'
|
|
path = self.root / key[0] / key[1] / f'{count // self.size:06d}{suffix}'
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.handles[handle_key] = path.open('w', encoding='utf-8')
|
|
self.files.append(path)
|
|
csv_path = path.with_suffix('.csv')
|
|
self.csv_handles[handle_key] = csv_path.open('w', encoding='utf-8-sig', newline='')
|
|
self.csv_writers[handle_key] = csv.DictWriter(self.csv_handles[handle_key], fieldnames=COLUMNS, delimiter=';')
|
|
self.csv_writers[handle_key].writeheader()
|
|
self.files.append(csv_path)
|
|
self.handles[handle_key].write(canonical(record) + '\n')
|
|
self.csv_writers[handle_key].writerow(csv_record(record))
|
|
self.counts[key] += 1
|
|
|
|
def close(self):
|
|
for file in (*self.handles.values(), *self.csv_handles.values()):
|
|
file.close()
|
|
|
|
|
|
def file_digest(path):
|
|
hasher = hashlib.sha256()
|
|
with path.open('rb') as file:
|
|
for block in iter(lambda: file.read(1024 * 1024), b''):
|
|
hasher.update(block)
|
|
return hasher.hexdigest()
|
|
|
|
|
|
def export(config, *, show_progress=False, allow_partial=False):
|
|
with storage(config) as db, ProgressDisplay('Экспорт', config.target_examples, enabled=show_progress) as progress:
|
|
capacity = db.execute('SELECT COALESCE(SUM(rows), 0) FROM episodes').fetchone()[0]
|
|
if capacity < config.target_examples and not allow_partial:
|
|
raise ValueError('Недостаточно сценариев: сначала выполните generate')
|
|
destination = config.output_dir / 'sft'
|
|
backup = config.output_dir / '.sft-previous'
|
|
# Recover an interrupted directory swap before making another export.
|
|
if backup.exists() and not destination.exists():
|
|
backup.rename(destination)
|
|
if destination.exists() and not (destination / 'manifest.json').exists():
|
|
raise ValueError('sft не является управляемым экспортом; выберите другой output_dir')
|
|
stage = Path(tempfile.mkdtemp(prefix='.sft-', dir=config.output_dir))
|
|
writer = Shards(stage, config.shard_size)
|
|
total, scenarios, failures = 0, 0, 0
|
|
domains, sources = Counter(), Counter()
|
|
try:
|
|
for (payload,) in db.execute('SELECT payload FROM episodes ORDER BY seq'):
|
|
if total >= config.target_examples:
|
|
break
|
|
episode = json.loads(payload)
|
|
scenarios += 1
|
|
domains[episode['scene']['domain']] += 1
|
|
sources[episode['source']['provider']] += 1
|
|
failures += int('teacher_error' in episode)
|
|
for record, structured_record in zip(training_rows(episode, config),
|
|
training_rows(episode, config, output_format='structured'), strict=True):
|
|
if total >= config.target_examples:
|
|
break
|
|
writer.write(record, structured_record)
|
|
total += 1
|
|
progress.update(total)
|
|
progress.update(status="контрольные суммы")
|
|
writer.close()
|
|
manifest = {'schema_version': VERSION, 'examples': total, 'scenarios': scenarios,
|
|
'partial': total < config.target_examples,
|
|
'examples_per_variant': {'text': total, 'structured': total},
|
|
'counts': {f'{split}/{mode}': n for (split, mode), n in sorted(writer.counts.items())},
|
|
'domains': dict(domains), 'sources': dict(sources), 'teacher_failures': failures,
|
|
'validation_level': 'logical_only',
|
|
'config': config.model_dump(mode='json'),
|
|
'files': {str(p.relative_to(stage)): file_digest(p) for p in writer.files}}
|
|
(stage / 'manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
if destination.exists():
|
|
destination.rename(backup)
|
|
stage.rename(destination)
|
|
if backup.exists():
|
|
shutil.rmtree(backup)
|
|
progress.update(total, status='')
|
|
return manifest
|
|
finally:
|
|
writer.close()
|
|
if stage.exists():
|
|
shutil.rmtree(stage)
|
|
|
|
|
|
def validate(config):
|
|
with storage(config) as db:
|
|
episodes, paths = 0, 0
|
|
for signature, payload, stored_rows in db.execute('SELECT signature,payload,rows FROM episodes'):
|
|
episode = json.loads(payload)
|
|
if signature != digest(episode['scene']) or episode['signature'] != signature or episode['scenario_id'] != signature:
|
|
raise ValueError('Неверная сигнатура сценария')
|
|
if 'step_texts' in episode:
|
|
validate_texts(episode, episode['step_texts'], config.prompts)
|
|
trace = validate_plan(episode, episode['plan'])
|
|
if trace != episode['trace']:
|
|
raise ValueError('Trace не соответствует плану')
|
|
paths += validate_branch(episode, episode['branch'])
|
|
if sum(1 for _ in training_rows(episode, config)) != stored_rows:
|
|
raise ValueError('Неверное число записей эпизода')
|
|
episodes += 1
|
|
root = config.output_dir / 'sft'
|
|
manifest_path = root / 'manifest.json'
|
|
if not manifest_path.exists():
|
|
raise ValueError('Нет экспорта SFT: выполните generate или export')
|
|
manifest = json.loads(manifest_path.read_text())
|
|
total = 0
|
|
expected_files = set(manifest['files'])
|
|
actual_files = {str(p.relative_to(root)) for pattern in ('*/*/*.jsonl', '*/*/*.csv') for p in root.glob(pattern)}
|
|
if expected_files != actual_files:
|
|
raise ValueError('Набор шардов не соответствует manifest')
|
|
for relative, checksum in manifest['files'].items():
|
|
path = root / relative
|
|
if path.resolve().is_relative_to(root.resolve()) is False or file_digest(path) != checksum:
|
|
raise ValueError(f'Повреждён шард {relative}')
|
|
for relative in manifest['files']:
|
|
path = root / relative
|
|
if path.suffix == '.jsonl' and '.structured.' not in path.name:
|
|
counterpart = path.with_suffix('.structured.jsonl')
|
|
if str(counterpart.relative_to(root)) not in expected_files:
|
|
raise ValueError('Отсутствует структурированная пара JSONL')
|
|
with path.open(encoding='utf-8') as file, counterpart.open(encoding='utf-8') as other:
|
|
for text_line, structured_line in zip_longest(file, other):
|
|
if text_line is None or structured_line is None:
|
|
raise ValueError('Разное число записей в паре JSONL')
|
|
text_row, structured_row = json.loads(text_line), json.loads(structured_line)
|
|
if (text_row['record_id'] != structured_row['record_id'] or
|
|
text_row['messages'][1] != structured_row['messages'][1]):
|
|
raise ValueError('Записи в паре JSONL не совпадают')
|
|
total += 1
|
|
if total != manifest['examples']:
|
|
raise ValueError('Неверное число экспортированных записей')
|
|
return {'episodes_checked': episodes, 'branch_paths_checked': paths, 'examples_checked': total}
|