100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
"""Agent 池服务：注册/注销/心跳/负载/离线剔除。"""
|
||
import logging
|
||
import time
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from app.constants import AgentStatus
|
||
from app.models.schemas import AgentHeartbeat, AgentInfo, AgentRegister
|
||
from app.repository.agent_repo import AgentRepo
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AgentService:
|
||
def __init__(self, redis: aioredis.Redis, repo: AgentRepo | None = None):
|
||
self.redis = redis
|
||
self.repo = repo or AgentRepo(redis)
|
||
|
||
async def register(self, reg: AgentRegister) -> AgentInfo:
|
||
agent = AgentInfo(
|
||
agent_id=reg.agent_id,
|
||
endpoint=reg.endpoint,
|
||
agent_tags=reg.agent_tags,
|
||
max_concurrent=reg.max_concurrent,
|
||
current_load=reg.current_load,
|
||
priority=reg.priority,
|
||
status=AgentStatus.READY,
|
||
create_time=time.time(),
|
||
last_heartbeat=time.time(),
|
||
)
|
||
created = await self.repo.upsert(agent)
|
||
logger.info("agent registered agent_id=%s created=%s", agent.agent_id, created)
|
||
return agent
|
||
|
||
async def unregister(self, agent_id: str) -> bool:
|
||
existed = await self.repo.get(agent_id)
|
||
if not existed:
|
||
return False
|
||
await self.repo.remove(agent_id)
|
||
logger.info("agent unregistered agent_id=%s", agent_id)
|
||
return True
|
||
|
||
async def heartbeat(self, hb: AgentHeartbeat) -> bool:
|
||
ok = await self.repo.beat(hb.agent_id, hb.current_load)
|
||
if not ok:
|
||
logger.warning("heartbeat from unknown agent agent_id=%s", hb.agent_id)
|
||
return ok
|
||
|
||
async def online(self) -> list[AgentInfo]:
|
||
return await self.repo.online()
|
||
|
||
async def all(self) -> list[AgentInfo]:
|
||
return await self.repo.all()
|
||
|
||
async def get(self, agent_id: str) -> AgentInfo | None:
|
||
return await self.repo.get(agent_id)
|
||
|
||
async def set_unavailable(self, agent_id: str) -> bool:
|
||
"""后台手动置为「不可用」（unavailable），心跳不会复活。"""
|
||
a = await self.repo.get(agent_id)
|
||
if not a:
|
||
return False
|
||
await self.repo.set_unavailable(agent_id)
|
||
logger.info("agent set unavailable agent_id=%s", agent_id)
|
||
return True
|
||
|
||
async def set_available(self, agent_id: str) -> bool:
|
||
"""后台手动恢复「可用」（ready），清除不可用状态。"""
|
||
a = await self.repo.get(agent_id)
|
||
if not a:
|
||
return False
|
||
await self.repo.set_available(agent_id)
|
||
logger.info("agent set available agent_id=%s", agent_id)
|
||
return True
|
||
|
||
async def set_priority(self, agent_id: str, priority: int) -> AgentInfo | None:
|
||
"""设置 Agent 调度优先级（1-5，越小越先分配）。"""
|
||
if priority < 1 or priority > 5:
|
||
raise ValueError("priority must be in [1, 5]")
|
||
ok = await self.repo.update_priority(agent_id, priority)
|
||
if not ok:
|
||
return None
|
||
logger.info("agent priority updated agent_id=%s priority=%s", agent_id, priority)
|
||
return await self.repo.get(agent_id)
|
||
|
||
async def purge_stale(self, timeout: float) -> list[str]:
|
||
"""剔除超时未心跳的 Agent：非 unavailable 的置为 offline（心跳失联）。
|
||
|
||
手动置为不可用的 agent 不参与心跳计时（已从心跳 ZSet 移除），保持 unavailable。
|
||
"""
|
||
stale = await self.repo.stale_agents(timeout)
|
||
purged = []
|
||
for aid in stale:
|
||
a = await self.repo.get(aid)
|
||
if a and a.status == AgentStatus.UNAVAILABLE:
|
||
continue
|
||
await self.repo.update_status(aid, AgentStatus.OFFLINE)
|
||
logger.info("agent purged (stale) agent_id=%s", aid)
|
||
purged.append(aid)
|
||
return purged |