Vladimir BryzgalovandClaude Opus 5
42a98bd703
Масштабирование ONNX зависит от процессора: замеры на Apple M4 не переносятся на Ryzen, а подбирать конфигурацию перезапусками мучительно. Теперь /v1/benchmark гоняет минуту записи на 1, 2, 4, 8 и 16 потоках и говорит, что поставить в config.toml. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
330 lines
16 KiB
Python
330 lines
16 KiB
Python
"""HTTP-сервис распознавания: приём файла, очередь, выдача результата."""
|
||
import logging
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from contextlib import asynccontextmanager
|
||
from pathlib import Path
|
||
|
||
from fastapi import Depends, FastAPI, File, HTTPException, Query, Request, UploadFile
|
||
from fastapi.openapi.docs import get_swagger_ui_html
|
||
from fastapi.openapi.utils import get_openapi
|
||
from fastapi.responses import JSONResponse
|
||
from fastapi.security import HTTPBearer
|
||
|
||
from app.config import Settings, load_settings
|
||
from app.pipeline import ModelsMissing, Pipeline, to_wav16k
|
||
from app.security import check_token, ip_allowed, parse_allowlist
|
||
from app.store import JobStatus, JobStore
|
||
from app.version import __version__
|
||
|
||
log = logging.getLogger("talkscore-asr")
|
||
|
||
settings: Settings = load_settings()
|
||
allowlist = parse_allowlist(settings.allow_ips)
|
||
store = JobStore(settings.data_dir / "jobs.db")
|
||
def make_pipeline() -> Pipeline:
|
||
"""Своя копия моделей на каждого воркера: они не рассчитаны на общий доступ."""
|
||
return Pipeline(
|
||
models_dir=settings.models_dir,
|
||
threads=settings.effective_threads(),
|
||
replacements_path=settings.replacements_path,
|
||
base_dir=settings.base_dir,
|
||
)
|
||
|
||
|
||
pipeline = make_pipeline()
|
||
|
||
_worker_stop = threading.Event()
|
||
_state: dict = {"ready": False, "error": None}
|
||
|
||
|
||
def _process(job_id: str, worker: Pipeline) -> None:
|
||
"""Задача уже помечена в работе тем, кто её забрал."""
|
||
job = store.get(job_id)
|
||
if job is None:
|
||
return
|
||
upload = settings.data_dir / "uploads" / job_id
|
||
try:
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
wav = Path(tmp) / "audio.wav"
|
||
to_wav16k(upload, wav, worker.ffmpeg)
|
||
speakers = int(job["options"].get("speakers", settings.speakers))
|
||
result = worker.transcribe(wav, num_speakers=speakers)
|
||
result["filename"] = job["filename"]
|
||
store.mark_done(job_id, result)
|
||
log.info("задача %s готова: %.1f с аудио, x%s", job_id,
|
||
result["duration_sec"], result["timing"]["realtime_factor"])
|
||
except Exception as exc: # noqa: BLE001 - в статус задачи должна попасть любая причина
|
||
log.exception("задача %s провалилась", job_id)
|
||
store.mark_failed(job_id, error=f"{type(exc).__name__}: {exc}")
|
||
finally:
|
||
upload.unlink(missing_ok=True)
|
||
|
||
|
||
def _worker_loop(index: int, worker: Pipeline) -> None:
|
||
"""Разбирает очередь. Задача захватывается атомарно, поэтому воркеров может быть много."""
|
||
last_cleanup = 0.0
|
||
while not _worker_stop.is_set():
|
||
if _state["ready"]:
|
||
job_id = store.claim_next()
|
||
if job_id:
|
||
log.info("воркер %d взял задачу %s", index, job_id)
|
||
_process(job_id, worker)
|
||
continue
|
||
# Уборкой занимается только первый воркер, чтобы не делать её хором.
|
||
if index == 0 and time.time() - last_cleanup > 3600:
|
||
removed = store.cleanup(settings.keep_results_hours)
|
||
if removed:
|
||
log.info("удалено старых задач: %d", removed)
|
||
last_cleanup = time.time()
|
||
_worker_stop.wait(1.0)
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
(settings.data_dir / "uploads").mkdir(parents=True, exist_ok=True)
|
||
try:
|
||
pipeline.warmup()
|
||
_state["ready"] = True
|
||
except ModelsMissing as exc:
|
||
_state["error"] = str(exc)
|
||
log.error("сервис запущен без моделей: %s", exc)
|
||
workers = []
|
||
for i in range(settings.effective_workers()):
|
||
# Первый воркер использует уже прогретый конвейер, остальные греются сами
|
||
# при первой задаче: держать копии моделей впустую незачем.
|
||
worker = pipeline if i == 0 else make_pipeline()
|
||
thread = threading.Thread(target=_worker_loop, args=(i, worker),
|
||
name=f"asr-worker-{i}", daemon=True)
|
||
thread.start()
|
||
workers.append(thread)
|
||
|
||
log.info("сервис слушает %s:%s, воркеров %d по %d потоков", settings.host,
|
||
settings.port, settings.effective_workers(), settings.effective_threads())
|
||
yield
|
||
_worker_stop.set()
|
||
for thread in workers:
|
||
thread.join(timeout=5)
|
||
|
||
|
||
# Штатные /docs и /openapi.json отключены: они не требуют токена. Вместо них
|
||
# ниже свои маршруты, закрытые тем же списком адресов, что и остальной сервис.
|
||
app = FastAPI(
|
||
title="talkscore-asr",
|
||
version=__version__,
|
||
description=(
|
||
"Локальное распознавание русской речи с разделением по говорящим.\n\n"
|
||
"Все методы, кроме `/health`, требуют заголовок `Authorization: Bearer <токен>`. "
|
||
"Нажмите **Authorize** и вставьте токен из `config.toml`."
|
||
),
|
||
lifespan=lifespan,
|
||
docs_url=None,
|
||
redoc_url=None,
|
||
openapi_url=None,
|
||
)
|
||
bearer_scheme = HTTPBearer(description="Токен из config.toml, раздел [security]")
|
||
|
||
|
||
def ip_guard(request: Request) -> None:
|
||
"""Только проверка адреса: нужна для страницы с описанием методов."""
|
||
client_ip = request.client.host if request.client else None
|
||
if not ip_allowed(client_ip, allowlist):
|
||
raise HTTPException(status_code=403, detail="адрес не в списке разрешённых")
|
||
|
||
|
||
def guard(request: Request) -> None:
|
||
"""Проверяет адрес и токен. Порядок важен: сначала сеть, потом секрет."""
|
||
client_ip = request.client.host if request.client else None
|
||
if not ip_allowed(client_ip, allowlist):
|
||
log.warning("отказано по адресу: %s", client_ip)
|
||
raise HTTPException(status_code=403, detail="адрес не в списке разрешённых")
|
||
if not check_token(request.headers.get("authorization"), settings.token):
|
||
raise HTTPException(status_code=401, detail="неверный или отсутствующий токен")
|
||
|
||
|
||
@app.get("/docs", include_in_schema=False)
|
||
def docs_page(request: Request):
|
||
"""Описание методов для браузера. Открывается только с разрешённых адресов."""
|
||
ip_guard(request)
|
||
return get_swagger_ui_html(openapi_url="openapi.json", title="talkscore-asr")
|
||
|
||
|
||
@app.get("/openapi.json", include_in_schema=False)
|
||
def openapi_schema(request: Request) -> JSONResponse:
|
||
ip_guard(request)
|
||
schema = get_openapi(title=app.title, version=app.version,
|
||
description=app.description, routes=app.routes)
|
||
schema["components"] = schema.get("components", {})
|
||
schema["components"]["securitySchemes"] = {
|
||
"bearer": {"type": "http", "scheme": "bearer"}}
|
||
schema["security"] = [{"bearer": []}]
|
||
return JSONResponse(schema)
|
||
|
||
|
||
@app.get("/health", summary="Состояние сервиса",
|
||
description="Единственный метод без токена - годится для мониторинга. "
|
||
"Показывает версию, очередь, число потоков и то, каким сервис "
|
||
"видит ваш адрес.")
|
||
def health(request: Request) -> JSONResponse:
|
||
"""Проверка живости - без токена, чтобы годилась для мониторинга.
|
||
|
||
Показывает адрес обратившегося и вердикт по списку доступа: без этого
|
||
отладка отказов превращается в гадание, а сам адрес клиенту и так известен.
|
||
Список разрешённых адресов при этом не раскрывается.
|
||
"""
|
||
client_ip = request.client.host if request.client else None
|
||
return JSONResponse({
|
||
"status": "ok" if _state["ready"] else "no_models",
|
||
"version": __version__,
|
||
"error": _state["error"],
|
||
"queue": store.stats(),
|
||
"threads": settings.effective_threads(),
|
||
"workers": settings.effective_workers(),
|
||
"your_ip": client_ip,
|
||
"your_ip_allowed": ip_allowed(client_ip, allowlist),
|
||
"ip_filter_active": bool(allowlist),
|
||
})
|
||
|
||
|
||
@app.post("/v1/jobs", dependencies=[Depends(guard)], summary="Отправить запись",
|
||
description="Принимает аудиофайл любого формата, который читает ffmpeg: "
|
||
"mp3, wav, m4a, ogg, opus, wma. Возвращает идентификатор задачи - "
|
||
"результат забирается отдельным запросом, потому что обработка "
|
||
"занимает минуты.")
|
||
async def create_job(
|
||
file: UploadFile = File(...),
|
||
speakers: int | None = Query(None, ge=0, le=10,
|
||
description="число говорящих, 0 = определить автоматически"),
|
||
) -> dict:
|
||
if not _state["ready"]:
|
||
raise HTTPException(status_code=503, detail=_state["error"] or "сервис ещё не готов")
|
||
|
||
job_id = store.create(filename=file.filename or "audio",
|
||
duration_sec=0.0,
|
||
options={"speakers": settings.speakers if speakers is None else speakers})
|
||
target = settings.data_dir / "uploads" / job_id
|
||
size = 0
|
||
try:
|
||
with target.open("wb") as out:
|
||
while chunk := await file.read(1 << 20):
|
||
size += len(chunk)
|
||
if size > settings.max_upload_bytes:
|
||
raise HTTPException(status_code=413,
|
||
detail=f"файл больше {settings.max_upload_mb} МБ")
|
||
out.write(chunk)
|
||
except HTTPException:
|
||
target.unlink(missing_ok=True)
|
||
store.mark_failed(job_id, error="файл слишком большой")
|
||
raise
|
||
if size == 0:
|
||
target.unlink(missing_ok=True)
|
||
store.mark_failed(job_id, error="пустой файл")
|
||
raise HTTPException(status_code=400, detail="пустой файл")
|
||
|
||
return {"job_id": job_id, "status": JobStatus.QUEUED,
|
||
"queue_position": store.queue_position(job_id)}
|
||
|
||
|
||
@app.post("/v1/benchmark", dependencies=[Depends(guard)], summary="Подобрать число потоков",
|
||
description="Прогоняет минуту из присланной записи на разном числе потоков "
|
||
"и показывает, где быстрее. Занимает пару минут и на это время "
|
||
"нагружает процессор.")
|
||
async def benchmark(file: UploadFile = File(...)) -> dict:
|
||
from app.bench import run_benchmark
|
||
from app.pipeline import read_wav
|
||
|
||
if not _state["ready"]:
|
||
raise HTTPException(status_code=503, detail=_state["error"] or "сервис ещё не готов")
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
src = Path(tmp) / (file.filename or "audio")
|
||
with src.open("wb") as out:
|
||
while chunk := await file.read(1 << 20):
|
||
out.write(chunk)
|
||
wav = Path(tmp) / "audio.wav"
|
||
to_wav16k(src, wav, pipeline.ffmpeg)
|
||
samples = read_wav(wav)
|
||
|
||
return run_benchmark(samples, models_dir=settings.models_dir,
|
||
base_dir=settings.base_dir,
|
||
replacements_path=settings.replacements_path)
|
||
|
||
|
||
@app.get("/v1/jobs/{job_id}", dependencies=[Depends(guard)], summary="Забрать результат",
|
||
description="Пока задача не готова, возвращает статус и место в очереди. "
|
||
"Готовый результат содержит реплики с говорящими и тайм-кодами, "
|
||
"статистику разговора и время обработки.")
|
||
def get_job(job_id: str) -> dict:
|
||
job = store.get(job_id)
|
||
if job is None:
|
||
raise HTTPException(status_code=404, detail="задача не найдена")
|
||
|
||
body = {"job_id": job_id, "status": job["status"], "filename": job["filename"]}
|
||
if job["status"] == JobStatus.QUEUED:
|
||
body["queue_position"] = store.queue_position(job_id)
|
||
if job["status"] == JobStatus.DONE:
|
||
body.update(job["result"])
|
||
if job["status"] == JobStatus.FAILED:
|
||
body["error"] = job["error"]
|
||
return body
|
||
|
||
|
||
@app.delete("/v1/jobs/{job_id}", dependencies=[Depends(guard)], summary="Удалить задачу",
|
||
description="Убирает задачу и её результат. Завершённые задачи и так "
|
||
"удаляются сами через срок из keep_results_hours.")
|
||
def delete_job(job_id: str) -> dict:
|
||
if store.get(job_id) is None:
|
||
raise HTTPException(status_code=404, detail="задача не найдена")
|
||
store.delete(job_id)
|
||
return {"deleted": job_id}
|
||
|
||
|
||
def _setup_console() -> None:
|
||
"""Windows-консоль по умолчанию не в UTF-8, иначе русский текст в логах рассыпается."""
|
||
for stream in (sys.stdout, sys.stderr):
|
||
try:
|
||
stream.reconfigure(encoding="utf-8", errors="replace")
|
||
except (AttributeError, ValueError):
|
||
pass
|
||
|
||
|
||
def run() -> None:
|
||
import uvicorn
|
||
|
||
_setup_console()
|
||
logging.basicConfig(level=logging.INFO,
|
||
format="%(asctime)s %(levelname)s %(name)s: %(message)s")
|
||
|
||
if not settings.token:
|
||
print("\n В config.toml пустой токен - сервис никого не пустит."
|
||
"\n Впишите значение в [security] token и перезапустите.\n")
|
||
|
||
# Модели проверяются здесь, а не по _state: lifespan отработает уже внутри
|
||
# uvicorn.run, и к тому моменту сообщение выводить поздно.
|
||
if not (settings.models_dir / "gigaam" / "config.json").is_file():
|
||
print("\n Модели не найдены. Сначала запустите download_models.bat"
|
||
"\n Сервис поднимется, но принимать записи не сможет.\n")
|
||
|
||
nets = len(allowlist)
|
||
access = f"только с {nets} адресов" if nets else "со ВСЕХ адресов"
|
||
print(f"\n Настройки: {settings.base_dir / 'config.toml'}"
|
||
f"\n Токен: {'задан' if settings.token else 'НЕ ЗАДАН, сервис никого не пустит'}"
|
||
f"\n Доступ: {access}"
|
||
f"\n Обработка: {settings.effective_workers()} задач одновременно"
|
||
f" по {settings.effective_threads()} потока")
|
||
print(f"\n talkscore-asr слушает http://{settings.host}:{settings.port}"
|
||
f"\n Потоков: {settings.effective_threads()}"
|
||
f"\n Проверка: curl http://localhost:{settings.port}/health"
|
||
"\n Остановить: Ctrl+C\n")
|
||
# Windows-консоль не разбирает ANSI-последовательности и печатает их как мусор
|
||
uvicorn.run(app, host=settings.host, port=settings.port, log_level="info",
|
||
use_colors=(os.name != "nt"))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run()
|