425 lines
17 KiB
Python
425 lines
17 KiB
Python
import asyncio
|
||
import os
|
||
import shutil
|
||
from datetime import datetime
|
||
from typing import Any, Callable, Dict, List, Optional
|
||
|
||
import aiofiles
|
||
import aiohttp
|
||
import uvicorn
|
||
from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, UploadFile
|
||
from fastapi.responses import StreamingResponse
|
||
from loguru import logger
|
||
from models import Task, TaskCreateRequest, TaskStatus
|
||
from processor import ProcessManager
|
||
|
||
|
||
# 定义任务管理器
|
||
class TaskManager:
|
||
def __init__(self):
|
||
self.tasks: Dict[str, Task] = {}
|
||
self.process_manager = ProcessManager(config_path="config/machines.json", task_manager=self)
|
||
self.callbacks: Dict[str, List[Callable[[Task], Any]]] = {}
|
||
self._task_lock = asyncio.Lock()
|
||
self._callback_lock = asyncio.Lock()
|
||
self._started = False
|
||
|
||
async def start(self):
|
||
"""启动任务管理器"""
|
||
if not self._started:
|
||
await self.process_manager.start()
|
||
self._started = True
|
||
logger.info("任务管理器已启动")
|
||
|
||
async def stop(self):
|
||
"""停止任务管理器"""
|
||
if self._started:
|
||
await self.process_manager.stop()
|
||
self._started = False
|
||
logger.info("任务管理器已停止")
|
||
|
||
async def add_task(
|
||
self, audio_file_path: str, video_file_path: str, output_dir: str, description: str, callback_url: Optional[str] = None
|
||
) -> Task:
|
||
"""添加新任务"""
|
||
async with self._task_lock:
|
||
task_id = f"task_{len(self.tasks) + 1}"
|
||
output_file_path = os.path.join(output_dir, f"{task_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.mp4")
|
||
task = Task(
|
||
task_id=task_id,
|
||
audio_file_path=audio_file_path,
|
||
video_file_path=video_file_path,
|
||
output_dir=output_dir,
|
||
description=description,
|
||
status=TaskStatus.PENDING,
|
||
created_at=datetime.now(),
|
||
callback_url=callback_url,
|
||
output_file_path=output_file_path,
|
||
)
|
||
self.tasks[task_id] = task
|
||
|
||
# 如果有回调URL，添加回调
|
||
if callback_url:
|
||
await self.add_callback(task_id, self._create_callback_handler(callback_url))
|
||
|
||
return task
|
||
|
||
def _create_callback_handler(self, callback_url: str) -> Callable[[Task], Any]:
|
||
"""创建回调处理器"""
|
||
|
||
async def callback_handler(task: Task):
|
||
try:
|
||
if callback_url.startswith("http"):
|
||
async with aiohttp.ClientSession() as session:
|
||
await session.post(callback_url, json=task.dict())
|
||
else:
|
||
logger.info(f"回调URL: {callback_url} 不是HTTP URL")
|
||
except Exception as e:
|
||
logger.error(f"Error sending callback to {callback_url}: {str(e)}")
|
||
|
||
return callback_handler
|
||
|
||
async def process_task_sync(self, task_id: str) -> Task:
|
||
"""同步处理任务"""
|
||
task = await self.get_task(task_id)
|
||
if not task:
|
||
raise ValueError(f"Task {task_id} not found")
|
||
|
||
try:
|
||
# 更新任务状态为处理中
|
||
await self.update_task_status(task_id, TaskStatus.PROCESSING)
|
||
|
||
# 生成唯一的文件ID
|
||
file_id = f"{task_id}_{int(datetime.now().timestamp())}"
|
||
|
||
# 处理视频
|
||
success = await self.process_manager.process_video(
|
||
task_id=task_id,
|
||
audio_file_path=task.audio_file_path,
|
||
video_file_path=task.video_file_path,
|
||
out_file_path=task.output_file_path,
|
||
)
|
||
|
||
if success:
|
||
# 更新任务状态为已入队
|
||
return await self.update_task_status(task_id, TaskStatus.QUEUED)
|
||
else:
|
||
return await self.update_task_status(task_id, TaskStatus.FAILED, error_message="Failed to queue video processing task")
|
||
except Exception as e:
|
||
return await self.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
|
||
|
||
async def update_task_status(
|
||
self, task_id: str, status: TaskStatus, machine_name: Optional[str] = None, error_message: Optional[str] = None
|
||
) -> Task:
|
||
"""更新任务状态"""
|
||
async with self._task_lock:
|
||
if task_id not in self.tasks:
|
||
raise ValueError(f"Task {task_id} not found")
|
||
|
||
task = self.tasks[task_id]
|
||
logger.info(f"更新任务状态: {task_id}, 从 {task.status} 到 {status}")
|
||
task.status = status
|
||
if machine_name:
|
||
task.machine_name = machine_name
|
||
if error_message:
|
||
task.error_message = error_message
|
||
if status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
|
||
task.completed_at = datetime.now()
|
||
logger.info(f"任务完成时间已设置: {task_id}, 状态: {status}")
|
||
|
||
# 触发回调
|
||
if task_id in self.callbacks:
|
||
logger.info(f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}")
|
||
for callback in self.callbacks[task_id]:
|
||
try:
|
||
await callback(task)
|
||
logger.info(f"任务回调执行成功: {task_id}")
|
||
except Exception as e:
|
||
logger.error(f"任务回调执行失败: {task_id}, 错误: {str(e)}")
|
||
|
||
return task
|
||
|
||
async def get_task(self, task_id: str) -> Optional[Task]:
|
||
"""获取任务信息"""
|
||
return self.tasks.get(task_id)
|
||
|
||
async def get_all_tasks(self) -> List[Task]:
|
||
"""获取所有任务"""
|
||
return list(self.tasks.values())
|
||
|
||
async def add_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None:
|
||
"""添加任务回调"""
|
||
async with self._callback_lock:
|
||
if task_id not in self.callbacks:
|
||
self.callbacks[task_id] = []
|
||
self.callbacks[task_id].append(callback)
|
||
|
||
async def remove_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None:
|
||
"""移除任务回调"""
|
||
async with self._callback_lock:
|
||
if task_id in self.callbacks:
|
||
self.callbacks[task_id].remove(callback)
|
||
|
||
|
||
# 创建FastAPI应用
|
||
app = FastAPI(title="视频处理API")
|
||
task_manager = TaskManager()
|
||
|
||
|
||
@app.on_event("startup")
|
||
async def startup_event():
|
||
"""应用启动时初始化"""
|
||
await task_manager.start()
|
||
|
||
|
||
@app.on_event("shutdown")
|
||
async def shutdown_event():
|
||
"""应用关闭时清理"""
|
||
await task_manager.stop()
|
||
|
||
|
||
# API路由
|
||
@app.post("/tasks/", response_model=Task)
|
||
async def create_task(
|
||
audio: UploadFile = File(...),
|
||
video: UploadFile = File(...),
|
||
description: str = Form(...),
|
||
callback_url: Optional[str] = Form(None),
|
||
background_tasks: BackgroundTasks = BackgroundTasks(),
|
||
):
|
||
"""创建新的视频处理任务"""
|
||
try:
|
||
# 创建临时目录
|
||
temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
try:
|
||
# 确保文件名不为None
|
||
if not audio.filename or not video.filename:
|
||
raise ValueError("文件名不能为空")
|
||
|
||
# 保存上传的文件
|
||
audio_path = os.path.join(temp_dir, audio.filename)
|
||
video_path = os.path.join(temp_dir, video.filename)
|
||
|
||
# 保存音频文件
|
||
logger.info(f"开始保存音频文件: {audio_path}")
|
||
async with aiofiles.open(audio_path, "wb") as f:
|
||
content = await audio.read()
|
||
await f.write(content)
|
||
logger.info(f"音频文件保存完成: {audio_path}")
|
||
|
||
# 保存视频文件
|
||
logger.info(f"开始保存视频文件: {video_path}")
|
||
async with aiofiles.open(video_path, "wb") as f:
|
||
content = await video.read()
|
||
await f.write(content)
|
||
logger.info(f"视频文件保存完成: {video_path}")
|
||
|
||
# 验证文件是否保存成功
|
||
if not os.path.exists(audio_path) or not os.path.exists(video_path):
|
||
raise ValueError("文件保存失败")
|
||
|
||
# 创建输出目录
|
||
output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 创建任务
|
||
task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
|
||
|
||
# 异步处理任务
|
||
background_tasks.add_task(process_video_task, task.task_id, audio_path, video_path, task.output_file_path)
|
||
|
||
return task
|
||
|
||
except Exception as e:
|
||
# 如果发生错误，立即清理临时文件
|
||
# if os.path.exists(temp_dir):
|
||
# shutil.rmtree(temp_dir)
|
||
logger.error(f"创建任务失败: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建任务失败: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@app.post("/tasks/sync")
|
||
async def create_task_sync(
|
||
audio: UploadFile = File(...),
|
||
video: UploadFile = File(...),
|
||
description: str = Form(...),
|
||
callback_url: Optional[str] = Form(None),
|
||
background_tasks: BackgroundTasks = BackgroundTasks(),
|
||
):
|
||
"""创建新的视频处理任务并返回处理结果"""
|
||
try:
|
||
logger.info("创建新的视频处理任务并返回处理结果")
|
||
# 创建临时目录
|
||
temp_dir = os.path.join(os.getcwd(), "temp", datetime.now().strftime("%Y%m%d_%H%M%S"))
|
||
logger.info(f"临时目录: {temp_dir}")
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
try:
|
||
# 确保文件名不为None
|
||
if not audio.filename or not video.filename:
|
||
raise ValueError("文件名不能为空")
|
||
logger.info(f"文件名: {audio.filename}, {video.filename}")
|
||
# 保存上传的文件
|
||
audio_path = os.path.join(temp_dir, audio.filename)
|
||
video_path = os.path.join(temp_dir, video.filename)
|
||
|
||
# 保存音频文件
|
||
logger.info(f"开始保存音频文件: {audio_path}")
|
||
async with aiofiles.open(audio_path, "wb") as f:
|
||
content = await audio.read()
|
||
await f.write(content)
|
||
logger.info(f"音频文件保存完成: {audio_path}")
|
||
|
||
# 保存视频文件
|
||
logger.info(f"开始保存视频文件: {video_path}")
|
||
async with aiofiles.open(video_path, "wb") as f:
|
||
content = await video.read()
|
||
await f.write(content)
|
||
logger.info(f"视频文件保存完成: {video_path}")
|
||
|
||
# 验证文件是否保存成功
|
||
if not os.path.exists(audio_path) or not os.path.exists(video_path):
|
||
raise ValueError("文件保存失败")
|
||
|
||
# 创建输出目录
|
||
output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d"))
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 创建任务
|
||
task = await task_manager.add_task(audio_path, video_path, output_dir, description, callback_url)
|
||
|
||
# 创建一个事件来等待任务完成
|
||
task_completed = asyncio.Event()
|
||
logger.info(f"创建任务完成事件: {task.task_id}")
|
||
|
||
async def task_completion_callback(completed_task: Task):
|
||
logger.info(f"任务完成回调被触发: {completed_task.task_id}, 状态: {completed_task.status}")
|
||
if completed_task.task_id == task.task_id and (
|
||
completed_task.status == TaskStatus.COMPLETED or completed_task.status == TaskStatus.FAILED
|
||
):
|
||
logger.info(f"设置任务完成事件: {task.task_id}")
|
||
task_completed.set()
|
||
|
||
# 添加回调
|
||
await task_manager.add_callback(task.task_id, task_completion_callback)
|
||
logger.info(f"已添加任务完成回调: {task.task_id}")
|
||
|
||
# 开始处理任务
|
||
await task_manager.process_task_sync(task.task_id)
|
||
logger.info(f"开始处理任务: {task.task_id}")
|
||
|
||
# 等待任务完成
|
||
try:
|
||
logger.info(f"等待任务完成: {task.task_id}")
|
||
await asyncio.wait_for(task_completed.wait(), timeout=1800) # 30分钟超时
|
||
##################sdsadsadsads
|
||
logger.info(f"任务完成事件已触发: {task.task_id}")
|
||
except asyncio.TimeoutError:
|
||
logger.error(f"任务处理超时: {task.task_id}")
|
||
raise HTTPException(status_code=500, detail="任务处理超时")
|
||
|
||
# 获取最终任务状态
|
||
result = await task_manager.get_task(task.task_id)
|
||
if not result:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
|
||
if result.status == TaskStatus.COMPLETED:
|
||
# 获取输出文件路径
|
||
output_file = result.output_file_path
|
||
if not os.path.exists(output_file):
|
||
logger.error(f"输出文件不存在: {output_file}")
|
||
raise HTTPException(status_code=404, detail="输出文件不存在")
|
||
|
||
# 获取文件名
|
||
filename = os.path.basename(output_file)
|
||
|
||
# 创建文件流
|
||
async def file_iterator():
|
||
async with aiofiles.open(output_file, "rb") as f:
|
||
while chunk := await f.read(8192): # 8KB chunks
|
||
yield chunk
|
||
|
||
return StreamingResponse(
|
||
file_iterator(), media_type="video/mp4", headers={"Content-Disposition": f'attachment; filename="{filename}"'}
|
||
)
|
||
else:
|
||
raise HTTPException(status_code=500, detail=result.error_message or "处理失败")
|
||
|
||
except Exception as e:
|
||
# 如果发生错误，立即清理临时文件
|
||
# if os.path.exists(temp_dir):
|
||
# shutil.rmtree(temp_dir)
|
||
logger.error(f"创建任务失败: {str(e)}", exc_info=True)
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建任务失败: {str(e)}", exc_info=True)
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
@app.get("/tasks/{task_id}", response_model=Task)
|
||
async def get_task(task_id: str):
|
||
"""获取任务信息"""
|
||
task = await task_manager.get_task(task_id)
|
||
if not task:
|
||
raise HTTPException(status_code=404, detail="Task not found")
|
||
return task
|
||
|
||
|
||
@app.get("/tasks/", response_model=List[Task])
|
||
async def get_all_tasks():
|
||
"""获取所有任务"""
|
||
return await task_manager.get_all_tasks()
|
||
|
||
|
||
@app.post("/tasks/{task_id}/callback")
|
||
async def add_task_callback(task_id: str, callback_url: str, background_tasks: BackgroundTasks):
|
||
"""添加任务完成回调"""
|
||
if not await task_manager.get_task(task_id):
|
||
raise HTTPException(status_code=404, detail="Task not found")
|
||
|
||
async def callback_handler(task: Task):
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
await session.post(callback_url, json=task.dict())
|
||
except Exception as e:
|
||
logger.error(f"Error sending callback to {callback_url}: {str(e)}")
|
||
|
||
await task_manager.add_callback(task_id, callback_handler)
|
||
return {"message": "Callback added successfully"}
|
||
|
||
|
||
# 后台任务处理函数
|
||
async def process_video_task(task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str):
|
||
"""处理视频任务的后台函数"""
|
||
try:
|
||
# 更新任务状态为处理中
|
||
await task_manager.update_task_status(task_id, TaskStatus.PROCESSING)
|
||
|
||
# 生成唯一的文件ID
|
||
file_id = f"{task_id}_{int(datetime.now().timestamp())}"
|
||
|
||
# 处理视频
|
||
success = await task_manager.process_manager.process_video(task_id, audio_file_path, video_file_path, out_file_path)
|
||
|
||
if success:
|
||
await task_manager.update_task_status(task_id, TaskStatus.COMPLETED)
|
||
else:
|
||
await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message="Video processing failed")
|
||
except Exception as e:
|
||
await task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
|
||
# 清理临时文件
|
||
# if os.path.exists(temp_dir):
|
||
# shutil.rmtree(temp_dir)
|
||
# logger.info(f"已清理临时目录: {temp_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run("api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info") # 开发模式下启用热重载 # 单进程模式
|