обновлен процесс обработки данных

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`. Одновременные писатели в одну папку
 блокируются. Неизменённый процедурный запуск воспроизводим побайтно.
 
 Полный экспорт пересобирается во временную папку и затем заменяет предыдущий.

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)

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(*(
+                tasks = [asyncio.create_task(save_episode(episode)) for episode in episodes]
-                        enrich(e, config.teacher, agent, config.prompts) for e 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,

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]}))

config.yaml

@@ -1,13 +1,14 @@
 # Число итоговых обучающих строк СУММАРНО по full/step/branch.
 # CLI --count имеет приоритет. Можно указать 0, 1, 100, 500000 и т.д.
 target_examples: 500000
-seed: 20261021
+seed: 20261032
-output_dir: dataset/assembly-ru
+output_dir: dataset/assembly-ru-branch-500k
 prompts_file: prompts.yaml  # инструкции LLM и обучающих записей
-modes: [full, step, branch]
+# modes: [full, step, branch]
-step_sampling: one  # one: один шаг на эпизод; all: каждый шаг отдельной строкой
+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

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})