From 6117538c75b37d0be9482ca4bbd16c07c98c3820 Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Mon, 28 Jul 2025 11:57:08 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E4=BA=86=E6=96=87=E4=BB=B6?= =?UTF-8?q?=E4=B8=8A=E4=BC=A0=E5=92=8Curl=E7=9A=84=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api.py | 440 +++++++++++++++++++++++++++++++++++---------------------- 1 file changed, 275 insertions(+), 165 deletions(-) diff --git a/api.py b/api.py index 7765373..16568e4 100644 --- a/api.py +++ b/api.py @@ -1,8 +1,8 @@ import asyncio import os -import shutil from datetime import datetime from typing import Any, Callable, Dict, List, Optional +import urllib.parse import aiofiles import aiohttp @@ -10,7 +10,7 @@ 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 models import Task, TaskStatus from processor import ProcessManager @@ -18,7 +18,9 @@ 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.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() @@ -39,12 +41,19 @@ class TaskManager: 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 + 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") + 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, @@ -60,7 +69,9 @@ class TaskManager: # 如果有回调URL，添加回调 if callback_url: - await self.add_callback(task_id, self._create_callback_handler(callback_url)) + await self.add_callback( + task_id, self._create_callback_handler(callback_url) + ) return task @@ -104,12 +115,22 @@ class TaskManager: # 更新任务状态为已入队 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") + 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)) + 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 + self, + task_id: str, + status: TaskStatus, + machine_name: Optional[str] = None, + error_message: Optional[str] = None, ) -> Task: """更新任务状态""" async with self._task_lock: @@ -129,7 +150,9 @@ class TaskManager: # 触发回调 if task_id in self.callbacks: - logger.info(f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}") + logger.info( + f"触发任务回调: {task_id}, 回调数量: {len(self.callbacks[task_id])}" + ) for callback in self.callbacks[task_id]: try: await callback(task) @@ -154,7 +177,9 @@ class TaskManager: self.callbacks[task_id] = [] self.callbacks[task_id].append(callback) - async def remove_callback(self, task_id: str, callback: Callable[[Task], Any]) -> None: + async def remove_callback( + self, task_id: str, callback: Callable[[Task], Any] + ) -> None: """移除任务回调""" async with self._callback_lock: if task_id in self.callbacks: @@ -181,62 +206,31 @@ async def shutdown_event(): # API路由 @app.post("/tasks/", response_model=Task) async def create_task( - audio: UploadFile = File(...), - video: UploadFile = File(...), + audio: Optional[UploadFile] = File(None), + video: Optional[UploadFile] = File(None), + audio_url: Optional[str] = Form(None), + video_url: Optional[str] = Form(None), description: str = Form(...), callback_url: Optional[str] = Form(None), background_tasks: BackgroundTasks = BackgroundTasks(), ): - """创建新的视频处理任务""" + """创建新的视频处理任务，支持文件上传或URL下载""" 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("文件名不能为空") + # 使用公共函数准备任务文件 + task, temp_dir = await prepare_task_files( + audio, video, audio_url, video_url, description, callback_url + ) - # 保存上传的文件 - audio_path = os.path.join(temp_dir, audio.filename) - video_path = os.path.join(temp_dir, video.filename) + # 异步处理任务 + background_tasks.add_task( + process_video_task, + task.task_id, + task.audio_file_path, + task.video_file_path, + task.output_file_path, + ) - # 保存音频文件 - 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)) + return task except Exception as e: logger.error(f"创建任务失败: {str(e)}") @@ -245,116 +239,88 @@ async def create_task( @app.post("/tasks/sync") async def create_task_sync( - audio: UploadFile = File(...), - video: UploadFile = File(...), + audio: Optional[UploadFile] = File(None), + video: Optional[UploadFile] = File(None), + audio_url: Optional[str] = Form(None), + video_url: Optional[str] = Form(None), description: str = Form(...), callback_url: Optional[str] = Form(None), - background_tasks: BackgroundTasks = BackgroundTasks(), ): - """创建新的视频处理任务并返回处理结果""" + """创建新的视频处理任务并返回处理结果，支持文件上传或URL下载""" 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) + + # 使用公共函数准备任务文件 + task, temp_dir = await prepare_task_files( + audio, video, audio_url, video_url, description, callback_url + ) + logger.info(f"任务创建成功: {task.task_id}") + logger.info(f"音频文件路径: {task.audio_file_path}") + logger.info(f"视频文件路径: {task.video_file_path}") + + # 创建一个事件来等待任务完成 + 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: - # 确保文件名不为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"等待任务完成: {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="任务处理超时") - # 保存音频文件 - 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}") + # 获取最终任务状态 + result = await task_manager.get_task(task.task_id) + if not result: + raise HTTPException(status_code=404, detail="任务不存在") - # 保存视频文件 - 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 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="输出文件不存在") - # 验证文件是否保存成功 - if not os.path.exists(audio_path) or not os.path.exists(video_path): - raise ValueError("文件保存失败") + # 获取文件名 + filename = os.path.basename(output_file) - # 创建输出目录 - output_dir = os.path.join("output", datetime.now().strftime("%Y%m%d")) - os.makedirs(output_dir, exist_ok=True) + # 创建文件流 + async def file_iterator(): + async with aiofiles.open(output_file, "rb") as f: + while chunk := await f.read(8192): # 8KB chunks + yield chunk - # 创建任务 - 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)) + 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: logger.error(f"创建任务失败: {str(e)}", exc_info=True) @@ -377,7 +343,9 @@ async def get_all_tasks(): @app.post("/tasks/{task_id}/callback") -async def add_task_callback(task_id: str, callback_url: str, background_tasks: BackgroundTasks): +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") @@ -394,7 +362,9 @@ async def add_task_callback(task_id: str, callback_url: str, background_tasks: B # 后台任务处理函数 -async def process_video_task(task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str): +async def process_video_task( + task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str +): """处理视频任务的后台函数""" try: # 更新任务状态为处理中 @@ -404,21 +374,161 @@ async def process_video_task(task_id: str, audio_file_path: str, video_file_path 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) + 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") + 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)) + 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}") +# 添加文件下载辅助函数 +async def download_file_from_url(url: str, destination_path: str) -> bool: + """从URL下载文件到指定路径""" + try: + logger.info(f"开始下载文件: {url} -> {destination_path}") + async with aiohttp.ClientSession() as session: + async with session.get( + url, timeout=aiohttp.ClientTimeout(total=300) + ) as response: + if response.status == 200: + # 获取Content-Length用于进度追踪 + content_length = response.headers.get("Content-Length") + if content_length: + logger.info(f"文件大小: {int(content_length)} bytes") + + async with aiofiles.open(destination_path, "wb") as f: + downloaded = 0 + async for chunk in response.content.iter_chunked(8192): + await f.write(chunk) + downloaded += len(chunk) + + logger.info( + f"文件下载完成: {destination_path} ({downloaded} bytes)" + ) + return True + else: + logger.error(f"下载失败，HTTP状态码: {response.status}") + return False + except Exception as e: + logger.error(f"下载文件时发生错误: {str(e)}") + return False + + +async def process_media_input( + media_file: Optional[UploadFile], + media_url: Optional[str], + media_type: str, + temp_dir: str, +) -> str: + """处理媒体输入（文件或URL），返回本地文件路径""" + if media_file and media_url: + raise ValueError(f"{media_type}不能同时提供文件和URL") + + if not media_file and not media_url: + raise ValueError(f"必须提供{media_type}文件或URL") + + if media_file: + # 处理上传文件 + if not media_file.filename: + raise ValueError(f"{media_type}文件名不能为空") + + file_path = os.path.join(temp_dir, media_file.filename) + logger.info(f"开始保存{media_type}文件: {file_path}") + async with aiofiles.open(file_path, "wb") as f: + content = await media_file.read() + await f.write(content) + logger.info(f"{media_type}文件保存完成: {file_path}") + return file_path + + else: + # 处理URL下载 + parsed_url = urllib.parse.urlparse(media_url) + filename = os.path.basename(parsed_url.path) + if not filename or "." not in filename: + # 如果URL没有文件名或没有扩展名，根据媒体类型生成 + if media_type == "音频": + extension = ".mp3" # 默认音频扩展名 + elif media_type == "视频": + extension = ".mp4" # 默认视频扩展名 + else: + extension = "" + + if not filename: + filename = f"{media_type}_{int(datetime.now().timestamp())}{extension}" + elif "." not in filename: + filename = f"{filename}{extension}" + + file_path = os.path.join(temp_dir, filename) + success = await download_file_from_url(media_url, file_path) + if not success: + raise ValueError(f"下载{media_type}文件失败: {media_url}") + return file_path + + +async def prepare_task_files( + audio: Optional[UploadFile], + video: Optional[UploadFile], + audio_url: Optional[str], + video_url: Optional[str], + description: str, + callback_url: Optional[str], +) -> tuple[Task, str]: + """ + 准备任务文件和创建任务的公共逻辑 + 返回创建的任务对象和临时目录路径 + """ + # 创建临时目录 + 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: + # 处理音频输入（文件或URL） + audio_path = await process_media_input(audio, audio_url, "音频", temp_dir) + + # 处理视频输入（文件或URL） + video_path = await process_media_input(video, video_url, "视频", temp_dir) + + # 验证文件是否存在 + 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 + ) + + return task, temp_dir + + except Exception as e: + # 如果发生错误，立即清理临时文件 + # if os.path.exists(temp_dir): + # shutil.rmtree(temp_dir) + logger.error(f"准备任务文件失败: {str(e)}") + raise + + if __name__ == "__main__": import uvicorn - uvicorn.run("api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info") # 开发模式下启用热重载 # 单进程模式 + uvicorn.run( + "api:app", host="0.0.0.0", port=8101, reload=True, workers=1, log_level="info" + ) # 开发模式下启用热重载 # 单进程模式