Coverage for agentlib/utils/multi_agent_system.py: 59%
158 statements
« prev ^ index » next coverage.py v7.4.4, created at 2026-08-13 10:24 +0000
1"""
2Module containing a local agency to test any LocalMASAgency system
3without the need of cloneMAP.
4"""
6import abc
7import json
8import logging
9import multiprocessing
10import threading
11from pathlib import Path
12from typing import List, Dict, Union, Any
13import time
15import pandas as pd
16from pydantic import (
17 field_validator,
18 ConfigDict,
19 BaseModel,
20 PrivateAttr,
21 Field,
22 FilePath,
23)
25from agentlib.core import Agent, Environment
26from agentlib.core.agent import AgentConfig
27from agentlib.utils.load_config import load_config
29logger = logging.getLogger(__name__)
32class MAS(BaseModel):
33 """Parent class for all MAS"""
35 model_config = ConfigDict(arbitrary_types_allowed=True)
37 agent_configs: List[Union[dict, FilePath, str]]
38 env: Union[Environment, dict, FilePath] = Field(
39 default_factory=Environment,
40 title="env",
41 description="The environment for the agents.",
42 )
43 variable_logging: bool = Field(
44 default=False,
45 title="variable_logging",
46 description="Enable variable logging in all agents with sampling rate of environment.",
47 )
48 use_direct_callback_databroker: bool = Field(
49 default=False,
50 description="If True, the `DirectCallbackDataBroker` will be used in all agents"
51 )
52 _agent_configs: Dict[str, AgentConfig] = PrivateAttr(default={})
54 def __init__(self, **data: Any) -> None:
55 """Add all agents as Agent object"""
56 super().__init__(**data)
57 for agent_config in self.agent_configs:
58 self.add_agent(config=agent_config)
60 @field_validator("agent_configs")
61 @classmethod
62 def setup_agents(cls, agent_configs):
63 """Load agent configs and add them."""
64 cfgs = []
65 for cfg in agent_configs:
66 cfgs.append(load_config(cfg, config_type=AgentConfig))
67 return cfgs
69 def add_agent(self, config: AgentConfig):
70 """
71 Add an agent to the local agency with the
72 given agent config.
74 Args:
75 config Dict: agent config
76 """
78 if self.variable_logging:
79 if isinstance(self.env, dict):
80 config = self.add_agent_logger(
81 config=config, sampling=self.env.get("t_sample", 60)
82 )
83 else:
84 config = self.add_agent_logger(
85 config=config, sampling=self.env.config.t_sample
86 )
87 if config.use_direct_callback_databroker and not self.use_direct_callback_databroker:
88 logger.warning(
89 "Agent %s explicitly sets use_direct_callback_databroker=True, "
90 "won't apply the MAS.use_direct_callback_databroker=False setting.",
91 config.id
92 )
93 else:
94 config.use_direct_callback_databroker = self.use_direct_callback_databroker
96 self._agent_configs[config.id] = config.model_copy()
97 logger.info("Registered agent %s in agency", config.id)
99 @staticmethod
100 def add_agent_logger(config: AgentConfig, sampling=60) -> AgentConfig:
101 """Adds the AgentLogger to the list of configs.
103 Args:
104 config dict: The config to be updated
105 sampling=
106 """
107 # Add Logger config
108 filename = f"variable_logs//Agent_{config.id}_Logger.log"
109 cfg = {
110 "module_id": "AgentLogger",
111 "type": "AgentLogger",
112 "t_sample": sampling,
113 "values_only": True,
114 "filename": filename,
115 "overwrite_log": True,
116 "clean_up": False,
117 }
118 config.modules.append(cfg)
119 return config
121 @abc.abstractmethod
122 def run(self, until):
123 """
124 Run the MAS.
125 Args:
126 until: The time until which the simulation should run.
128 Returns:
130 """
131 raise NotImplementedError("'run' is not implemented by the parent class MAS.")
134class LocalMASAgency(MAS):
135 """
136 Local LocalMASAgency agency class which holds the agents in a common environment,
137 executes and terminates them.
138 """
140 _agents: Dict[str, Agent] = PrivateAttr(default={})
142 @field_validator("env")
143 @classmethod
144 def setup_env(cls, env):
145 """Setup the env if a config is given."""
146 if isinstance(env, Environment):
147 return env
148 if isinstance(env, (Path, str)):
149 if Path(env).exists():
150 with open(env, "r") as f:
151 env = json.load(f)
152 return Environment(config=env)
154 def add_agent(self, config: AgentConfig):
155 """Also setup the agent directly"""
156 super().add_agent(config=config)
157 self.setup_agent(id=config.id)
159 def stop_agency(self):
160 """Stop all threads"""
161 logger.info("Stopping agency")
162 self.terminate_agents()
164 _gui_process: multiprocessing.Process = PrivateAttr(default=None)
166 def show_gui(self):
167 """
168 Interactively visualizes the dependencies between the agents,
169 their modules and variables using Dash Cytoscape.
170 """
171 from agentlib.utils.plotting.dependency_graph import show_dependency_graph
172 self._gui_process = show_dependency_graph(self)
174 def stop_gui(self, prompt="--- Press Enter to stop the GUI and finish the script ---"):
175 """
176 Blocks the script until the user presses Enter, then safely
177 terminates the Dash GUI process.
178 """
179 try:
180 if prompt:
181 time.sleep(3)
182 input(f"\n{prompt}\n")
183 except KeyboardInterrupt:
184 pass
185 finally:
186 self._gui_process.terminate()
187 self._gui_process.join()
188 self._gui_process = None
189 print("Terminated GUI.")
191 def run(self, until):
192 """Execute the LocalMASAgency and terminate it after run is finished"""
193 self.env.run(until=until)
194 self.stop_agency()
196 def __enter__(self):
197 """Enable 'with' statement"""
198 return self
200 def __exit__(self, exc_type, exc_val, exc_tb):
201 """On exit in 'with' statement, stop the agency"""
202 self.stop_agency()
204 def terminate_agents(self):
205 """Terminate all agents modules."""
206 logger.info("Terminating all agent modules")
207 for agent in self._agents.values():
208 agent.terminate()
210 def setup_agent(self, id: str) -> Agent:
211 """Setup the agent matching the given id"""
212 # pylint: disable=redefined-builtin
213 agent = Agent(env=self.env, config=self._agent_configs[id])
214 self._agents[agent.id] = agent
215 return agent
217 def get_agent(self, id: str) -> Agent:
218 """Get the agent matching the given id"""
219 # pylint: disable=redefined-builtin, inconsistent-return-statements
220 try:
221 return self._agents[id]
222 except KeyError:
223 KeyError(f"Given id '{id}' is not in the set of agents.")
225 def get_results(self, cleanup: bool = False) -> Dict[str, pd.DataFrame]:
226 """
227 Get all results of the agentLogger
228 Args:
229 cleanup: If true, read files are deleted.
231 Returns:
232 Dict[str, pd.DataFrame]: key is the agent_id, value the dataframe
233 """
234 results = {}
235 for agent in self._agents.values():
236 new_res = agent.get_results(cleanup=cleanup)
237 results[agent.id] = new_res
238 return results
243class LocalCloneMAPAgency(LocalMASAgency):
244 """
245 Local LocalMASAgency agency class which tries to mimic cloneMAP
246 behaviour for the local execution.
247 """
249 def run(self, until=None):
250 pass # Already running
252 def terminate_agents(self):
253 """Terminate all agents modules."""
254 logger.info("Can't terminate agents yet in this MAS")
256 def setup_agent(self, id: str):
257 """Setup the agent matching the given id"""
259 # pylint: disable=redefined-builtin
260 def _get_ag(env, ag_config):
261 ag = Agent(env=Environment(config=env), config=ag_config)
262 ag.env.run()
263 return ag
265 thread = threading.Thread(
266 target=_get_ag,
267 kwargs={
268 "env": self.env.config.model_copy(),
269 "ag_config": self._agent_configs[id].copy(),
270 },
271 )
272 thread.start()
273 self._agents[id] = thread
276def agent_process(
277 agent_config: Union[dict, FilePath],
278 until: float,
279 env: Union[dict, FilePath],
280 results_dict: dict,
281 cleanup=True,
282 log_level=logging.ERROR,
283):
284 """
285 Function to initialize and start an agent in its own process.
286 Collects results from the agent and stores them
287 in the passed results_dict.
288 Args:
289 cleanup:
290 agent_config: Config for an agent.
291 until: Simulation runtime
292 env: config for an environment
293 results_dict: dict from process manager
294 log_level: the log level for this process
296 Returns:
298 """
299 logging.basicConfig(level=log_level)
300 env = Environment(config=env)
301 agent = Agent(config=agent_config, env=env)
302 agent.env.run(until=until)
303 results = agent.get_results(cleanup)
304 for mod in agent.modules:
305 mod.terminate()
306 results_dict[agent.id] = results
309class MultiProcessingMAS(MAS):
310 """
311 Helper class to conveniently run multi-agent-systems in separate processes.
312 """
314 env: Union[dict, FilePath] = Field(
315 default_factory=lambda: Environment(config={"rt": True}),
316 title="env",
317 description="The environment for the agents.",
318 )
319 cleanup: bool = Field(
320 default=False,
321 description="Whether agents should clean the results files after " "running.",
322 )
323 log_level: int = Field(
324 default=logging.ERROR, description="Loglevel to set for the processes."
325 )
327 _processes: List[multiprocessing.Process] = PrivateAttr(default=[])
328 _results_dict: Dict[str, pd.DataFrame] = PrivateAttr(default={})
330 @field_validator("env")
331 @classmethod
332 def setup_env(cls, env):
333 """Setup the env if a config is given."""
334 if isinstance(env, Environment):
335 env = env.config.model_dump()
336 elif isinstance(env, (Path, str)):
337 if Path(env).exists():
338 with open(env, "r") as f:
339 env = json.load(f)
340 assert env.setdefault("rt", True), (
341 "Synchronization between processes relies on time, RealTimeEnvironment "
342 "is required."
343 )
344 return env
346 def __init__(self, **data: Any) -> None:
347 super().__init__(**data)
348 manager = multiprocessing.Manager()
349 self._results_dict = manager.dict()
351 def run(self, until):
352 """Execute the multi-agent-system in parallel and terminate it after
353 run is finished"""
354 for agent in self._agent_configs.values():
355 kwargs = {
356 "agent_config": agent,
357 "until": until,
358 "env": self.env,
359 "results_dict": self._results_dict,
360 "cleanup": self.cleanup,
361 "log_level": self.log_level,
362 }
363 process = multiprocessing.Process(
364 target=agent_process, name=agent.id, kwargs=kwargs
365 )
366 self._processes.append(process)
367 process.start()
368 for process in self._processes:
369 process.join()
371 def get_results(self) -> Dict[str, pd.DataFrame]:
372 """
373 Get all results of the agentLogger
374 Returns:
375 Dict[str, pd.DataFrame]: key is the agent_id, value the dataframe
376 """
377 return dict(self._results_dict)