You've already forked FrameTour-RenderWorker
feat(task): 优化 GPU 调度以支持特定任务类型
- 添加 GPU_REQUIRED_TASK_TYPES 集合定义需要 GPU 加速的任务类型 - 修改任务执行逻辑仅对需要 GPU 的任务类型获取 GPU 设备 - 更新 GPU 设备释放逻辑确保仅在实际分配设备时进行释放 - 改进日志记录和资源管理流程
This commit is contained in:
@@ -12,6 +12,12 @@ from typing import Dict, Optional, TYPE_CHECKING
|
|||||||
|
|
||||||
from domain.task import Task, TaskType
|
from domain.task import Task, TaskType
|
||||||
from domain.result import TaskResult, ErrorCode
|
from domain.result import TaskResult, ErrorCode
|
||||||
|
|
||||||
|
# 需要 GPU 加速的任务类型
|
||||||
|
GPU_REQUIRED_TASK_TYPES = {
|
||||||
|
TaskType.RENDER_SEGMENT_VIDEO,
|
||||||
|
TaskType.COMPOSE_TRANSITION,
|
||||||
|
}
|
||||||
from domain.config import WorkerConfig
|
from domain.config import WorkerConfig
|
||||||
from core.handler import TaskHandler
|
from core.handler import TaskHandler
|
||||||
from services.lease_service import LeaseService
|
from services.lease_service import LeaseService
|
||||||
@@ -179,9 +185,10 @@ class TaskExecutor:
|
|||||||
)
|
)
|
||||||
lease_service.start()
|
lease_service.start()
|
||||||
|
|
||||||
# 获取 GPU 设备
|
# 获取 GPU 设备(仅对需要 GPU 的任务类型)
|
||||||
device_index = None
|
device_index = None
|
||||||
if self.gpu_scheduler.enabled:
|
needs_gpu = task.task_type in GPU_REQUIRED_TASK_TYPES
|
||||||
|
if needs_gpu and self.gpu_scheduler.enabled:
|
||||||
device_index = self.gpu_scheduler.acquire()
|
device_index = self.gpu_scheduler.acquire()
|
||||||
if device_index is not None:
|
if device_index is not None:
|
||||||
logger.info(f"[task:{task_id}] Assigned to GPU device {device_index}")
|
logger.info(f"[task:{task_id}] Assigned to GPU device {device_index}")
|
||||||
@@ -227,8 +234,8 @@ class TaskExecutor:
|
|||||||
if handler:
|
if handler:
|
||||||
handler.clear_gpu_device()
|
handler.clear_gpu_device()
|
||||||
|
|
||||||
# 释放 GPU 设备
|
# 释放 GPU 设备(仅当实际分配了设备时)
|
||||||
if self.gpu_scheduler.enabled:
|
if device_index is not None:
|
||||||
self.gpu_scheduler.release(device_index)
|
self.gpu_scheduler.release(device_index)
|
||||||
|
|
||||||
# 停止租约续期
|
# 停止租约续期
|
||||||
|
|||||||
Reference in New Issue
Block a user