74 lines
3.9 KiB
Python
74 lines
3.9 KiB
Python
"""CPU-only EmbeddingGemma worker; independent of GPU mode and chat leases."""
|
|
import http.client
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import secrets
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
|
|
REPO='ggml-org/embeddinggemma-300m-qat-q8_0-GGUF'
|
|
FILE='embeddinggemma-300m-qat-Q8_0.gguf'
|
|
def supported(model):
|
|
return bool(model and model.get('kind')=='embeddings' and model.get('repo')==REPO and model.get('file')==FILE)
|
|
|
|
class Embeddings:
|
|
def __init__(self,root,catalog,worker):
|
|
self.root=Path(root);self.catalog=catalog;self.worker=worker;self.lock=threading.RLock();self.process=None;self.identity=None;self.port=None;self.key=None
|
|
def blockers(self,p):
|
|
if not supported(p.get('model')):return ['Noch kein geprüftes Embedding-Rezept für diese Datei. Unterstützt: EmbeddingGemma 300M QAT Q8_0.']
|
|
try:self.worker.build()
|
|
except ValueError as exc:return [str(exc)]
|
|
return []
|
|
def stop(self):
|
|
with self.lock:
|
|
process=self.process;self.process=None;self.identity=None
|
|
if process and process.poll() is None:
|
|
os.killpg(process.pid,signal.SIGTERM)
|
|
try:process.wait(timeout=10)
|
|
except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
|
|
(self.root/'key').unlink(missing_ok=True)
|
|
def ensure(self,p):
|
|
identity=(p['id'],p['revision'],self.worker.runtime.status()['active'])
|
|
if self.process and self.process.poll() is None and self.identity==identity:return
|
|
self.stop();errors=self.blockers(p)
|
|
if errors:raise ValueError(errors[0])
|
|
entry=self.catalog.entry(p['model_id']);model=self.catalog.root/entry['id']/'model.gguf';directory=self.worker.build()
|
|
self.root.mkdir(parents=True,exist_ok=True,mode=0o700)
|
|
with socket.socket() as sock:sock.bind(('127.0.0.1',0));self.port=sock.getsockname()[1]
|
|
self.key=secrets.token_urlsafe(32);keyfile=self.root/'key'
|
|
fd=os.open(keyfile,os.O_WRONLY|os.O_CREAT|os.O_TRUNC,0o600)
|
|
with os.fdopen(fd,'w') as out:out.write(self.key+'\n')
|
|
env=dict(os.environ,CUDA_VISIBLE_DEVICES='')
|
|
args=[str(directory/'build/bin/llama-server'),'--model',str(model),'--alias',p['name'],'--embedding','--pooling','mean','--ctx-size','4096','--batch-size','4096','--ubatch-size','1024','--parallel','4','--threads','8','--threads-batch','8','--n-gpu-layers','0','--host','127.0.0.1','--port',str(self.port),'--api-key-file',str(keyfile),'--no-ui']
|
|
self.process=subprocess.Popen(args,env=env,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL,start_new_session=True)
|
|
try:
|
|
deadline=time.monotonic()+120
|
|
while time.monotonic()<deadline:
|
|
if self.process.poll() is not None:raise ValueError('EmbeddingGemma konnte nicht gestartet werden. llama.cpp-Build und Modell prüfen.')
|
|
conn=http.client.HTTPConnection('127.0.0.1',self.port,timeout=2)
|
|
try:
|
|
conn.request('GET','/health',headers={'Authorization':'Bearer '+self.key})
|
|
if conn.getresponse().status==200:self.identity=identity;keyfile.unlink(missing_ok=True);return
|
|
except OSError:pass
|
|
finally:conn.close()
|
|
time.sleep(.25)
|
|
raise ValueError('Embedding-Start hat das Zeitlimit überschritten.')
|
|
except Exception:self.stop();raise
|
|
def generate(self,p,data):
|
|
with self.lock:
|
|
self.ensure(p);conn=http.client.HTTPConnection('127.0.0.1',self.port,timeout=120)
|
|
try:
|
|
conn.request('POST','/v1/embeddings',body=json.dumps(data).encode(),headers={'Content-Type':'application/json','Authorization':'Bearer '+self.key})
|
|
response=conn.getresponse();raw=response.read(4*1024*1024+1)
|
|
if response.status!=200:raise ValueError('Embedding-Anfrage abgelehnt. Textlänge und Eingabe prüfen (maximal ca. 1024 Token je Text).')
|
|
if len(raw)>4*1024*1024:raise ValueError('Embedding-Antwort zu groß.')
|
|
result=json.loads(raw)
|
|
if not result.get('data'):raise ValueError('Keine Embedding-Vektoren zurückgegeben.')
|
|
return result
|
|
except (OSError,ValueError):raise ValueError('Embedding-Anfrage fehlgeschlagen. Eingabelänge und Laufzeit prüfen.') from None
|
|
finally:conn.close()
|