79 lines
3.6 KiB
Python
79 lines
3.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 → 失败回退）。"""
|
||
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:
|
||
return False # 暂无可用 Agent，保持 pending，等待下次调度
|
||
# 绑定 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)
|
||
# 标记 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 = 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)
|
||
await self.task_repo.mark_pending(request_id)
|
||
await self.agent_repo.adjust_load(agent.agent_id, -1)
|
||
await self.agent_repo.update_status(agent.agent_id, AgentStatus.READY)
|
||
logger.warning("task dispatch push failed, reverted request_id=%s agent=%s", request_id, agent.agent_id)
|
||
return False
|
||
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 |