zikai
00e2a95fb7
- 音频/视频转双语(英/中)SRT 字幕，Docker 容器化，CPU 开发/GPU 生产同一份代码 - faster-whisper ASR(词级时间戳) + 断句时间戳重算 + NLLB 翻译(模型不共驻) - 分片上传(断点续传) + SQLite 持久化 + 主页/历史/日志页面 - 历史页文件名搜索；缓存定时清理(默认保留7天，可配置) - 双 Dockerfile(cpu/gpu) + setup/start/stop 脚本
250 lines
8.9 KiB
Python
250 lines
8.9 KiB
Python
"""分片上传服务：会话管理 + 分片落盘 + 拼接 + 创建转写任务。
|
||
|
||
存储布局::
|
||
|
||
<work_dir>/<upload_id>/0.part 分片暂存
|
||
<work_dir>/<upload_id>/1.part
|
||
...
|
||
<upload_dir>/<yyyy>/<mm>/<uuid>.<ext> complete 后的正式视频
|
||
|
||
与 server 的区别：视频无需 sha256 去重（每个视频都转写），complete 直接创建 Task。
|
||
管线触发由 controller 调用 pipeline.enqueue_task，本服务不依赖 pipeline。
|
||
"""
|
||
|
||
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)
|
||
|
||
# ---------------- 拼接 + 创建任务 ----------------
|
||
|
||
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 = Task(
|
||
filename=session.filename,
|
||
source_path=rel,
|
||
status="queued",
|
||
progress=0.0,
|
||
)
|
||
self.db.add(task)
|
||
session.status = "completed"
|
||
session.final_path = rel
|
||
session.updated_at = datetime.now(timezone.utc)
|
||
self.db.commit()
|
||
self.db.refresh(task)
|
||
# 正向关联：session → task（替代旧的 source_path 反向查找）
|
||
session.task_id = task.id
|
||
self.db.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 顺序拼接全部分片为正式视频文件。"""
|
||
ext = Path(session.filename).suffix or ".mp4"
|
||
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)
|