Skip to content

Introduce _HardKeyAgentSet in agents.py#3219

Merged
quaquel merged 7 commits intomesa:mainfrom
codebreaker32:refactor/agentset#2
Jan 29, 2026
Merged

Introduce _HardKeyAgentSet in agents.py#3219
quaquel merged 7 commits intomesa:mainfrom
codebreaker32:refactor/agentset#2

Conversation

@codebreaker32
Copy link
Copy Markdown
Collaborator

@codebreaker32 codebreaker32 commented Jan 28, 2026

Summary

This PR introduces _HardKeyAgentSet, that inherits from AbstractAgentSet using strong references (standard dict) instead of weak references. This eliminates WeakKeyDictionary overhead for Model-managed collections where lifecycle is explicitly controlled.

Motive

See #3128 and #3160

Implementation:

  • Inherits from AbstractAgentSet
  • Uses dict[A, None] for storage (vs WeakKeyDictionary)
  • All abstract methods implemented with dict-based optimizations
  • copying strong-ref sets creates memory leak risk. By downgrading to AgentSet, we get best of both worlds

Changes

  1. Introduced _HardKeyAgentSet
  2. Moved to_list() from AgentSet to AbstractAgentSet

✅ Fully backward compatible - no changes to existing APIs. This is a new internal class.
Ticks (2) from #3209

@github-actions
Copy link
Copy Markdown

Performance benchmarks:

Model Size Init time [95% CI] Run time [95% CI]
BoltzmannWealth small 🟢 -6.0% [-6.5%, -5.6%] 🟢 -5.1% [-5.4%, -4.9%]
BoltzmannWealth large 🟢 -3.9% [-4.5%, -3.4%] 🔵 -2.8% [-4.3%, -1.5%]
Schelling small 🔵 +0.4% [+0.0%, +0.8%] 🔵 -1.3% [-1.5%, -1.1%]
Schelling large 🔵 -1.9% [-2.6%, -1.3%] 🔵 -0.6% [-2.5%, +1.1%]
WolfSheep small 🟢 -6.2% [-6.5%, -5.9%] 🟢 -4.5% [-4.7%, -4.2%]
WolfSheep large 🟢 -6.1% [-7.0%, -5.1%] 🟢 -5.9% [-7.6%, -4.3%]
BoidFlockers small 🟢 -7.7% [-8.0%, -7.3%] 🔵 -0.2% [-0.4%, -0.0%]
BoidFlockers large 🟢 -6.2% [-6.5%, -5.8%] 🔵 -0.1% [-0.4%, +0.1%]

@codebreaker32
Copy link
Copy Markdown
Collaborator Author

As _HardKeyAgent doesn't inherit from Sequence, DeffuantWeisbuchModel is raising TypeError

        if not isinstance(population, _Sequence):
>           raise TypeError("Population must be a sequence.  "
                            "For dicts or sets, use sorted(d).")
E           TypeError: Population must be a sequence.  For dicts or sets, use sorted(d).

@quaquel
Copy link
Copy Markdown
Member

quaquel commented Jan 28, 2026

As _HardKeyAgent doesn't inherit from Sequence, DeffuantWeisbuchModel is raising TypeError

We have to fix the model anyway because of the deprecation. The source seems to be self.random.sample(self.agents, 2). Which is easily replaced with self.random.sample(self.agents.tolist(), 2).

@codebreaker32
Copy link
Copy Markdown
Collaborator Author

codebreaker32 commented Jan 28, 2026

I think to_list() is quite good improvement
For eg. In this DeffuantWeisbuchModel

# Looks like O(N), but is actually O(N²)
for _ in range(self.n):
  # random.sample calls __getitem__ repeatedly.
  # Since our __getitem__ rebuilds the list every time (O(N)), this line is expensive.
  agent_a, agent_b = self.random.sample(self.agents, 2)

Initially user was doing this, unaware of what was going under the hood
Now this API explicitly tells them what they are doing giving them chance to optimise

agents_list = self.agents.to_list()
for _ in range(self.n):
        agent_a, agent_b = self.random.sample(agents_list, 2)

# Truly O(N)

However this was only possible because loop didn't involve any addition or deletion but I think it is still
cool

@quaquel
Copy link
Copy Markdown
Member

quaquel commented Jan 28, 2026

Initially user was doing this, unaware of what was going under the hood

And this was exactly the reason why I wanted to deprecate it. I saw other models doing exactly this without realizing the overhead.

@quaquel
Copy link
Copy Markdown
Member

quaquel commented Jan 29, 2026

Coverage sits at 87%. Can you check this and see if you can add relevant tests to bring this up? I'll try to review later today.

@codebreaker32
Copy link
Copy Markdown
Collaborator Author

Can you check this and see if you can add relevant tests to bring this up?

Done

mesa/model.py Outdated

@property
def agents(self) -> AgentSet[A]:
def agents(self) -> _HardKeyAgentSet[A]:
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might it be a good idea to update the docstrings for agents and agents_by_type to highlight that they return an agentset with hard refs and should not be messed with, since they are the key agent registration system in Mesa?

Copy link
Copy Markdown
Member

@quaquel quaquel left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One remaining question regarding how to properly document this. Codewise, I am fine with merging this.

@codebreaker32
Copy link
Copy Markdown
Collaborator Author

One remaining question regarding how to properly document this. Codewise, I am fine with merging this.

I see two options:

  1. Update it in this PR only,
  2. We can update documentation in next PR where I am planning to remove the self._agents and updating documentation

If you prefer option 2, I will revert the changes in model.py

@codebreaker32
Copy link
Copy Markdown
Collaborator Author

I prefer option 2, we can finalize the storage refactor there

@quaquel
Copy link
Copy Markdown
Member

quaquel commented Jan 29, 2026

Option 2 is fine with me as well. I'll take a last look at this and merge it hopefully this afternoon.

@quaquel quaquel merged commit b0f5c81 into mesa:main Jan 29, 2026
14 checks passed
@codebreaker32 codebreaker32 deleted the refactor/agentset#2 branch January 29, 2026 14:24
@codebreaker32 codebreaker32 changed the title Introduce _HardKeyAgentSet and refactor Model storage to use it Introduce _HardKeyAgentSet in agents.py Jan 29, 2026
@codebreaker32
Copy link
Copy Markdown
Collaborator Author

PR title and description updated to reflect reverted model.py changes

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement Release notes label

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants