60 lines
4.3 KiB
Python
60 lines
4.3 KiB
Python
import json
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
from chat_test import ChatTests
|
|
from inference import Scheduler,InferenceError,LlamaWorker
|
|
|
|
class Reply:
|
|
status=200
|
|
def __init__(self,events):self.events=iter(events)
|
|
def read1(self,_):return next(self.events,b'')
|
|
|
|
class Worker:
|
|
def __init__(self):self.error=None;self.loaded=False;self.conn=Mock();self.conn.getresponse.return_value=Reply([b'data: {"choices":[{"delta":{"content":"OK"}}]}\n\n',b'data: [DONE]\n\n'])
|
|
def stop(self):self.loaded=False
|
|
def ensure(self,p,cancel=lambda:False):
|
|
if self.error:raise InferenceError(self.error)
|
|
self.loaded=True
|
|
def status(self):return dict(state='ready' if self.loaded else 'stopped',error=self.error)
|
|
def connect(self):return self.conn,'synthetic-only'
|
|
|
|
class ChatTestsTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.profile=dict(id='p',name='test',revision=1,kind='chat',runnable=True,parameters=dict(slots=2,temperature=.2,top_p=.8,top_k=20))
|
|
self.worker=Worker();self.manager=ChatTests(SimpleNamespace(status=lambda:dict(profiles=[self.profile])),self.worker,Scheduler(self.worker))
|
|
def request(self):return dict(profile_id='p',messages=[dict(role='user',content='synthetic test')],max_tokens=8)
|
|
def finish(self):
|
|
deadline=time.monotonic()+3
|
|
while self.manager.status()['job']['state']=='running' and time.monotonic()<deadline:time.sleep(.01)
|
|
return self.manager.status()['job']
|
|
def test_answer_stream_and_no_prompt_in_status(self):
|
|
self.manager.start(self.request());job=self.finish();self.assertEqual(job['state'],'complete');self.assertEqual(job['answer'],'OK');self.assertNotIn('synthetic test',json.dumps(job));self.assertNotIn('messages',job)
|
|
self.assertTrue(self.worker.loaded);self.manager.unload();self.assertFalse(self.worker.loaded)
|
|
def test_memory_rejection_reaches_user_and_releases_lease(self):
|
|
self.worker.error='GPU-Split passt nicht in den Speicher.';self.manager.start(self.request());job=self.finish();self.assertEqual(job['state'],'failed');self.assertIn('Speicher',job['error']);self.assertEqual(self.manager.scheduler.active,0)
|
|
def test_queue_cancel_does_not_stop_other_request(self):
|
|
with self.manager.scheduler.lease(('other',)):
|
|
self.worker.loaded=True;self.manager.start(self.request());self.manager.stop();job=self.finish();self.assertEqual(job['state'],'cancelled');self.assertTrue(self.worker.loaded)
|
|
with self.assertRaises(ValueError):self.manager.unload()
|
|
def test_truncated_stream_is_not_success(self):
|
|
self.worker.conn.getresponse.return_value=Reply([b'data: {"choices":[]}\n\n']);self.manager.start(self.request());self.assertEqual(self.finish()['state'],'failed')
|
|
def test_native_timings_and_usage_from_final_empty_choices(self):
|
|
event=dict(choices=[],timings=dict(prompt_n=100,prompt_ms=200,prompt_per_second=500,predicted_n=20,predicted_ms=1000,predicted_per_second=20,secret='excluded'),usage=dict(prompt_tokens=150,completion_tokens=20,total_tokens=170))
|
|
self.worker.conn.getresponse.return_value=Reply([('data: '+json.dumps(event)+'\n\n').encode(),b'data: [DONE]\n\n'])
|
|
self.manager.start(self.request());job=self.finish()
|
|
self.assertEqual(job['timings']['prompt_per_second'],500);self.assertEqual(job['usage']['prompt_tokens'],150);self.assertNotIn('secret',job['timings'])
|
|
body=json.loads(self.worker.conn.request.call_args.args[2]);self.assertTrue(body['stream_options']['include_usage'])
|
|
def test_missing_and_invalid_timings_are_not_invented(self):
|
|
self.worker.conn.getresponse.return_value=Reply([b'data: {"choices":[],"timings":{"prompt_n":true,"prompt_ms":-1,"predicted_per_second":"42"},"usage":null}\n\n',b'data: [DONE]\n\n'])
|
|
self.manager.start(self.request());job=self.finish();self.assertEqual(job['timings'],{});self.assertEqual(job['usage'],{})
|
|
def test_validation(self):
|
|
for data in [dict(self.request(),max_tokens=True),dict(self.request(),messages=[dict(role='user',content='')]),dict(self.request(),profile_id='missing')]:
|
|
with self.assertRaises(ValueError):self.manager.start(data)
|
|
def test_oom_claim_requires_counter_evidence(self):
|
|
worker=object.__new__(LlamaWorker);worker.oom_before=4;worker.oom_count=lambda:4;self.assertIn('nicht eindeutig',worker.crash_message());worker.oom_count=lambda:5;self.assertIn('OOM-Kill im Deck-RAM',worker.crash_message())
|
|
|
|
if __name__=='__main__':unittest.main()
|