You've already forked FrameTour-RenderWorker
feat(render): 实现渲染系统v2核心架构
- 添加v2支持的任务类型常量定义 - 更新软件版本至0.0.9 - 定义v2统一音视频编码参数 - 实现系统信息工具get_sys_info_v2方法 - 新增get_capabilities和_get_gpu_info功能 - 创建core模块及TaskHandler抽象基类 - 添加渲染系统设计文档包括集群架构、v2 PRD和Worker PRD - 实现任务处理器抽象基类及接口规范
This commit is contained in:
@@ -1,3 +1,9 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
常量定义
|
||||
"""
|
||||
|
||||
# v1 支持的特性
|
||||
SUPPORT_FEATURE = (
|
||||
'simple_render_algo',
|
||||
'gpu_accelerate',
|
||||
@@ -6,4 +12,35 @@ SUPPORT_FEATURE = (
|
||||
'rclone_upload',
|
||||
'custom_re_encode',
|
||||
)
|
||||
SOFTWARE_VERSION = '0.0.8'
|
||||
|
||||
# 软件版本
|
||||
SOFTWARE_VERSION = '0.0.9'
|
||||
|
||||
# v2 支持的任务类型
|
||||
V2_TASK_TYPES = (
|
||||
'RENDER_SEGMENT_VIDEO',
|
||||
'PREPARE_JOB_AUDIO',
|
||||
'PACKAGE_SEGMENT_TS',
|
||||
'FINALIZE_MP4',
|
||||
)
|
||||
|
||||
# v2 默认能力
|
||||
V2_DEFAULT_CAPABILITIES = list(V2_TASK_TYPES)
|
||||
|
||||
# v2 统一视频编码参数
|
||||
V2_VIDEO_ENCODE_PARAMS = {
|
||||
'codec': 'libx264',
|
||||
'preset': 'medium',
|
||||
'profile': 'main',
|
||||
'level': '4.0',
|
||||
'crf': '23',
|
||||
'pix_fmt': 'yuv420p',
|
||||
}
|
||||
|
||||
# v2 统一音频编码参数
|
||||
V2_AUDIO_ENCODE_PARAMS = {
|
||||
'codec': 'aac',
|
||||
'bitrate': '128k',
|
||||
'sample_rate': '48000',
|
||||
'channels': '2',
|
||||
}
|
||||
|
||||
12
core/__init__.py
Normal file
12
core/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
核心抽象层
|
||||
|
||||
包含任务处理器抽象基类等核心接口定义。
|
||||
"""
|
||||
|
||||
from core.handler import TaskHandler
|
||||
|
||||
__all__ = [
|
||||
'TaskHandler',
|
||||
]
|
||||
79
core/handler.py
Normal file
79
core/handler.py
Normal file
@@ -0,0 +1,79 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务处理器抽象基类
|
||||
|
||||
定义任务处理器的接口规范。
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from domain.task import Task, TaskType
|
||||
from domain.result import TaskResult
|
||||
|
||||
|
||||
class TaskHandler(ABC):
|
||||
"""
|
||||
任务处理器抽象基类
|
||||
|
||||
所有任务处理器都必须继承此类并实现相应方法。
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def handle(self, task: 'Task') -> 'TaskResult':
|
||||
"""
|
||||
处理任务的主方法
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
|
||||
Returns:
|
||||
TaskResult: 任务结果(成功或失败)
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_supported_type(self) -> 'TaskType':
|
||||
"""
|
||||
返回此处理器支持的任务类型
|
||||
|
||||
Returns:
|
||||
TaskType: 支持的任务类型枚举值
|
||||
"""
|
||||
pass
|
||||
|
||||
def before_handle(self, task: 'Task') -> None:
|
||||
"""
|
||||
处理前钩子(可选重写)
|
||||
|
||||
用于任务执行前的准备工作,如日志记录、资源检查等。
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
"""
|
||||
pass
|
||||
|
||||
def after_handle(self, task: 'Task', result: 'TaskResult') -> None:
|
||||
"""
|
||||
处理后钩子(可选重写)
|
||||
|
||||
用于任务执行后的清理工作,如资源释放、统计记录等。
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
result: 任务结果
|
||||
"""
|
||||
pass
|
||||
|
||||
def validate_task(self, task: 'Task') -> bool:
|
||||
"""
|
||||
验证任务是否有效(可选重写)
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
|
||||
Returns:
|
||||
bool: 任务是否有效
|
||||
"""
|
||||
return True
|
||||
24
domain/__init__.py
Normal file
24
domain/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
领域模型层
|
||||
|
||||
包含任务实体、结果、配置等核心数据结构。
|
||||
"""
|
||||
|
||||
from domain.task import Task, TaskType, TaskStatus, RenderSpec, OutputSpec, AudioSpec, AudioProfile
|
||||
from domain.result import TaskResult, ErrorCode, RETRY_CONFIG
|
||||
from domain.config import WorkerConfig
|
||||
|
||||
__all__ = [
|
||||
'Task',
|
||||
'TaskType',
|
||||
'TaskStatus',
|
||||
'RenderSpec',
|
||||
'OutputSpec',
|
||||
'AudioSpec',
|
||||
'AudioProfile',
|
||||
'TaskResult',
|
||||
'ErrorCode',
|
||||
'RETRY_CONFIG',
|
||||
'WorkerConfig',
|
||||
]
|
||||
122
domain/config.py
Normal file
122
domain/config.py
Normal file
@@ -0,0 +1,122 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Worker 配置模型
|
||||
|
||||
定义 Worker 运行时的配置参数。
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
# 默认支持的任务类型
|
||||
DEFAULT_CAPABILITIES = [
|
||||
"RENDER_SEGMENT_VIDEO",
|
||||
"PREPARE_JOB_AUDIO",
|
||||
"PACKAGE_SEGMENT_TS",
|
||||
"FINALIZE_MP4"
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkerConfig:
|
||||
"""
|
||||
Worker 配置
|
||||
|
||||
包含 Worker 运行所需的所有配置参数。
|
||||
"""
|
||||
# API 配置
|
||||
api_endpoint: str
|
||||
access_key: str
|
||||
worker_id: str
|
||||
|
||||
# 并发控制
|
||||
max_concurrency: int = 4
|
||||
|
||||
# 心跳配置
|
||||
heartbeat_interval: int = 5 # 秒
|
||||
|
||||
# 租约配置
|
||||
lease_extension_threshold: int = 60 # 秒,提前多久续期
|
||||
lease_extension_duration: int = 300 # 秒,每次续期时长
|
||||
|
||||
# 目录配置
|
||||
temp_dir: str = "/tmp/render_worker"
|
||||
|
||||
# 能力配置
|
||||
capabilities: List[str] = field(default_factory=lambda: DEFAULT_CAPABILITIES.copy())
|
||||
|
||||
# FFmpeg 配置
|
||||
ffmpeg_timeout: int = 3600 # 秒,FFmpeg 执行超时
|
||||
|
||||
# 下载/上传配置
|
||||
download_timeout: int = 300 # 秒,下载超时
|
||||
upload_timeout: int = 600 # 秒,上传超时
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> 'WorkerConfig':
|
||||
"""从环境变量创建配置"""
|
||||
|
||||
# API 端点,优先使用 V2 版本
|
||||
api_endpoint = os.getenv('API_ENDPOINT_V2') or os.getenv('API_ENDPOINT', '')
|
||||
if not api_endpoint:
|
||||
raise ValueError("API_ENDPOINT_V2 or API_ENDPOINT environment variable is required")
|
||||
|
||||
# Access Key
|
||||
access_key = os.getenv('ACCESS_KEY', '')
|
||||
if not access_key:
|
||||
raise ValueError("ACCESS_KEY environment variable is required")
|
||||
|
||||
# Worker ID
|
||||
worker_id = os.getenv('WORKER_ID', '100001')
|
||||
|
||||
# 并发数
|
||||
max_concurrency = int(os.getenv('MAX_CONCURRENCY', '4'))
|
||||
|
||||
# 心跳间隔
|
||||
heartbeat_interval = int(os.getenv('HEARTBEAT_INTERVAL', '5'))
|
||||
|
||||
# 租约配置
|
||||
lease_extension_threshold = int(os.getenv('LEASE_EXTENSION_THRESHOLD', '60'))
|
||||
lease_extension_duration = int(os.getenv('LEASE_EXTENSION_DURATION', '300'))
|
||||
|
||||
# 临时目录
|
||||
temp_dir = os.getenv('TEMP_DIR', os.getenv('TEMP', '/tmp/render_worker'))
|
||||
|
||||
# 能力列表
|
||||
capabilities_str = os.getenv('CAPABILITIES', '')
|
||||
if capabilities_str:
|
||||
capabilities = [c.strip() for c in capabilities_str.split(',') if c.strip()]
|
||||
else:
|
||||
capabilities = DEFAULT_CAPABILITIES.copy()
|
||||
|
||||
# FFmpeg 超时
|
||||
ffmpeg_timeout = int(os.getenv('FFMPEG_TIMEOUT', '3600'))
|
||||
|
||||
# 下载/上传超时
|
||||
download_timeout = int(os.getenv('DOWNLOAD_TIMEOUT', '300'))
|
||||
upload_timeout = int(os.getenv('UPLOAD_TIMEOUT', '600'))
|
||||
|
||||
return cls(
|
||||
api_endpoint=api_endpoint,
|
||||
access_key=access_key,
|
||||
worker_id=worker_id,
|
||||
max_concurrency=max_concurrency,
|
||||
heartbeat_interval=heartbeat_interval,
|
||||
lease_extension_threshold=lease_extension_threshold,
|
||||
lease_extension_duration=lease_extension_duration,
|
||||
temp_dir=temp_dir,
|
||||
capabilities=capabilities,
|
||||
ffmpeg_timeout=ffmpeg_timeout,
|
||||
download_timeout=download_timeout,
|
||||
upload_timeout=upload_timeout
|
||||
)
|
||||
|
||||
def get_work_dir_path(self, task_id: str) -> str:
|
||||
"""获取任务工作目录路径"""
|
||||
return os.path.join(self.temp_dir, f"task_{task_id}")
|
||||
|
||||
def ensure_temp_dir(self) -> None:
|
||||
"""确保临时目录存在"""
|
||||
os.makedirs(self.temp_dir, exist_ok=True)
|
||||
105
domain/result.py
Normal file
105
domain/result.py
Normal file
@@ -0,0 +1,105 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务结果模型
|
||||
|
||||
定义错误码、重试配置、任务结果等数据结构。
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Dict, Any, List
|
||||
|
||||
|
||||
class ErrorCode(Enum):
|
||||
"""错误码枚举"""
|
||||
E_INPUT_UNAVAILABLE = "E_INPUT_UNAVAILABLE" # 素材不可访问/404
|
||||
E_FFMPEG_FAILED = "E_FFMPEG_FAILED" # FFmpeg 执行失败
|
||||
E_UPLOAD_FAILED = "E_UPLOAD_FAILED" # 上传失败
|
||||
E_SPEC_INVALID = "E_SPEC_INVALID" # renderSpec 非法
|
||||
E_TIMEOUT = "E_TIMEOUT" # 执行超时
|
||||
E_UNKNOWN = "E_UNKNOWN" # 未知错误
|
||||
|
||||
|
||||
# 重试配置
|
||||
RETRY_CONFIG: Dict[ErrorCode, Dict[str, Any]] = {
|
||||
ErrorCode.E_INPUT_UNAVAILABLE: {
|
||||
'max_retries': 3,
|
||||
'backoff': [1, 2, 5] # 重试间隔(秒)
|
||||
},
|
||||
ErrorCode.E_FFMPEG_FAILED: {
|
||||
'max_retries': 2,
|
||||
'backoff': [1, 3]
|
||||
},
|
||||
ErrorCode.E_UPLOAD_FAILED: {
|
||||
'max_retries': 3,
|
||||
'backoff': [1, 2, 5]
|
||||
},
|
||||
ErrorCode.E_SPEC_INVALID: {
|
||||
'max_retries': 0, # 不重试
|
||||
'backoff': []
|
||||
},
|
||||
ErrorCode.E_TIMEOUT: {
|
||||
'max_retries': 2,
|
||||
'backoff': [5, 10]
|
||||
},
|
||||
ErrorCode.E_UNKNOWN: {
|
||||
'max_retries': 1,
|
||||
'backoff': [2]
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskResult:
|
||||
"""
|
||||
任务结果
|
||||
|
||||
封装任务执行的结果,包括成功数据或失败信息。
|
||||
"""
|
||||
success: bool
|
||||
data: Optional[Dict[str, Any]] = None
|
||||
error_code: Optional[ErrorCode] = None
|
||||
error_message: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def ok(cls, data: Dict[str, Any]) -> 'TaskResult':
|
||||
"""创建成功结果"""
|
||||
return cls(success=True, data=data)
|
||||
|
||||
@classmethod
|
||||
def fail(cls, error_code: ErrorCode, error_message: str) -> 'TaskResult':
|
||||
"""创建失败结果"""
|
||||
return cls(
|
||||
success=False,
|
||||
error_code=error_code,
|
||||
error_message=error_message
|
||||
)
|
||||
|
||||
def to_report_dict(self) -> Dict[str, Any]:
|
||||
"""
|
||||
转换为上报格式
|
||||
|
||||
用于 API 上报时的数据格式转换。
|
||||
"""
|
||||
if self.success:
|
||||
return {'result': self.data}
|
||||
else:
|
||||
return {
|
||||
'errorCode': self.error_code.value if self.error_code else 'E_UNKNOWN',
|
||||
'errorMessage': self.error_message or 'Unknown error'
|
||||
}
|
||||
|
||||
def can_retry(self) -> bool:
|
||||
"""是否可以重试"""
|
||||
if self.success:
|
||||
return False
|
||||
if not self.error_code:
|
||||
return True
|
||||
config = RETRY_CONFIG.get(self.error_code, {})
|
||||
return config.get('max_retries', 0) > 0
|
||||
|
||||
def get_retry_config(self) -> Dict[str, Any]:
|
||||
"""获取重试配置"""
|
||||
if not self.error_code:
|
||||
return {'max_retries': 1, 'backoff': [2]}
|
||||
return RETRY_CONFIG.get(self.error_code, {'max_retries': 1, 'backoff': [2]})
|
||||
249
domain/task.py
Normal file
249
domain/task.py
Normal file
@@ -0,0 +1,249 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务领域模型
|
||||
|
||||
定义任务类型、任务实体、渲染规格、输出规格等数据结构。
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, Any, Optional, List
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class TaskType(Enum):
|
||||
"""任务类型枚举"""
|
||||
RENDER_SEGMENT_VIDEO = "RENDER_SEGMENT_VIDEO" # 渲染视频片段
|
||||
PREPARE_JOB_AUDIO = "PREPARE_JOB_AUDIO" # 生成全局音频
|
||||
PACKAGE_SEGMENT_TS = "PACKAGE_SEGMENT_TS" # 封装 TS 分片
|
||||
FINALIZE_MP4 = "FINALIZE_MP4" # 产出最终 MP4
|
||||
|
||||
|
||||
class TaskStatus(Enum):
|
||||
"""任务状态枚举"""
|
||||
PENDING = "PENDING"
|
||||
RUNNING = "RUNNING"
|
||||
SUCCESS = "SUCCESS"
|
||||
FAILED = "FAILED"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RenderSpec:
|
||||
"""
|
||||
渲染规格
|
||||
|
||||
用于 RENDER_SEGMENT_VIDEO 任务,定义视频渲染参数。
|
||||
"""
|
||||
crop_enable: bool = False
|
||||
crop_size: Optional[str] = None
|
||||
speed: str = "1.0"
|
||||
lut_url: Optional[str] = None
|
||||
overlay_url: Optional[str] = None
|
||||
effects: Optional[str] = None
|
||||
zoom_cut: bool = False
|
||||
video_crop: Optional[str] = None
|
||||
face_pos: Optional[str] = None
|
||||
transitions: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict]) -> 'RenderSpec':
|
||||
"""从字典创建 RenderSpec"""
|
||||
if not data:
|
||||
return cls()
|
||||
return cls(
|
||||
crop_enable=data.get('cropEnable', False),
|
||||
crop_size=data.get('cropSize'),
|
||||
speed=str(data.get('speed', '1.0')),
|
||||
lut_url=data.get('lutUrl'),
|
||||
overlay_url=data.get('overlayUrl'),
|
||||
effects=data.get('effects'),
|
||||
zoom_cut=data.get('zoomCut', False),
|
||||
video_crop=data.get('videoCrop'),
|
||||
face_pos=data.get('facePos'),
|
||||
transitions=data.get('transitions')
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputSpec:
|
||||
"""
|
||||
输出规格
|
||||
|
||||
用于 RENDER_SEGMENT_VIDEO 任务,定义视频输出参数。
|
||||
"""
|
||||
width: int = 1080
|
||||
height: int = 1920
|
||||
fps: int = 30
|
||||
bitrate: int = 4000000
|
||||
codec: str = "h264"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict]) -> 'OutputSpec':
|
||||
"""从字典创建 OutputSpec"""
|
||||
if not data:
|
||||
return cls()
|
||||
return cls(
|
||||
width=data.get('width', 1080),
|
||||
height=data.get('height', 1920),
|
||||
fps=data.get('fps', 30),
|
||||
bitrate=data.get('bitrate', 4000000),
|
||||
codec=data.get('codec', 'h264')
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioSpec:
|
||||
"""
|
||||
音频规格
|
||||
|
||||
用于 PREPARE_JOB_AUDIO 任务中的片段叠加音效。
|
||||
"""
|
||||
audio_url: Optional[str] = None
|
||||
volume: float = 1.0
|
||||
fade_in_ms: int = 10
|
||||
fade_out_ms: int = 10
|
||||
start_ms: int = 0
|
||||
delay_ms: int = 0
|
||||
loop_enable: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict]) -> Optional['AudioSpec']:
|
||||
"""从字典创建 AudioSpec"""
|
||||
if not data:
|
||||
return None
|
||||
return cls(
|
||||
audio_url=data.get('audioUrl'),
|
||||
volume=float(data.get('volume', 1.0)),
|
||||
fade_in_ms=int(data.get('fadeInMs', 10)),
|
||||
fade_out_ms=int(data.get('fadeOutMs', 10)),
|
||||
start_ms=int(data.get('startMs', 0)),
|
||||
delay_ms=int(data.get('delayMs', 0)),
|
||||
loop_enable=data.get('loopEnable', False)
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AudioProfile:
|
||||
"""
|
||||
音频配置
|
||||
|
||||
用于 PREPARE_JOB_AUDIO 任务的全局音频参数。
|
||||
"""
|
||||
sample_rate: int = 48000
|
||||
channels: int = 2
|
||||
codec: str = "aac"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict]) -> 'AudioProfile':
|
||||
"""从字典创建 AudioProfile"""
|
||||
if not data:
|
||||
return cls()
|
||||
return cls(
|
||||
sample_rate=data.get('sampleRate', 48000),
|
||||
channels=data.get('channels', 2),
|
||||
codec=data.get('codec', 'aac')
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
"""
|
||||
任务实体
|
||||
|
||||
表示一个待执行的渲染任务。
|
||||
"""
|
||||
task_id: str
|
||||
task_type: TaskType
|
||||
priority: int
|
||||
lease_expire_time: datetime
|
||||
payload: Dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict) -> 'Task':
|
||||
"""从 API 响应字典创建 Task"""
|
||||
lease_time_str = data.get('leaseExpireTime', '')
|
||||
|
||||
# 解析 ISO 8601 时间格式
|
||||
if lease_time_str:
|
||||
if lease_time_str.endswith('Z'):
|
||||
lease_time_str = lease_time_str[:-1] + '+00:00'
|
||||
try:
|
||||
lease_expire_time = datetime.fromisoformat(lease_time_str)
|
||||
except ValueError:
|
||||
# 解析失败时使用当前时间 + 5分钟
|
||||
lease_expire_time = datetime.now()
|
||||
else:
|
||||
lease_expire_time = datetime.now()
|
||||
|
||||
return cls(
|
||||
task_id=str(data['taskId']),
|
||||
task_type=TaskType(data['taskType']),
|
||||
priority=data.get('priority', 0),
|
||||
lease_expire_time=lease_expire_time,
|
||||
payload=data.get('payload', {})
|
||||
)
|
||||
|
||||
def get_job_id(self) -> str:
|
||||
"""获取作业 ID"""
|
||||
return str(self.payload.get('jobId', ''))
|
||||
|
||||
def get_segment_id(self) -> Optional[str]:
|
||||
"""获取片段 ID(如果有)"""
|
||||
segment_id = self.payload.get('segmentId')
|
||||
return str(segment_id) if segment_id else None
|
||||
|
||||
def get_plan_segment_index(self) -> int:
|
||||
"""获取计划片段索引"""
|
||||
return int(self.payload.get('planSegmentIndex', 0))
|
||||
|
||||
def get_duration_ms(self) -> int:
|
||||
"""获取时长(毫秒)"""
|
||||
return int(self.payload.get('durationMs', 5000))
|
||||
|
||||
def get_material_url(self) -> Optional[str]:
|
||||
"""获取素材 URL"""
|
||||
return self.payload.get('boundMaterialUrl') or self.payload.get('sourceRef')
|
||||
|
||||
def get_render_spec(self) -> RenderSpec:
|
||||
"""获取渲染规格"""
|
||||
return RenderSpec.from_dict(self.payload.get('renderSpec'))
|
||||
|
||||
def get_output_spec(self) -> OutputSpec:
|
||||
"""获取输出规格"""
|
||||
return OutputSpec.from_dict(self.payload.get('output'))
|
||||
|
||||
def get_bgm_url(self) -> Optional[str]:
|
||||
"""获取 BGM URL"""
|
||||
return self.payload.get('bgmUrl')
|
||||
|
||||
def get_total_duration_ms(self) -> int:
|
||||
"""获取总时长(毫秒)"""
|
||||
return int(self.payload.get('totalDurationMs', 0))
|
||||
|
||||
def get_segments(self) -> List[Dict]:
|
||||
"""获取片段列表"""
|
||||
return self.payload.get('segments', [])
|
||||
|
||||
def get_audio_profile(self) -> AudioProfile:
|
||||
"""获取音频配置"""
|
||||
return AudioProfile.from_dict(self.payload.get('audioProfile'))
|
||||
|
||||
def get_video_url(self) -> Optional[str]:
|
||||
"""获取视频 URL(用于 PACKAGE_SEGMENT_TS)"""
|
||||
return self.payload.get('videoUrl')
|
||||
|
||||
def get_audio_url(self) -> Optional[str]:
|
||||
"""获取音频 URL(用于 PACKAGE_SEGMENT_TS)"""
|
||||
return self.payload.get('audioUrl')
|
||||
|
||||
def get_start_time_ms(self) -> int:
|
||||
"""获取开始时间(毫秒)"""
|
||||
return int(self.payload.get('startTimeMs', 0))
|
||||
|
||||
def get_m3u8_url(self) -> Optional[str]:
|
||||
"""获取 m3u8 URL(用于 FINALIZE_MP4)"""
|
||||
return self.payload.get('m3u8Url')
|
||||
|
||||
def get_ts_list(self) -> List[str]:
|
||||
"""获取 TS 列表(用于 FINALIZE_MP4)"""
|
||||
return self.payload.get('tsList', [])
|
||||
20
handlers/__init__.py
Normal file
20
handlers/__init__.py
Normal file
@@ -0,0 +1,20 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务处理器层
|
||||
|
||||
包含各种任务类型的具体处理器实现。
|
||||
"""
|
||||
|
||||
from handlers.base import BaseHandler
|
||||
from handlers.render_video import RenderSegmentVideoHandler
|
||||
from handlers.prepare_audio import PrepareJobAudioHandler
|
||||
from handlers.package_ts import PackageSegmentTsHandler
|
||||
from handlers.finalize_mp4 import FinalizeMp4Handler
|
||||
|
||||
__all__ = [
|
||||
'BaseHandler',
|
||||
'RenderSegmentVideoHandler',
|
||||
'PrepareJobAudioHandler',
|
||||
'PackageSegmentTsHandler',
|
||||
'FinalizeMp4Handler',
|
||||
]
|
||||
280
handlers/base.py
Normal file
280
handlers/base.py
Normal file
@@ -0,0 +1,280 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务处理器基类
|
||||
|
||||
提供所有处理器共用的基础功能。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
import shutil
|
||||
import tempfile
|
||||
import subprocess
|
||||
from abc import ABC
|
||||
from typing import Optional, List, TYPE_CHECKING
|
||||
|
||||
from core.handler import TaskHandler
|
||||
from domain.task import Task
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
from domain.config import WorkerConfig
|
||||
from util import oss
|
||||
from util.ffmpeg import subprocess_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from services.api_client import APIClientV2
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# v2 统一视频编码参数(来自集成文档)
|
||||
VIDEO_ENCODE_ARGS = [
|
||||
'-c:v', 'libx264',
|
||||
'-preset', 'medium',
|
||||
'-profile:v', 'main',
|
||||
'-level', '4.0',
|
||||
'-crf', '23',
|
||||
'-pix_fmt', 'yuv420p',
|
||||
]
|
||||
|
||||
# v2 统一音频编码参数
|
||||
AUDIO_ENCODE_ARGS = [
|
||||
'-c:a', 'aac',
|
||||
'-b:a', '128k',
|
||||
'-ar', '48000',
|
||||
'-ac', '2',
|
||||
]
|
||||
|
||||
|
||||
class BaseHandler(TaskHandler, ABC):
|
||||
"""
|
||||
任务处理器基类
|
||||
|
||||
提供所有处理器共用的基础功能,包括:
|
||||
- 临时目录管理
|
||||
- 文件下载/上传
|
||||
- FFmpeg 命令执行
|
||||
- 日志记录
|
||||
"""
|
||||
|
||||
def __init__(self, config: WorkerConfig, api_client: 'APIClientV2'):
|
||||
"""
|
||||
初始化处理器
|
||||
|
||||
Args:
|
||||
config: Worker 配置
|
||||
api_client: API 客户端
|
||||
"""
|
||||
self.config = config
|
||||
self.api_client = api_client
|
||||
|
||||
def before_handle(self, task: Task) -> None:
|
||||
"""处理前钩子"""
|
||||
logger.debug(f"[task:{task.task_id}] Before handle: {task.task_type.value}")
|
||||
|
||||
def after_handle(self, task: Task, result: TaskResult) -> None:
|
||||
"""处理后钩子"""
|
||||
status = "success" if result.success else "failed"
|
||||
logger.debug(f"[task:{task.task_id}] After handle: {status}")
|
||||
|
||||
def create_work_dir(self, task_id: str = None) -> str:
|
||||
"""
|
||||
创建临时工作目录
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID(用于目录命名)
|
||||
|
||||
Returns:
|
||||
工作目录路径
|
||||
"""
|
||||
# 确保临时根目录存在
|
||||
os.makedirs(self.config.temp_dir, exist_ok=True)
|
||||
|
||||
# 创建唯一的工作目录
|
||||
prefix = f"task_{task_id}_" if task_id else "task_"
|
||||
work_dir = tempfile.mkdtemp(dir=self.config.temp_dir, prefix=prefix)
|
||||
|
||||
logger.debug(f"Created work directory: {work_dir}")
|
||||
return work_dir
|
||||
|
||||
def cleanup_work_dir(self, work_dir: str) -> None:
|
||||
"""
|
||||
清理临时工作目录
|
||||
|
||||
Args:
|
||||
work_dir: 工作目录路径
|
||||
"""
|
||||
if not work_dir or not os.path.exists(work_dir):
|
||||
return
|
||||
|
||||
try:
|
||||
shutil.rmtree(work_dir)
|
||||
logger.debug(f"Cleaned up work directory: {work_dir}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to cleanup work directory {work_dir}: {e}")
|
||||
|
||||
def download_file(self, url: str, dest: str, timeout: int = None) -> bool:
|
||||
"""
|
||||
下载文件
|
||||
|
||||
Args:
|
||||
url: 文件 URL
|
||||
dest: 目标路径
|
||||
timeout: 超时时间(秒)
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
if timeout is None:
|
||||
timeout = self.config.download_timeout
|
||||
|
||||
try:
|
||||
result = oss.download_from_oss(url, dest)
|
||||
if result:
|
||||
file_size = os.path.getsize(dest) if os.path.exists(dest) else 0
|
||||
logger.debug(f"Downloaded: {url} -> {dest} ({file_size} bytes)")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Download failed: {url} -> {e}")
|
||||
return False
|
||||
|
||||
def upload_file(
|
||||
self,
|
||||
task_id: str,
|
||||
file_type: str,
|
||||
file_path: str,
|
||||
file_name: str = None
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
上传文件并返回访问 URL
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
file_type: 文件类型(video/audio/ts/mp4)
|
||||
file_path: 本地文件路径
|
||||
file_name: 文件名(可选)
|
||||
|
||||
Returns:
|
||||
访问 URL,失败返回 None
|
||||
"""
|
||||
# 获取上传 URL
|
||||
upload_info = self.api_client.get_upload_url(task_id, file_type, file_name)
|
||||
if not upload_info:
|
||||
logger.error(f"[task:{task_id}] Failed to get upload URL")
|
||||
return None
|
||||
|
||||
upload_url = upload_info.get('uploadUrl')
|
||||
access_url = upload_info.get('accessUrl')
|
||||
|
||||
if not upload_url:
|
||||
logger.error(f"[task:{task_id}] Invalid upload URL response")
|
||||
return None
|
||||
|
||||
# 上传文件
|
||||
try:
|
||||
result = oss.upload_to_oss(upload_url, file_path)
|
||||
if result:
|
||||
file_size = os.path.getsize(file_path)
|
||||
logger.info(f"[task:{task_id}] Uploaded: {file_path} ({file_size} bytes)")
|
||||
return access_url
|
||||
else:
|
||||
logger.error(f"[task:{task_id}] Upload failed: {file_path}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Upload error: {e}")
|
||||
return None
|
||||
|
||||
def run_ffmpeg(
|
||||
self,
|
||||
cmd: List[str],
|
||||
task_id: str,
|
||||
timeout: int = None
|
||||
) -> bool:
|
||||
"""
|
||||
执行 FFmpeg 命令
|
||||
|
||||
Args:
|
||||
cmd: FFmpeg 命令参数列表
|
||||
task_id: 任务 ID(用于日志)
|
||||
timeout: 超时时间(秒)
|
||||
|
||||
Returns:
|
||||
是否成功
|
||||
"""
|
||||
if timeout is None:
|
||||
timeout = self.config.ffmpeg_timeout
|
||||
|
||||
# 日志记录命令(限制长度)
|
||||
cmd_str = ' '.join(cmd)
|
||||
if len(cmd_str) > 500:
|
||||
cmd_str = cmd_str[:500] + '...'
|
||||
logger.info(f"[task:{task_id}] FFmpeg: {cmd_str}")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
timeout=timeout,
|
||||
**subprocess_args(False)
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
stderr = result.stderr.decode('utf-8', errors='replace')[:1000]
|
||||
logger.error(f"[task:{task_id}] FFmpeg failed (code={result.returncode}): {stderr}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.error(f"[task:{task_id}] FFmpeg timeout after {timeout}s")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] FFmpeg error: {e}")
|
||||
return False
|
||||
|
||||
def probe_duration(self, file_path: str) -> Optional[float]:
|
||||
"""
|
||||
探测媒体文件时长
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
|
||||
Returns:
|
||||
时长(秒),失败返回 None
|
||||
"""
|
||||
try:
|
||||
from util.ffmpeg import probe_video_info
|
||||
_, _, duration = probe_video_info(file_path)
|
||||
return float(duration) if duration else None
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to probe duration: {file_path} -> {e}")
|
||||
return None
|
||||
|
||||
def get_file_size(self, file_path: str) -> int:
|
||||
"""
|
||||
获取文件大小
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
|
||||
Returns:
|
||||
文件大小(字节)
|
||||
"""
|
||||
try:
|
||||
return os.path.getsize(file_path)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def ensure_file_exists(self, file_path: str, min_size: int = 0) -> bool:
|
||||
"""
|
||||
确保文件存在且大小满足要求
|
||||
|
||||
Args:
|
||||
file_path: 文件路径
|
||||
min_size: 最小大小(字节)
|
||||
|
||||
Returns:
|
||||
是否满足要求
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
return False
|
||||
return os.path.getsize(file_path) >= min_size
|
||||
190
handlers/finalize_mp4.py
Normal file
190
handlers/finalize_mp4.py
Normal file
@@ -0,0 +1,190 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
最终 MP4 合并处理器
|
||||
|
||||
处理 FINALIZE_MP4 任务,将所有 TS 分片合并为最终可下载的 MP4 文件。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from handlers.base import BaseHandler
|
||||
from domain.task import Task, TaskType
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FinalizeMp4Handler(BaseHandler):
|
||||
"""
|
||||
最终 MP4 合并处理器
|
||||
|
||||
职责:
|
||||
- 下载所有 TS 分片
|
||||
- 使用 concat demuxer 合并
|
||||
- 产出最终 MP4(remux,不重编码)
|
||||
- 上传 MP4 产物
|
||||
|
||||
关键约束:
|
||||
- 优先使用 remux(复制流,不重新编码)
|
||||
- 使用 aac_adtstoasc bitstream filter 处理音频
|
||||
"""
|
||||
|
||||
def get_supported_type(self) -> TaskType:
|
||||
return TaskType.FINALIZE_MP4
|
||||
|
||||
def handle(self, task: Task) -> TaskResult:
|
||||
"""处理 MP4 合并任务"""
|
||||
work_dir = self.create_work_dir(task.task_id)
|
||||
|
||||
try:
|
||||
# 获取 TS 列表
|
||||
ts_list = task.get_ts_list()
|
||||
m3u8_url = task.get_m3u8_url()
|
||||
|
||||
if not ts_list and not m3u8_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_SPEC_INVALID,
|
||||
"Missing tsList or m3u8Url"
|
||||
)
|
||||
|
||||
output_file = os.path.join(work_dir, 'final.mp4')
|
||||
|
||||
if ts_list:
|
||||
# 方式1:使用 TS 列表
|
||||
result = self._process_ts_list(task, work_dir, ts_list, output_file)
|
||||
else:
|
||||
# 方式2:使用 m3u8 URL
|
||||
result = self._process_m3u8(task, work_dir, m3u8_url, output_file)
|
||||
|
||||
if not result.success:
|
||||
return result
|
||||
|
||||
# 验证输出文件
|
||||
if not self.ensure_file_exists(output_file, min_size=4096):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"MP4 output file is missing or too small"
|
||||
)
|
||||
|
||||
# 获取文件大小
|
||||
file_size = self.get_file_size(output_file)
|
||||
|
||||
# 上传产物
|
||||
mp4_url = self.upload_file(task.task_id, 'mp4', output_file)
|
||||
if not mp4_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_UPLOAD_FAILED,
|
||||
"Failed to upload MP4"
|
||||
)
|
||||
|
||||
return TaskResult.ok({
|
||||
'mp4Url': mp4_url,
|
||||
'fileSizeBytes': file_size
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task.task_id}] Unexpected error: {e}", exc_info=True)
|
||||
return TaskResult.fail(ErrorCode.E_UNKNOWN, str(e))
|
||||
|
||||
finally:
|
||||
self.cleanup_work_dir(work_dir)
|
||||
|
||||
def _process_ts_list(
|
||||
self,
|
||||
task: Task,
|
||||
work_dir: str,
|
||||
ts_list: List[str],
|
||||
output_file: str
|
||||
) -> TaskResult:
|
||||
"""
|
||||
使用 TS 列表处理
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
work_dir: 工作目录
|
||||
ts_list: TS URL 列表
|
||||
output_file: 输出文件路径
|
||||
|
||||
Returns:
|
||||
TaskResult
|
||||
"""
|
||||
# 1. 下载所有 TS 分片
|
||||
ts_files = []
|
||||
for i, ts_url in enumerate(ts_list):
|
||||
ts_file = os.path.join(work_dir, f'seg_{i}.ts')
|
||||
if not self.download_file(ts_url, ts_file):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_INPUT_UNAVAILABLE,
|
||||
f"Failed to download TS segment {i}: {ts_url}"
|
||||
)
|
||||
ts_files.append(ts_file)
|
||||
|
||||
logger.info(f"[task:{task.task_id}] Downloaded {len(ts_files)} TS segments")
|
||||
|
||||
# 2. 创建 concat 文件列表
|
||||
concat_file = os.path.join(work_dir, 'concat.txt')
|
||||
with open(concat_file, 'w', encoding='utf-8') as f:
|
||||
for ts_file in ts_files:
|
||||
# 路径中的反斜杠需要转义或使用正斜杠
|
||||
ts_path = ts_file.replace('\\', '/')
|
||||
f.write(f"file '{ts_path}'\n")
|
||||
|
||||
# 3. 构建合并命令(remux,不重编码)
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-hide_banner',
|
||||
'-f', 'concat',
|
||||
'-safe', '0',
|
||||
'-i', concat_file,
|
||||
'-c', 'copy', # 复制流,不重编码
|
||||
'-bsf:a', 'aac_adtstoasc', # 音频 bitstream filter
|
||||
output_file
|
||||
]
|
||||
|
||||
# 4. 执行 FFmpeg
|
||||
if not self.run_ffmpeg(cmd, task.task_id):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"MP4 concatenation failed"
|
||||
)
|
||||
|
||||
return TaskResult.ok({})
|
||||
|
||||
def _process_m3u8(
|
||||
self,
|
||||
task: Task,
|
||||
work_dir: str,
|
||||
m3u8_url: str,
|
||||
output_file: str
|
||||
) -> TaskResult:
|
||||
"""
|
||||
使用 m3u8 URL 处理
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
work_dir: 工作目录
|
||||
m3u8_url: m3u8 URL
|
||||
output_file: 输出文件路径
|
||||
|
||||
Returns:
|
||||
TaskResult
|
||||
"""
|
||||
# 构建命令
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-hide_banner',
|
||||
'-protocol_whitelist', 'file,http,https,tcp,tls',
|
||||
'-i', m3u8_url,
|
||||
'-c', 'copy',
|
||||
'-bsf:a', 'aac_adtstoasc',
|
||||
output_file
|
||||
]
|
||||
|
||||
# 执行 FFmpeg
|
||||
if not self.run_ffmpeg(cmd, task.task_id):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"MP4 conversion from m3u8 failed"
|
||||
)
|
||||
|
||||
return TaskResult.ok({})
|
||||
175
handlers/package_ts.py
Normal file
175
handlers/package_ts.py
Normal file
@@ -0,0 +1,175 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
TS 分片封装处理器
|
||||
|
||||
处理 PACKAGE_SEGMENT_TS 任务,将视频片段和对应时间区间的音频封装为 TS 分片。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from handlers.base import BaseHandler
|
||||
from domain.task import Task, TaskType
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PackageSegmentTsHandler(BaseHandler):
|
||||
"""
|
||||
TS 分片封装处理器
|
||||
|
||||
职责:
|
||||
- 下载视频片段
|
||||
- 下载全局音频
|
||||
- 截取对应时间区间的音频
|
||||
- 封装为 TS 分片
|
||||
- 上传 TS 产物
|
||||
|
||||
关键约束:
|
||||
- TS 必须包含音视频同轨
|
||||
- 使用 output_ts_offset 保证时间戳连续
|
||||
- 输出 extinfDurationSec 供 m3u8 使用
|
||||
"""
|
||||
|
||||
def get_supported_type(self) -> TaskType:
|
||||
return TaskType.PACKAGE_SEGMENT_TS
|
||||
|
||||
def handle(self, task: Task) -> TaskResult:
|
||||
"""处理 TS 封装任务"""
|
||||
work_dir = self.create_work_dir(task.task_id)
|
||||
|
||||
try:
|
||||
# 解析参数
|
||||
video_url = task.get_video_url()
|
||||
audio_url = task.get_audio_url()
|
||||
start_time_ms = task.get_start_time_ms()
|
||||
duration_ms = task.get_duration_ms()
|
||||
|
||||
if not video_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_SPEC_INVALID,
|
||||
"Missing videoUrl"
|
||||
)
|
||||
|
||||
if not audio_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_SPEC_INVALID,
|
||||
"Missing audioUrl"
|
||||
)
|
||||
|
||||
# 计算时间参数
|
||||
start_sec = start_time_ms / 1000.0
|
||||
duration_sec = duration_ms / 1000.0
|
||||
|
||||
# 1. 下载视频片段
|
||||
video_file = os.path.join(work_dir, 'video.mp4')
|
||||
if not self.download_file(video_url, video_file):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_INPUT_UNAVAILABLE,
|
||||
f"Failed to download video: {video_url}"
|
||||
)
|
||||
|
||||
# 2. 下载全局音频
|
||||
audio_file = os.path.join(work_dir, 'audio.aac')
|
||||
if not self.download_file(audio_url, audio_file):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_INPUT_UNAVAILABLE,
|
||||
f"Failed to download audio: {audio_url}"
|
||||
)
|
||||
|
||||
# 3. 构建 TS 封装命令
|
||||
output_file = os.path.join(work_dir, 'segment.ts')
|
||||
cmd = self._build_command(
|
||||
video_file=video_file,
|
||||
audio_file=audio_file,
|
||||
output_file=output_file,
|
||||
start_sec=start_sec,
|
||||
duration_sec=duration_sec
|
||||
)
|
||||
|
||||
# 4. 执行 FFmpeg
|
||||
if not self.run_ffmpeg(cmd, task.task_id):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"TS packaging failed"
|
||||
)
|
||||
|
||||
# 5. 验证输出文件
|
||||
if not self.ensure_file_exists(output_file, min_size=1024):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"TS output file is missing or too small"
|
||||
)
|
||||
|
||||
# 6. 获取实际时长(用于 EXTINF)
|
||||
actual_duration = self.probe_duration(output_file)
|
||||
extinf_duration = actual_duration if actual_duration else duration_sec
|
||||
|
||||
# 7. 上传产物
|
||||
ts_url = self.upload_file(task.task_id, 'ts', output_file)
|
||||
if not ts_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_UPLOAD_FAILED,
|
||||
"Failed to upload TS"
|
||||
)
|
||||
|
||||
return TaskResult.ok({
|
||||
'tsUrl': ts_url,
|
||||
'extinfDurationSec': extinf_duration
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task.task_id}] Unexpected error: {e}", exc_info=True)
|
||||
return TaskResult.fail(ErrorCode.E_UNKNOWN, str(e))
|
||||
|
||||
finally:
|
||||
self.cleanup_work_dir(work_dir)
|
||||
|
||||
def _build_command(
|
||||
self,
|
||||
video_file: str,
|
||||
audio_file: str,
|
||||
output_file: str,
|
||||
start_sec: float,
|
||||
duration_sec: float
|
||||
) -> List[str]:
|
||||
"""
|
||||
构建 TS 封装命令
|
||||
|
||||
Args:
|
||||
video_file: 视频文件路径
|
||||
audio_file: 音频文件路径
|
||||
output_file: 输出文件路径
|
||||
start_sec: 开始时间(秒)
|
||||
duration_sec: 时长(秒)
|
||||
|
||||
Returns:
|
||||
FFmpeg 命令参数列表
|
||||
"""
|
||||
cmd = [
|
||||
'ffmpeg', '-y', '-hide_banner',
|
||||
# 视频输入
|
||||
'-i', video_file,
|
||||
# 音频输入(从 start_sec 开始截取 duration_sec)
|
||||
'-ss', str(start_sec),
|
||||
'-t', str(duration_sec),
|
||||
'-i', audio_file,
|
||||
# 映射流
|
||||
'-map', '0:v:0', # 使用第一个输入的视频流
|
||||
'-map', '1:a:0', # 使用第二个输入的音频流
|
||||
# 复制编码(不重新编码)
|
||||
'-c:v', 'copy',
|
||||
'-c:a', 'copy',
|
||||
# 关键:时间戳偏移,保证整体连续
|
||||
'-output_ts_offset', str(start_sec),
|
||||
# 复用参数
|
||||
'-muxdelay', '0',
|
||||
'-muxpreload', '0',
|
||||
# 输出格式
|
||||
'-f', 'mpegts',
|
||||
output_file
|
||||
]
|
||||
|
||||
return cmd
|
||||
251
handlers/prepare_audio.py
Normal file
251
handlers/prepare_audio.py
Normal file
@@ -0,0 +1,251 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
全局音频准备处理器
|
||||
|
||||
处理 PREPARE_JOB_AUDIO 任务,生成整个视频的连续音频轨道。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
from handlers.base import BaseHandler, AUDIO_ENCODE_ARGS
|
||||
from domain.task import Task, TaskType, AudioSpec, AudioProfile
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PrepareJobAudioHandler(BaseHandler):
|
||||
"""
|
||||
全局音频准备处理器
|
||||
|
||||
职责:
|
||||
- 下载全局 BGM
|
||||
- 下载各片段叠加音效
|
||||
- 构建复杂混音命令
|
||||
- 执行混音
|
||||
- 上传音频产物
|
||||
|
||||
关键约束:
|
||||
- 全局 BGM 连续生成一次,贯穿整个时长
|
||||
- 禁止使用 amix normalize=1
|
||||
- 只对叠加音轨做极短淡入淡出(5-20ms)
|
||||
- 不对 BGM 做边界 fade
|
||||
"""
|
||||
|
||||
def get_supported_type(self) -> TaskType:
|
||||
return TaskType.PREPARE_JOB_AUDIO
|
||||
|
||||
def handle(self, task: Task) -> TaskResult:
|
||||
"""处理音频准备任务"""
|
||||
work_dir = self.create_work_dir(task.task_id)
|
||||
|
||||
try:
|
||||
# 解析参数
|
||||
total_duration_ms = task.get_total_duration_ms()
|
||||
if total_duration_ms <= 0:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_SPEC_INVALID,
|
||||
"Invalid totalDurationMs"
|
||||
)
|
||||
|
||||
total_duration_sec = total_duration_ms / 1000.0
|
||||
audio_profile = task.get_audio_profile()
|
||||
bgm_url = task.get_bgm_url()
|
||||
segments = task.get_segments()
|
||||
|
||||
# 1. 下载 BGM(如有)
|
||||
bgm_file = None
|
||||
if bgm_url:
|
||||
bgm_file = os.path.join(work_dir, 'bgm.mp3')
|
||||
if not self.download_file(bgm_url, bgm_file):
|
||||
logger.warning(f"[task:{task.task_id}] Failed to download BGM")
|
||||
bgm_file = None
|
||||
|
||||
# 2. 下载叠加音效
|
||||
sfx_files = []
|
||||
for i, seg in enumerate(segments):
|
||||
audio_spec_data = seg.get('audioSpecJson')
|
||||
if audio_spec_data:
|
||||
audio_spec = AudioSpec.from_dict(audio_spec_data)
|
||||
if audio_spec and audio_spec.audio_url:
|
||||
sfx_file = os.path.join(work_dir, f'sfx_{i}.mp3')
|
||||
if self.download_file(audio_spec.audio_url, sfx_file):
|
||||
sfx_files.append({
|
||||
'file': sfx_file,
|
||||
'spec': audio_spec,
|
||||
'segment': seg
|
||||
})
|
||||
else:
|
||||
logger.warning(f"[task:{task.task_id}] Failed to download SFX {i}")
|
||||
|
||||
# 3. 构建音频混音命令
|
||||
output_file = os.path.join(work_dir, 'audio_full.aac')
|
||||
cmd = self._build_audio_command(
|
||||
bgm_file=bgm_file,
|
||||
sfx_files=sfx_files,
|
||||
output_file=output_file,
|
||||
total_duration_sec=total_duration_sec,
|
||||
audio_profile=audio_profile
|
||||
)
|
||||
|
||||
# 4. 执行 FFmpeg
|
||||
if not self.run_ffmpeg(cmd, task.task_id):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"Audio mixing failed"
|
||||
)
|
||||
|
||||
# 5. 验证输出文件
|
||||
if not self.ensure_file_exists(output_file, min_size=1024):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"Audio output file is missing or too small"
|
||||
)
|
||||
|
||||
# 6. 上传产物
|
||||
audio_url = self.upload_file(task.task_id, 'audio', output_file)
|
||||
if not audio_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_UPLOAD_FAILED,
|
||||
"Failed to upload audio"
|
||||
)
|
||||
|
||||
return TaskResult.ok({
|
||||
'audioUrl': audio_url
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task.task_id}] Unexpected error: {e}", exc_info=True)
|
||||
return TaskResult.fail(ErrorCode.E_UNKNOWN, str(e))
|
||||
|
||||
finally:
|
||||
self.cleanup_work_dir(work_dir)
|
||||
|
||||
def _build_audio_command(
|
||||
self,
|
||||
bgm_file: Optional[str],
|
||||
sfx_files: List[Dict],
|
||||
output_file: str,
|
||||
total_duration_sec: float,
|
||||
audio_profile: AudioProfile
|
||||
) -> List[str]:
|
||||
"""
|
||||
构建音频混音命令
|
||||
|
||||
Args:
|
||||
bgm_file: BGM 文件路径(可选)
|
||||
sfx_files: 叠加音效列表
|
||||
output_file: 输出文件路径
|
||||
total_duration_sec: 总时长(秒)
|
||||
audio_profile: 音频配置
|
||||
|
||||
Returns:
|
||||
FFmpeg 命令参数列表
|
||||
"""
|
||||
sample_rate = audio_profile.sample_rate
|
||||
channels = audio_profile.channels
|
||||
|
||||
# 情况1:无 BGM 也无叠加音效 -> 生成静音
|
||||
if not bgm_file and not sfx_files:
|
||||
return [
|
||||
'ffmpeg', '-y', '-hide_banner',
|
||||
'-f', 'lavfi',
|
||||
'-i', f'anullsrc=r={sample_rate}:cl=stereo',
|
||||
'-t', str(total_duration_sec),
|
||||
'-c:a', 'aac', '-b:a', '128k',
|
||||
output_file
|
||||
]
|
||||
|
||||
# 情况2:仅 BGM,无叠加音效
|
||||
if not sfx_files:
|
||||
return [
|
||||
'ffmpeg', '-y', '-hide_banner',
|
||||
'-i', bgm_file,
|
||||
'-t', str(total_duration_sec),
|
||||
'-c:a', 'aac', '-b:a', '128k',
|
||||
'-ar', str(sample_rate), '-ac', str(channels),
|
||||
output_file
|
||||
]
|
||||
|
||||
# 情况3:BGM + 叠加音效 -> 复杂滤镜
|
||||
inputs = []
|
||||
if bgm_file:
|
||||
inputs.extend(['-i', bgm_file])
|
||||
for sfx in sfx_files:
|
||||
inputs.extend(['-i', sfx['file']])
|
||||
|
||||
filter_parts = []
|
||||
input_idx = 0
|
||||
|
||||
# BGM 处理(或生成静音底轨)
|
||||
if bgm_file:
|
||||
filter_parts.append(
|
||||
f"[0:a]atrim=0:{total_duration_sec},asetpts=PTS-STARTPTS,"
|
||||
f"apad=whole_dur={total_duration_sec}[bgm]"
|
||||
)
|
||||
input_idx = 1
|
||||
else:
|
||||
filter_parts.append(
|
||||
f"anullsrc=r={sample_rate}:cl=stereo,"
|
||||
f"atrim=0:{total_duration_sec}[bgm]"
|
||||
)
|
||||
input_idx = 0
|
||||
|
||||
# 叠加音效处理
|
||||
sfx_labels = []
|
||||
for i, sfx in enumerate(sfx_files):
|
||||
idx = input_idx + i
|
||||
spec = sfx['spec']
|
||||
seg = sfx['segment']
|
||||
|
||||
# 计算时间参数
|
||||
start_time_ms = seg.get('startTimeMs', 0)
|
||||
duration_ms = seg.get('durationMs', 5000)
|
||||
delay_ms = start_time_ms + spec.delay_ms
|
||||
delay_sec = delay_ms / 1000.0
|
||||
duration_sec = duration_ms / 1000.0
|
||||
|
||||
# 淡入淡出参数(极短,5-20ms)
|
||||
fade_in_sec = spec.fade_in_ms / 1000.0
|
||||
fade_out_sec = spec.fade_out_ms / 1000.0
|
||||
|
||||
# 音量
|
||||
volume = spec.volume
|
||||
|
||||
label = f"sfx{i}"
|
||||
sfx_labels.append(f"[{label}]")
|
||||
|
||||
# 构建滤镜:延迟 + 淡入淡出 + 音量
|
||||
# 注意:只对叠加音轨做淡入淡出,不对 BGM 做
|
||||
sfx_filter = (
|
||||
f"[{idx}:a]"
|
||||
f"adelay={int(delay_ms)}|{int(delay_ms)},"
|
||||
f"afade=t=in:st={delay_sec}:d={fade_in_sec},"
|
||||
f"afade=t=out:st={delay_sec + duration_sec - fade_out_sec}:d={fade_out_sec},"
|
||||
f"volume={volume}"
|
||||
f"[{label}]"
|
||||
)
|
||||
filter_parts.append(sfx_filter)
|
||||
|
||||
# 混音(关键:normalize=0,禁止归一化)
|
||||
# dropout_transition=0 表示输入结束时不做渐变
|
||||
mix_inputs = "[bgm]" + "".join(sfx_labels)
|
||||
num_inputs = 1 + len(sfx_files)
|
||||
filter_parts.append(
|
||||
f"{mix_inputs}amix=inputs={num_inputs}:duration=first:"
|
||||
f"dropout_transition=0:normalize=0[out]"
|
||||
)
|
||||
|
||||
filter_complex = ';'.join(filter_parts)
|
||||
|
||||
cmd = ['ffmpeg', '-y', '-hide_banner'] + inputs + [
|
||||
'-filter_complex', filter_complex,
|
||||
'-map', '[out]',
|
||||
'-c:a', 'aac', '-b:a', '128k',
|
||||
'-ar', str(sample_rate), '-ac', str(channels),
|
||||
output_file
|
||||
]
|
||||
|
||||
return cmd
|
||||
274
handlers/render_video.py
Normal file
274
handlers/render_video.py
Normal file
@@ -0,0 +1,274 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
视频片段渲染处理器
|
||||
|
||||
处理 RENDER_SEGMENT_VIDEO 任务,将原素材渲染为符合输出规格的视频片段。
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from handlers.base import BaseHandler, VIDEO_ENCODE_ARGS
|
||||
from domain.task import Task, TaskType, RenderSpec, OutputSpec
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RenderSegmentVideoHandler(BaseHandler):
|
||||
"""
|
||||
视频片段渲染处理器
|
||||
|
||||
职责:
|
||||
- 下载素材文件
|
||||
- 下载 LUT 文件(如有)
|
||||
- 下载叠加层(如有)
|
||||
- 构建 FFmpeg 渲染命令
|
||||
- 执行渲染
|
||||
- 上传产物
|
||||
"""
|
||||
|
||||
def get_supported_type(self) -> TaskType:
|
||||
return TaskType.RENDER_SEGMENT_VIDEO
|
||||
|
||||
def handle(self, task: Task) -> TaskResult:
|
||||
"""处理视频渲染任务"""
|
||||
work_dir = self.create_work_dir(task.task_id)
|
||||
|
||||
try:
|
||||
# 解析参数
|
||||
material_url = task.get_material_url()
|
||||
if not material_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_SPEC_INVALID,
|
||||
"Missing material URL (boundMaterialUrl or sourceRef)"
|
||||
)
|
||||
|
||||
render_spec = task.get_render_spec()
|
||||
output_spec = task.get_output_spec()
|
||||
duration_ms = task.get_duration_ms()
|
||||
|
||||
# 1. 下载素材
|
||||
input_file = os.path.join(work_dir, 'input.mp4')
|
||||
if not self.download_file(material_url, input_file):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_INPUT_UNAVAILABLE,
|
||||
f"Failed to download material: {material_url}"
|
||||
)
|
||||
|
||||
# 2. 下载 LUT(如有)
|
||||
lut_file = None
|
||||
if render_spec.lut_url:
|
||||
lut_file = os.path.join(work_dir, 'lut.cube')
|
||||
if not self.download_file(render_spec.lut_url, lut_file):
|
||||
logger.warning(f"[task:{task.task_id}] Failed to download LUT, continuing without it")
|
||||
lut_file = None
|
||||
|
||||
# 3. 下载叠加层(如有)
|
||||
overlay_file = None
|
||||
if render_spec.overlay_url:
|
||||
# 根据 URL 后缀确定文件扩展名
|
||||
ext = '.png'
|
||||
if render_spec.overlay_url.lower().endswith('.jpg') or render_spec.overlay_url.lower().endswith('.jpeg'):
|
||||
ext = '.jpg'
|
||||
overlay_file = os.path.join(work_dir, f'overlay{ext}')
|
||||
if not self.download_file(render_spec.overlay_url, overlay_file):
|
||||
logger.warning(f"[task:{task.task_id}] Failed to download overlay, continuing without it")
|
||||
overlay_file = None
|
||||
|
||||
# 4. 构建 FFmpeg 命令
|
||||
output_file = os.path.join(work_dir, 'output.mp4')
|
||||
cmd = self._build_command(
|
||||
input_file=input_file,
|
||||
output_file=output_file,
|
||||
render_spec=render_spec,
|
||||
output_spec=output_spec,
|
||||
duration_ms=duration_ms,
|
||||
lut_file=lut_file,
|
||||
overlay_file=overlay_file
|
||||
)
|
||||
|
||||
# 5. 执行 FFmpeg
|
||||
if not self.run_ffmpeg(cmd, task.task_id):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"FFmpeg rendering failed"
|
||||
)
|
||||
|
||||
# 6. 验证输出文件
|
||||
if not self.ensure_file_exists(output_file, min_size=4096):
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_FFMPEG_FAILED,
|
||||
"Output file is missing or too small"
|
||||
)
|
||||
|
||||
# 7. 获取实际时长
|
||||
actual_duration = self.probe_duration(output_file)
|
||||
actual_duration_ms = int(actual_duration * 1000) if actual_duration else duration_ms
|
||||
|
||||
# 8. 上传产物
|
||||
video_url = self.upload_file(task.task_id, 'video', output_file)
|
||||
if not video_url:
|
||||
return TaskResult.fail(
|
||||
ErrorCode.E_UPLOAD_FAILED,
|
||||
"Failed to upload video"
|
||||
)
|
||||
|
||||
return TaskResult.ok({
|
||||
'videoUrl': video_url,
|
||||
'actualDurationMs': actual_duration_ms
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task.task_id}] Unexpected error: {e}", exc_info=True)
|
||||
return TaskResult.fail(ErrorCode.E_UNKNOWN, str(e))
|
||||
|
||||
finally:
|
||||
self.cleanup_work_dir(work_dir)
|
||||
|
||||
def _build_command(
|
||||
self,
|
||||
input_file: str,
|
||||
output_file: str,
|
||||
render_spec: RenderSpec,
|
||||
output_spec: OutputSpec,
|
||||
duration_ms: int,
|
||||
lut_file: Optional[str] = None,
|
||||
overlay_file: Optional[str] = None
|
||||
) -> List[str]:
|
||||
"""
|
||||
构建 FFmpeg 渲染命令
|
||||
|
||||
Args:
|
||||
input_file: 输入文件路径
|
||||
output_file: 输出文件路径
|
||||
render_spec: 渲染规格
|
||||
output_spec: 输出规格
|
||||
duration_ms: 目标时长(毫秒)
|
||||
lut_file: LUT 文件路径(可选)
|
||||
overlay_file: 叠加层文件路径(可选)
|
||||
|
||||
Returns:
|
||||
FFmpeg 命令参数列表
|
||||
"""
|
||||
cmd = ['ffmpeg', '-y', '-hide_banner']
|
||||
|
||||
# 输入文件
|
||||
cmd.extend(['-i', input_file])
|
||||
|
||||
# 叠加层输入
|
||||
if overlay_file:
|
||||
cmd.extend(['-i', overlay_file])
|
||||
|
||||
# 构建视频滤镜链
|
||||
filters = self._build_video_filters(
|
||||
render_spec=render_spec,
|
||||
output_spec=output_spec,
|
||||
lut_file=lut_file,
|
||||
has_overlay=overlay_file is not None
|
||||
)
|
||||
|
||||
# 应用滤镜
|
||||
if overlay_file:
|
||||
# 使用 filter_complex 处理叠加
|
||||
cmd.extend(['-filter_complex', filters])
|
||||
elif filters:
|
||||
cmd.extend(['-vf', filters])
|
||||
|
||||
# 编码参数(v2 统一参数)
|
||||
cmd.extend(VIDEO_ENCODE_ARGS)
|
||||
|
||||
# 帧率
|
||||
fps = output_spec.fps
|
||||
cmd.extend(['-r', str(fps)])
|
||||
|
||||
# GOP 大小(关键帧间隔)
|
||||
gop_size = fps * 2 # 2秒一个关键帧
|
||||
cmd.extend(['-g', str(gop_size)])
|
||||
cmd.extend(['-keyint_min', str(gop_size)])
|
||||
|
||||
# 时长
|
||||
duration_sec = duration_ms / 1000.0
|
||||
cmd.extend(['-t', str(duration_sec)])
|
||||
|
||||
# 无音频(视频片段不包含音频)
|
||||
cmd.append('-an')
|
||||
|
||||
# 输出文件
|
||||
cmd.append(output_file)
|
||||
|
||||
return cmd
|
||||
|
||||
def _build_video_filters(
|
||||
self,
|
||||
render_spec: RenderSpec,
|
||||
output_spec: OutputSpec,
|
||||
lut_file: Optional[str] = None,
|
||||
has_overlay: bool = False
|
||||
) -> str:
|
||||
"""
|
||||
构建视频滤镜链
|
||||
|
||||
Args:
|
||||
render_spec: 渲染规格
|
||||
output_spec: 输出规格
|
||||
lut_file: LUT 文件路径
|
||||
has_overlay: 是否有叠加层
|
||||
|
||||
Returns:
|
||||
滤镜字符串
|
||||
"""
|
||||
filters = []
|
||||
width = output_spec.width
|
||||
height = output_spec.height
|
||||
|
||||
# 1. 变速处理
|
||||
speed = float(render_spec.speed) if render_spec.speed else 1.0
|
||||
if speed != 1.0 and speed > 0:
|
||||
# setpts 公式:PTS / speed
|
||||
pts_factor = 1.0 / speed
|
||||
filters.append(f"setpts={pts_factor}*PTS")
|
||||
|
||||
# 2. LUT 调色
|
||||
if lut_file:
|
||||
# 路径中的反斜杠需要转义
|
||||
lut_path = lut_file.replace('\\', '/')
|
||||
filters.append(f"lut3d='{lut_path}'")
|
||||
|
||||
# 3. 裁切处理
|
||||
if render_spec.crop_enable and render_spec.face_pos:
|
||||
# 根据人脸位置进行智能裁切
|
||||
try:
|
||||
fx, fy = map(float, render_spec.face_pos.split(','))
|
||||
# 计算裁切区域(保持输出比例)
|
||||
target_ratio = width / height
|
||||
# 假设裁切到目标比例
|
||||
filters.append(
|
||||
f"crop='min(iw,ih*{target_ratio})':'min(ih,iw/{target_ratio})':"
|
||||
f"'(iw-min(iw,ih*{target_ratio}))*{fx}':"
|
||||
f"'(ih-min(ih,iw/{target_ratio}))*{fy}'"
|
||||
)
|
||||
except (ValueError, ZeroDivisionError):
|
||||
logger.warning(f"Invalid face position: {render_spec.face_pos}")
|
||||
elif render_spec.zoom_cut:
|
||||
# 中心缩放裁切
|
||||
target_ratio = width / height
|
||||
filters.append(
|
||||
f"crop='min(iw,ih*{target_ratio})':'min(ih,iw/{target_ratio})'"
|
||||
)
|
||||
|
||||
# 4. 缩放和填充
|
||||
scale_filter = (
|
||||
f"scale={width}:{height}:force_original_aspect_ratio=decrease,"
|
||||
f"pad={width}:{height}:(ow-iw)/2:(oh-ih)/2:black"
|
||||
)
|
||||
filters.append(scale_filter)
|
||||
|
||||
# 5. 构建最终滤镜
|
||||
if has_overlay:
|
||||
# 使用 filter_complex 格式
|
||||
base_filters = ','.join(filters) if filters else 'copy'
|
||||
return f"[0:v]{base_filters}[base];[base][1:v]overlay=0:0"
|
||||
else:
|
||||
return ','.join(filters) if filters else ''
|
||||
188
index_v2.py
Normal file
188
index_v2.py
Normal file
@@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RenderWorker v2 入口
|
||||
|
||||
支持 v2 API 协议的渲染 Worker,处理以下任务类型:
|
||||
- RENDER_SEGMENT_VIDEO: 渲染视频片段
|
||||
- PREPARE_JOB_AUDIO: 生成全局音频
|
||||
- PACKAGE_SEGMENT_TS: 封装 TS 分片
|
||||
- FINALIZE_MP4: 产出最终 MP4
|
||||
|
||||
使用方法:
|
||||
python index_v2.py
|
||||
|
||||
环境变量:
|
||||
API_ENDPOINT_V2: v2 API 端点(或使用 API_ENDPOINT)
|
||||
ACCESS_KEY: Worker 认证密钥
|
||||
WORKER_ID: Worker ID(默认 100001)
|
||||
MAX_CONCURRENCY: 最大并发数(默认 4)
|
||||
HEARTBEAT_INTERVAL: 心跳间隔秒数(默认 5)
|
||||
TEMP_DIR: 临时文件目录
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import signal
|
||||
import logging
|
||||
|
||||
# 加载配置(必须在其他模块之前)
|
||||
import config
|
||||
|
||||
from domain.config import WorkerConfig
|
||||
from services.api_client import APIClientV2
|
||||
from services.task_executor import TaskExecutor
|
||||
|
||||
# 日志配置
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s',
|
||||
datefmt='%Y-%m-%d %H:%M:%S'
|
||||
)
|
||||
logger = logging.getLogger('worker_v2')
|
||||
|
||||
|
||||
class WorkerV2:
|
||||
"""
|
||||
v2 渲染 Worker 主类
|
||||
|
||||
负责:
|
||||
- 配置加载
|
||||
- API 客户端初始化
|
||||
- 任务执行器管理
|
||||
- 主循环运行
|
||||
- 优雅退出处理
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化 Worker"""
|
||||
# 加载配置
|
||||
try:
|
||||
self.config = WorkerConfig.from_env()
|
||||
except ValueError as e:
|
||||
logger.error(f"Configuration error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# 初始化 API 客户端
|
||||
self.api_client = APIClientV2(self.config)
|
||||
|
||||
# 初始化任务执行器
|
||||
self.task_executor = TaskExecutor(self.config, self.api_client)
|
||||
|
||||
# 运行状态
|
||||
self.running = True
|
||||
|
||||
# 确保临时目录存在
|
||||
self.config.ensure_temp_dir()
|
||||
|
||||
# 注册信号处理器
|
||||
self._setup_signal_handlers()
|
||||
|
||||
def _setup_signal_handlers(self):
|
||||
"""设置信号处理器"""
|
||||
# Windows 不支持 SIGTERM
|
||||
signal.signal(signal.SIGINT, self._signal_handler)
|
||||
if hasattr(signal, 'SIGTERM'):
|
||||
signal.signal(signal.SIGTERM, self._signal_handler)
|
||||
|
||||
def _signal_handler(self, signum, frame):
|
||||
"""
|
||||
信号处理,优雅退出
|
||||
|
||||
Args:
|
||||
signum: 信号编号
|
||||
frame: 当前栈帧
|
||||
"""
|
||||
signal_name = signal.Signals(signum).name
|
||||
logger.info(f"Received signal {signal_name}, initiating shutdown...")
|
||||
self.running = False
|
||||
|
||||
def run(self):
|
||||
"""主循环"""
|
||||
logger.info("=" * 60)
|
||||
logger.info("RenderWorker v2 Starting")
|
||||
logger.info("=" * 60)
|
||||
logger.info(f"Worker ID: {self.config.worker_id}")
|
||||
logger.info(f"API Endpoint: {self.config.api_endpoint}")
|
||||
logger.info(f"Max Concurrency: {self.config.max_concurrency}")
|
||||
logger.info(f"Heartbeat Interval: {self.config.heartbeat_interval}s")
|
||||
logger.info(f"Capabilities: {', '.join(self.config.capabilities)}")
|
||||
logger.info(f"Temp Directory: {self.config.temp_dir}")
|
||||
logger.info("=" * 60)
|
||||
|
||||
consecutive_errors = 0
|
||||
max_consecutive_errors = 10
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
# 心跳同步并拉取任务
|
||||
current_task_ids = self.task_executor.get_current_task_ids()
|
||||
tasks = self.api_client.sync(current_task_ids)
|
||||
|
||||
# 提交新任务
|
||||
for task in tasks:
|
||||
if self.task_executor.submit_task(task):
|
||||
logger.info(f"Submitted task: {task.task_id} ({task.task_type.value})")
|
||||
|
||||
# 重置错误计数
|
||||
consecutive_errors = 0
|
||||
|
||||
# 等待下次心跳
|
||||
time.sleep(self.config.heartbeat_interval)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Keyboard interrupt received")
|
||||
self.running = False
|
||||
except Exception as e:
|
||||
consecutive_errors += 1
|
||||
logger.error(f"Worker loop error ({consecutive_errors}/{max_consecutive_errors}): {e}", exc_info=True)
|
||||
|
||||
# 连续错误过多,增加等待时间
|
||||
if consecutive_errors >= max_consecutive_errors:
|
||||
logger.error("Too many consecutive errors, waiting 30 seconds...")
|
||||
time.sleep(30)
|
||||
consecutive_errors = 0
|
||||
else:
|
||||
time.sleep(5)
|
||||
|
||||
# 优雅关闭
|
||||
self._shutdown()
|
||||
|
||||
def _shutdown(self):
|
||||
"""优雅关闭"""
|
||||
logger.info("Shutting down...")
|
||||
|
||||
# 等待当前任务完成
|
||||
current_count = self.task_executor.get_current_task_count()
|
||||
if current_count > 0:
|
||||
logger.info(f"Waiting for {current_count} running task(s) to complete...")
|
||||
|
||||
# 关闭执行器
|
||||
self.task_executor.shutdown(wait=True)
|
||||
|
||||
# 关闭 API 客户端
|
||||
self.api_client.close()
|
||||
|
||||
logger.info("Worker stopped")
|
||||
|
||||
|
||||
def main():
|
||||
"""主函数"""
|
||||
# 初始化 OpenTelemetry(如果可用)
|
||||
try:
|
||||
from telemetry import init_opentelemetry
|
||||
init_opentelemetry()
|
||||
logger.info("OpenTelemetry initialized")
|
||||
except ImportError:
|
||||
logger.debug("OpenTelemetry not available, skipping initialization")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize OpenTelemetry: {e}")
|
||||
|
||||
# 创建并运行 Worker
|
||||
worker = WorkerV2()
|
||||
worker.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
16
services/__init__.py
Normal file
16
services/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
服务层
|
||||
|
||||
包含 API 客户端、任务执行器、租约服务等服务组件。
|
||||
"""
|
||||
|
||||
from services.api_client import APIClientV2
|
||||
from services.lease_service import LeaseService
|
||||
from services.task_executor import TaskExecutor
|
||||
|
||||
__all__ = [
|
||||
'APIClientV2',
|
||||
'LeaseService',
|
||||
'TaskExecutor',
|
||||
]
|
||||
371
services/api_client.py
Normal file
371
services/api_client.py
Normal file
@@ -0,0 +1,371 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
v2 API 客户端
|
||||
|
||||
实现与渲染服务端 v2 接口的通信。
|
||||
"""
|
||||
|
||||
import logging
|
||||
import subprocess
|
||||
import requests
|
||||
from typing import Dict, List, Optional, Any
|
||||
|
||||
from domain.task import Task
|
||||
from domain.config import WorkerConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class APIClientV2:
|
||||
"""
|
||||
v2 API 客户端
|
||||
|
||||
负责与渲染服务端的所有 HTTP 通信。
|
||||
"""
|
||||
|
||||
def __init__(self, config: WorkerConfig):
|
||||
"""
|
||||
初始化 API 客户端
|
||||
|
||||
Args:
|
||||
config: Worker 配置
|
||||
"""
|
||||
self.config = config
|
||||
self.base_url = config.api_endpoint.rstrip('/')
|
||||
self.access_key = config.access_key
|
||||
self.worker_id = config.worker_id
|
||||
self.session = requests.Session()
|
||||
|
||||
# 设置默认请求头
|
||||
self.session.headers.update({
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json'
|
||||
})
|
||||
|
||||
def sync(self, current_task_ids: List[str]) -> List[Task]:
|
||||
"""
|
||||
心跳同步并拉取任务
|
||||
|
||||
Args:
|
||||
current_task_ids: 当前正在执行的任务 ID 列表
|
||||
|
||||
Returns:
|
||||
List[Task]: 新分配的任务列表
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/worker/sync"
|
||||
|
||||
# 将 task_id 转换为整数(服务端期望 []int64)
|
||||
task_ids_int = [int(tid) for tid in current_task_ids if tid.isdigit()]
|
||||
|
||||
payload = {
|
||||
'accessKey': self.access_key,
|
||||
'workerId': self.worker_id,
|
||||
'capabilities': self.config.capabilities,
|
||||
'maxConcurrency': self.config.max_concurrency,
|
||||
'currentTaskCount': len(current_task_ids),
|
||||
'currentTaskIds': task_ids_int,
|
||||
'ffmpegVersion': self._get_ffmpeg_version(),
|
||||
'codecInfo': self._get_codec_info(),
|
||||
'systemInfo': self._get_system_info()
|
||||
}
|
||||
|
||||
try:
|
||||
resp = self.session.post(url, json=payload, timeout=10)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
if data.get('code') != 200:
|
||||
logger.warning(f"Sync failed: {data.get('message')}")
|
||||
return []
|
||||
|
||||
# 解析任务列表
|
||||
tasks = []
|
||||
for task_data in data.get('data', {}).get('tasks', []):
|
||||
try:
|
||||
task = Task.from_dict(task_data)
|
||||
tasks.append(task)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse task: {e}")
|
||||
|
||||
if tasks:
|
||||
logger.info(f"Received {len(tasks)} new tasks")
|
||||
|
||||
return tasks
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
logger.warning("Sync timeout")
|
||||
return []
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Sync request error: {e}")
|
||||
return []
|
||||
except Exception as e:
|
||||
logger.error(f"Sync error: {e}")
|
||||
return []
|
||||
|
||||
def report_start(self, task_id: str) -> bool:
|
||||
"""
|
||||
报告任务开始
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}/start"
|
||||
|
||||
try:
|
||||
resp = self.session.post(
|
||||
url,
|
||||
json={'workerId': self.worker_id},
|
||||
timeout=10
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.debug(f"[task:{task_id}] Start reported")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"[task:{task_id}] Report start failed: {resp.status_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Report start error: {e}")
|
||||
return False
|
||||
|
||||
def report_success(self, task_id: str, result: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
报告任务成功
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
result: 任务结果数据
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}/success"
|
||||
|
||||
try:
|
||||
resp = self.session.post(
|
||||
url,
|
||||
json={
|
||||
'workerId': self.worker_id,
|
||||
'result': result
|
||||
},
|
||||
timeout=10
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.debug(f"[task:{task_id}] Success reported")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"[task:{task_id}] Report success failed: {resp.status_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Report success error: {e}")
|
||||
return False
|
||||
|
||||
def report_fail(self, task_id: str, error_code: str, error_message: str) -> bool:
|
||||
"""
|
||||
报告任务失败
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
error_code: 错误码
|
||||
error_message: 错误信息
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}/fail"
|
||||
|
||||
try:
|
||||
resp = self.session.post(
|
||||
url,
|
||||
json={
|
||||
'workerId': self.worker_id,
|
||||
'errorCode': error_code,
|
||||
'errorMessage': error_message[:1000] # 限制长度
|
||||
},
|
||||
timeout=10
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.debug(f"[task:{task_id}] Failure reported")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"[task:{task_id}] Report fail failed: {resp.status_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Report fail error: {e}")
|
||||
return False
|
||||
|
||||
def get_upload_url(self, task_id: str, file_type: str, file_name: str = None) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
获取上传 URL
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
file_type: 文件类型(video/audio/ts/mp4)
|
||||
file_name: 文件名(可选)
|
||||
|
||||
Returns:
|
||||
Dict 包含 uploadUrl 和 accessUrl,失败返回 None
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}/uploadUrl"
|
||||
|
||||
payload = {'fileType': file_type}
|
||||
if file_name:
|
||||
payload['fileName'] = file_name
|
||||
|
||||
try:
|
||||
resp = self.session.post(url, json=payload, timeout=10)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
if data.get('code') == 200:
|
||||
return data.get('data')
|
||||
logger.warning(f"[task:{task_id}] Get upload URL failed: {resp.status_code}")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Get upload URL error: {e}")
|
||||
return None
|
||||
|
||||
def extend_lease(self, task_id: str, extension: int = None) -> bool:
|
||||
"""
|
||||
延长租约
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
extension: 延长秒数(默认使用配置值)
|
||||
|
||||
Returns:
|
||||
bool: 是否成功
|
||||
"""
|
||||
if extension is None:
|
||||
extension = self.config.lease_extension_duration
|
||||
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}/extend-lease"
|
||||
|
||||
try:
|
||||
resp = self.session.post(
|
||||
url,
|
||||
params={
|
||||
'workerId': self.worker_id,
|
||||
'extension': extension
|
||||
},
|
||||
timeout=10
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
logger.debug(f"[task:{task_id}] Lease extended by {extension}s")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"[task:{task_id}] Extend lease failed: {resp.status_code}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Extend lease error: {e}")
|
||||
return False
|
||||
|
||||
def get_task_info(self, task_id: str) -> Optional[Dict]:
|
||||
"""
|
||||
获取任务详情
|
||||
|
||||
Args:
|
||||
task_id: 任务 ID
|
||||
|
||||
Returns:
|
||||
任务详情字典,失败返回 None
|
||||
"""
|
||||
url = f"{self.base_url}/render/v2/task/{task_id}"
|
||||
|
||||
try:
|
||||
resp = self.session.get(url, timeout=10)
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
if data.get('code') == 200:
|
||||
return data.get('data')
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Get task info error: {e}")
|
||||
return None
|
||||
|
||||
def _get_ffmpeg_version(self) -> str:
|
||||
"""获取 FFmpeg 版本"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['ffmpeg', '-version'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
first_line = result.stdout.split('\n')[0]
|
||||
if 'version' in first_line:
|
||||
parts = first_line.split()
|
||||
for i, part in enumerate(parts):
|
||||
if part == 'version' and i + 1 < len(parts):
|
||||
return parts[i + 1]
|
||||
return 'unknown'
|
||||
except Exception:
|
||||
return 'unknown'
|
||||
|
||||
def _get_codec_info(self) -> str:
|
||||
"""获取支持的编解码器信息"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['ffmpeg', '-codecs'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
# 检查常用编解码器
|
||||
codecs = []
|
||||
output = result.stdout
|
||||
if 'libx264' in output:
|
||||
codecs.append('libx264')
|
||||
if 'libx265' in output or 'hevc' in output:
|
||||
codecs.append('libx265')
|
||||
if 'aac' in output:
|
||||
codecs.append('aac')
|
||||
if 'libfdk_aac' in output:
|
||||
codecs.append('libfdk_aac')
|
||||
return ', '.join(codecs) if codecs else 'unknown'
|
||||
except Exception:
|
||||
return 'unknown'
|
||||
|
||||
def _get_system_info(self) -> Dict[str, Any]:
|
||||
"""获取系统信息"""
|
||||
try:
|
||||
import platform
|
||||
import psutil
|
||||
|
||||
info = {
|
||||
'os': platform.system(),
|
||||
'cpu': f"{psutil.cpu_count()} cores",
|
||||
'memory': f"{psutil.virtual_memory().total // (1024**3)}GB",
|
||||
'cpuUsage': f"{psutil.cpu_percent()}%",
|
||||
'memoryAvailable': f"{psutil.virtual_memory().available // (1024**3)}GB"
|
||||
}
|
||||
|
||||
# 尝试获取 GPU 信息
|
||||
gpu_info = self._get_gpu_info()
|
||||
if gpu_info:
|
||||
info['gpu'] = gpu_info
|
||||
|
||||
return info
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _get_gpu_info(self) -> Optional[str]:
|
||||
"""获取 GPU 信息"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
if result.returncode == 0:
|
||||
gpu_name = result.stdout.strip().split('\n')[0]
|
||||
return gpu_name
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
def close(self):
|
||||
"""关闭会话"""
|
||||
self.session.close()
|
||||
110
services/lease_service.py
Normal file
110
services/lease_service.py
Normal file
@@ -0,0 +1,110 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
租约续期服务
|
||||
|
||||
后台线程定期为正在执行的任务续期租约。
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from services.api_client import APIClientV2
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LeaseService:
|
||||
"""
|
||||
租约续期服务
|
||||
|
||||
在后台线程中定期调用 API 延长任务租约,
|
||||
防止长时间任务因租约过期被回收。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_client: 'APIClientV2',
|
||||
task_id: str,
|
||||
interval: int = 60,
|
||||
extension: int = 300
|
||||
):
|
||||
"""
|
||||
初始化租约服务
|
||||
|
||||
Args:
|
||||
api_client: API 客户端
|
||||
task_id: 任务 ID
|
||||
interval: 续期间隔(秒),默认 60 秒
|
||||
extension: 每次续期时长(秒),默认 300 秒
|
||||
"""
|
||||
self.api_client = api_client
|
||||
self.task_id = task_id
|
||||
self.interval = interval
|
||||
self.extension = extension
|
||||
self.running = False
|
||||
self.thread: threading.Thread = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
def start(self):
|
||||
"""启动租约续期线程"""
|
||||
if self.running:
|
||||
logger.warning(f"[task:{self.task_id}] Lease service already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
self._stop_event.clear()
|
||||
self.thread = threading.Thread(
|
||||
target=self._run,
|
||||
name=f"LeaseService-{self.task_id}",
|
||||
daemon=True
|
||||
)
|
||||
self.thread.start()
|
||||
logger.debug(f"[task:{self.task_id}] Lease service started (interval={self.interval}s)")
|
||||
|
||||
def stop(self):
|
||||
"""停止租约续期线程"""
|
||||
if not self.running:
|
||||
return
|
||||
|
||||
self.running = False
|
||||
self._stop_event.set()
|
||||
|
||||
if self.thread and self.thread.is_alive():
|
||||
self.thread.join(timeout=5)
|
||||
|
||||
logger.debug(f"[task:{self.task_id}] Lease service stopped")
|
||||
|
||||
def _run(self):
|
||||
"""续期线程主循环"""
|
||||
while self.running:
|
||||
# 等待指定间隔或收到停止信号
|
||||
if self._stop_event.wait(timeout=self.interval):
|
||||
# 收到停止信号
|
||||
break
|
||||
|
||||
if self.running:
|
||||
self._extend_lease()
|
||||
|
||||
def _extend_lease(self):
|
||||
"""执行租约续期"""
|
||||
try:
|
||||
success = self.api_client.extend_lease(self.task_id, self.extension)
|
||||
if success:
|
||||
logger.debug(f"[task:{self.task_id}] Lease extended by {self.extension}s")
|
||||
else:
|
||||
logger.warning(f"[task:{self.task_id}] Failed to extend lease")
|
||||
except Exception as e:
|
||||
logger.warning(f"[task:{self.task_id}] Lease extension error: {e}")
|
||||
|
||||
def __enter__(self):
|
||||
"""上下文管理器入口"""
|
||||
self.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""上下文管理器出口"""
|
||||
self.stop()
|
||||
return False
|
||||
234
services/task_executor.py
Normal file
234
services/task_executor.py
Normal file
@@ -0,0 +1,234 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
任务执行器
|
||||
|
||||
管理任务的并发执行,协调处理器、租约服务等组件。
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor, Future
|
||||
from typing import Dict, Optional, TYPE_CHECKING
|
||||
|
||||
from domain.task import Task, TaskType
|
||||
from domain.result import TaskResult, ErrorCode
|
||||
from domain.config import WorkerConfig
|
||||
from core.handler import TaskHandler
|
||||
from services.lease_service import LeaseService
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from services.api_client import APIClientV2
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TaskExecutor:
|
||||
"""
|
||||
任务执行器
|
||||
|
||||
负责任务的并发调度和执行,包括:
|
||||
- 注册和管理任务处理器
|
||||
- 维护任务执行状态
|
||||
- 协调租约续期
|
||||
- 上报执行结果
|
||||
"""
|
||||
|
||||
def __init__(self, config: WorkerConfig, api_client: 'APIClientV2'):
|
||||
"""
|
||||
初始化任务执行器
|
||||
|
||||
Args:
|
||||
config: Worker 配置
|
||||
api_client: API 客户端
|
||||
"""
|
||||
self.config = config
|
||||
self.api_client = api_client
|
||||
|
||||
# 任务处理器注册表
|
||||
self.handlers: Dict[TaskType, TaskHandler] = {}
|
||||
|
||||
# 当前任务跟踪
|
||||
self.current_tasks: Dict[str, Task] = {}
|
||||
self.current_futures: Dict[str, Future] = {}
|
||||
|
||||
# 线程池
|
||||
self.executor = ThreadPoolExecutor(
|
||||
max_workers=config.max_concurrency,
|
||||
thread_name_prefix="TaskWorker"
|
||||
)
|
||||
|
||||
# 线程安全锁
|
||||
self.lock = threading.Lock()
|
||||
|
||||
# 注册处理器
|
||||
self._register_handlers()
|
||||
|
||||
def _register_handlers(self):
|
||||
"""注册所有任务处理器"""
|
||||
# 延迟导入以避免循环依赖
|
||||
from handlers.render_video import RenderSegmentVideoHandler
|
||||
from handlers.prepare_audio import PrepareJobAudioHandler
|
||||
from handlers.package_ts import PackageSegmentTsHandler
|
||||
from handlers.finalize_mp4 import FinalizeMp4Handler
|
||||
|
||||
handlers = [
|
||||
RenderSegmentVideoHandler(self.config, self.api_client),
|
||||
PrepareJobAudioHandler(self.config, self.api_client),
|
||||
PackageSegmentTsHandler(self.config, self.api_client),
|
||||
FinalizeMp4Handler(self.config, self.api_client),
|
||||
]
|
||||
|
||||
for handler in handlers:
|
||||
task_type = handler.get_supported_type()
|
||||
self.handlers[task_type] = handler
|
||||
logger.debug(f"Registered handler for {task_type.value}")
|
||||
|
||||
def get_current_task_ids(self) -> list:
|
||||
"""
|
||||
获取当前正在执行的任务 ID 列表
|
||||
|
||||
Returns:
|
||||
任务 ID 列表
|
||||
"""
|
||||
with self.lock:
|
||||
return list(self.current_tasks.keys())
|
||||
|
||||
def get_current_task_count(self) -> int:
|
||||
"""
|
||||
获取当前正在执行的任务数量
|
||||
|
||||
Returns:
|
||||
任务数量
|
||||
"""
|
||||
with self.lock:
|
||||
return len(self.current_tasks)
|
||||
|
||||
def can_accept_task(self) -> bool:
|
||||
"""
|
||||
检查是否可以接受新任务
|
||||
|
||||
Returns:
|
||||
是否可以接受
|
||||
"""
|
||||
return self.get_current_task_count() < self.config.max_concurrency
|
||||
|
||||
def submit_task(self, task: Task) -> bool:
|
||||
"""
|
||||
提交任务到线程池
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
|
||||
Returns:
|
||||
是否提交成功
|
||||
"""
|
||||
with self.lock:
|
||||
# 检查任务是否已在执行
|
||||
if task.task_id in self.current_tasks:
|
||||
logger.warning(f"[task:{task.task_id}] Task already running, skipping")
|
||||
return False
|
||||
|
||||
# 检查是否有对应的处理器
|
||||
if task.task_type not in self.handlers:
|
||||
logger.error(f"[task:{task.task_id}] No handler for type: {task.task_type.value}")
|
||||
return False
|
||||
|
||||
# 记录任务
|
||||
self.current_tasks[task.task_id] = task
|
||||
|
||||
# 提交到线程池
|
||||
future = self.executor.submit(self._process_task, task)
|
||||
self.current_futures[task.task_id] = future
|
||||
|
||||
logger.info(f"[task:{task.task_id}] Submitted ({task.task_type.value})")
|
||||
return True
|
||||
|
||||
def _process_task(self, task: Task):
|
||||
"""
|
||||
处理单个任务(在线程池中执行)
|
||||
|
||||
Args:
|
||||
task: 任务实体
|
||||
"""
|
||||
task_id = task.task_id
|
||||
logger.info(f"[task:{task_id}] Starting {task.task_type.value}")
|
||||
|
||||
# 启动租约续期服务
|
||||
lease_service = LeaseService(
|
||||
self.api_client,
|
||||
task_id,
|
||||
interval=self.config.lease_extension_threshold,
|
||||
extension=self.config.lease_extension_duration
|
||||
)
|
||||
lease_service.start()
|
||||
|
||||
try:
|
||||
# 报告任务开始
|
||||
self.api_client.report_start(task_id)
|
||||
|
||||
# 获取处理器
|
||||
handler = self.handlers.get(task.task_type)
|
||||
if not handler:
|
||||
raise ValueError(f"No handler for task type: {task.task_type}")
|
||||
|
||||
# 执行前钩子
|
||||
handler.before_handle(task)
|
||||
|
||||
# 执行任务
|
||||
result = handler.handle(task)
|
||||
|
||||
# 执行后钩子
|
||||
handler.after_handle(task, result)
|
||||
|
||||
# 上报结果
|
||||
if result.success:
|
||||
self.api_client.report_success(task_id, result.data)
|
||||
logger.info(f"[task:{task_id}] Completed successfully")
|
||||
else:
|
||||
error_code = result.error_code.value if result.error_code else 'E_UNKNOWN'
|
||||
self.api_client.report_fail(task_id, error_code, result.error_message or '')
|
||||
logger.error(f"[task:{task_id}] Failed: {result.error_message}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[task:{task_id}] Exception: {e}", exc_info=True)
|
||||
self.api_client.report_fail(task_id, 'E_UNKNOWN', str(e))
|
||||
|
||||
finally:
|
||||
# 停止租约续期
|
||||
lease_service.stop()
|
||||
|
||||
# 从当前任务中移除
|
||||
with self.lock:
|
||||
self.current_tasks.pop(task_id, None)
|
||||
self.current_futures.pop(task_id, None)
|
||||
|
||||
def shutdown(self, wait: bool = True):
|
||||
"""
|
||||
关闭执行器
|
||||
|
||||
Args:
|
||||
wait: 是否等待所有任务完成
|
||||
"""
|
||||
logger.info("Shutting down task executor...")
|
||||
|
||||
# 关闭线程池
|
||||
self.executor.shutdown(wait=wait)
|
||||
|
||||
# 清理状态
|
||||
with self.lock:
|
||||
self.current_tasks.clear()
|
||||
self.current_futures.clear()
|
||||
|
||||
logger.info("Task executor shutdown complete")
|
||||
|
||||
def get_handler(self, task_type: TaskType) -> Optional[TaskHandler]:
|
||||
"""
|
||||
获取指定类型的处理器
|
||||
|
||||
Args:
|
||||
task_type: 任务类型
|
||||
|
||||
Returns:
|
||||
处理器实例,不存在则返回 None
|
||||
"""
|
||||
return self.handlers.get(task_type)
|
||||
@@ -1,14 +1,24 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
系统信息工具
|
||||
|
||||
提供系统信息采集功能。
|
||||
"""
|
||||
|
||||
import os
|
||||
import platform
|
||||
from datetime import datetime
|
||||
|
||||
import psutil
|
||||
from constant import SUPPORT_FEATURE, SOFTWARE_VERSION
|
||||
from constant import SUPPORT_FEATURE, SOFTWARE_VERSION, V2_DEFAULT_CAPABILITIES
|
||||
|
||||
|
||||
def get_sys_info():
|
||||
"""
|
||||
Returns a dictionary with system information.
|
||||
获取系统信息(v1 格式)
|
||||
|
||||
Returns:
|
||||
dict: 系统信息字典
|
||||
"""
|
||||
info = {
|
||||
'version': SOFTWARE_VERSION,
|
||||
@@ -22,3 +32,65 @@ def get_sys_info():
|
||||
'support_feature': SUPPORT_FEATURE
|
||||
}
|
||||
return info
|
||||
|
||||
|
||||
def get_sys_info_v2():
|
||||
"""
|
||||
获取系统信息(v2 格式)
|
||||
|
||||
Returns:
|
||||
dict: v2 API 所需的系统信息字典
|
||||
"""
|
||||
mem = psutil.virtual_memory()
|
||||
|
||||
info = {
|
||||
'os': platform.system(),
|
||||
'cpu': f"{os.cpu_count()} cores",
|
||||
'memory': f"{mem.total // (1024**3)}GB",
|
||||
'cpuUsage': f"{psutil.cpu_percent()}%",
|
||||
'memoryAvailable': f"{mem.available // (1024**3)}GB",
|
||||
'platform': platform.system(),
|
||||
'pythonVersion': platform.python_version(),
|
||||
}
|
||||
|
||||
# 尝试获取 GPU 信息
|
||||
gpu_info = _get_gpu_info()
|
||||
if gpu_info:
|
||||
info['gpu'] = gpu_info
|
||||
|
||||
return info
|
||||
|
||||
|
||||
def get_capabilities():
|
||||
"""
|
||||
获取 Worker 支持的能力列表
|
||||
|
||||
Returns:
|
||||
list: 能力列表
|
||||
"""
|
||||
return V2_DEFAULT_CAPABILITIES.copy()
|
||||
|
||||
|
||||
def _get_gpu_info():
|
||||
"""
|
||||
尝试获取 GPU 信息
|
||||
|
||||
Returns:
|
||||
str: GPU 信息,失败返回 None
|
||||
"""
|
||||
try:
|
||||
import subprocess
|
||||
# 尝试使用 nvidia-smi
|
||||
result = subprocess.run(
|
||||
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5
|
||||
)
|
||||
if result.returncode == 0:
|
||||
gpu_name = result.stdout.strip().split('\n')[0]
|
||||
return gpu_name
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
Reference in New Issue
Block a user