Performance benchmarks:

Would you mind separating this into 2 PRs? The agent creation one will be straightforward to merge.

The agent set stuff is more difficult. I don't want to change the internal weakkeydict to a dict while still under mesa 3.X. It might cause all kinds of unexpected results in existing models that rely on agentset having a weakkeydict. I'm also currently thinking of keeping the existing AgentSet as a weakkeydict, adding a separate dict version used only internally for model.agents and model.agents_by_type, and removing model._agents (because it's redundant from there on out). But even that might not be fully backward compatible because people might use a copy operation on their agentset. What is also poorly defined at the moment is the behavior of all methods that return an agentset.

I have reverted the creation change.

I'm also currently thinking of keeping the existing AgentSet as a weakkeydict, adding a separate dict version used only internally for model.agents and model.agents_by_type

What if we can keep AgentSet as it is and inherit a new subclass StrongAgentSet(or some better name) for dict-only purposes. It'd be more cleaner and all the methods of AgentSet would be still available to it

What if we can keep AgentSet as it is and inherit a new subclass StrongAgentSet(or some better name) for dict-only purposes. It'd be more cleaner and all the methods of AgentSet would be still available to it

Yes that is a logical next step.

The main issue however is how to properly integrate this into mesa for model.agents and model.agents_by_type. Clearly, for performance, we want to use normal dicts here. And they are controlled by model.register_agent and model.deregister_agent, which in turn is tied to agent.remove. So, memory leaks here are not a concern.

However, what happens if someone makes a copy of these hard ref agent sets?. This can happen via any of the methods with an inplace keyword argument. If you say inplace=False, you'll get a copy. And now we create the potential for a memory leak because these copies are not controlled by model.deregister_agent.

I see two solutions:

1. Don't allow inplace on any of the methods of the hard ref agent set. This avoids the implict copy and forces a user to use copy.copy. Next, we can make it clear in the docstring that making a copy is a bad idea due to the potential for memory leaks.
2. Allow inplace, but always return the current AgentSet (so with weakkeydict). It keeps the API the same between both; however, it implicitly changes the datatype of the object you are dealing with.

Allow inplace, but always return the current AgentSet (so with weakkeydict). It keeps the API the same between both; however, it implicitly changes the datatype of the object you are dealing with.

The implementation I suggested already addresses this naturally. We don't need to write any extra logic to enforce it; the inheritance structure handles it for us.

We define _StrongAgentSet to strictly override the storage mechanism:

class _StrongAgentSet(AgentSet):
    """Internal optimization for Model.agents. Uses strong refs for speed."""
    
    def _storage(self, agents):
        return dict.fromkeys(agents)

In the base AgentSet class, methods like select, shuffle, and sort explicitly call the AgentSet constructor for their return values

def select(self, ...):
    # ...
    return AgentSet(agents, self.random)  # <--- Explicitly returns a WEAK AgentSet

1. Input: You call model.agents.select(...) on the optimized _StrongAgentSet (Internal/Strong).
2. Output: The method runs, filters the agents, and returns a standard AgentSet (Public/Weak).

This automatically creates the safety boundary you described:

• Internal Engine (model.agents): Fast, Strong Refs.
• External Views (return values): Safe, Weak Refs.

For extra safety, we can explicitly override __copy__, to ensure that Python's standard copy operations also respect this "downgrade" to weak references.

Thanks for the clarification. This PR currently subclasses the weakkey version from the hard key version. I suggest changing this around.

Moreover, there are many optimizations inside the current agent set when we iterate. These really need to stay in place because they have a profound impact on performance. From memory, the one in shuffle is the most critical, but the others also matter. Basically, we cannot just change the type of the dict used internally and keep everything else the same.

Basically, we cannot just change the type of the dict used internally and keep everything else the same.

You are completely right, we cannot just swap the storage backend and expect it to work. The current AgentSet implementation is tightly coupled to WeakKeyDictionary specific APIs (like .keyrefs() and .data), which don't exist on a standard dict.

To make _StrongAgentSet work, We have to do more than just change _storage; Instead have to override every method that relies on those weakref internals.

I’ve implemented a working version where _StrongAgentSet passes all safety checks and correctly downgrades copies to weak sets. Happy to leave the final merge decision to you

Can you update the title to reflect what this PR does.

I'll try to review this tomorrow. A first quick look suggests most if fine.

What are your thoughts on updating the internals of Mesa to start using this?

What are your thoughts on updating the internals of Mesa to start using this?

In my opinion, this change isn't introducing strong ownership; it is optimizing the model.agents to match the strong ownership that already exists, removing the redundant weak-reference overhead.

I think updating model.agents to use _StrongAgentSet is the right move because:

1. Conceptually, the Model is the "source of truth" and owner of the agents. In the current implementation, model.py already enforces this ownership via the internal _agents dict. However, the public model.agents (currently a standard AgentSet) wraps these in WeakKeyDictionary, adding unnecessary overhead to check for liveness that the Model already guarantees. Switching to _StrongAgentSet removes this redundancy and optimizes the hot loop.

2. Since the Model already holds strong references (in _agents), "Lazy Deletion" (relying on GC to remove agents) is already disabled in practice. Adopting _StrongAgentSet formalizes this: agents persist until agent.remove() is explicitly called.

My view: This is a positive standard. Relying on Python's Garbage Collector for simulation logic is non-deterministic. Explicit remove() calls ensure the simulation state is stable and predictable.

Integration Plan
I recommend we apply this specifically to Model.agents

• Change: Initialize self._all_agents = _StrongAgentSet([], self.random) and self._agents_by_type inside Model.__init__.
• Constraint: Keep _StrongAgentSet strictly internal. It should never be instantiated by users directly.

I also believe this is be gonna a critical change as AgentSet directly interacts/will interact with AgentDataSet

For Schelling, you should also update agents_by_type to use the dict based agentset. Since this one is already managed by Model.deregister_agent it should be a trivial change.

Some cons that I could think for this approach are:-

1. Copying Costs
   AgentSet.copy() is usually fast (shallow copy of weakrefs). _StrongAgentSet.copy() forces the creation of a new AgentSet, which requires generating new weak references for every agent in the set. This makes operations like copy() or non-inplace shuffle() significantly heavier.

2. Polymorphism Issue
   _StrongAgentSet shadows AgentSet. If a future contributor adds a new method to AgentSet that relies on internal weakref structures (like .data or .keyrefs()) and forgets to override it in _StrongAgentSet, it will cause runtime crashes (AttributeError).

On the GitHub side, it’s running on shared, virtual runners. Performance is influenced by other workloads. Not the most stable platform for benchmarks (but free).

With 3.4.2 released, we can move towards merging this.

Small note: We have maintenance branches (or can create ones), which to we can backport specific fixes. main always points to at least a new minor version, sometimes a major, but never a patch. So don’t let that slow you down :)

it’s running on shared, virtual runners. Performance is influenced by other workloads

Agree. But I found some points that maybe we could fix from our side

1. We currently measure from the very first run. This captures the cold start penalty. Standard practice(On netlogo and ASV) is to run the model 3+ times to "warm up" the system before starting the timer.
2. gc.enable() is currently on, meaning the Garbage Collector fires randomly during runs. If GC hits during a specific run, that run looks artificially slow. The standard approach is gc.disable() during the timed loop and gc.collect() manually between runs.
3. We currently track the minimum(fastest_init). While common, this is extremely noisy on shared CI runners. The industry standard (used by NetLogo and ASV) is to calculate the median and Interquartile Range (IQR), which filters out outliers caused by system spikes.
4. We use timeit.default_timer(). We should use time.perf_counter(), which provides nanosecond resolution and is monotonic (unaffected by system clock updates).

and

5. We don't properly clean up after a run. As indicated in Memory leak in models #3179, there is still some memory issue. It would be good to call model.remove_agents() at the end of run_model.

@codebreaker32, If you have time, a PR would be welcome. Otherwise, I'll take a look hopefully tonight.

@codebreaker32, If you have time, a PR would be welcome

Sure, I am on it

3. We currently track the minimum(fastest_init). While common, this is extremely noisy on shared CI runners. The industry standard (used by NetLogo and ASV) is to calculate the median and Interquartile Range (IQR), which filters out outliers caused by system spikes.

To give some context, there are good arguments for this: While compute execution is often slower, is rarely faster. The distribution tends to be heavily skewed. When originally developing these benchmarks, using the minimum was a conscious decision with more stable results.

That said, improvements to our benchmarking are certainly welcome.

Polymorphism Issue
_StrongAgentSet shadows AgentSet. If a future contributor adds a new method to AgentSet that relies on internal weakref structures (like .data or .keyrefs()) and forgets to override it in _StrongAgentSet, it will cause runtime crashes (AttributeError).

Did we address this yet?

To give some context, there are good arguments for this: While compute execution is often slower, is rarely faster. The distribution tends to be heavily skewed. When originally developing these benchmarks, using the minimum was a conscious decision with more stable results.

After testing both approaches(results shared on #3203), I strongly agree with you :)

Polymorphism Issue
_StrongAgentSet shadows AgentSet. If a future contributor adds a new method to AgentSet that relies on internal weakref structures (like .data or .keyrefs()) and forgets to override it in _StrongAgentSet, it will cause runtime crashes (AttributeError).

Did we address this yet?

I thought it is okay for now after what @quaquel said

However, I doubt there will be major changes to AgentSet. The API has proven very stable after its introduction and early optimizations.

Could you check if there's a simple test we could add to check for this?

Could you check if there's a simple test we could add to check for this?

I think you are overthinking the issue. First, the API has been very stable. Second, if new stuff is added to agentset, we are likely to remember this and test accordingly. Third, I expect to go in and clean this up in a follow-up PR by adding an ABC AgentSet class that both weakkeydict and normal dict inherit from.

Could you check if there's a simple test we could add to check for this?

Also, DW There are enough tests for _StrongAgentSet to save us from unsafe method definition or overrides :)

Performance benchmarks:

Above benchmarks ran using old global_benchmarks.py on this PR and updated on main, Now I have merged the main into this branch. You can run again for the updated behavior

Performance benchmarks:

I am reviewing the code locally at the moment. I have a few minor comments still:

1. The typing on model for model.agents and model.agents_by_type is not entirely correct anymore since these return the new _StrongAgentSet instead of AgentSet. I guess because of inheritance it still works.
2. It might also be good to slightly expand the documentation on model.agents and model.agents_by_type, warning people to use this with care.
3. It would be good to open an issue indicating that model._agent and model._all_agents can be removed. I know this will be a messy job.

You are right since _StrongAgentSet is a subclass, the typing holds up via inheritance, but it is subtle. Also, I acknowledge that the current structure (with _agent, _all_agents, and the inheritance mix) is due for a cleanup.

I see this complete process as a 3-step roadmap to maintain backward compatibility and easy to review:

1. Refactor AgentSet: Introduce an AbstractAgentSet and make the standard AgentSet inherit from it.
2. Formalize Strong Storage: Introduce HardKeyAgentSet (inheriting from AbstractAgentSet) and formally swap it into the Model.
3. Cleanup: Remove redundant instances like model._all_agents once the new hierarchy is stable.

We can open a tracking issue for this roadmap. If you are fine with this approach, I will proceed

Sounds like a plan. Does this mean that this PR will be superseded and so does not need to be merged, or do you want to start from this PR?

Does this mean that this PR will be superseded and so does not need to be merged, or do you want to start from this PR?

I am fine either way!

If you want to get an ABC in immediately, we can merge this PR by tomorrow( I will finalize it by tonight).
Alternatively, if you prefer to have the clean architecture in place first, we can consider this PR superseded. I would then close this and start the new PR series (beginning with AbstractAgentSet) from scratch and tracking issue.

Which approach do you prefer?

Let's start from the ABC instead of merging this. The git history will be cleaner and easier to reconstruct in the future, if needed, what changed and why.