67 lines
4.9 KiB
Python
67 lines
4.9 KiB
Python
import tempfile,unittest,json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock,patch
|
|
from auto_test import AutoTests,model_layers,fine_splits
|
|
from inference import Scheduler,InferenceError
|
|
|
|
class AutoTestsTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp=tempfile.TemporaryDirectory();self.worker=Mock();self.worker.memory_plan={'gpu_layers':999};self.worker.catalog.entry.return_value=dict(id='m',kind='chat',file='model.gguf',profile_eligible=True)
|
|
self.profiles=Mock();self.scheduler=Scheduler(self.worker);self.auto=AutoTests(Path(self.tmp.name)/'auto.json',self.profiles,self.worker,self.scheduler)
|
|
self.gpus=[dict(name='RTX 5080',uuid='first'),dict(name='RTX 3060',uuid='second')]
|
|
def tearDown(self):self.tmp.cleanup()
|
|
def start(self):
|
|
with patch('auto_test.probe',return_value=self.gpus):self.auto.start(dict(model_id='m',max_context=4096,mtp=False))
|
|
self.auto.thread.join(3);self.assertFalse(self.auto.thread.is_alive())
|
|
def test_primary_first_results_persist_and_save_explicitly(self):
|
|
self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':40})
|
|
self.start();job=self.auto.status()['job'];self.assertEqual(job['state'],'complete');self.assertEqual(len(job['results']),6)
|
|
self.assertTrue(all(r['parameters']['gpu_devices']==['first'] for r in job['results']));self.profiles.save.assert_not_called();self.assertEqual(self.scheduler.active,0)
|
|
self.auto.save(dict(job_id=job['id'],index=0,name='saved'));self.profiles.save.assert_called_once()
|
|
self.assertEqual(AutoTests(self.auto.path,self.profiles,self.worker,self.scheduler).job['state'],'complete')
|
|
def test_secondary_before_cpu(self):
|
|
def plan(p,cancel):
|
|
if p['parameters']['gpu_offload']=='full':raise InferenceError('does not fit')
|
|
self.worker.plan.side_effect=plan;self.auto.request=Mock(return_value={});self.auto.benchmark=Mock(return_value={'predicted_per_second':10});self.start()
|
|
rows=self.auto.job['results'];first_success=next(i for i,r in enumerate(rows) if r['success']);self.assertEqual(first_success,18);self.assertEqual(rows[first_success]['parameters']['gpu_offload'],'auto')
|
|
def test_busy_lease_does_not_stop_another_worker(self):
|
|
with self.scheduler.lease(('other',)):
|
|
self.worker.reset_mock();self.start();self.assertEqual(self.auto.job['state'],'failed');self.worker.stop.assert_not_called()
|
|
def test_cancel_releases_lease_and_does_not_save_profile(self):
|
|
self.worker.plan.side_effect=lambda *a,**k:self.auto.cancel.set();self.auto.request=Mock(side_effect=InterruptedError());self.start();self.assertEqual(self.auto.job['state'],'cancelled');self.assertEqual(self.scheduler.active,0);self.profiles.save.assert_not_called()
|
|
def test_validation_and_failed_result_save(self):
|
|
with self.assertRaises(ValueError):self.auto.start(dict(model_id='m',max_context=True,mtp=False))
|
|
self.auto.job={'id':'a','model_id':'m','results':[{'success':False}]}
|
|
with self.assertRaises(ValueError):self.auto.save(dict(job_id='a',index=0,name='x'))
|
|
def test_prompt_filled_by_token_count(self):
|
|
self.auto.request=Mock(side_effect=[{'tokens':list(range(4000))},{'timings':{'prompt_n':3072,'predicted_n':64,'predicted_per_second':12}}]);r=self.auto.benchmark(4096);self.assertEqual(r['prompt_n'],3072);self.assertEqual(len(self.auto.request.call_args.args[1]['prompt']),3072)
|
|
class FineTests(unittest.TestCase):
|
|
def test_single_layer_steps_between_coarse_candidates(self):
|
|
import math
|
|
values=fine_splits(66,[85,15],90)
|
|
self.assertEqual([k for k,_ in values],[58,59])
|
|
for k,split in values:self.assertEqual(math.ceil(66*split[0]/sum(split)),k)
|
|
self.assertEqual(fine_splits(66,[95,5],100)[-1][0],65)
|
|
def test_gguf_metadata_mtp_and_output(self):
|
|
import struct
|
|
def text(s):b=s.encode();return struct.pack('<Q',len(b))+b
|
|
blob=b'GGUF'+struct.pack('<IQQ',3,0,3)+text('general.architecture')+struct.pack('<I',8)+text('qwen35')
|
|
blob+=text('qwen35.block_count')+struct.pack('<II',4,65)+text('qwen35.nextn_predict_layers')+struct.pack('<II',4,1)
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p=Path(d)/'model.gguf';p.write_bytes(blob);self.assertEqual(model_layers(p,True),66);self.assertEqual(model_layers(p,False),65)
|
|
p.write_bytes(blob[:-2])
|
|
with self.assertRaises(ValueError):model_layers(p,True)
|
|
def test_refinement_integrated_after_first_fit(self):
|
|
fixture=AutoTestsTests();fixture.setUp()
|
|
try:
|
|
a=fixture.auto;a.job=dict(id='j',model_id='m',max_context=4096,mtp=True,attempts=0,results=[],started_at=0)
|
|
seen=[]
|
|
def candidate(context,tier,devices,split,ratio,offload,layer_count=None):
|
|
seen.append((context,ratio,layer_count));return [True] if ratio and ratio[0]<=90 else []
|
|
a.test_candidate=candidate;fixture.worker.catalog.root=Path('/unused')
|
|
with patch('auto_test.model_layers',return_value=66):a.run('first','second')
|
|
self.assertEqual([v[2] for v in seen if v[2] is not None],[61,61])
|
|
finally:fixture.tearDown()
|
|
if __name__=='__main__':unittest.main()
|