Support for asynchronous step() functions

[wang-boyu]

Motivation

Currently agents' step() functions are executed sequentially through AgentSet's do() and shuffle_do() functions. If there are I/O operations involved in step() (e.g., disk or network I/O operations), they can potentially reduce performance and reponsiveness of the model. For example, in the upcoming Mesa-LLM project, agents will most likely need to query large language models at each step, through either local deployment or API calls.

Note that sequential execution of agents' step() functions may be necessary, depending on specific model logic. For example in Schelling's segregation model, each agent's decision depends on other agents' earlier actions, so step() functions need to be called one after another. Conversely, in models where agents can observe and act simultaneously, or if I/O operations can be performed independently, these expensive I/Os might be a bottleneck.

Potential Solution?

I'm wondering whether these cases can benefit from asynchronous execution of agents' step() functions. There could be a new async step function (or with any other function name) defined by users such as:

class MyAgent(mesa.Agent):
    ...

    # new method for async io
    async def async_step(self):
        # expensive I/O operations
        ....

    ...

But AgentSet will need to be updated accordingly, for instance with new async versions of do() and shuffle_do():

class AgentSet(MutableSet, Sequence):
    ...

    async def do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
        # use asyncio module to gather results from agents' `async_step()` functions
        await ...

    # and similarly for `shuffle_do()`
    async def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
        await ...

However, this means that model.step() functions will also need to be async:

class MyModel(mesa.Model):
    ...

    async def async_step(self):
        await self.agents.async_do("async_step")

As a result, SolaraViz will not work unless we update it to use async model step function too.

I also thought about letting users create their own AgentSet class:

class CustomAgentSet(mesa.AgentSet):
    ...

    async def do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
        # use asyncio module to gather results from agents' `async_step()` functions
        await ...

    # and similarly for `shuffle_do()`
    async def shuffle_do(self, method: str | Callable, *args, **kwargs) -> AgentSet:
        await ...

But mesa.Model uses our own mesa.AgentSet and there's no way of inserting a custom AgentSet. Even if it can be done, there is still the same problem of converting model.step() to async and making it work with SolaraViz.

Any better ideas?

[Spartan-71]

Here is what I found while researching about async support in Solara:

1. Task decorator:

• @solara.lab.task decorator turns a function or coroutine function into a task object.
• Allows you to run heavy and time-consuming tasks (API calls, I/O operations) in the background while UI is responsive.
• The task's state can be monitored through these attributes: .not_called, .pending, .finished, .cancelled, and '.error`.
• Reference: Docs Link

2. Use_task hook:

• hook that runs a function or coroutine function as a task and returns a result.
• unlike @task decorator the result is not shared globally.
• Reference: Docs Link

So the decorator and hook can be useful in integrating async step function support with Solara visualization.