114 lines
4.9 KiB
Python
114 lines
4.9 KiB
Python
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
import tempfile
|
|
from types import SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from assembly_dataset.config import Config, TeacherConfig
|
|
from assembly_dataset.pipeline import generate, export, validate
|
|
from assembly_dataset.teacher import Plan
|
|
from assembly_dataset.world import solve
|
|
from assembly_dataset.language import action_text
|
|
from llm.config import OpenAiConfig
|
|
|
|
|
|
class TestAgent:
|
|
def __init__(self, invalid=False):
|
|
self.active = 0
|
|
self.peak = 0
|
|
self.invalid = invalid
|
|
|
|
async def run(self, prompt, **kwargs):
|
|
self.active += 1
|
|
self.peak = max(self.peak, self.active)
|
|
try:
|
|
await asyncio.sleep(0.001)
|
|
plan = {'status': 'ready', 'actions': []} if self.invalid else solve(json.loads(prompt)['scene'])[0]
|
|
if plan['actions']:
|
|
scene = json.loads(prompt)['scene']
|
|
plan['actions'] = [{**a, 'text': action_text(a, scene, Config().prompts)} for a in plan['actions']]
|
|
return SimpleNamespace(output=Plan.model_validate(plan))
|
|
finally:
|
|
self.active -= 1
|
|
|
|
|
|
class PipelineTests(unittest.TestCase):
|
|
def teacher_config(self, path):
|
|
return Config(output_dir=Path(path), target_examples=30,
|
|
teacher=TeacherConfig(enabled=True, fraction=1, concurrency=2, retries=1,
|
|
llm=OpenAiConfig(provider='openai', model='test')))
|
|
|
|
def test_teacher_concurrency_and_validated_provenance(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = self.teacher_config(tmp)
|
|
agent = TestAgent()
|
|
asyncio.run(generate(config, agent))
|
|
report = export(config)
|
|
self.assertEqual(agent.peak, 2)
|
|
self.assertEqual(report['sources'], {'openai': 10})
|
|
self.assertEqual(report['teacher_failures'], 0)
|
|
self.assertEqual(validate(config)['examples_checked'], 30)
|
|
|
|
def test_invalid_teacher_cannot_poison_corpus(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = self.teacher_config(tmp)
|
|
asyncio.run(generate(config, TestAgent(invalid=True)))
|
|
report = export(config)
|
|
self.assertEqual(report['teacher_failures'], 10)
|
|
self.assertEqual(report['sources'], {'symbolic': 10})
|
|
self.assertEqual(validate(config)['examples_checked'], 30)
|
|
|
|
def test_restart_after_interrupted_generation_equals_clean_run(self):
|
|
import assembly_dataset.pipeline as pipeline
|
|
original = pipeline.make_episode
|
|
def interrupted(index, config):
|
|
if index >= 40:
|
|
raise KeyboardInterrupt()
|
|
return original(index, config)
|
|
with tempfile.TemporaryDirectory() as first, tempfile.TemporaryDirectory() as second:
|
|
config = Config(output_dir=Path(first), target_examples=150)
|
|
with patch.object(pipeline, 'make_episode', side_effect=interrupted):
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
asyncio.run(generate(config))
|
|
asyncio.run(generate(config))
|
|
resumed = export(config)
|
|
clean = Config(output_dir=Path(second), target_examples=150)
|
|
asyncio.run(generate(clean))
|
|
uninterrupted = export(clean)
|
|
self.assertEqual(resumed['files'], uninterrupted['files'])
|
|
self.assertEqual(validate(config)['examples_checked'], 150)
|
|
|
|
def test_corrupted_shard_is_detected(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=4)
|
|
asyncio.run(generate(config))
|
|
export(config)
|
|
next(Path(tmp).glob('sft/*/*/*.jsonl')).write_text('broken\n')
|
|
with self.assertRaisesRegex(ValueError, 'Повреждён'):
|
|
validate(config)
|
|
|
|
def test_all_steps_and_one_mode_obey_total(self):
|
|
for modes in (['step'], ['branch'], ['full', 'step', 'branch']):
|
|
with self.subTest(modes=modes), tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=37, modes=modes, step_sampling='all')
|
|
asyncio.run(generate(config))
|
|
report = export(config)
|
|
self.assertEqual(report['examples'], 37)
|
|
self.assertEqual(validate(config)['examples_checked'], 37)
|
|
|
|
def test_smaller_export_does_not_remove_canonical_episodes(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config = Config(output_dir=Path(tmp), target_examples=100)
|
|
asyncio.run(generate(config))
|
|
original = export(config)
|
|
config.target_examples = 1
|
|
self.assertEqual(export(config)['examples'], 1)
|
|
config.target_examples = 100
|
|
self.assertEqual(export(config)['files'], original['files'])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|