Coverage for agentos/orchestration/distributed.py: 46%

1""" 

2AgentOS v1.14.2 — Distributed Orchestration (Ray-based Agent Swarm). 

3 

4受 Ray Serve / Ray Core 启发，为 AgentOS 增加分布式编排层。 

5Agent 不再局限于单进程，可以在多台机器上组成 Swarm， 

6自动负载均衡、容错恢复、跨节点通信。 

7 

8Core features: 

9- RayAgentActor: Ray Actor 封装的 Agent 实例 

10- DistSwarmCoordinator: 分布式 Swarm 协调器 

11- AgentPlacementStrategy: 智能 Agent 放置（CPU/GPU/内存感知） 

12- DistTaskQueue: 分布式任务队列（Ray 原生） 

13- CrossNodeBus: 跨节点消息总线 

14- FaultTolerance: Actor 重启/状态恢复 

15 

16Architecture: 

17 DistSwarmCoordinator (head node) 

18 ├── RayAgentActor[0] (worker node 1) 

19 │ ├── ToolAgent instance 

20 │ └── Local memory store 

21 ├── RayAgentActor[1] (worker node 2) 

22 ├── ... 

23 └── DistTaskQueue → automatic load balancing 

24 

25Usage: 

26 coordinator = DistSwarmCoordinator(num_workers=4) 

27 coordinator.start() 

28 result = await coordinator.submit(task="Summarize all PDFs in /data/") 

29 coordinator.shutdown() 

30""" 

31 

32from __future__ import annotations 

33 

34import time 

35import uuid 

36from collections.abc import Callable 

37from dataclasses import dataclass 

38from enum import StrEnum 

39from typing import ( 

40 Any, 

41) 

42 

43# ── Ray Optional Import ──────────────────── 

44 

45_HAS_RAY = False 

46try: 

47 import ray 

48 

49 _HAS_RAY = True 

50except ImportError: 

51 ray = None # type: ignore 

52 

53 

54def _require_ray(): 

55 """Raise helpful error if ray is not installed.""" 

56 if not _HAS_RAY: 

57 raise ImportError( 

58 "The distributed orchestration module requires 'ray'. " 

59 "Install it with: pip install ray" 

60 ) 

61 

62 

63# ── Data Models ───────────────────────────── 

64 

65 

66@dataclass 

67class AgentPlacementSpec: 

68 """Agent placement specification for distributed deployment.""" 

69 

70 cpu: float = 1.0 

71 gpu: float = 0.0 

72 memory_mb: int = 512 

73 node_affinity: str | None = None 

74 strategy: str = "spread" # spread | pack | custom 

75 

76 

77class DistTaskStatus(StrEnum): 

78 PENDING = "pending" 

79 RUNNING = "running" 

80 COMPLETED = "completed" 

81 FAILED = "failed" 

82 CANCELLED = "cancelled" 

83 

84 

85class DistAgentStatus(StrEnum): 

86 IDLE = "idle" 

87 BUSY = "busy" 

88 OFFLINE = "offline" 

89 

90 

91class PlacementStrategy(StrEnum): 

92 SPREAD = "spread" 

93 PACK = "pack" 

94 RANDOM = "random" 

95 CUSTOM = "custom" 

96 

97 

98@dataclass 

99class DistSwarmConfig: 

100 """Configuration for distributed swarm orchestrator.""" 

101 

102 num_workers: int = 4 

103 cpus_per_worker: float = 1.0 

104 gpus_per_worker: float = 0.0 

105 memory_per_worker_mb: int = 1024 

106 placement: PlacementStrategy = PlacementStrategy.SPREAD 

107 heartbeat_interval: float = 5.0 

108 task_timeout: float = 300.0 

109 

110 

111@dataclass 

112class DistTaskRecord: 

113 """Record of a distributed task.""" 

114 

115 task_id: str 

116 status: DistTaskStatus = DistTaskStatus.PENDING 

117 assigned_actor: str | None = None 

118 created_at: float = 0.0 

119 started_at: float | None = None 

120 completed_at: float | None = None 

121 result: Any = None 

122 error: str | None = None 

123 

124 

125class CrossNodeMailbox: 

126 """Mailbox for cross-node message passing.""" 

127 

128 def __init__(self, mailbox_id: str = ""): 

129 self.mailbox_id = mailbox_id or uuid.uuid4().hex[:8] 

130 self._messages: list[dict[str, Any]] = [] 

131 

132 async def send(self, message: dict[str, Any]) -> None: 

133 self._messages.append(message) 

134 

135 async def receive(self, timeout: float = 5.0) -> dict[str, Any] | None: 

136 if self._messages: 

137 return self._messages.pop(0) 

138 return None 

139 

140 async def receive_all(self) -> list[dict[str, Any]]: 

141 msgs = list(self._messages) 

142 self._messages.clear() 

143 return msgs 

144 

145 

146class CrossNodeBus: 

147 """Cross-node message bus for distributed communication.""" 

148 

149 def __init__(self, bus_id: str = ""): 

150 self.bus_id = bus_id or uuid.uuid4().hex[:8] 

151 self._mailboxes: dict[str, CrossNodeMailbox] = {} 

152 self._subscribers: dict[str, list[Callable]] = {} 

153 

154 def create_mailbox(self, name: str = "") -> CrossNodeMailbox: 

155 mbox = CrossNodeMailbox(name) 

156 self._mailboxes[mbox.mailbox_id] = mbox 

157 return mbox 

158 

159 def get_mailbox(self, mailbox_id: str) -> CrossNodeMailbox | None: 

160 return self._mailboxes.get(mailbox_id) 

161 

162 async def broadcast(self, topic: str, payload: dict[str, Any]) -> None: 

163 for cb in self._subscribers.get(topic, []): 

164 try: 

165 await cb(payload) 

166 except Exception: 

167 pass 

168 

169 def subscribe(self, topic: str, callback: Callable) -> None: 

170 self._subscribers.setdefault(topic, []).append(callback) 

171 

172 

173# ── Ray Agent Actor (only if ray is available) ── 

174 

175if _HAS_RAY: 

176 

177 @ray.remote 

178 class RayAgentActor: 

179 """Ray Actor wrapping an Agent instance for distributed execution.""" 

180 

181 def __init__(self, actor_name: str = "", node_id: str = ""): 

182 self.name = actor_name or uuid.uuid4().hex[:8] 

183 self.node_id = node_id or ray.get_runtime_context().get_node_id() 

184 self.status = DistAgentStatus.IDLE 

185 self.tasks_completed: int = 0 

186 self.tasks_failed: int = 0 

187 self._shutdown: bool = False 

188 

189 def get_status(self) -> dict[str, Any]: 

190 return { 

191 "name": self.name, 

192 "node_id": self.node_id, 

193 "status": self.status.value, 

194 "tasks_completed": self.tasks_completed, 

195 "tasks_failed": self.tasks_failed, 

196 } 

197 

198 async def execute(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]: 

199 self.status = DistAgentStatus.BUSY 

200 try: 

201 result = {"task_id": task_id, "status": "ok", "data": payload} 

202 self.tasks_completed += 1 

203 return result 

204 except Exception as e: 

205 self.tasks_failed += 1 

206 return {"task_id": task_id, "status": "error", "error": str(e)} 

207 finally: 

208 self.status = DistAgentStatus.IDLE 

209 

210 def shutdown(self) -> None: 

211 self._shutdown = True 

212 self.status = DistAgentStatus.OFFLINE 

213 

214else: 

215 # Placeholder when ray is not installed 

216 class RayAgentActor: 

217 def __init__(self, *args, **kwargs): 

218 _require_ray() 

219 

220 

221# ── Distributed Task Queue ────────────────── 

222 

223 

224class DistTaskQueue: 

225 """Distributed task queue with load balancing.""" 

226 

227 def __init__(self, max_size: int = 1000): 

228 self.max_size = max_size 

229 self._queue: list[DistTaskRecord] = [] 

230 self._results: dict[str, Any] = {} 

231 

232 def submit(self, payload: dict[str, Any], timeout: float = 300.0) -> DistTaskRecord: 

233 task_id = uuid.uuid4().hex[:16] 

234 record = DistTaskRecord( 

235 task_id=task_id, 

236 created_at=time.time(), 

237 ) 

238 self._queue.append(record) 

239 return record 

240 

241 def get_result(self, task_id: str, timeout: float = 30.0) -> Any: 

242 return self._results.get(task_id) 

243 

244 def mark_complete(self, task_id: str, result: Any) -> None: 

245 self._results[task_id] = result 

246 for rec in self._queue: 

247 if rec.task_id == task_id: 

248 rec.status = DistTaskStatus.COMPLETED 

249 rec.result = result 

250 rec.completed_at = time.time() 

251 

252 def list_pending(self) -> list[DistTaskRecord]: 

253 return [r for r in self._queue if r.status == DistTaskStatus.PENDING] 

254 

255 

256# ── Distributed Swarm Coordinator ─────────── 

257 

258 

259class DistSwarmCoordinator: 

260 """Coordinates a distributed swarm of agent actors.""" 

261 

262 def __init__( 

263 self, 

264 config: DistSwarmConfig | None = None, 

265 num_workers: int = 4, 

266 ): 

267 self.config = config or DistSwarmConfig(num_workers=num_workers) 

268 self._actors: list[Any] = [] 

269 self.bus = CrossNodeBus() 

270 self._started: bool = False 

271 

272 async def start(self) -> None: 

273 _require_ray() 

274 if not ray.is_initialized(): 

275 ray.init(ignore_reinit_error=True) 

276 for i in range(self.config.num_workers): 

277 actor = RayAgentActor.remote( # type: ignore[union-attr] 

278 actor_name=f"worker-{i}", 

279 ) 

280 self._actors.append(actor) 

281 self._started = True 

282 

283 async def stop(self) -> None: 

284 for actor in self._actors: 

285 try: 

286 actor.shutdown.remote() # type: ignore[union-attr] 

287 except Exception: 

288 pass 

289 self._actors.clear() 

290 self._started = False 

291 

292 async def submit(self, payload: dict[str, Any], timeout: float = 300.0) -> Any: 

293 _require_ray() 

294 if not self._actors: 

295 await self.start() 

296 actor = self._actors[0] # Simple round-robin 

297 result_ref = actor.execute.remote( # type: ignore[union-attr] 

298 uuid.uuid4().hex[:16], payload 

299 ) 

300 try: 

301 return ray.get(result_ref, timeout=timeout) 

302 except Exception as e: 

303 return {"error": str(e)} 

304 

305 def is_running(self) -> bool: 

306 return self._started 

307 

308 def actor_count(self) -> int: 

309 return len(self._actors) 

310 

311 

312def quick_start(num_workers: int = 4) -> DistSwarmCoordinator: 

313 """Quickly create and start a distributed swarm coordinator.""" 

314 return DistSwarmCoordinator(num_workers=num_workers)