102 lines
4.6 KiB
Python
102 lines
4.6 KiB
Python
"""规则调度服务：标签匹配→负载过滤→最优选择→任务绑定→推送下发。"""
|
||
import logging
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from app.constants import AgentStatus, TaskStatus
|
||
from app.models.schemas import AgentInfo
|
||
from app.repository.agent_repo import AgentRepo
|
||
from app.repository.task_repo import TaskRepo
|
||
from app.services.relay import RelayService
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class Scheduler:
|
||
def __init__(
|
||
self,
|
||
redis: aioredis.Redis,
|
||
task_repo: TaskRepo | None = None,
|
||
agent_repo: AgentRepo | None = None,
|
||
relay: RelayService | None = None,
|
||
):
|
||
self.redis = redis
|
||
self.task_repo = task_repo or TaskRepo(redis)
|
||
self.agent_repo = agent_repo or AgentRepo(redis)
|
||
self.relay = relay or RelayService(redis)
|
||
|
||
async def pick_candidate(self, task_tags: list[str]) -> AgentInfo | None:
|
||
"""按规则挑选最优 Agent：标签匹配→负载过滤→优先级最小且空闲。"""
|
||
candidates = await self.agent_repo.by_tags(task_tags)
|
||
# 负载过滤：剔除满载
|
||
candidates = [a for a in candidates if a.current_load < a.max_concurrent]
|
||
if not candidates:
|
||
return None
|
||
# 最优：优先级最小（越先分配），同优先级取负载最低、在线时间最长（稳定）
|
||
candidates.sort(key=lambda a: (a.priority, a.current_load, -a.create_time))
|
||
return candidates[0]
|
||
|
||
async def dispatch(self, request_id: str) -> bool:
|
||
"""将单个 pending 任务下发给匹配 Agent（绑定 → 推送 payload → 失败回退）。
|
||
|
||
成功 / 无可用 Agent / 推送失败时，都会在任务上记录 error_info（成功时清空），
|
||
便于前端与日志定位“为什么没分配”。
|
||
"""
|
||
task = await self.task_repo.get(request_id)
|
||
if not task or task.status != TaskStatus.PENDING:
|
||
return False
|
||
agent = await self.pick_candidate(task.task_tags)
|
||
if not agent:
|
||
# 暂无可用 Agent（标签不匹配 / 离线 / 满载），保持 pending，等待下次调度
|
||
await self.task_repo.update(
|
||
request_id,
|
||
error_info="无可用 Agent：标签不匹配或全部离线/满载",
|
||
)
|
||
logger.info("task no candidate request_id=%s tags=%s", request_id, task.task_tags)
|
||
return False
|
||
# 绑定 Agent，标记 running
|
||
await self.task_repo.update(
|
||
request_id,
|
||
status=TaskStatus.RUNNING,
|
||
agent_id=agent.agent_id,
|
||
progress=0,
|
||
)
|
||
await self.task_repo.mark_running(request_id)
|
||
# 记录调度前状态，供推送失败回退时恢复
|
||
prev_status = agent.status
|
||
# 标记 Agent 为处理中
|
||
await self.agent_repo.update_status(agent.agent_id, AgentStatus.PROCESSING)
|
||
# 更新 Agent 负载
|
||
await self.agent_repo.adjust_load(agent.agent_id, 1)
|
||
# 真实推送 payload 到 Agent 端点
|
||
pushed, reason = await self.relay.dispatch_command(agent.agent_id, request_id, task.payload or {})
|
||
if not pushed:
|
||
# 推送失败：回退 pending 并释放负载，按原状态恢复 Agent，等待下次调度
|
||
await self.task_repo.update(
|
||
request_id,
|
||
status=TaskStatus.PENDING,
|
||
agent_id="",
|
||
progress=0,
|
||
error_info=f"推送失败: {reason}",
|
||
)
|
||
await self.task_repo.mark_pending(request_id)
|
||
await self.agent_repo.adjust_load(agent.agent_id, -1)
|
||
# 若该 Agent 本为 PROCESSING（还有并发容量），保持 PROCESSING；否则置回 READY
|
||
restore = prev_status if prev_status == AgentStatus.PROCESSING else AgentStatus.READY
|
||
await self.agent_repo.update_status(agent.agent_id, restore)
|
||
logger.warning("task dispatch push failed, reverted request_id=%s agent=%s reason=%s",
|
||
request_id, agent.agent_id, reason)
|
||
return False
|
||
# 推送成功：清空 error_info
|
||
await self.task_repo.update(request_id, error_info="")
|
||
logger.info("task dispatched request_id=%s -> agent=%s", request_id, agent.agent_id)
|
||
return True
|
||
|
||
async def dispatch_pending(self, batch_size: int = 20) -> int:
|
||
"""调度所有可调度的 pending 任务，返回成功下发数。"""
|
||
pending = await self.task_repo.pending_tasks()
|
||
dispatched = 0
|
||
for rid in pending[:batch_size]:
|
||
if await self.dispatch(rid):
|
||
dispatched += 1
|
||
return dispatched |