98 lines
6.1 KiB
Python
98 lines
6.1 KiB
Python
"""Ephemeral admin chat tests using the same worker and leases as the API."""
|
||
import json
|
||
import math
|
||
import socket
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from inference import InferenceError
|
||
|
||
class ChatTests:
|
||
def __init__(self,profiles,worker,scheduler):
|
||
self.profiles=profiles;self.worker=worker;self.scheduler=scheduler
|
||
self.lock=threading.RLock();self.job=None;self.cancel=threading.Event();self.socket=None
|
||
def status(self):
|
||
with self.lock:job=dict(self.job) if self.job else None
|
||
return dict(job=job,worker=self.worker.status(),scheduler=self.scheduler.status())
|
||
def start(self,data):
|
||
if set(data)!={'profile_id','messages','max_tokens'}:raise ValueError('Profil, Nachrichten und Antwortlimit erforderlich.')
|
||
messages=data['messages'];limit=data['max_tokens']
|
||
if type(limit)!=int or not 1<=limit<=4096:raise ValueError('Antwortlimit: 1–4096 Token.')
|
||
if not isinstance(messages,list) or not 1<=len(messages)<=32:raise ValueError('1–32 Textnachrichten erforderlich.')
|
||
for m in messages:
|
||
if not isinstance(m,dict) or set(m)!={'role','content'} or m['role'] not in ('user','assistant','system') or not isinstance(m['content'],str) or not m['content'].strip():raise ValueError('Nur Textnachrichten mit Rolle user, assistant oder system unterstützt.')
|
||
if sum(len(m['content']) for m in messages)>24000:raise ValueError('Testverlauf zu lang (maximal 24.000 Zeichen). Neuen Chat beginnen.')
|
||
profile=next((p for p in self.profiles.status()['profiles'] if p['id']==data['profile_id'] and p['kind']=='chat'),None)
|
||
if not profile or not profile['runnable']:raise ValueError('Kein ausführbares Chatprofil. Modell und CUDA-Build prüfen.')
|
||
with self.lock:
|
||
if self.job and self.job['state']=='running':raise ValueError('Ein Chat-Test läuft bereits.')
|
||
self.cancel.clear();self.job=dict(id=uuid.uuid4().hex,state='running',phase='Wartet auf freie Modellreservierung',profile_id=profile['id'],profile_name=profile['name'],started_at=time.time(),answer='',reasoning='',error=None,timings={},usage={})
|
||
threading.Thread(target=self._run,args=(profile,messages,limit),daemon=True).start()
|
||
return dict(self.job)
|
||
def phase(self,text):
|
||
with self.lock:self.job['phase']=text
|
||
def stop(self):
|
||
self.cancel.set()
|
||
with self.lock:sock=self.socket
|
||
if sock:
|
||
try:sock.shutdown(socket.SHUT_RDWR)
|
||
except OSError:pass
|
||
return {'cancellation_requested':True}
|
||
def unload(self):
|
||
if not self.scheduler.unload_idle():raise ValueError('Es laufen noch Anfragen. Erst deren Ende abwarten; andere Antworten werden nicht abgebrochen.')
|
||
return self.status()
|
||
def _run(self,profile,messages,limit):
|
||
conn=None;outcome='failed';message='Chat-Test fehlgeschlagen.';error=None
|
||
try:
|
||
key=('chat',profile['id'],profile['revision'])
|
||
def allowed():return not self.cancel.is_set() and any(p['id']==profile['id'] and p['revision']==profile['revision'] for p in self.profiles.status()['profiles'])
|
||
def prepare():
|
||
self.phase('Speicher wird geprüft · Profil wird geladen')
|
||
self.worker.ensure(profile,cancel=self.cancel.is_set)
|
||
with self.scheduler.lease(key,profile['parameters']['slots'],prepare,allowed=allowed):
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
self.phase('Antwort wird erzeugt')
|
||
body=dict(model=profile['name'],messages=messages,max_tokens=limit,stream=True,stream_options={"include_usage":True},**{k:profile['parameters'][k] for k in ('temperature','top_p','top_k')})
|
||
conn,token=self.worker.connect()
|
||
conn.request('POST','/v1/chat/completions',json.dumps(body).encode(),headers={'Content-Type':'application/json','Authorization':'Bearer '+token})
|
||
with self.lock:self.socket=conn.sock
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
response=conn.getresponse()
|
||
if response.status!=200:
|
||
raise InferenceError('llama.cpp hat den Test abgelehnt. Kontextbudget, Nachrichtenlänge und Profilparameter prüfen.')
|
||
buffer=b'';total=0;done=False;deadline=time.monotonic()+600
|
||
while time.monotonic()<deadline:
|
||
if self.cancel.is_set():raise InterruptedError()
|
||
chunk=response.read1(8192)
|
||
if not chunk:break
|
||
buffer+=chunk;total+=len(chunk)
|
||
if total>4*1024**2:raise InferenceError('Testantwort überschreitet 4 MiB.')
|
||
while b'\n' in buffer:
|
||
line,buffer=buffer.split(b'\n',1);line=line.strip()
|
||
if not line.startswith(b'data:'):continue
|
||
payload=line[5:].strip()
|
||
if payload==b'[DONE]':done=True;break
|
||
event=json.loads(payload)
|
||
if 'error' in event:raise InferenceError('Der Modellworker meldet einen Fehler während der Antwort.')
|
||
with self.lock:
|
||
for section,fields in [('timings',('prompt_n','prompt_ms','prompt_per_second','predicted_n','predicted_ms','predicted_per_second')),('usage',('prompt_tokens','completion_tokens','total_tokens'))]:
|
||
values=event.get(section)
|
||
if isinstance(values,dict):
|
||
self.job[section]={k:v for k,v in values.items() if k in fields and type(v) in (int,float) and math.isfinite(v) and v>=0}
|
||
for choice in event.get('choices',[]):
|
||
delta=choice.get('delta',{})
|
||
with self.lock:
|
||
for source,target in [('content','answer'),('reasoning_content','reasoning')]:
|
||
if isinstance(delta.get(source),str):self.job[target]+=delta[source]
|
||
if choice.get('finish_reason'):self.job['finish_reason']=choice['finish_reason']
|
||
if done:break
|
||
if not done:raise InferenceError('Antwort unterbrochen oder Zeitlimit erreicht. '+(self.worker.status().get('error') or 'Ein GPU-OOM ist ohne eindeutigen Nachweis nicht bestätigt.'))
|
||
outcome='complete';message='Antwort fertig · Modell bleibt für weitere Anfragen geladen'
|
||
except Exception as exc:
|
||
cancelled=self.cancel.is_set()
|
||
message='Chat-Test abgebrochen.' if cancelled else (str(exc) if isinstance(exc,InferenceError) else 'Verbindung zum Modell unterbrochen. '+(self.worker.status().get('error') or 'Speicherfehler, Modellabsturz oder Zeitlimit möglich; Ursache nicht eindeutig.'))
|
||
outcome='cancelled' if cancelled else 'failed';error=None if cancelled else message
|
||
finally:
|
||
if conn:conn.close()
|
||
with self.lock:self.socket=None;self.job.update(state=outcome,error=error,phase=message,finished_at=time.time())
|