199 lines
8.7 KiB
Python
199 lines
8.7 KiB
Python
"""Bounded, cancellable model subprocesses with CPU as the default device."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from contextlib import closing
|
|
from contextvars import ContextVar
|
|
from pathlib import Path
|
|
from typing import Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from app.config import BACKEND_DIR
|
|
from app.database.db import connect
|
|
from app.errors import ApiError
|
|
from app.local_models.catalog import CATALOG
|
|
from app.local_models.manager import model_path, read_state
|
|
from app.providers.base import ProviderError
|
|
|
|
|
|
class RuntimeConfig(BaseModel):
|
|
device: Literal["cpu", "cuda"] = "cpu"
|
|
cpu_threads: int = Field(default=2, ge=1, le=32)
|
|
memory_limit_mb: int = Field(default=8192, ge=1024, le=131072)
|
|
gpu_memory_limit_mb: int = Field(default=4096, ge=512, le=65536)
|
|
timeout_seconds: int = Field(default=1800, ge=30, le=14400)
|
|
embedding_model: Literal["bekko", "granite"] = "bekko"
|
|
version: int = Field(default=1, ge=1)
|
|
|
|
|
|
runtime_context = ContextVar("runtime_config", default=None)
|
|
runtime_progress = ContextVar("runtime_progress", default=None)
|
|
|
|
|
|
def configuration():
|
|
if runtime_context.get() is not None:
|
|
return runtime_context.get()
|
|
with closing(connect()) as conn:
|
|
conn.execute("CREATE TABLE IF NOT EXISTS local_runtime_config (id INTEGER PRIMARY KEY CHECK(id=1), config_json TEXT NOT NULL)")
|
|
row = conn.execute("SELECT config_json FROM local_runtime_config WHERE id=1").fetchone()
|
|
return RuntimeConfig.model_validate_json(row[0]) if row else RuntimeConfig()
|
|
|
|
|
|
def configure(request):
|
|
from app.database.db import transaction
|
|
configuration()
|
|
with closing(connect()) as conn, transaction(conn):
|
|
row = conn.execute("SELECT config_json FROM local_runtime_config WHERE id=1").fetchone()
|
|
previous = RuntimeConfig.model_validate_json(row[0]) if row else RuntimeConfig()
|
|
if request.version != previous.version:
|
|
raise ApiError(409, "VERSION_CONFLICT", "Local runtime settings changed; reload first.")
|
|
request = request.model_copy(update={"version": request.version + 1})
|
|
conn.execute("INSERT OR REPLACE INTO local_runtime_config VALUES (1,?)", (request.model_dump_json(),))
|
|
return request
|
|
|
|
|
|
def interpreter():
|
|
return Path(os.getenv("APP_MODEL_PYTHON", str(BACKEND_DIR / ".venv-models" / ("Scripts/python.exe" if os.name == "nt" else "bin/python"))))
|
|
|
|
|
|
class Runtime:
|
|
def __init__(self):
|
|
self.active = {}
|
|
self.active_files = {}
|
|
self.waiters = []
|
|
self.counter = 0
|
|
self.diagnostics = []
|
|
|
|
def in_use(self, key):
|
|
return key in self.active.values()
|
|
|
|
def media_in_use(self, path):
|
|
target = str(Path(path).resolve())
|
|
return any(target in paths for paths in self.active_files.values())
|
|
|
|
async def infer(self, key, operation, payload, *, priority=10):
|
|
if read_state(key)["status"] != "installed":
|
|
raise ProviderError("LOCAL_MODEL_NOT_INSTALLED", "请先在模型配置中下载本地模型。")
|
|
if not interpreter().is_file():
|
|
raise ProviderError("LOCAL_RUNTIME_NOT_INSTALLED", "请先运行本地模型 CPU/CUDA 安装脚本。")
|
|
config = configuration()
|
|
self.counter += 1
|
|
ticket = (priority, self.counter)
|
|
self.waiters.append(ticket)
|
|
process = None
|
|
attempt = None
|
|
try:
|
|
# One resident model at a time prevents overlapping CPU/GPU allocations.
|
|
while self.active or ticket != min(self.waiters):
|
|
await asyncio.sleep(0.05)
|
|
self.waiters.remove(ticket)
|
|
self.active[ticket] = key
|
|
self.active_files[ticket] = {str(Path(payload[name]).resolve()) for name in ("source", "reference") if payload.get(name)}
|
|
# Deletion may have occurred while this request was queued.
|
|
if read_state(key)["status"] != "installed":
|
|
raise ProviderError("LOCAL_MODEL_NOT_INSTALLED", "模型文件已被删除。")
|
|
from app.services.usage_service import UsageAttempt
|
|
attempt = UsageAttempt("local-models", CATALOG[key].repository, "local", operation, source="local")
|
|
env = {**os.environ, "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1",
|
|
"HF_HUB_DISABLE_TELEMETRY": "1", "OMP_NUM_THREADS": str(config.cpu_threads),
|
|
"PYTHONIOENCODING": "utf-8"}
|
|
process = await asyncio.create_subprocess_exec(str(interpreter()), str(Path(__file__).with_name("worker.py")),
|
|
stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL,
|
|
env=env, limit=16 * 1024 * 1024, **({"creationflags": 0x08000000} if os.name == "nt" else {}))
|
|
request = {"key": key, "operation": operation, "model_path": str(model_path(key).resolve()),
|
|
"config": config.model_dump(), "payload": payload}
|
|
async def receive():
|
|
process.stdin.write(json.dumps(request).encode())
|
|
await process.stdin.drain()
|
|
process.stdin.close()
|
|
final = None
|
|
while line := await process.stdout.readline():
|
|
if len(line) > 16 * 1024 * 1024:
|
|
raise ProviderError("LOCAL_MODEL_INVALID_RESPONSE", "本地模型输出超限。")
|
|
message = json.loads(line)
|
|
if "progress" in message:
|
|
callback = runtime_progress.get()
|
|
if callback:
|
|
callback(message)
|
|
else:
|
|
final = message
|
|
await process.wait()
|
|
return final
|
|
try:
|
|
result = await asyncio.wait_for(receive(), config.timeout_seconds)
|
|
except TimeoutError as exc:
|
|
raise ProviderError("LOCAL_MODEL_TIMEOUT", "本地模型处理超时。") from exc
|
|
if process.returncode != 0:
|
|
raise ProviderError("LOCAL_MODEL_PROCESS_FAILED", "本地模型进程退出，请检查依赖与资源预算。")
|
|
if not isinstance(result, dict):
|
|
raise ProviderError("LOCAL_MODEL_INVALID_RESPONSE", "本地模型进程未返回有效结果。")
|
|
if "error_code" in result:
|
|
raise ProviderError(result["error_code"], result.get("message", "本地推理失败。"))
|
|
attempt.observe(result)
|
|
attempt.completed = True
|
|
self.diagnostics.append({"model": CATALOG[key].repository, "revision": CATALOG[key].revision,
|
|
**result.get("diagnostics", {})})
|
|
self.diagnostics = self.diagnostics[-100:]
|
|
return result["result"]
|
|
finally:
|
|
if ticket in self.waiters:
|
|
self.waiters.remove(ticket)
|
|
if process is not None and process.returncode is None:
|
|
process.kill()
|
|
await process.wait()
|
|
self.active.pop(ticket, None)
|
|
self.active_files.pop(ticket, None)
|
|
if attempt:
|
|
attempt.persist()
|
|
|
|
|
|
runtime = Runtime()
|
|
|
|
|
|
class LocalEmbedding:
|
|
dim = 384
|
|
|
|
@property
|
|
def model_id(self):
|
|
spec = CATALOG[configuration().embedding_model]
|
|
return f"{spec.repository}@{spec.revision}"
|
|
|
|
@property
|
|
def version(self):
|
|
return CATALOG[configuration().embedding_model].revision
|
|
|
|
@property
|
|
def available(self):
|
|
return read_state(configuration().embedding_model)["status"] == "installed" and interpreter().is_file()
|
|
|
|
async def embed_documents(self, texts):
|
|
return await runtime.infer(configuration().embedding_model, "embedding", {"texts": texts}, priority=0)
|
|
|
|
async def embed_query(self, query):
|
|
return (await self.embed_documents([query]))[0]
|
|
|
|
|
|
class LocalSpeech:
|
|
@property
|
|
def available(self):
|
|
return self.available_for("transcription")
|
|
|
|
def available_for(self, capability):
|
|
key = "qwen3-asr" if capability == "transcription" else "eres2netv2"
|
|
return read_state(key)["status"] == "installed" and interpreter().is_file()
|
|
|
|
async def transcribe(self, source, language):
|
|
from app.providers.routing import RoutedTranscript
|
|
from app.contracts import TranscriptSegment
|
|
result = await runtime.infer("qwen3-asr", "transcription", {"source": str(source.resolve()), "language": language})
|
|
return RoutedTranscript(text=result["text"], source="local",
|
|
segments=[TranscriptSegment(**s) for s in result["segments"]])
|
|
|
|
async def match(self, source, reference):
|
|
result = await runtime.infer("eres2netv2", "speaker_matching",
|
|
{"source": str(source.resolve()), "reference": str(reference.resolve())}, priority=0)
|
|
return result["score"]
|