diff --git a/README.md b/README.md index 914db7b..91dc26a 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ uv run python -m assembly_dataset validate процент выполнения, скорость, прошедшее и приблизительное оставшееся время. При продолжении корпуса генерация начинается с уже сохранённого количества. -Счётчик генерации обновляется после сохранения пакета проверенных эпизодов, +Счётчик генерации обновляется после сохранения каждого проверенного эпизода, а не после каждого API-запроса. Во время вызовов модели отображается «ожидание LLM». Экспорт отдельно показывает запись строк и подготовку контрольных сумм. Оценка времени появляется после первых обновлений и может меняться из-за @@ -228,9 +228,20 @@ dataset/assembly-ru/ SQLite служит транзакционным checkpoint и индексом уникальности. Он хранит полные эпизоды, включая trace и teacher-ошибки. Пакет в памяти ограничен размером очереди; -SFT экспортируется потоково по шардам. После прерывания повторите команду: -принятые эпизоды сохраняются, незавершённый пакет повторяется. При teacher-режиме -запросы незавершённого пакета могут повториться. Одновременные писатели в одну папку +Каждый готовый ответ сохраняется сразу, не ожидая остальных запросов пакета. +При Ctrl+C, SIGTERM или ошибке генерации CLI экспортирует накопленные данные +в JSONL/CSV; неполный экспорт отмечается `partial: true` в manifest. +После исчерпания повторных попыток ошибки API останавливают генерацию. +Невалидные планы, как прежде, заменяются проверенным символическим планом с отметкой ошибки. +Для корпуса только из ответов LLM запускайте `generate --teacher-only`: при +невалидном ответе генерация остановится с экспортом накопленных записей, +а символическая замена не попадёт в базу. Флаг действует на новые записи; +ранее сохранённые эпизоды он не удаляет. +Для продолжения повторите команду с теми же настройками: сохранённые эпизоды +не запрашиваются повторно, незавершённые запросы выполняются заново. +После принудительного завершения (SIGKILL) данные остаются в SQLite; +экспортировать их без запросов к LLM можно командой +`uv run python -m assembly_dataset export --allow-partial`. Одновременные писатели в одну папку блокируются. Неизменённый процедурный запуск воспроизводим побайтно. Полный экспорт пересобирается во временную папку и затем заменяет предыдущий. diff --git a/assembly_dataset/__main__.py b/assembly_dataset/__main__.py index 483ab76..2c95ffc 100644 --- a/assembly_dataset/__main__.py +++ b/assembly_dataset/__main__.py @@ -1,6 +1,7 @@ import argparse import asyncio import json +import signal from pathlib import Path import sys @@ -8,6 +9,18 @@ from .config import load_config from .pipeline import generate, export, validate +async def run_generation(config, *, show_progress, teacher_only=False): + loop = asyncio.get_running_loop() + task = asyncio.current_task() + previous = signal.getsignal(signal.SIGTERM) + loop.add_signal_handler(signal.SIGTERM, task.cancel) + try: + return await generate(config, show_progress=show_progress, teacher_only=teacher_only) + finally: + loop.remove_signal_handler(signal.SIGTERM) + signal.signal(signal.SIGTERM, previous) + + def main(): parser = argparse.ArgumentParser(description='Датасет высокоуровневых планов сборки: full / step / branch') parser.add_argument('--config', type=Path, default=Path(__file__).resolve().parents[1] / 'config.yaml') @@ -15,6 +28,8 @@ def main(): parser.add_argument('--count', type=int, help='Точное суммарное число SFT-записей; 0 допустим') parser.add_argument('--output-dir', type=Path, help='Папка корпуса (относительно текущего каталога)') parser.add_argument('--dry-run', action='store_true', help='Показать настройки без записи и запросов') + parser.add_argument('--teacher-only', action='store_true', help='Сохранять только проверенные ответы LLM; останавливать генерацию при символической замене') + parser.add_argument('--allow-partial', action='store_true', help='Экспортировать накопленные записи, даже если цель ещё не достигнута') parser.add_argument('--no-progress', action='store_true', help='Отключить индикаторы прогресса') args = parser.parse_args() try: @@ -32,11 +47,20 @@ def main(): if config.teacher.enabled: from dotenv import load_dotenv load_dotenv(Path(__file__).resolve().parents[1] / '.env') - asyncio.run(generate(config, show_progress=not args.no_progress)) + try: + asyncio.run(run_generation(config, show_progress=not args.no_progress, teacher_only=args.teacher_only)) + except (Exception, KeyboardInterrupt, asyncio.CancelledError) as exc: + print(f'Генерация остановлена ({type(exc).__name__}). Экспорт накопленных данных…', file=sys.stderr) + try: + saved = export(config, show_progress=not args.no_progress, allow_partial=True) + print(f"Сохранено {saved['examples']} записей: {config.output_dir / 'sft'}", file=sys.stderr) + except Exception as export_error: + print(f'Экспорт не выполнен ({type(export_error).__name__}); контрольные точки остаются в {config.output_dir / "episodes.sqlite3"}', file=sys.stderr) + return 130 if isinstance(exc, (KeyboardInterrupt, asyncio.CancelledError)) else 1 result = export(config, show_progress=not args.no_progress) result = {k: result[k] for k in ('examples', 'scenarios', 'counts', 'teacher_failures')} elif args.command == 'export': - result = export(config, show_progress=not args.no_progress) + result = export(config, show_progress=not args.no_progress, allow_partial=args.allow_partial) result = {k: result[k] for k in ('examples', 'scenarios', 'counts', 'teacher_failures')} else: result = validate(config) diff --git a/assembly_dataset/pipeline.py b/assembly_dataset/pipeline.py index c38d45d..ad54f45 100644 --- a/assembly_dataset/pipeline.py +++ b/assembly_dataset/pipeline.py @@ -53,12 +53,13 @@ def storage(config): db.close() -async def generate(config: Config, agent=None, *, show_progress=False): +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: + 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 @@ -66,28 +67,49 @@ async def generate(config: Config, agent=None, *, show_progress=False): 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: + 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: - from .teacher import enrich progress.update(status='ожидание LLM') - episodes = [item[0] for item in await asyncio.gather(*( - enrich(e, config.teacher, agent, config.prompts) for e in episodes))] + 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: - for episode in episodes: - rows = sum(1 for _ in training_rows(episode, config)) - cursor = db.execute('INSERT OR IGNORE INTO episodes(signature,payload,rows) VALUES (?,?,?)', - (episode['signature'], canonical(episode), rows)) - if cursor.rowcount: - count += rows - duplicates = 0 - else: - duplicates += 1 - index = episode['candidate_index'] + 1 - if count >= config.target_examples: - break db.execute("UPDATE meta SET value=? WHERE key='next_index'", (str(index),)) if duplicates > 10000: raise ValueError('Пространство сценариев исчерпано: расширьте диапазон компонентов или domains') @@ -140,10 +162,10 @@ def file_digest(path): return hasher.hexdigest() -def export(config, *, show_progress=False): +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: + if capacity < config.target_examples and not allow_partial: raise ValueError('Недостаточно сценариев: сначала выполните generate') destination = config.output_dir / 'sft' backup = config.output_dir / '.sft-previous' @@ -175,6 +197,7 @@ def export(config, *, show_progress=False): 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, diff --git a/assembly_dataset/teacher.py b/assembly_dataset/teacher.py index 0e48e1c..e576e8a 100644 --- a/assembly_dataset/teacher.py +++ b/assembly_dataset/teacher.py @@ -4,6 +4,7 @@ import re from typing import Literal from pydantic import Field, field_validator +from pydantic_ai.exceptions import UnexpectedModelBehavior from llm.config import ConfigModel from .prompts import render from .language import validate_texts, action_text @@ -63,6 +64,8 @@ async def enrich(episode, config, agent, prompts): except Exception as exc: # API exceptions may contain request headers. Persist the class only. error = {'scenario_id': episode['scenario_id'], 'error_type': type(exc).__name__, 'attempts': attempt + 1} + if not isinstance(exc, (ValueError, UnexpectedModelBehavior)) and attempt == config.retries: + raise RuntimeError(f'Запрос LLM не выполнен: {type(exc).__name__}. Генерация остановлена.') from None if isinstance(exc, ValueError): prompt = render(prompts.teacher_retry, canonical({'instruction': episode['instruction'], 'scene': teacher_scene, 'validation_error': str(exc)[:500]})) diff --git a/config.yaml b/config.yaml index 69e03eb..839bf1e 100644 --- a/config.yaml +++ b/config.yaml @@ -1,13 +1,14 @@ # Число итоговых обучающих строк СУММАРНО по full/step/branch. # CLI --count имеет приоритет. Можно указать 0, 1, 100, 500000 и т.д. target_examples: 500000 -seed: 20261021 -output_dir: dataset/assembly-ru +seed: 20261032 +output_dir: dataset/assembly-ru-branch-500k prompts_file: prompts.yaml # инструкции LLM и обучающих записей -modes: [full, step, branch] -step_sampling: one # one: один шаг на эпизод; all: каждый шаг отдельной строкой +# modes: [full, step, branch] +modes: [branch] +step_sampling: all # one: один шаг на эпизод; all: каждый шаг отдельной строкой shard_size: 1000 -split: [1.0, 0.0, 0.0] # разделение по сценарию; малые наборы не гарантируют все split +split: [0.8, 0.2, 0] # разделение по сценарию; малые наборы не гарантируют все split min_components: 1 max_components: 10 blocked_fraction: 0.2 @@ -37,7 +38,7 @@ teacher: # Для существующей интеграции Codex замените llm выше: llm: provider: codex - model: gpt-6-astra + model: gpt-6-luna credential_path: .config/codex-credentials.json generation: temperature: null diff --git a/tests/test_interrupted_generation.py b/tests/test_interrupted_generation.py new file mode 100644 index 0000000..b222cbd --- /dev/null +++ b/tests/test_interrupted_generation.py @@ -0,0 +1,125 @@ +import asyncio +from contextlib import closing, redirect_stderr, redirect_stdout +from io import StringIO +import json +import os +import signal +from pathlib import Path +import sqlite3 +import tempfile +import unittest +from unittest.mock import patch + +from assembly_dataset.config import Config +from assembly_dataset.pipeline import generate, export, validate +from tests import test_assembly_pipeline as helpers + + +class InterruptedGenerationTests(unittest.TestCase): + def test_completed_response_survives_cancel_and_resume(self): + with tempfile.TemporaryDirectory() as tmp: + config = helpers.PipelineTests().teacher_config(tmp) + config.modes = ['branch'] + config.target_examples = 4 + class SlowAgent(helpers.TestAgent): + calls = 0 + async def run(self, prompt, **kwargs): + self.calls += 1 + if self.calls == 1: + await asyncio.Event().wait() + return await super().run(prompt, **kwargs) + async def run(): + task = asyncio.create_task(generate(config, SlowAgent())) + for _ in range(200): + await asyncio.sleep(0.01) + if (config.output_dir / 'episodes.sqlite3').exists(): + with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db: + if db.execute('select count(*) from episodes').fetchone()[0]: + break + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + asyncio.run(run()) + with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db: + self.assertEqual(db.execute('select count(*) from episodes').fetchone()[0], 1) + self.assertEqual(export(config, allow_partial=True)['examples'], 1) + self.assertEqual(validate(config)['examples_checked'], 1) + agent = helpers.TestAgent() + asyncio.run(generate(config, agent)) + self.assertEqual(export(config)['examples'], 4) + self.assertEqual(validate(config)['episodes_checked'], 4) + + def test_api_failure_stops_without_symbolic_replacement(self): + with tempfile.TemporaryDirectory() as tmp: + config = helpers.PipelineTests().teacher_config(tmp) + class FailedAgent: + async def run(self, *args, **kwargs): + raise RuntimeError('quota exhausted secret') + with self.assertRaisesRegex(RuntimeError, 'RuntimeError') as caught: + asyncio.run(generate(config, FailedAgent())) + self.assertNotIn('secret', str(caught.exception)) + with closing(sqlite3.connect(config.output_dir / 'episodes.sqlite3')) as db: + self.assertEqual(db.execute('select count(*) from episodes').fetchone()[0], 0) + + def test_cli_exports_checkpoint_on_interrupt_or_error(self): + from assembly_dataset.__main__ import main + for failure, code in [(KeyboardInterrupt, 130), (RuntimeError, 1), (signal.SIGTERM, 130)]: + with self.subTest(failure=failure), tempfile.TemporaryDirectory() as tmp: + config = Config(output_dir=Path(tmp), target_examples=3, modes=["branch"]) + asyncio.run(generate(config)) + config.target_examples = 30 + async def interrupted(*args, **kwargs): + if failure == signal.SIGTERM: + os.kill(os.getpid(), signal.SIGTERM) + await asyncio.Event().wait() + else: + raise failure() + with patch('sys.argv', ['assembly_dataset', '--no-progress']), patch('assembly_dataset.__main__.load_config', return_value=config), patch('assembly_dataset.__main__.generate', interrupted), redirect_stderr(StringIO()), redirect_stdout(StringIO()): + self.assertEqual(main(), code) + manifest = json.loads((config.output_dir / 'sft/manifest.json').read_text()) + self.assertEqual(manifest['examples'], 3) + self.assertTrue(manifest['partial']) + self.assertEqual(validate(config)['examples_checked'], 3) + + def test_enough_rows_do_not_skip_interrupted_batch_holes(self): + from assembly_dataset.pipeline import storage + from assembly_dataset.world import make_episode, canonical + with tempfile.TemporaryDirectory() as tmp: + config = Config(output_dir=Path(tmp), target_examples=1, modes=['branch']) + # Response for candidate 1 arrived; candidate 0 is still pending. + episode = make_episode(1, config) + with storage(config) as db, db: + db.execute('INSERT INTO episodes VALUES (?,?,?,?)', + (2, episode['signature'], canonical(episode), 1)) + asyncio.run(generate(config)) + with storage(config) as db: + indices = [json.loads(p)['candidate_index'] for (p,) in db.execute('select payload from episodes order by seq')] + self.assertEqual(indices, [0, 1]) + + def test_invalid_structured_response_keeps_validation_fallback(self): + from pydantic_ai.exceptions import UnexpectedModelBehavior + with tempfile.TemporaryDirectory() as tmp: + config = helpers.PipelineTests().teacher_config(tmp) + config.target_examples = 1 + class InvalidAgent: + async def run(self, *args, **kwargs): + raise UnexpectedModelBehavior('Invalid structured output') + asyncio.run(generate(config, InvalidAgent())) + self.assertEqual(export(config)['teacher_failures'], 1) + + def test_teacher_only_rejects_symbolic_fallback(self): + from assembly_dataset.pipeline import storage + with tempfile.TemporaryDirectory() as tmp: + config = helpers.PipelineTests().teacher_config(tmp) + config.target_examples = 1 + with self.assertRaisesRegex(ValueError, 'teacher-only'): + asyncio.run(generate(config, helpers.TestAgent(invalid=True), teacher_only=True)) + with storage(config) as db: + self.assertEqual(db.execute('SELECT COUNT(*) FROM episodes').fetchone()[0], 0) + + def test_teacher_only_keeps_valid_teacher_response(self): + with tempfile.TemporaryDirectory() as tmp: + config = helpers.PipelineTests().teacher_config(tmp) + config.target_examples = 1 + asyncio.run(generate(config, helpers.TestAgent(), teacher_only=True)) + self.assertEqual(export(config)['sources'], {'openai': 1})