126 lines
6.6 KiB
Python
126 lines
6.6 KiB
Python
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})
|