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 dataclasses import dataclass 

37from enum import Enum 

38from typing import ( 

39 Any, Callable, Dict, List, Optional, 

40) 

41 

42# ── Ray Optional Import ──────────────────── 

43 

44_HAS_RAY = False 

45try: 

46 import ray 

47 _HAS_RAY = True 

48except ImportError: 

49 ray = None # type: ignore 

50 

51 

52def _require_ray(): 

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

54 if not _HAS_RAY: 

55 raise ImportError( 

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

57 "Install it with: pip install ray" 

58 ) 

59 

60 

61# ── Data Models ───────────────────────────── 

62 

63@dataclass 

64class AgentPlacementSpec: 

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

66 cpu: float = 1.0 

67 gpu: float = 0.0 

68 memory_mb: int = 512 

69 node_affinity: Optional[str] = None 

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

71 

72 

73class DistTaskStatus(str, Enum): 

74 PENDING = "pending" 

75 RUNNING = "running" 

76 COMPLETED = "completed" 

77 FAILED = "failed" 

78 CANCELLED = "cancelled" 

79 

80 

81class DistAgentStatus(str, Enum): 

82 IDLE = "idle" 

83 BUSY = "busy" 

84 OFFLINE = "offline" 

85 

86 

87class PlacementStrategy(str, Enum): 

88 SPREAD = "spread" 

89 PACK = "pack" 

90 RANDOM = "random" 

91 CUSTOM = "custom" 

92 

93 

94@dataclass 

95class DistSwarmConfig: 

96 """Configuration for distributed swarm orchestrator.""" 

97 num_workers: int = 4 

98 cpus_per_worker: float = 1.0 

99 gpus_per_worker: float = 0.0 

100 memory_per_worker_mb: int = 1024 

101 placement: PlacementStrategy = PlacementStrategy.SPREAD 

102 heartbeat_interval: float = 5.0 

103 task_timeout: float = 300.0 

104 

105 

106@dataclass 

107class DistTaskRecord: 

108 """Record of a distributed task.""" 

109 task_id: str 

110 status: DistTaskStatus = DistTaskStatus.PENDING 

111 assigned_actor: Optional[str] = None 

112 created_at: float = 0.0 

113 started_at: Optional[float] = None 

114 completed_at: Optional[float] = None 

115 result: Any = None 

116 error: Optional[str] = None 

117 

118 

119class CrossNodeMailbox: 

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

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

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

123 self._messages: List[Dict[str, Any]] = [] 

124 

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

126 self._messages.append(message) 

127 

128 async def receive(self, timeout: float = 5.0) -> Optional[Dict[str, Any]]: 

129 if self._messages: 

130 return self._messages.pop(0) 

131 return None 

132 

133 async def receive_all(self) -> List[Dict[str, Any]]: 

134 msgs = list(self._messages) 

135 self._messages.clear() 

136 return msgs 

137 

138 

139class CrossNodeBus: 

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

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

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

143 self._mailboxes: Dict[str, CrossNodeMailbox] = {} 

144 self._subscribers: Dict[str, List[Callable]] = {} 

145 

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

147 mbox = CrossNodeMailbox(name) 

148 self._mailboxes[mbox.mailbox_id] = mbox 

149 return mbox 

150 

151 def get_mailbox(self, mailbox_id: str) -> Optional[CrossNodeMailbox]: 

152 return self._mailboxes.get(mailbox_id) 

153 

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

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

156 try: 

157 await cb(payload) 

158 except Exception: 

159 pass 

160 

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

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

163 

164 

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

166 

167if _HAS_RAY: 

168 

169 @ray.remote 

170 class RayAgentActor: 

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

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

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

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

175 self.status = DistAgentStatus.IDLE 

176 self.tasks_completed: int = 0 

177 self.tasks_failed: int = 0 

178 self._shutdown: bool = False 

179 

180 def get_status(self) -> Dict[str, Any]: 

181 return { 

182 "name": self.name, 

183 "node_id": self.node_id, 

184 "status": self.status.value, 

185 "tasks_completed": self.tasks_completed, 

186 "tasks_failed": self.tasks_failed, 

187 } 

188 

189 async def execute( 

190 self, task_id: str, payload: Dict[str, Any] 

191 ) -> Dict[str, Any]: 

192 self.status = DistAgentStatus.BUSY 

193 try: 

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

195 self.tasks_completed += 1 

196 return result 

197 except Exception as e: 

198 self.tasks_failed += 1 

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

200 finally: 

201 self.status = DistAgentStatus.IDLE 

202 

203 def shutdown(self) -> None: 

204 self._shutdown = True 

205 self.status = DistAgentStatus.OFFLINE 

206 

207else: 

208 # Placeholder when ray is not installed 

209 class RayAgentActor: 

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

211 _require_ray() 

212 

213 

214# ── Distributed Task Queue ────────────────── 

215 

216class DistTaskQueue: 

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

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

219 self.max_size = max_size 

220 self._queue: List[DistTaskRecord] = [] 

221 self._results: Dict[str, Any] = {} 

222 

223 def submit( 

224 self, payload: Dict[str, Any], timeout: float = 300.0 

225 ) -> DistTaskRecord: 

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

227 record = DistTaskRecord( 

228 task_id=task_id, 

229 created_at=time.time(), 

230 ) 

231 self._queue.append(record) 

232 return record 

233 

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

235 return self._results.get(task_id) 

236 

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

238 self._results[task_id] = result 

239 for rec in self._queue: 

240 if rec.task_id == task_id: 

241 rec.status = DistTaskStatus.COMPLETED 

242 rec.result = result 

243 rec.completed_at = time.time() 

244 

245 def list_pending(self) -> List[DistTaskRecord]: 

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

247 

248 

249# ── Distributed Swarm Coordinator ─────────── 

250 

251class DistSwarmCoordinator: 

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

253 def __init__( 

254 self, 

255 config: Optional[DistSwarmConfig] = None, 

256 num_workers: int = 4, 

257 ): 

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

259 self._actors: List[Any] = [] 

260 self.bus = CrossNodeBus() 

261 self._started: bool = False 

262 

263 async def start(self) -> None: 

264 _require_ray() 

265 if not ray.is_initialized(): 

266 ray.init(ignore_reinit_error=True) 

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

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

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

270 ) 

271 self._actors.append(actor) 

272 self._started = True 

273 

274 async def stop(self) -> None: 

275 for actor in self._actors: 

276 try: 

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

278 except Exception: 

279 pass 

280 self._actors.clear() 

281 self._started = False 

282 

283 async def submit( 

284 self, payload: Dict[str, Any], timeout: float = 300.0 

285 ) -> Any: 

286 _require_ray() 

287 if not self._actors: 

288 await self.start() 

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

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

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

292 ) 

293 try: 

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

295 except Exception as e: 

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

297 

298 def is_running(self) -> bool: 

299 return self._started 

300 

301 def actor_count(self) -> int: 

302 return len(self._actors) 

303 

304 

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

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

307 return DistSwarmCoordinator(num_workers=num_workers)