73110848f4
- 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: 安装时预下载模型权重
263 lines
9.6 KiB
Python
263 lines
9.6 KiB
Python
"""分片上传服务：会话管理 + 分片落盘 + 拼接 + 创建转写任务。
|
||
|
||
存储布局::
|
||
|
||
<work_dir>/<upload_id>/0.part 分片暂存
|
||
<work_dir>/<upload_id>/1.part
|
||
...
|
||
<upload_dir>/<yyyy>/<mm>/<uuid>.<ext> complete 后的正式视频
|
||
|
||
complete 创建转写 Task，管线触发由 controller 调用 scheduler.enqueue_task，
|
||
本服务不依赖 scheduler（避免循环依赖）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
import shutil
|
||
import uuid
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
|
||
from fastapi import HTTPException
|
||
from sqlalchemy.orm import Session
|
||
|
||
from ..config import get_settings
|
||
from ..models.task import Task
|
||
from ..models.upload_session import UploadSession
|
||
from ..schemas.task import (
|
||
ChunkUploadResponse,
|
||
CompleteResponse,
|
||
CreateSessionRequest,
|
||
CreateSessionResponse,
|
||
SessionStatusResponse,
|
||
)
|
||
|
||
logger = logging.getLogger("audio2text.upload")
|
||
|
||
|
||
class UploadService:
|
||
def __init__(self, db: Session) -> None:
|
||
self.db = db
|
||
s = get_settings()
|
||
self.upload_root = s.upload_dir()
|
||
self.work_root = s.work_dir()
|
||
self.chunk_bytes = s.storage.chunk_bytes
|
||
self.session_ttl = s.storage.chunk_session_ttl_seconds
|
||
|
||
# ---------------- 会话生命周期 ----------------
|
||
|
||
def create_session(self, body: CreateSessionRequest) -> CreateSessionResponse:
|
||
upload_id = uuid.uuid4().hex
|
||
session = UploadSession(
|
||
upload_id=upload_id,
|
||
filename=body.filename,
|
||
size_bytes=body.size_bytes,
|
||
chunk_size=body.chunk_size,
|
||
total_chunks=body.total_chunks,
|
||
uploaded_chunks=[],
|
||
status="pending",
|
||
)
|
||
self.db.add(session)
|
||
self.db.commit()
|
||
self._session_dir(upload_id).mkdir(parents=True, exist_ok=True)
|
||
return CreateSessionResponse(
|
||
upload_id=upload_id,
|
||
filename=body.filename,
|
||
size_bytes=body.size_bytes,
|
||
chunk_size=body.chunk_size,
|
||
total_chunks=body.total_chunks,
|
||
)
|
||
|
||
def get_status(self, upload_id: str) -> SessionStatusResponse:
|
||
session = self._require_session(upload_id)
|
||
return SessionStatusResponse(
|
||
upload_id=session.upload_id,
|
||
filename=session.filename,
|
||
size_bytes=session.size_bytes,
|
||
chunk_size=session.chunk_size,
|
||
total_chunks=session.total_chunks,
|
||
uploaded_chunks=list(session.uploaded_chunks or []),
|
||
completed=(session.status == "completed"),
|
||
task_id=session.task_id,
|
||
)
|
||
|
||
# ---------------- 分片写入 ----------------
|
||
|
||
def write_chunk(self, upload_id: str, index: int, data: bytes) -> list[int]:
|
||
session = self._require_session(upload_id)
|
||
self._validate_index(session, index)
|
||
|
||
session_dir = self._session_dir(upload_id)
|
||
session_dir.mkdir(parents=True, exist_ok=True)
|
||
chunk_path = session_dir / f"{index}.part"
|
||
|
||
try:
|
||
with chunk_path.open("wb") as out:
|
||
out.write(data)
|
||
out.flush()
|
||
os.fsync(out.fileno())
|
||
except Exception:
|
||
chunk_path.unlink(missing_ok=True)
|
||
raise
|
||
|
||
uploaded = list(session.uploaded_chunks or [])
|
||
if index not in uploaded:
|
||
uploaded.append(index)
|
||
session.uploaded_chunks = uploaded
|
||
session.updated_at = datetime.now(timezone.utc)
|
||
self.db.commit()
|
||
return sorted(uploaded)
|
||
|
||
# ---------------- 拼接 + 创建任务 ----------------
|
||
|
||
# 允许的音视频扩展名白名单（防可执行文件落盘到上传目录）
|
||
_ALLOWED_EXTS = frozenset({
|
||
".mp4", ".mkv", ".avi", ".mov", ".webm", ".flv",
|
||
".mp3", ".wav", ".flac", ".aac", ".m4a", ".ogg", ".wma",
|
||
})
|
||
|
||
# ---------------- 拼接 + 创建任务 ----------------
|
||
|
||
def complete(self, upload_id: str) -> CompleteResponse:
|
||
session = self._require_session(upload_id)
|
||
|
||
# 幂等：已 complete 则返回已建任务
|
||
if session.status == "completed" and session.task_id is not None:
|
||
task = self.db.get(Task, session.task_id)
|
||
if task is not None:
|
||
return CompleteResponse(
|
||
task_id=task.id, filename=task.filename,
|
||
size_bytes=session.size_bytes, status=task.status,
|
||
)
|
||
|
||
uploaded = set(session.uploaded_chunks or [])
|
||
missing = [i for i in range(session.total_chunks) if i not in uploaded]
|
||
if missing:
|
||
raise HTTPException(
|
||
409,
|
||
f"分片不齐全：缺失 {len(missing)} 个，例如 {sorted(missing)[:10]}",
|
||
)
|
||
|
||
final_path = self._assemble(session)
|
||
rel = str(final_path.relative_to(self.upload_root))
|
||
|
||
# 单事务：建 Task + 更新 session 状态 + 关联 task_id 一次 commit
|
||
# 避免双 commit 之间崩溃产生孤儿 Task（Task 已建但 session.task_id 为空）
|
||
task = Task(
|
||
filename=session.filename,
|
||
source_path=rel,
|
||
status="queued",
|
||
progress=0.0,
|
||
)
|
||
self.db.add(task)
|
||
session.status = "completed"
|
||
session.final_path = rel
|
||
session.task_id = None # 占位，flush 后用 task.id 赋值
|
||
session.updated_at = datetime.now(timezone.utc)
|
||
self.db.flush() # 拿到 task.id（不 commit，仍在事务内）
|
||
session.task_id = task.id
|
||
self.db.commit()
|
||
self.db.refresh(task)
|
||
|
||
# 清理分片暂存（commit 后，即使清理失败也不影响已建任务）
|
||
self._cleanup_session_dir(upload_id)
|
||
|
||
logger.info("上传完成 task_id=%s file=%s size=%d", task.id, session.filename, session.size_bytes)
|
||
return CompleteResponse(
|
||
task_id=task.id, filename=task.filename,
|
||
size_bytes=session.size_bytes, status=task.status,
|
||
)
|
||
|
||
# ---------------- 过期会话清理 ----------------
|
||
|
||
def reap_stale_sessions(self) -> int:
|
||
"""清理被放弃的会话：pending 且 updated_at 超过 ttl。
|
||
|
||
时间比较统一用 aware UTC datetime，避免 naive datetime 的 .timestamp()
|
||
按本地时区算导致的偏移。
|
||
"""
|
||
cutoff = datetime.now(timezone.utc) - timedelta(seconds=self.session_ttl)
|
||
sessions = (
|
||
self.db.query(UploadSession)
|
||
.filter(UploadSession.status == "pending")
|
||
.all()
|
||
)
|
||
n = 0
|
||
for session in sessions:
|
||
updated = session.updated_at
|
||
if updated is None:
|
||
continue
|
||
# SQLite 存 naive datetime，统一补 UTC 后比较
|
||
if updated.tzinfo is None:
|
||
updated = updated.replace(tzinfo=timezone.utc)
|
||
if updated < cutoff:
|
||
self._cleanup_session_dir(session.upload_id)
|
||
self.db.delete(session)
|
||
n += 1
|
||
logger.info("清理过期分片会话 upload_id=%s file=%s", session.upload_id, session.filename)
|
||
if n:
|
||
self.db.commit()
|
||
return n
|
||
|
||
# ---------------- 内部 ----------------
|
||
|
||
def _require_session(self, upload_id: str) -> UploadSession:
|
||
if not upload_id:
|
||
raise HTTPException(400, "upload_id 不能为空")
|
||
session = (
|
||
self.db.query(UploadSession)
|
||
.filter(UploadSession.upload_id == upload_id)
|
||
.first()
|
||
)
|
||
if session is None:
|
||
raise HTTPException(404, f"会话不存在或已过期：{upload_id}")
|
||
return session
|
||
|
||
@staticmethod
|
||
def _validate_index(session: UploadSession, index: int) -> None:
|
||
if index < 0 or index >= session.total_chunks:
|
||
raise HTTPException(400, f"分片下标越界：{index} 不在 [0, {session.total_chunks})")
|
||
|
||
def _session_dir(self, upload_id: str) -> Path:
|
||
return self.work_root / upload_id
|
||
|
||
def _assemble(self, session: UploadSession) -> Path:
|
||
"""按 index 顺序拼接全部分片为正式视频文件。"""
|
||
# 扩展名取自客户端 filename，但做白名单净化：不在允许列表内则回退 .bin
|
||
ext = Path(session.filename).suffix.lower()
|
||
if ext not in self._ALLOWED_EXTS:
|
||
ext = ".bin"
|
||
now = datetime.now(timezone.utc)
|
||
sub = self.upload_root / f"{now:%Y}" / f"{now:%m}"
|
||
sub.mkdir(parents=True, exist_ok=True)
|
||
final = sub / f"{uuid.uuid4().hex}{ext}"
|
||
part_path = final.with_suffix(final.suffix + ".part")
|
||
|
||
session_dir = self._session_dir(session.upload_id)
|
||
try:
|
||
with part_path.open("wb") as out:
|
||
for index in range(session.total_chunks):
|
||
chunk_path = session_dir / f"{index}.part"
|
||
if not chunk_path.is_file():
|
||
raise HTTPException(409, f"拼接时发现分片缺失：{index}.part")
|
||
with chunk_path.open("rb") as src:
|
||
while buf := src.read(self.chunk_bytes):
|
||
out.write(buf)
|
||
out.flush()
|
||
os.fsync(out.fileno())
|
||
os.replace(part_path, final)
|
||
except Exception:
|
||
part_path.unlink(missing_ok=True)
|
||
raise
|
||
return final
|
||
|
||
def _cleanup_session_dir(self, upload_id: str) -> None:
|
||
session_dir = self._session_dir(upload_id)
|
||
try:
|
||
if session_dir.exists():
|
||
shutil.rmtree(session_dir, ignore_errors=True)
|
||
except Exception as exc: # pragma: no cover
|
||
logger.warning("清理会话目录失败 upload_id=%s: %s", upload_id, exc)
|