216 lines
8.8 KiB
Python
216 lines
8.8 KiB
Python
"""Agent 池 Redis 存储层。"""
|
||
import time
|
||
from typing import Any
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from app import constants as C
|
||
from app.constants import AgentStatus
|
||
from app.models.schemas import AgentInfo
|
||
|
||
|
||
class AgentRepo:
|
||
def __init__(self, redis: aioredis.Redis):
|
||
self.redis = redis
|
||
|
||
@staticmethod
|
||
def _info_key(agent_id: str) -> str:
|
||
return C.AGENT_INFO_KEY.format(agent_id=agent_id)
|
||
|
||
@staticmethod
|
||
def _tag_key(tag: str) -> str:
|
||
return C.AGENT_TAG_KEY.format(tag=tag)
|
||
|
||
# ---------- 写入 ----------
|
||
async def upsert(self, agent: AgentInfo, *, heartbeat: bool = False) -> bool:
|
||
"""写入 Agent 信息，返回是否新建。非心跳时为全量注册。"""
|
||
key = self._info_key(agent.agent_id)
|
||
existed = await self.redis.exists(key)
|
||
if not existed:
|
||
await self.redis.sadd(C.AGENT_ALL_SET, agent.agent_id)
|
||
# 建立标签索引
|
||
for tag in agent.agent_tags:
|
||
await self.redis.sadd(self._tag_key(tag), agent.agent_id)
|
||
mapping: dict[str, Any] = {
|
||
"agent_id": agent.agent_id,
|
||
"endpoint": agent.endpoint,
|
||
"agent_tags": ",".join(agent.agent_tags),
|
||
"max_concurrent": str(agent.max_concurrent),
|
||
"current_load": str(agent.current_load),
|
||
"priority": str(agent.priority),
|
||
"last_heartbeat": str(agent.last_heartbeat),
|
||
"status": agent.status.value,
|
||
"create_time": str(agent.create_time),
|
||
}
|
||
await self.redis.hset(key, mapping=mapping)
|
||
# 心跳 ZSet 记录时间戳
|
||
await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent.agent_id: agent.last_heartbeat})
|
||
return not existed
|
||
|
||
async def update_status(self, agent_id: str, status: AgentStatus) -> bool:
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
await self.redis.hset(key, "status", status.value)
|
||
# 离线（心跳失联）与不可用均退出心跳 ZSet，避免被误判/干扰
|
||
if status in (AgentStatus.OFFLINE, AgentStatus.UNAVAILABLE):
|
||
await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
|
||
return True
|
||
|
||
async def set_unavailable(self, agent_id: str) -> bool:
|
||
"""后台手动置为不可用（unavailable），心跳不会复活。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
await self.redis.hset(key, "status", AgentStatus.UNAVAILABLE.value)
|
||
await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
|
||
return True
|
||
|
||
async def set_available(self, agent_id: str) -> bool:
|
||
"""后台手动恢复可用（ready），清除不可用状态。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
await self.redis.hset(key, "status", AgentStatus.READY.value)
|
||
# 恢复就绪时补齐心跳时间戳，避免被误判为超时
|
||
await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: time.time()})
|
||
return True
|
||
|
||
async def mark_idle(self, agent_id: str) -> bool:
|
||
"""任务结束/取消后标记 Agent 空闲。
|
||
|
||
若 agent 已被手动置为不可用（unavailable），保持不可用、不复活；
|
||
否则置回 ready。
|
||
"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
cur = await self.get(agent_id)
|
||
if cur is None:
|
||
return False
|
||
if cur.status == AgentStatus.UNAVAILABLE:
|
||
return True
|
||
await self.redis.hset(key, "status", AgentStatus.READY.value)
|
||
return True
|
||
|
||
async def beat(self, agent_id: str, current_load: int) -> bool:
|
||
"""更新心跳时间与负载，按状态机规则流转状态。
|
||
|
||
- 不可用（unavailable）：保持不可用，不因心跳复活。
|
||
- 离线（offline，心跳失联）：恢复 ready。
|
||
- ready/processing/stopping：保持原状态。
|
||
"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
now = time.time()
|
||
cur = await self.get(agent_id)
|
||
if cur is None:
|
||
return False
|
||
if cur.status == AgentStatus.UNAVAILABLE:
|
||
status = AgentStatus.UNAVAILABLE
|
||
elif cur.status == AgentStatus.OFFLINE:
|
||
status = AgentStatus.READY
|
||
else:
|
||
status = cur.status
|
||
await self.redis.hset(
|
||
key,
|
||
mapping={
|
||
"last_heartbeat": str(now),
|
||
"current_load": str(current_load),
|
||
"status": status.value,
|
||
},
|
||
)
|
||
await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: now})
|
||
return True
|
||
|
||
async def update_priority(self, agent_id: str, priority: int) -> bool:
|
||
"""更新 Agent 调度优先级。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
await self.redis.hset(key, "priority", str(priority))
|
||
return True
|
||
|
||
async def adjust_load(self, agent_id: str, delta: int) -> bool:
|
||
"""原子调整 Agent 负载（delta 可为正/负），负载下限为 0。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
# hincrby 为原子增量操作，兼容 fakeredis（不支持 Lua eval）
|
||
new = await self.redis.hincrby(key, "current_load", delta)
|
||
if new < 0:
|
||
await self.redis.hset(key, "current_load", "0")
|
||
return True
|
||
|
||
# ---------- 读取 ----------
|
||
async def get(self, agent_id: str) -> AgentInfo | None:
|
||
raw = await self.redis.hgetall(self._info_key(agent_id))
|
||
if not raw:
|
||
return None
|
||
return AgentInfo(
|
||
agent_id=raw.get("agent_id", ""),
|
||
endpoint=raw.get("endpoint", ""),
|
||
agent_tags=[t for t in raw.get("agent_tags", "").split(",") if t],
|
||
max_concurrent=int(raw.get("max_concurrent", 1) or 1),
|
||
current_load=int(raw.get("current_load", 0) or 0),
|
||
priority=int(raw.get("priority", 3) or 3),
|
||
last_heartbeat=float(raw.get("last_heartbeat", 0) or 0),
|
||
status=AgentStatus(raw.get("status", AgentStatus.READY.value)),
|
||
create_time=float(raw.get("create_time", 0) or 0),
|
||
)
|
||
|
||
async def online(self) -> list[AgentInfo]:
|
||
# 可参与调度的状态：ready（空闲）与 processing（忙碌但仍有并发容量）。
|
||
# offline（心跳失联）、unavailable（手动禁用）与 stopping（正在停止）不参与调度。
|
||
return [a for a in await self.all() if a.status in (AgentStatus.READY, AgentStatus.PROCESSING)]
|
||
|
||
async def all(self) -> list[AgentInfo]:
|
||
ids = list(await self.redis.smembers(C.AGENT_ALL_SET))
|
||
out = []
|
||
for aid in ids:
|
||
a = await self.get(aid)
|
||
if a:
|
||
out.append(a)
|
||
return out
|
||
|
||
async def by_tags(self, tags: list[str]) -> list[AgentInfo]:
|
||
"""按标签索引取候选 Agent 池（取所有标签的交集，无标签则返回全部可调度）。
|
||
|
||
只返回可调度的 agent（ready/processing），offline/unavailable/stopping 被排除。
|
||
"""
|
||
if not tags:
|
||
return await self.online()
|
||
keys = [self._tag_key(t) for t in tags]
|
||
if len(keys) == 1:
|
||
ids = list(await self.redis.smembers(keys[0]))
|
||
else:
|
||
await self.redis.sinterstore("agent:tmp:intersect", keys)
|
||
ids = list(await self.redis.smembers("agent:tmp:intersect"))
|
||
await self.redis.delete("agent:tmp:intersect")
|
||
out = []
|
||
for aid in ids:
|
||
a = await self.get(aid)
|
||
if a and a.status in (AgentStatus.READY, AgentStatus.PROCESSING):
|
||
out.append(a)
|
||
return out
|
||
|
||
async def remove(self, agent_id: str) -> None:
|
||
"""注销：删除信息、心跳 ZSet、标签索引、全量索引。"""
|
||
a = await self.get(agent_id)
|
||
if a:
|
||
for tag in a.agent_tags:
|
||
await self.redis.srem(self._tag_key(tag), agent_id)
|
||
await self.redis.delete(self._info_key(agent_id))
|
||
await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
|
||
await self.redis.srem(C.AGENT_ALL_SET, agent_id)
|
||
|
||
async def stale_agents(self, timeout: float) -> list[str]:
|
||
"""返回超过 timeout 秒未心跳的 Agent ID（基于 ZSet score）。"""
|
||
cutoff = time.time() - timeout
|
||
scored = await self.redis.zrangebyscore(C.AGENT_HEARTBEAT_ZSET, 0, cutoff)
|
||
return list(scored)
|
||
|
||
async def heartbeat_map(self) -> dict[str, float]:
|
||
"""返回 {agent_id: last_heartbeat}。"""
|
||
return {k: float(v) for k, v in await self.redis.zrange(C.AGENT_HEARTBEAT_ZSET, 0, -1, withscores=True)} |