You've already forked FrameTour-RenderWorker
Jerry Yan
24de32e6bb
- 添加v2支持的任务类型常量定义 - 更新软件版本至0.0.9 - 定义v2统一音视频编码参数 - 实现系统信息工具get_sys_info_v2方法 - 新增get_capabilities和_get_gpu_info功能 - 创建core模块及TaskHandler抽象基类 - 添加渲染系统设计文档包括集群架构、v2 PRD和Worker PRD - 实现任务处理器抽象基类及接口规范
372 lines
11 KiB
Python
372 lines
11 KiB
Python
# -*- 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()
|