You've already forked FrameTour-RenderWorker
feat(gpu): 添加多显卡调度支持
- 新增 GPUDevice 数据类定义 GPU 设备信息 - 扩展 WorkerConfig 添加 gpu_devices 配置项 - 从环境变量 GPU_DEVICES 读取多显卡设备配置 - 实现 GPUScheduler 提供轮询调度功能 - 修改 FFmpeg 参数生成支持设备指定 - 添加线程本地存储管理当前 GPU 设备 - 更新任务执行器集成 GPU 设备分配 - 实现 GPU 设备自动检测和验证功能 - 添加相关日志记录和状态监控
This commit is contained in:
@@ -11,6 +11,7 @@ import logging
|
||||
import shutil
|
||||
import tempfile
|
||||
import subprocess
|
||||
import threading
|
||||
from abc import ABC
|
||||
from typing import Optional, List, Dict, Any, Tuple, TYPE_CHECKING
|
||||
|
||||
@@ -75,23 +76,33 @@ def get_video_encode_args(hw_accel: str = HW_ACCEL_NONE) -> List[str]:
|
||||
]
|
||||
|
||||
|
||||
def get_hwaccel_decode_args(hw_accel: str = HW_ACCEL_NONE) -> List[str]:
|
||||
def get_hwaccel_decode_args(hw_accel: str = HW_ACCEL_NONE, device_index: Optional[int] = None) -> List[str]:
|
||||
"""
|
||||
获取硬件加速解码参数(输入文件之前使用)
|
||||
|
||||
Args:
|
||||
hw_accel: 硬件加速类型 (none, qsv, cuda)
|
||||
device_index: GPU 设备索引,用于多显卡调度
|
||||
|
||||
Returns:
|
||||
FFmpeg 硬件加速解码参数列表
|
||||
"""
|
||||
if hw_accel == HW_ACCEL_CUDA:
|
||||
# CUDA 硬件加速解码
|
||||
# 注意:使用 cuda 作为 hwaccel,但输出到系统内存以便 CPU 滤镜处理
|
||||
return ['-hwaccel', 'cuda', '-hwaccel_output_format', 'cuda']
|
||||
args = ['-hwaccel', 'cuda']
|
||||
# 多显卡模式下指定设备
|
||||
if device_index is not None:
|
||||
args.extend(['-hwaccel_device', str(device_index)])
|
||||
args.extend(['-hwaccel_output_format', 'cuda'])
|
||||
return args
|
||||
elif hw_accel == HW_ACCEL_QSV:
|
||||
# QSV 硬件加速解码
|
||||
return ['-hwaccel', 'qsv', '-hwaccel_output_format', 'qsv']
|
||||
args = ['-hwaccel', 'qsv']
|
||||
# QSV 在 Windows 上使用 -qsv_device
|
||||
if device_index is not None:
|
||||
args.extend(['-qsv_device', str(device_index)])
|
||||
args.extend(['-hwaccel_output_format', 'qsv'])
|
||||
return args
|
||||
else:
|
||||
return []
|
||||
|
||||
@@ -248,9 +259,13 @@ class BaseHandler(TaskHandler, ABC):
|
||||
- 临时目录管理
|
||||
- 文件下载/上传
|
||||
- FFmpeg 命令执行
|
||||
- GPU 设备管理(多显卡调度)
|
||||
- 日志记录
|
||||
"""
|
||||
|
||||
# 线程本地存储:用于存储当前线程的 GPU 设备索引
|
||||
_thread_local = threading.local()
|
||||
|
||||
def __init__(self, config: WorkerConfig, api_client: 'APIClientV2'):
|
||||
"""
|
||||
初始化处理器
|
||||
@@ -267,6 +282,39 @@ class BaseHandler(TaskHandler, ABC):
|
||||
max_size_gb=config.cache_max_size_gb
|
||||
)
|
||||
|
||||
# ========== GPU 设备管理 ==========
|
||||
|
||||
def set_gpu_device(self, device_index: int) -> None:
|
||||
"""
|
||||
设置当前线程的 GPU 设备索引
|
||||
|
||||
由 TaskExecutor 在任务执行前调用。
|
||||
|
||||
Args:
|
||||
device_index: GPU 设备索引
|
||||
"""
|
||||
self._thread_local.gpu_device = device_index
|
||||
|
||||
def get_gpu_device(self) -> Optional[int]:
|
||||
"""
|
||||
获取当前线程的 GPU 设备索引
|
||||
|
||||
Returns:
|
||||
GPU 设备索引,未设置则返回 None
|
||||
"""
|
||||
return getattr(self._thread_local, 'gpu_device', None)
|
||||
|
||||
def clear_gpu_device(self) -> None:
|
||||
"""
|
||||
清除当前线程的 GPU 设备索引
|
||||
|
||||
由 TaskExecutor 在任务执行后调用。
|
||||
"""
|
||||
if hasattr(self._thread_local, 'gpu_device'):
|
||||
del self._thread_local.gpu_device
|
||||
|
||||
# ========== FFmpeg 参数生成 ==========
|
||||
|
||||
def get_video_encode_args(self) -> List[str]:
|
||||
"""
|
||||
获取当前配置的视频编码参数
|
||||
@@ -278,12 +326,13 @@ class BaseHandler(TaskHandler, ABC):
|
||||
|
||||
def get_hwaccel_decode_args(self) -> List[str]:
|
||||
"""
|
||||
获取硬件加速解码参数(在输入文件之前使用)
|
||||
获取硬件加速解码参数(支持设备指定)
|
||||
|
||||
Returns:
|
||||
FFmpeg 硬件加速解码参数列表
|
||||
"""
|
||||
return get_hwaccel_decode_args(self.config.hw_accel)
|
||||
device_index = self.get_gpu_device()
|
||||
return get_hwaccel_decode_args(self.config.hw_accel, device_index)
|
||||
|
||||
def get_hwaccel_filter_prefix(self) -> str:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user