43 lines
3.4 KiB
Python
43 lines
3.4 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch,Mock
|
|
from inference import LlamaWorker,InferenceError
|
|
from profiles import SCHEMAS,CHAT_GPU_DEFAULTS
|
|
|
|
class WorkerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp=tempfile.TemporaryDirectory();self.root=Path(self.tmp.name);directory=self.root/'runtime'/'build-id'/'build/bin';directory.mkdir(parents=True);(directory/'llama-server').touch()
|
|
self.entry=dict(id='model',file='test.gguf',size=1024**3)
|
|
self.runtime=SimpleNamespace(root=self.root/'runtime',status=lambda:dict(active='build-id',builds=[dict(id='build-id',backend='CUDA',fit_tool=True)]))
|
|
self.worker=LlamaWorker(self.root/'worker',SimpleNamespace(root=self.root/'models',entry=lambda _:self.entry),self.runtime)
|
|
self.profile=dict(id='p',name='medium',revision=1,model_id='model',parameters={**{k:v[2] for k,v in SCHEMAS['chat'].items()},**CHAT_GPU_DEFAULTS,'gpu_devices':['gpu-first','gpu-second'],'split_mode':'layer','tensor_split':[85,15]})
|
|
self.gpus=[dict(uuid='gpu-first',processes=0,total_mib=16000,free_mib=15000),dict(uuid='gpu-second',processes=0,total_mib=12000,free_mib=11000)]
|
|
def tearDown(self):self.tmp.cleanup()
|
|
def patches(self):
|
|
real=Path.read_text
|
|
return patch.object(Path,'read_text',lambda path,*a,**kw:'MemAvailable: 50000000 kB\n' if str(path)=='/proc/meminfo' else real(path,*a,**kw))
|
|
def test_fixed_split_uses_prediction_and_preserves_device_order(self):
|
|
process=Mock();process.poll.return_value=None
|
|
http=Mock();http.getresponse.return_value.status=200
|
|
with self.patches(),patch('inference.probe',return_value=self.gpus),patch('inference.cgroup_headroom',return_value=32*1024**3),patch.object(self.worker,'_fit_command',return_value='CUDA0 9000 1000 1000\nCUDA1 2000 1000 1000\nHost 512 0 256\n') as fit,patch('inference.subprocess.Popen',return_value=process) as launch,patch('inference.http.client.HTTPConnection',return_value=http),patch('inference.threading.Thread'):
|
|
self.worker.ensure(self.profile)
|
|
args=launch.call_args.args[0];env=launch.call_args.kwargs['env']
|
|
self.assertEqual(env['CUDA_VISIBLE_DEVICES'],'gpu-first,gpu-second');self.assertEqual(args[args.index('--tensor-split')+1],'85,15');self.assertIn('--kv-unified',args);self.assertIn('--fit-print',fit.call_args.args[0]);self.assertEqual(self.worker.status()['state'],'ready')
|
|
def test_busy_gpu_is_rejected_before_fit_or_spawn(self):
|
|
self.gpus[0]['processes']=1
|
|
with patch('inference.probe',return_value=self.gpus),patch('inference.subprocess.Popen') as launch:
|
|
with self.assertRaises(InferenceError):self.worker.ensure(self.profile)
|
|
launch.assert_not_called()
|
|
def test_fixed_split_exhaustion_does_not_load_model(self):
|
|
with self.patches(),patch('inference.probe',return_value=self.gpus),patch('inference.cgroup_headroom',return_value=32*1024**3),patch.object(self.worker,'_fit_command',return_value='CUDA0 50000 1000 1000\nCUDA1 2000 1000 1000\nHost 512 0 256\n') as fit,patch('inference.subprocess.Popen') as launch:
|
|
with self.assertRaises(InferenceError):self.worker.ensure(self.profile)
|
|
launch.assert_not_called();self.assertLessEqual(fit.call_count,11)
|
|
def test_build_symlink_cannot_escape_deck_state(self):
|
|
self.runtime.status=lambda:dict(active='outside',builds=[dict(id='outside',backend='CUDA',fit_tool=True)])
|
|
(self.root/'runtime'/'outside').symlink_to('/usr')
|
|
with self.assertRaises(InferenceError):self.worker.build()
|
|
|
|
if __name__=='__main__':unittest.main()
|