Skip to content

Commit

Permalink
time: Use weak references to agents in schedulers
Browse files Browse the repository at this point in the history
This commit modifies the schedulers to use weak references for agent management. The update aims to improve memory efficiency and garbage collection, especially in scenarios where agents are dynamically added and removed during model execution.

It also ensures that the Agent remove() method also removes the Agent from a schedule.

Key Changes:
- Agents within schedulers are now stored as weak references (`weakref.ref[Agent]`), reducing the memory footprint and allowing for more efficient garbage collection.
- The `add` method in `BaseScheduler` and derived classes has been updated to store agents as weak references.
- The `remove` method has been modified to handle both direct agent objects and weak references, improving flexibility and robustness.
- Updated `do_each` method to iterate over a copy of agent keys and check the existence of agents before calling methods, ensuring stability during dynamic agent removal.
- Added a `agents` property to return a list of live agent instances, maintaining backward compatibility for user code that iterates over agents.
- Ensured all existing tests pass, confirming the stability and correctness of the changes.

The changes result in a more memory-efficient scheduler implementation in Mesa, while preserving existing functionality and ensuring backward compatibility.

All 27 current tests pass without any modifications, so there shouldn't be changes breaking compatibility.
  • Loading branch information
EwoutH committed Dec 17, 2023
1 parent a74b120 commit 89b61ec
Showing 1 changed file with 24 additions and 11 deletions.
35 changes: 24 additions & 11 deletions mesa/time.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from __future__ import annotations

import heapq
import weakref
from collections import defaultdict

# mypy
Expand Down Expand Up @@ -64,7 +65,7 @@ def __init__(self, model: Model) -> None:
self.model = model
self.steps = 0
self.time: TimeT = 0
self._agents: dict[int, Agent] = {}
self._agents: dict[int, weakref.ref[Agent]] = {}

def add(self, agent: Agent) -> None:
"""Add an Agent object to the schedule.
Expand All @@ -78,15 +79,17 @@ def add(self, agent: Agent) -> None:
f"Agent with unique id {agent.unique_id!r} already added to scheduler"
)

self._agents[agent.unique_id] = agent
self._agents[agent.unique_id] = weakref.ref(agent)

def remove(self, agent: Agent) -> None:
"""Remove all instances of a given agent from the schedule.
def remove(self, agent: Agent | weakref.ref[Agent]) -> None:
"""Remove an agent from the schedule.
Args:
agent: An agent object.
agent: An agent object or a weak reference to an agent.
"""
del self._agents[agent.unique_id]
agent_id = (agent() if isinstance(agent, weakref.ref) else agent).unique_id
if agent_id in self._agents:
del self._agents[agent_id]

def step(self) -> None:
"""Execute the step of all the agents, one at a time."""
Expand All @@ -102,7 +105,12 @@ def get_agent_count(self) -> int:

@property
def agents(self) -> list[Agent]:
return list(self._agents.values())
"""Return a list of live agent instances."""
return [
agent_ref()
for agent_ref in self._agents.values()
if agent_ref() is not None
]

def get_agent_keys(self, shuffle: bool = False) -> list[int]:
# To be able to remove and/or add agents during stepping
Expand All @@ -113,13 +121,18 @@ def get_agent_keys(self, shuffle: bool = False) -> list[int]:
return agent_keys

def do_each(self, method, agent_keys=None, shuffle=False):
"""Performs a method on each agent, managing weak references."""
if agent_keys is None:
agent_keys = self.get_agent_keys()
agent_keys = list(self._agents.keys())
if shuffle:
self.model.random.shuffle(agent_keys)
for agent_key in agent_keys:
if agent_key in self._agents:
getattr(self._agents[agent_key], method)()

for agent_key in list(agent_keys): # Create a copy of the list to iterate over
agent_ref = self._agents.get(agent_key)
if agent_ref is not None:
agent = agent_ref()
if agent is not None:
getattr(agent, method)()


class RandomActivation(BaseScheduler):
Expand Down

0 comments on commit 89b61ec

Please sign in to comment.