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

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