init: heygen 队列

This commit is contained in:
2025-07-25 15:20:42 +08:00
commit 8301a19003
7 changed files with 1423 additions and 0 deletions

60
.gitignore vendored Normal file
View File

@@ -0,0 +1,60 @@
# IDE 和编辑器文件
.idea/
.vscode/
*.swp
*.swo
.DS_Store
# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
venv/
env/
.env
.venv/
.venv
# 日志文件
*.log
logs/
# 临时文件
*.tmp
*.temp
*.bak
# 数据库文件
*.db
*.sqlite3
# 系统文件
Thumbs.db
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
output/
temp/
config/machines.json

168
README.md Normal file
View File

@@ -0,0 +1,168 @@
# Multi-HeyGem 视频处理服务
Multi-HeyGem 是一个分布式视频处理服务,支持多机器并行处理视频和音频合成任务。该服务提供了 RESTful API 接口,可以方便地集成到其他系统中。
## 功能特点
- 支持多机器分布式处理
- 异步任务处理
- RESTful API 接口
- 任务状态追踪
- 支持回调通知
- 详细的日志记录
- 自动负载均衡
## 系统要求
配置机器信息:
编辑 `config/machines.json` 文件,添加您的处理机器信息:
```json
{
"machines": [
{
"name": "机器名称",
"host": "机器IP地址",
"port": 8000,
"max_concurrent": 1,
"enabled": true
}
]
}
```
## 启动服务
```bash
python api.py
```
服务将在 `http://0.0.0.0:8101` 启动。
## API 接口
通过 `http://localhost:8101/docs` 可查看 api 接口
### 创建任务
创建任务接口需要上传文件,以下是使用示例:
#### 同步生成示例(直接返回结果)
```python
import requests
import json
# API 服务地址
API_URL = "http://localhost:8101"
# 准备文件
files = {
'audio': ('audio.mp3', open('./samples/input/audio.mp3', 'rb')),
'video': ('video.mp4', open('./samples/input/video.mp4', 'rb'))
}
# 准备请求数据
data = {
'description': '测试任务'
}
# 发送同步生成请求
response = requests.post(
f"{API_URL}/tasks/sync",
files=files,
data=data
)
# 处理响应
if response.status_code == 200:
# 直接保存返回的视频流
output_path = './output/result.mp4'
with open(output_path, 'wb') as f:
f.write(response.content)
print(f"视频生成成功!已保存到: {output_path}")
else:
print(f"视频生成失败: {response.text}")
```
#### 异步调用示例
```python
import aiohttp
import aiofiles
import asyncio
async def create_task():
API_URL = "http://localhost:8101"
# 准备文件
audio_file_path = './samples/input/audio.mp3'
video_file_path = './samples/input/video.mp4'
# 创建 FormData
data = aiohttp.FormData()
# 添加音频文件
async with aiofiles.open(audio_file_path, 'rb') as f:
audio_content = await f.read()
data.add_field(
'audio_file',
audio_content,
filename='audio.mp3',
content_type='audio/mpeg'
)
# 添加视频文件
async with aiofiles.open(video_file_path, 'rb') as f:
video_content = await f.read()
data.add_field(
'video_file',
video_content,
filename='video.mp4',
content_type='video/mp4'
)
# 添加其他参数
data.add_field('description', '测试任务')
data.add_field('callback_url', 'http://your-callback-url.com/webhook') # 可选的回调地址
# 发送请求
async with aiohttp.ClientSession() as session:
async with session.post(f"{API_URL}/tasks", data=data) as response:
if response.status == 200:
task = await response.json()
print(f"任务创建成功任务ID: {task['task_id']}")
print(f"任务状态: {task['status']}")
# 查询任务状态
task_id = task['task_id']
async with session.get(f"{API_URL}/tasks/{task_id}") as status_response:
task_status = await status_response.json()
print(f"当前任务状态: {task_status['status']}")
else:
error_text = await response.text()
print(f"任务创建失败: {error_text}")
# 运行异步函数
asyncio.run(create_task())
```
### 任务状态说明
- `pending`: 任务等待处理
- `processing`: 任务正在处理中
- `completed`: 任务处理完成
- `failed`: 任务处理失败
### 回调通知
当任务状态发生变化时,如果设置了 `callback_url`,系统会向该地址发送 POST 请求,请求体包含:
```json
{
"task_id": "任务ID",
"status": "任务状态",
"completed_at": "完成时间",
"error_message": "错误信息(如果有)"
}
```

424
api.py Normal file
View File

@@ -0,0 +1,424 @@
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") # 开发模式下启用热重载 # 单进程模式

View File

@@ -0,0 +1,11 @@
{
"machines": [
{
"name": "1",
"host": "192.168.0.131",
"port": 8000,
"max_concurrent": 1,
"enabled": true
}
]
}

36
models.py Normal file
View File

@@ -0,0 +1,36 @@
from datetime import datetime
from enum import Enum
from typing import Optional
from pydantic import BaseModel
# 定义任务状态枚举
class TaskStatus(Enum):
PENDING = "pending"
QUEUED = "queued"
PROCESSING = "processing"
COMPLETED = "completed"
FAILED = "failed"
# 定义任务模型
class Task(BaseModel):
task_id: str
audio_file_path: str
video_file_path: str
output_dir: str
output_file_path: str
description: str
status: TaskStatus
created_at: datetime
completed_at: Optional[datetime] = None
error_message: Optional[str] = None
machine_name: Optional[str] = None
callback_url: Optional[str] = None
# 定义任务创建请求模型
class TaskCreateRequest(BaseModel):
description: str
callback_url: Optional[str] = None

716
processor.py Normal file
View File

@@ -0,0 +1,716 @@
import asyncio
import asyncio.exceptions
import json
import os
import random
import subprocess
import time
import traceback
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, List, Optional
import aiofiles
import aiohttp
from loguru import logger
from models import TaskStatus # 从models.py导入TaskStatus
from moviepy.editor import VideoFileClip
from pydub import AudioSegment
# 配置日志输出
# 创建logs目录
os.makedirs("logs", exist_ok=True)
# 添加错误日志文件输出只输出ERROR级别的日志
logger.add(
"logs/error.log",
rotation="500 MB",
retention="10 days",
level="ERROR",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
)
@dataclass
class Machine:
name: str
host: str
port: int
max_concurrent: int
enabled: bool
current_tasks: int = 0
lock: Optional[asyncio.Lock] = None
class ProcessManager:
def __init__(self, config_path: str, log_dir: str = "logs", task_manager=None):
self.machines: Dict[str, Machine] = {}
self.config_path = config_path
self.log_dir = log_dir
self.request_log_file = None
self.task_manager = task_manager # 添加task_manager引用
self._load_config()
# 添加文件锁字典
self.file_locks: Dict[str, asyncio.Lock] = {}
self.output_dir_locks: Dict[str, asyncio.Lock] = {}
# 添加任务队列
self.task_queue: asyncio.Queue = asyncio.Queue()
# 添加状态汇报任务
self.status_task: Optional[asyncio.Task] = None
# 添加运行状态标志
self.is_running = False
# 添加处理任务列表
self.process_tasks: List[asyncio.Task] = []
# 初始化日志文件
self._init_log_file()
# 添加机器使用统计
self.machine_stats: Dict[str, Dict] = {}
for machine_name in self.machines:
self.machine_stats[machine_name] = {
"total_tasks": 0,
"current_tasks": 0,
"last_used": None,
"average_processing_time": 0,
"success_rate": 1.0,
"total_processing_time": 0,
}
def _init_log_file(self):
"""初始化日志文件"""
# 确保日志目录存在
os.makedirs(self.log_dir, exist_ok=True)
# 创建新的日志文件
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
self.request_log_file = os.path.join(self.log_dir, f"video_generation_requests_{timestamp}.md")
# 创建表头
with open(self.request_log_file, "w", encoding="utf-8") as f:
f.write(
f"""# 视频生成请求记录 - {timestamp}
| 开始时间 | 结束时间 | 用时(秒) | 设备 | 视频大小(MB) | 音频大小(MB) | 视频时长(秒) | 音频时长(秒) | 视频名称 | 音频名称 | 输出路径 | 状态 | 错误信息 |
|----------|----------|----------|------|-------------|-------------|-------------|-------------|----------|----------|----------|------|----------|
"""
)
def write_request_log(
self,
start_time: float,
end_time: float,
machine_name: str,
video_size: float,
audio_size: float,
video_duration: float,
audio_duration: float,
video_name: str,
audio_name: str,
output_path: str,
status: bool,
error_msg: str = "",
):
"""记录视频生成请求的详细信息"""
if not self.request_log_file:
logger.error("日志文件未初始化")
return
duration = end_time - start_time
# 追加新的记录
with open(self.request_log_file, "a", encoding="utf-8") as f:
f.write(
f"| {datetime.fromtimestamp(start_time).strftime('%Y-%m-%d %H:%M:%S')} | "
f"{datetime.fromtimestamp(end_time).strftime('%Y-%m-%d %H:%M:%S')} | "
f"{duration:.2f} | {machine_name} | {video_size:.2f} | {audio_size:.2f} | "
f"{video_duration:.2f} | {audio_duration:.2f} | {video_name} | {audio_name} | "
f"{output_path} | {'成功' if status else '失败'} | {error_msg} |\n"
)
def _load_config(self):
"""加载机器配置"""
config_path = os.path.abspath(self.config_path)
with open(config_path, "r", encoding="utf-8") as f:
config = json.load(f)
for machine_config in config["machines"]:
if machine_config["enabled"]:
machine = Machine(
name=machine_config["name"],
host=machine_config["host"],
port=machine_config["port"],
max_concurrent=machine_config["max_concurrent"],
enabled=machine_config["enabled"],
lock=asyncio.Lock(),
)
self.machines[machine.name] = machine
async def get_available_machine(self) -> Optional[Machine]:
"""获取一个可用的机器"""
for machine in self.machines.values():
async with machine.lock: # type: ignore
if machine.current_tasks < machine.max_concurrent:
machine.current_tasks += 1
return machine
return None
async def release_machine(self, machine: Machine):
"""释放机器资源"""
async with machine.lock: # type: ignore
machine.current_tasks -= 1
def _get_file_lock(self, file_path: str) -> asyncio.Lock:
"""获取文件锁"""
if file_path not in self.file_locks:
self.file_locks[file_path] = asyncio.Lock()
return self.file_locks[file_path]
def _get_output_dir_lock(self, output_dir: str) -> asyncio.Lock:
"""获取输出目录锁"""
if output_dir not in self.output_dir_locks:
self.output_dir_locks[output_dir] = asyncio.Lock()
return self.output_dir_locks[output_dir]
async def _safe_copy_file(self, src: str, dst: str, timeout: float = 300) -> bool:
"""安全地复制文件,带超时控制"""
try:
# 获取文件锁
file_lock = self._get_file_lock(src)
# 尝试获取锁,带超时控制
try:
async with asyncio.timeout(timeout):
await file_lock.acquire()
except asyncio.exceptions.TimeoutError:
logger.warning(f"等待文件锁超时: {src}")
return False
except Exception as e:
logger.error(f"获取文件锁失败: {src}, 错误: {str(e)}")
return False
try:
# 确保源文件存在
if not os.path.exists(src):
logger.error(f"源文件不存在: {src}")
return False
# 确保目标目录存在
os.makedirs(os.path.dirname(dst), exist_ok=True)
# 复制文件
async with aiofiles.open(src, "rb") as fsrc:
content = await fsrc.read()
async with aiofiles.open(dst, "wb") as fdst:
await fdst.write(content)
return True
finally:
file_lock.release()
except Exception as e:
logger.error(f"复制文件失败 {src} -> {dst}: {str(e)}")
return False
async def start(self):
"""启动进程管理器"""
self.is_running = True
# 启动状态汇报任务
self.status_task = asyncio.create_task(self._status_report())
# 启动任务处理循环
# 计算总的最大并发数
total_max_concurrent = sum(machine.max_concurrent for machine in self.machines.values())
self.process_tasks = [asyncio.create_task(self._process_tasks()) for _ in range(total_max_concurrent)]
async def stop(self):
"""停止进程管理器"""
self.is_running = False
if self.status_task:
self.status_task.cancel()
try:
await self.status_task
except asyncio.CancelledError:
pass
# 取消所有处理任务
for task in self.process_tasks:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def _status_report(self):
"""定时打印状态汇报"""
while self.is_running:
report: List[str] = []
# 汇报等待队列中的任务数量
report.append(f"\n状态: 队列等待任务数量: {self.task_queue.qsize()}")
# 更新机器统计信息
for machine_name, machine in self.machines.items():
stats = self.machine_stats[machine_name]
stats["current_tasks"] = machine.current_tasks
# 计算使用率
usage_rate = (machine.current_tasks / machine.max_concurrent) * 100
report.append(f" - {machine_name}: 当前负载: {machine.current_tasks}/{machine.max_concurrent} ({usage_rate:.1f}%)")
logger.info("\n".join(report))
await asyncio.sleep(5) # 每5秒打印一次状态
async def _process_tasks(self):
"""处理任务队列中的任务"""
while self.is_running:
try:
# 从队列中获取任务
task = await self.task_queue.get()
logger.info(f"获取任务: {task}")
task_id, audio_file_path, video_file_path, output_file_path = task
# 等待可用机器
machine = await self._wait_for_available_machine()
if not machine:
logger.error("无法获取可用机器,任务将重新入队")
# 将任务重新放回队列
await self.task_queue.put(task)
# 更新任务状态为失败
if self.task_manager and task_id:
await self.task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message="无法获取可用机器")
self.task_queue.task_done()
await asyncio.sleep(5) # 等待5秒后重试
continue
# 处理任务
try:
success = await self._process_single_task(task_id, machine, audio_file_path, video_file_path, output_file_path)
if not success:
logger.error(f"任务处理失败: \n音频文件: {audio_file_path}\n视频文件: {video_file_path}\n输出文件: {output_file_path}")
# 更新任务状态为失败
if self.task_manager and task_id:
await self.task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message="任务处理失败")
finally:
self.task_queue.task_done()
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"任务处理出错: {str(e)} {traceback.format_exc()}")
# 更新任务状态为失败
if self.task_manager and task_id:
await self.task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message=str(e))
self.task_queue.task_done()
async def _wait_for_available_machine(self, timeout: float = 3600) -> Optional[Machine]:
"""等待获取可用机器,只关注负载"""
start_time = asyncio.get_event_loop().time()
while asyncio.get_event_loop().time() - start_time < timeout:
# 获取所有未满的机器
available_machines = []
for machine in self.machines.values():
async with machine.lock: # type: ignore
if machine.current_tasks < machine.max_concurrent:
available_machines.append(machine)
if available_machines:
# 选择负载最低的机器
selected_machine = min(available_machines, key=lambda m: m.current_tasks / m.max_concurrent)
# 再次检查机器是否可用
async with selected_machine.lock: # type: ignore
if selected_machine.current_tasks < selected_machine.max_concurrent:
selected_machine.current_tasks += 1
logger.info(
f"选择机器 {selected_machine.name} 处理任务,当前负载: {selected_machine.current_tasks}/{selected_machine.max_concurrent}"
)
return selected_machine
# 所有机器都满了,等待后重试
logger.warning("所有机器都已满,等待资源释放")
await asyncio.sleep(5)
async def _process_single_task(self, task_id: str, machine: Machine, audio_file_path: str, video_file_path: str, output_file_path: str) -> bool:
"""处理单个任务"""
start_time = time.time()
success = False
try:
logger.info(f"开始处理视频 {output_file_path}")
logger.info(f"任务ID: {task_id}")
logger.info(f"音频文件路径: {audio_file_path} 是否存在: {os.path.exists(audio_file_path)}")
logger.info(f"视频文件路径: {video_file_path} 是否存在: {os.path.exists(video_file_path)}")
logger.info(f"使用机器: {machine.name}")
url = f"http://{machine.host}:{machine.port}/process"
async with aiohttp.ClientSession() as session:
# 准备文件数据
data = aiohttp.FormData()
# 直接使用原文件
async with aiofiles.open(audio_file_path, "rb") as f:
audio_content = await f.read()
data.add_field("audio", audio_content, filename=os.path.basename(audio_file_path), content_type="audio/mpeg")
async with aiofiles.open(video_file_path, "rb") as f:
video_content = await f.read()
data.add_field("video", video_content, filename=os.path.basename(video_file_path), content_type="video/mp4")
# 获取音频和视频时长
try:
# 获取音频时长
audio = AudioSegment.from_mp3(audio_file_path)
audio_duration = len(audio) / 1000.0 # 转换为秒
# 获取视频时长
video = VideoFileClip(video_file_path)
video_duration = video.duration
video.close()
logger.info(f"音频时长: {audio_duration:.2f}")
logger.info(f"视频时长: {video_duration:.2f}")
except Exception as e:
logger.warning(f"获取音频或视频时长失败: {str(e)}")
# 如果视频时长小于音频时长+1s则使用ffmpeg通过循环的形式延长视频时长到音频时长+1s
if video_duration < audio_duration + 1:
logger.info(f"视频时长小于音频时长+1s通过循环的形式延长视频时长到音频时长+1s")
# 计算需要循环的次数
target_duration = audio_duration + 1
loop_count = int(target_duration / video_duration) + 1
# 使用ffmpeg通过循环的形式延长视频时长
command = [
"ffmpeg",
"-stream_loop",
str(loop_count - 1), # 循环次数减1因为原始视频也算一次
"-i",
video_file_path,
"-c",
"copy",
"-t",
str(target_duration), # 设置目标时长
"-f",
"mp4", # 明确指定输出格式
"-y",
video_file_path + ".tmp", # 使用临时文件避免覆盖原文件
]
try:
# 使用 subprocess.run 替代 asyncio.create_subprocess_exec
logger.info(f"执行FFmpeg命令: {' '.join(command)}")
result = subprocess.run(command, capture_output=True, text=True, check=False)
if result.returncode != 0:
error_msg = result.stderr if result.stderr else "未知错误"
logger.error(f"扩展视频时长失败FFmpeg返回码: {result.returncode}, 错误信息: {error_msg}")
return False
logger.info(f"扩展视频时长完成,扩展后视频时长: {target_duration:.2f}")
# 替换原文件
os.replace(video_file_path + ".tmp", video_file_path)
except Exception as e:
import traceback
error_trace = traceback.format_exc()
logger.error(f"扩展视频时长时发生错误:\n命令: {' '.join(command)}\n错误信息: {str(e)}\n完整堆栈:\n{error_trace}")
return False
# 如果视频时长大于音频时长+1s则随机裁剪到音频时长+1s
elif video_duration > audio_duration + 1:
logger.info(f"视频时长大于音频时长+1s随机裁剪到音频时长+1s")
target_duration = audio_duration + 1
# 计算随机起始时间,确保剩余时长足够
max_start_time = video_duration - target_duration
start_time = random.uniform(0, max_start_time)
# 使用ffmpeg随机裁剪视频
command = [
"ffmpeg",
"-ss",
str(start_time),
"-i",
video_file_path,
"-t",
str(target_duration),
"-c",
"copy",
"-f",
"mp4", # 明确指定输出格式
"-y",
video_file_path + ".tmp",
]
logger.info(f"实际请求音视频时长: {target_duration:.2f}")
try:
# 使用 subprocess.run 替代 asyncio.create_subprocess_exec
logger.info(f"执行FFmpeg命令: {' '.join(command)}")
result = subprocess.run(command, capture_output=True, text=True, check=False)
if result.returncode != 0:
error_msg = result.stderr if result.stderr else "未知错误"
logger.error(f"裁剪视频时长失败FFmpeg返回码: {result.returncode}, 错误信息: {error_msg}")
return False
logger.info(f"裁剪视频时长完成,裁剪后视频时长: {target_duration:.2f}")
# 替换原文件
os.replace(video_file_path + ".tmp", video_file_path)
logger.info(f"替换原文件完成")
except Exception as e:
import traceback
error_trace = traceback.format_exc()
logger.error(f"裁剪视频时长时发生错误:\n命令: {' '.join(command)}\n错误信息: {str(e)}\n完整堆栈:\n{error_trace}")
return False
# 打印文件大小
audio_size = len(audio_content) / (1024 * 1024) # 转换为MB
video_size = len(video_content) / (1024 * 1024) # 转换为MB
logger.info(f"音频文件大小: {audio_size:.2f}MB")
logger.info(f"视频文件大小: {video_size:.2f}MB")
# 发送请求
logger.info(f"开始进行视频生成,视频名称: {os.path.basename(video_file_path)}, 音频名称: {os.path.basename(audio_file_path)}")
start_time = time.time()
# 设置30分钟超时
timeout = aiohttp.ClientTimeout(total=1800) # 30分钟 = 1800秒
async with session.post(url, data=data, timeout=timeout) as response:
if response.status == 200:
# 检查响应的 Content-Type
content_type = response.headers.get("Content-Type", "").lower()
# 如果是 JSON 响应
if "application/json" in content_type:
try:
response_json = await response.json()
if not response_json.get("status", True): # 如果status为false或不存在
error_msg = response_json.get("msg", "未知错误")
logger.error(f"服务器处理失败: {error_msg}")
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
error_msg,
)
return False
except json.JSONDecodeError:
logger.error("解析JSON响应失败")
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
"解析JSON响应失败",
)
return False
# 如果是视频文件响应
elif "video/mp4" in content_type:
# 确保输出目录存在
output_dir = os.path.dirname(output_file_path)
output_lock = self._get_output_dir_lock(output_dir)
# 尝试获取输出目录锁,带超时控制
try:
async with asyncio.timeout(300): # 5分钟超时
await output_lock.acquire()
except asyncio.exceptions.TimeoutError:
logger.error(f"等待输出目录锁超时: {output_dir}")
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
f"等待输出目录锁超时: {output_dir}",
)
return False
try:
os.makedirs(output_dir, exist_ok=True)
# 保存响应内容,添加重试机制
max_retries = 3
retry_delay = 5 # 秒
last_error = None
for attempt in range(max_retries):
try:
# 读取响应内容
content = await response.read()
if not content:
raise ValueError("响应内容为空")
# 保存文件
async with aiofiles.open(output_file_path, "wb") as f:
await f.write(content)
# 验证文件是否成功保存
if not os.path.exists(output_file_path) or os.path.getsize(output_file_path) == 0:
raise ValueError("文件保存失败或文件大小为0")
logger.info(f"成功保存视频文件: {output_file_path}")
# 更新任务状态为完成
if self.task_manager and task_id:
await self.task_manager.update_task_status(task_id, TaskStatus.COMPLETED)
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
True,
)
return True
except Exception as e:
last_error = e
if attempt < max_retries - 1:
logger.warning(f"保存视频文件失败,{retry_delay}秒后重试 ({attempt + 1}/{max_retries}): {str(e)}")
await asyncio.sleep(retry_delay)
else:
logger.error(f"保存视频文件失败,达到最大重试次数: {str(last_error)}")
# 更新任务状态为失败
if self.task_manager and task_id:
await self.task_manager.update_task_status(task_id, TaskStatus.FAILED, error_message=str(last_error))
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
f"保存视频文件失败: {str(last_error)}",
)
return False
finally:
output_lock.release()
else:
logger.error(f"未知的响应类型: {content_type}")
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
f"未知的响应类型: {content_type}",
)
return False
else:
error_text = await response.text()
try:
error_json = json.loads(error_text)
error_msg = error_json.get("msg", error_text)
logger.error(f"处理视频失败: HTTP {response.status} - {error_msg}")
except json.JSONDecodeError:
logger.error(f"处理视频失败: HTTP {response.status} - {error_text}")
self.write_request_log(
start_time,
time.time(),
machine.name,
video_size,
audio_size,
target_duration,
audio_duration,
os.path.basename(video_file_path),
os.path.basename(audio_file_path),
output_file_path,
False,
f"HTTP {response.status} - {error_text}",
)
return False
return success
except Exception as e:
logger.error(f"处理视频时发生错误: {str(e)}")
return False
finally:
await self.release_machine(machine)
async def process_video(self, task_id: str, audio_file_path: str, video_file_path: str, out_file_path: str) -> Optional[str]:
"""添加视频处理任务到队列"""
try:
# 生成输出文件路径
# 添加任务到队列
await self.task_queue.put((task_id, audio_file_path, video_file_path, out_file_path))
return out_file_path
except Exception as e:
logger.error(f"添加任务到队列失败: {str(e)}")
return None
def get_machine_status(self) -> List[Dict]:
"""获取所有机器的状态"""
return [
{
"name": machine.name,
"host": machine.host,
"current_tasks": machine.current_tasks,
"max_concurrent": machine.max_concurrent,
"enabled": machine.enabled,
}
for machine in self.machines.values()
]
async def main():
"""
命令行入口函数
"""
import argparse
parser = argparse.ArgumentParser(description="视频处理管理器")
args = parser.parse_args()
logger.info("启动视频处理管理器...")
# 初始化进程管理器
process_manager = ProcessManager(config_path="config/machines.json")
await process_manager.start()
try:
# 保持程序运行
while True:
await asyncio.sleep(1)
except KeyboardInterrupt:
logger.info("收到停止信号,正在关闭...")
finally:
await process_manager.stop()
if __name__ == "__main__":
exit(asyncio.run(main()))

8
requirements.txt Normal file
View File

@@ -0,0 +1,8 @@
aiofiles>=23.2.1
aiohttp>=3.9.1
fastapi>=0.104.1
loguru>=0.7.2
moviepy>=1.0.3
pydantic>=2.5.2
pydub>=0.25.1
uvicorn>=0.24.0