86 lines
4.7 KiB
Python
86 lines
4.7 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
from runtime import Runtime
|
|
|
|
class RuntimeTests(unittest.TestCase):
|
|
def test_untrusted_build_arguments(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d)
|
|
for ref in ('main; reboot','--upload-pack=x','../main','https://evil/repo'):
|
|
with self.assertRaises(ValueError):r.start(ref)
|
|
with self.assertRaises(ValueError):r.start('b9000',jobs=32)
|
|
def test_missing_tools_does_not_spawn(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d)
|
|
with patch.object(r,'prerequisites',return_value={'cuda_ready':False}),patch('runtime.subprocess.Popen') as spawn:
|
|
with self.assertRaises(ValueError):r.start('b9000')
|
|
with self.assertRaises(ValueError):r.start('v0.5.0')
|
|
self.assertEqual(r.prerequisites.call_count,2)
|
|
spawn.assert_not_called()
|
|
def test_selection_rollback_persistence(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d);r.state['builds']=[{'id':'one'},{'id':'two'}]
|
|
with self.assertRaises(ValueError):r.activate('../unknown')
|
|
r.activate('one');r.activate('two');r.rollback()
|
|
self.assertEqual(Runtime(d).status()['active'],'one')
|
|
def test_restart_marks_interrupted(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d);r.state['job']={'state':'running'};r.save()
|
|
self.assertEqual(Runtime(d).status()['job']['state'],'interrupted')
|
|
def test_cancel_before_command(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d);r.stop()
|
|
with patch('runtime.subprocess.Popen') as spawn:
|
|
with self.assertRaises(InterruptedError):r.run(['git'],Path(d),'test')
|
|
spawn.assert_not_called()
|
|
def test_fit_requires_selected_build_and_bounded_context(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d)
|
|
for context,slots in [(0,1),(160000,99),(True,1)]:
|
|
with self.assertRaises(ValueError):r.fit('medium',context,slots)
|
|
with patch.object(r,'references',return_value={'profiles':[dict(id='medium',available=True)]}):
|
|
with self.assertRaises(ValueError):r.fit('medium',160000,2)
|
|
def test_official_prereleases_are_not_hidden(self):
|
|
row=dict(tag_name='b11229',name='build',published_at='date',body='notes',html_url='https://github.com/ggml-org/llama.cpp/releases',prerelease=True,draft=False)
|
|
with tempfile.TemporaryDirectory() as d,patch('runtime.github',side_effect=[[row],dict(row,tag_name='v0.5.0',prerelease=False)]):
|
|
releases=Runtime(d).releases()['releases'];self.assertEqual(releases[0]['tag'],'v0.5.0');self.assertTrue(releases[1]['prerelease'])
|
|
def test_fit_excludes_busy_gpu_and_preserves_context(self):
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d);r.state.update(active='build',builds=[dict(id='build',backend='CUDA',fit_tool=True)])
|
|
gpus=[dict(uuid='GPU-free',name='3060',total_mib=12000,free_mib=2600),dict(uuid='GPU-busy',name='5080',total_mib=16000,free_mib=173)]
|
|
ref=dict(id='medium',available=True,file='model.gguf',ubatch=256)
|
|
result=type('Result',(),dict(returncode=0,stdout='-c 160000 -ngl 4',stderr=''))()
|
|
with patch.object(r,'references',return_value={'profiles':[ref]}),patch.object(r,'prerequisites',return_value={'gpus':gpus}),patch('runtime.subprocess.run',return_value=result) as run:
|
|
value=r.fit('medium',160000,2)
|
|
self.assertTrue(value['success']);self.assertFalse(value['model_loaded'])
|
|
self.assertEqual(value['excluded_gpus'],['5080'])
|
|
self.assertEqual(run.call_args.kwargs['env']['CUDA_VISIBLE_DEVICES'],'GPU-free')
|
|
args=run.call_args.args[0];self.assertEqual(args[args.index('--ctx-size')+1],'160000');self.assertEqual(args[args.index('--parallel')+1],'2')
|
|
def test_fit_source_adapter_is_explicit_and_idempotent(self):
|
|
from runtime import prepare_fit_source
|
|
with tempfile.TemporaryDirectory() as d:
|
|
path=Path(d)/'tools/fit-params/fit-params.cpp';path.parent.mkdir(parents=True)
|
|
path.write_text(' llama_backend_init();\n')
|
|
prepare_fit_source(d);first=path.read_text();prepare_fit_source(d)
|
|
self.assertEqual(first,path.read_text());self.assertIn('params.kv_unified = true',first)
|
|
path.write_text('unknown upstream layout')
|
|
with self.assertRaises(ValueError):prepare_fit_source(d)
|
|
def test_cancels_only_owned_real_process(self):
|
|
import sys,threading,time
|
|
with tempfile.TemporaryDirectory() as d:
|
|
r=Runtime(d);r.state['job']={'state':'running'};errors=[]
|
|
def work():
|
|
try:r.run([sys.executable,'-c','import time; time.sleep(60)'],Path(d),'test',70)
|
|
except InterruptedError:pass
|
|
except Exception as exc:errors.append(exc)
|
|
t=threading.Thread(target=work);t.start()
|
|
for _ in range(100):
|
|
if r.process is not None:break
|
|
time.sleep(.01)
|
|
child=r.process
|
|
self.assertIsNotNone(child)
|
|
r.stop();t.join(3)
|
|
self.assertFalse(t.is_alive());self.assertIsNotNone(child.poll());self.assertEqual(errors,[])
|