feat: 调度器+并发管线+GPU优化+日志分级+前端修复
- scheduler: ffmpeg 异步线程 + GPU 串行调度 + 模型复用（2N→2 次加载） - pipeline: 阶段拆分（extract/asr/translate），中间数据存 Task 字段 - translate_service: 长度排序批处理，padding 浪费减少 91% - model_manager: ASR/翻译不共驻，BatchedInferencePipeline 批量解码 - 日志分级: INFO=任务流转里程碑，DEBUG=进度详情；默认 INFO - 前端: 日志最新在上+滚动感知+退避轮询；24h 时间；上传中状态显示 - /health: 返回完整 Whisper/NLLB 配置 - upload_service: 单事务 complete + 扩展名白名单 - task_router: 合并 UploadSession 虚拟任务到列表 - Dockerfile: CPU/GPU 独立构建链，deps 缓存稳定 - prefetch_models: 安装时预下载模型权重
This commit is contained in:
audio2text dev
236
app/services/scheduler.py
Normal file
236
app/services/scheduler.py
Normal file
@@ -0,0 +1,236 @@
|
||||
"""任务调度器：ffmpeg 异步提取 + GPU 阶段串行 + 模型复用。
|
||||
|
||||
设计动机：多任务时不应串行等一个任务全跑完才下一个。ffmpeg 是纯 CPU，可与 GPU 阶段
|
||||
并行；GPU 阶段（ASR + 翻译）串行化（共享显存），但卸载模型前查队列，有同类待处理
|
||||
任务就继续用当前模型，减少重复加载/卸载。
|
||||
|
||||
数据流：
|
||||
enqueue_task ──► ffmpeg 线程（每任务一个，异步，CPU）
|
||||
│ 提取音频 → task.wav_path → status=transcribing
|
||||
▼（唤醒 GPU 线程）
|
||||
GPU 调度线程（单线程，常驻）
|
||||
① 取 status=transcribing 的任务，get_asr()
|
||||
while 还有 transcribing 任务: asr_phase → status=translating
|
||||
（ASR 队列空，切翻译）
|
||||
② 取 status=translating 的任务，get_translator()
|
||||
while 还有 translating 任务: translate_phase → status=done
|
||||
（翻译队列空，回到 ① 等待）
|
||||
|
||||
模型复用：N 个任务的模型切换次数从 2N 降到最优 2 次（一批 ASR 全做完 → 切翻译 → 一批翻译全做完）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from ..database import get_session_local
|
||||
from ..models.task import (
|
||||
Task, STATUS_EXTRACTING, STATUS_TRANSCRIBING, STATUS_SEGMENTING,
|
||||
STATUS_TRANSLATING, STATUS_QUEUED, STATUS_FAILED,
|
||||
)
|
||||
from . import pipeline
|
||||
from .model_manager import get_model_manager
|
||||
|
||||
logger = logging.getLogger("audio2text.scheduler")
|
||||
|
||||
# 唤醒 GPU 调度线程的事件（新任务入队或 ffmpeg 完成时 set）
|
||||
_wake_event = threading.Event()
|
||||
# GPU 调度线程单例
|
||||
_scheduler_thread: threading.Thread | None = None
|
||||
_scheduler_started = False
|
||||
|
||||
|
||||
def enqueue_task(task_id: int) -> None:
|
||||
"""任务入队：起 ffmpeg 线程提取音频 + 唤醒 GPU 调度线程。
|
||||
|
||||
替代旧 pipeline.enqueue_task（每任务一个线程跑完整管线）。
|
||||
ffmpeg 在独立线程跑（CPU，与 GPU 并行），完成后 GPU 调度线程接管 ASR+翻译。
|
||||
"""
|
||||
_ensure_scheduler_running()
|
||||
t = threading.Thread(
|
||||
target=_extract_audio_async, args=(task_id,),
|
||||
name=f"ffmpeg-{task_id}", daemon=True,
|
||||
)
|
||||
t.start()
|
||||
logger.info("任务 %d 已入队，开始音频提取。", task_id)
|
||||
|
||||
|
||||
def start_scheduler() -> None:
|
||||
"""启动 GPU 调度线程（应用启动时调一次，幂等）。"""
|
||||
global _scheduler_thread, _scheduler_started
|
||||
if _scheduler_started:
|
||||
return
|
||||
_scheduler_started = True
|
||||
_reset_stuck_tasks()
|
||||
_scheduler_thread = threading.Thread(
|
||||
target=_gpu_scheduler, name="gpu-scheduler", daemon=True,
|
||||
)
|
||||
_scheduler_thread.start()
|
||||
logger.info("GPU 调度线程已启动。")
|
||||
|
||||
|
||||
def _reset_stuck_tasks() -> None:
|
||||
"""启动时清理卡在中间状态的任务（进程上次崩溃残留）。
|
||||
|
||||
transcribing 但无 wav_path、translating 但无 segments_json 的任务，
|
||||
是上次进程异常退出留下的孤儿。标记为 failed 避免调度线程反复尝试。
|
||||
"""
|
||||
db = get_session_local()()
|
||||
try:
|
||||
stuck = (
|
||||
db.query(Task)
|
||||
.filter(Task.status.in_([
|
||||
STATUS_EXTRACTING, STATUS_TRANSCRIBING, STATUS_SEGMENTING, STATUS_TRANSLATING,
|
||||
]))
|
||||
.all()
|
||||
)
|
||||
n = 0
|
||||
for task in stuck:
|
||||
reason = ""
|
||||
if task.status in (STATUS_TRANSCRIBING, STATUS_SEGMENTING) and not task.wav_path:
|
||||
reason = f"重启时发现 {task.status} 状态但无 wav_path"
|
||||
elif task.status == STATUS_TRANSLATING and not task.segments_json:
|
||||
reason = f"重启时发现 translating 状态但无 segments_json"
|
||||
elif task.status == STATUS_EXTRACTING:
|
||||
reason = "重启时发现 extracting 状态（ffmpeg 未完成）"
|
||||
if reason:
|
||||
task.status = STATUS_FAILED
|
||||
task.error = reason[:2000]
|
||||
task.updated_at = datetime.now(timezone.utc)
|
||||
n += 1
|
||||
logger.warning("清理卡住的任务 %d：%s", task.id, reason)
|
||||
if n:
|
||||
db.commit()
|
||||
logger.info("共清理 %d 个卡住的任务。", n)
|
||||
except Exception as exc: # pragma: no cover
|
||||
logger.error("清理卡住任务时出错：%s", exc)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _ensure_scheduler_running() -> None:
|
||||
"""确保 GPU 调度线程在跑（enqueue 时调，防止 lifespan 未启动的边界情况）。"""
|
||||
if not _scheduler_started:
|
||||
start_scheduler()
|
||||
|
||||
|
||||
# ---------------- ffmpeg 异步提取 ----------------
|
||||
|
||||
def _extract_audio_async(task_id: int) -> None:
|
||||
"""独立线程跑 ffmpeg 提取（CPU），完成后唤醒 GPU 调度线程。"""
|
||||
db = get_session_local()()
|
||||
try:
|
||||
task = db.get(Task, task_id)
|
||||
if task is None:
|
||||
logger.error("任务 %d 不存在，ffmpeg 线程退出。", task_id)
|
||||
return
|
||||
pipeline.extract_phase(db, task)
|
||||
_wake_event.set() # 通知 GPU 线程有新任务
|
||||
except Exception as exc:
|
||||
logger.exception("任务 %d ffmpeg 提取失败：%s", task_id, exc)
|
||||
pipeline.mark_failed(db, task_id, str(exc))
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
# ---------------- GPU 调度线程 ----------------
|
||||
|
||||
def _gpu_scheduler() -> None:
|
||||
"""常驻 GPU 调度线程：串行处理 ASR + 翻译，切换模型前查队列复用。
|
||||
|
||||
循环逻辑：
|
||||
1. 处理所有待 ASR 任务（Whisper 只加载一次）
|
||||
2. 处理所有待翻译任务（NLLB 只加载一次）
|
||||
3. 都空了 → 等待唤醒
|
||||
每个阶段失败的任务标记 failed，不影响其他任务。
|
||||
"""
|
||||
logger.info("GPU 调度线程开始运行。")
|
||||
while True:
|
||||
try:
|
||||
# 优先处理 ASR 队列：把所有待 ASR 的任务一次性做完（模型复用）
|
||||
asr_count = _drain_asr_queue()
|
||||
# 再处理翻译队列：把所有待翻译的任务一次性做完（模型复用）
|
||||
trans_count = _drain_translate_queue()
|
||||
|
||||
if asr_count == 0 and trans_count == 0:
|
||||
# 两队列都空，等待新任务唤醒
|
||||
_wake_event.wait(timeout=60)
|
||||
_wake_event.clear()
|
||||
except Exception as exc: # pragma: no cover
|
||||
# 调度线程不能死，任何异常都捕获后继续
|
||||
logger.exception("GPU 调度线程异常（已恢复）：%s", exc)
|
||||
|
||||
|
||||
def _drain_asr_queue() -> int:
|
||||
"""连续处理所有 status=transcribing 的任务，Whisper 只加载一次。
|
||||
|
||||
Returns: 本轮处理的任务数
|
||||
"""
|
||||
n = 0
|
||||
mm = get_model_manager()
|
||||
while True:
|
||||
db = get_session_local()()
|
||||
task = None # 预定义，防 query 抛异常时 except 引用未定义变量
|
||||
try:
|
||||
# 取最早一个待 ASR 的任务（按 id 升序，FIFO）
|
||||
task = (
|
||||
db.query(Task)
|
||||
.filter(Task.status == STATUS_TRANSCRIBING)
|
||||
.order_by(Task.id.asc())
|
||||
.first()
|
||||
)
|
||||
if task is None:
|
||||
break # ASR 队列空
|
||||
# 加载 ASR 模型（若翻译器在内存，model_manager 自动卸载它，并记 INFO）
|
||||
mm.get_asr()
|
||||
# 执行 ASR + 断句
|
||||
pipeline.asr_phase(db, task)
|
||||
n += 1
|
||||
except Exception as exc:
|
||||
tid = task.id if task is not None else -1
|
||||
logger.exception("任务 %d ASR 阶段失败：%s", tid, exc)
|
||||
if task is not None:
|
||||
pipeline.mark_failed(db, task.id, str(exc))
|
||||
finally:
|
||||
db.close()
|
||||
if n > 0:
|
||||
logger.info("ASR 批次完成：处理 %d 个任务。", n)
|
||||
return n
|
||||
|
||||
|
||||
def _drain_translate_queue() -> int:
|
||||
"""连续处理所有 status=translating 的任务，NLLB 只加载一次。
|
||||
|
||||
Returns: 本轮处理的任务数
|
||||
"""
|
||||
n = 0
|
||||
mm = get_model_manager()
|
||||
while True:
|
||||
db = get_session_local()()
|
||||
task = None # 预定义，防 query 抛异常时 except 引用未定义变量
|
||||
try:
|
||||
task = (
|
||||
db.query(Task)
|
||||
.filter(Task.status == STATUS_TRANSLATING)
|
||||
.order_by(Task.id.asc())
|
||||
.first()
|
||||
)
|
||||
if task is None:
|
||||
break # 翻译队列空
|
||||
# 加载翻译模型（若 ASR 在内存，model_manager 自动卸载它，并记 INFO）
|
||||
mm.get_translator()
|
||||
# 执行翻译 + 写 SRT
|
||||
pipeline.translate_phase(db, task)
|
||||
n += 1
|
||||
except Exception as exc:
|
||||
tid = task.id if task is not None else -1
|
||||
logger.exception("任务 %d 翻译阶段失败：%s", tid, exc)
|
||||
if task is not None:
|
||||
pipeline.mark_failed(db, task.id, str(exc))
|
||||
finally:
|
||||
db.close()
|
||||
if n > 0:
|
||||
logger.info("翻译批次完成：处理 %d 个任务。", n)
|
||||
return n
|
||||
Reference in New Issue
Block a user