academy.manager¶

academy/manager.py

Manager ¶

Manager(
    exchange_client: ExchangeClient[ExchangeTransportT],
    executors: (
        Executor
        | MutableMapping[str, Executor | None]
        | None
    ) = None,
    *,
    default_executor: str = "event_loop",
    max_retries: int = 0
)

Bases: Generic[ExchangeTransportT], NoPickleMixin

Launch and manage running agents.

A manager is used to launch agents in the current event loop or in Executors and interact with/manage those agents.

Tip

This class can be used as a context manager. Upon exiting the context, running agents will be shutdown, any agent handles created by the manager will be closed, and the executors will be shutdown.

Tip

When using ProcessPoolExecutors, use the initializer argument to configure logging in the worker processes that will execute agents.

Note

The manager takes ownership of the exchange client and executors, meaning the manager will clean up those resources when the manager is closed.

Parameters:

exchange_client (ExchangeClient[ExchangeTransportT]) –

Exchange client.
executors (Executor | MutableMapping[str, Executor | None] | None, default: None ) –

An executor instance or mapping of names to executors to use to run agents. If a single executor is provided, it is set as the default executor with name 'default', overriding any value of default_executor. Setting executor=None (the default) or passing {'name': None, ...} can be used to launch agents in the current event loop.
default_executor (str, default: 'event_loop' ) –

Specify the name of the default executor to use when not specified in launch().
max_retries (int, default: 0 ) –

Maximum number of times to retry running an agent if it exits with an error.

Raises:

ValueError –

If default_executor is specified but does not exist in executors.

Source code in academy/manager.py

def __init__(
    self,
    exchange_client: ae.ExchangeClient[ExchangeTransportT],
    executors: Executor
    | MutableMapping[str, Executor | None]
    | None = None,
    *,
    default_executor: str = 'event_loop',
    max_retries: int = 0,
) -> None:
    self._executors: dict[str, Executor | None] = {'event_loop': None}
    if isinstance(executors, Executor):
        self._executors.update({'default': executors})
        default_executor = 'default'
    elif isinstance(executors, MutableMapping):
        self._executors.update(executors)

    if default_executor not in self._executors:
        raise ValueError(
            f'No executor named "{default_executor}" was provided to '
            'use as the default.',
        )

    self._exchange_client = exchange_client
    self._exchange_factory = exchange_client.factory()
    self._user_id = self._exchange_client.client_id
    self._default_executor = default_executor
    self._max_retries = max_retries
    self._closed = asyncio.Event()

    self._handles: dict[AgentId[Any], Handle[Any]] = {}
    self._acbs: dict[AgentId[Any], _ACB[Any]] = {}

    logger.info(
        'Initialized manager (%s)',
        self.user_id,
        extra={'academy.user_id': self.user_id},
    )

exchange_client `property` ¶

exchange_client: ExchangeClient[ExchangeTransportT]

User client for the exchange.

exchange_factory `property` ¶

exchange_factory: ExchangeFactory[ExchangeTransportT]

Client factory for the exchange.

user_id `property` ¶

user_id: EntityId

Exchange client user ID of this manager.

from_exchange_factory `async` `classmethod` ¶

from_exchange_factory(
    factory: ExchangeFactory[ExchangeTransportT],
    executors: (
        Executor
        | MutableMapping[str, Executor | None]
        | None
    ) = None,
    *,
    default_executor: str = "event_loop",
    max_retries: int = 0
) -> Self

Instantiate a new exchange client and manager from a factory.

Source code in academy/manager.py

@classmethod
async def from_exchange_factory(
    cls,
    factory: ae.ExchangeFactory[ExchangeTransportT],
    executors: Executor
    | MutableMapping[str, Executor | None]
    | None = None,
    *,
    default_executor: str = 'event_loop',
    max_retries: int = 0,
) -> Self:
    """Instantiate a new exchange client and manager from a factory."""
    client = await factory.create_user_client()
    return cls(
        client,
        executors,
        default_executor=default_executor,
        max_retries=max_retries,
    )

close `async` ¶

close(close_exchange: bool = True) -> None

Shutdown the manager and cleanup resources.

Request all running agents to shut down.
Wait for all running agents to shut down.
Close the exchange client. (if close_exchange)
Shutdown the executors.
Raise an exceptions returned by agents.

Parameters:

close_exchange (bool, default: True ) –

If the exchange_client should be closed.

Raises:

Exception –

Any exceptions raised by agents.

Source code in academy/manager.py

async def close(self, close_exchange: bool = True) -> None:
    """Shutdown the manager and cleanup resources.

    1. Request all running agents to shut down.
    1. Wait for all running agents to shut down.
    1. Close the exchange client. (if close_exchange)
    1. Shutdown the executors.
    1. Raise an exceptions returned by agents.

    Args:
        close_exchange: If the exchange_client should be closed.

    Raises:
        Exception: Any exceptions raised by agents.
    """
    self._closed.set()  # Stop from launching retries

    for acb in self._acbs.values():
        if not acb.task.done():
            handle = self.get_handle(acb.agent_id)
            with contextlib.suppress(AgentTerminatedError):
                await handle.shutdown()
    logger.debug('Requested shutdown from all agents')

    for acb in self._acbs.values():
        await acb.task
    logger.debug('All agents have completed')

    if close_exchange:
        await self.exchange_client.close()

    for executor in self._executors.values():
        if executor is not None:
            executor.shutdown(wait=True, cancel_futures=True)

    exceptions = (acb.task.exception() for acb in self._acbs.values())
    exceptions_only = tuple(e for e in exceptions if e is not None)
    raise_exceptions(
        exceptions_only,
        message='Caught failures in agent while shutting down.',
    )

    logger.info(
        'Closed manager (%s)',
        self.user_id,
        extra={'academy.user_id': self.user_id},
    )

add_executor ¶

add_executor(name: str, executor: Executor) -> Self

Add an executor to the manager.

Note

It is not possible to remove an executor as this could create complications if an agent is already running in that executor.

Parameters:

name (str) –

Name of the executor used when launching agents.
executor (Executor) –

Executor instance.

Returns:

Self –

Self for chaining.

Raises:

ValueError –

If an executor with name already exists.

Source code in academy/manager.py

def add_executor(self, name: str, executor: Executor) -> Self:
    """Add an executor to the manager.

    Note:
        It is not possible to remove an executor as this could create
        complications if an agent is already running in that executor.

    Args:
        name: Name of the executor used when launching agents.
        executor: Executor instance.

    Returns:
        Self for chaining.

    Raises:
        ValueError: If an executor with `name` already exists.
    """
    if name in self._executors:
        raise ValueError(f'Executor named "{name}" already exists.')
    self._executors[name] = executor
    return self

set_default_executor ¶

set_default_executor(name: str) -> Self

Set the default executor by name.

Parameters:

name (str) –

Name of the executor to use as default.

Returns:

Self –

Self for chaining.

Raises:

ValueError –

If no executor with name exists.

Source code in academy/manager.py

def set_default_executor(self, name: str) -> Self:
    """Set the default executor by name.

    Args:
        name: Name of the executor to use as default.

    Returns:
        Self for chaining.

    Raises:
        ValueError: If no executor with `name` exists.
    """
    if name not in self._executors:
        raise ValueError(f'An executor named "{name}" does not exist.')
    self._default_executor = name
    return self

launch `async` ¶

launch(
    agent: AgentT | type[AgentT],
    *,
    args: tuple[Any, ...] | None = None,
    kwargs: dict[str, Any] | None = None,
    config: RuntimeConfig | None = None,
    executor: str | None = None,
    submit_kwargs: dict[str, Any] | None = None,
    name: str | None = None,
    registration: AgentRegistration[AgentT] | None = None,
    init_logging: bool = False,
    loglevel: int | str = INFO,
    logfile: str | None = None
) -> Handle[AgentT]

Launch a new agent with a specified agent.

Parameters:

agent (AgentT | type[AgentT]) –

Agent instance the agent will implement or the agent type that will be initialized on the worker using args and kwargs.
args (tuple[Any, ...] | None, default: None ) –

Positional arguments used to initialize the agent. Ignored if agent is already an instance.
kwargs (dict[str, Any] | None, default: None ) –

Keyword arguments used to initialize the agent. Ignored if agent is already an instance.
config (RuntimeConfig | None, default: None ) –

Agent run configuration.
executor (str | None, default: None ) –

Name of the executor instance to use. If None, uses the default executor.
submit_kwargs (dict[str, Any] | None, default: None ) –

Additional arguments to pass to the submit function of the executor (i.e. a resource specification).
name (str | None, default: None ) –

Readable name of the agent used when registering a new agent.
registration (AgentRegistration[AgentT] | None, default: None ) –

If None, a new agent will be registered with the exchange.
init_logging (bool, default: False ) –

Initialize logging before running agent.
loglevel (int | str, default: INFO ) –

Level of logging.
logfile (str | None, default: None ) –

Location to write logs.

Returns:

Handle[AgentT] –

Handle (client bound) used to interact with the agent.

Raises:

RuntimeError –

If registration is provided and an agent with that ID has already been executed.

Source code in academy/manager.py

async def launch(  # noqa: PLR0913
    self,
    agent: AgentT | type[AgentT],
    *,
    args: tuple[Any, ...] | None = None,
    kwargs: dict[str, Any] | None = None,
    config: RuntimeConfig | None = None,
    executor: str | None = None,
    submit_kwargs: dict[str, Any] | None = None,
    name: str | None = None,
    registration: AgentRegistration[AgentT] | None = None,
    init_logging: bool = False,
    loglevel: int | str = logging.INFO,
    logfile: str | None = None,
) -> Handle[AgentT]:
    """Launch a new agent with a specified agent.

    Args:
        agent: Agent instance the agent will implement or the
            agent type that will be initialized on the worker using
            `args` and `kwargs`.
        args: Positional arguments used to initialize the agent.
            Ignored if `agent` is already an instance.
        kwargs: Keyword arguments used to initialize the agent.
            Ignored if `agent` is already an instance.
        config: Agent run configuration.
        executor: Name of the executor instance to use. If `None`, uses
            the default executor.
        submit_kwargs: Additional arguments to pass to the submit function
            of the executor (i.e. a resource specification).
        name: Readable name of the agent used when registering a new agent.
        registration: If `None`, a new agent will be registered with
            the exchange.
        init_logging: Initialize logging before running agent.
        loglevel: Level of logging.
        logfile: Location to write logs.

    Returns:
        Handle (client bound) used to interact with the agent.

    Raises:
        RuntimeError: If `registration` is provided and an agent with
            that ID has already been executed.
    """
    executor = executor if executor is not None else self._default_executor
    assert executor is not None
    executor_instance = self._executors[executor]

    if registration is None:
        agent_type = agent if isinstance(agent, type) else type(agent)
        registration = await self.register_agent(agent_type, name=name)
    elif registration.agent_id in self._acbs:
        raise RuntimeError(
            f'{registration.agent_id} has already been executed.',
        )

    if init_logging and (
        executor_instance is None
        or isinstance(executor_instance, ThreadPoolExecutor)
    ):
        warnings.warn(
            f'`init_logging` was specified for agent '
            f'{registration.agent_id} running in the same process as the '
            f'Manager. `init_logging` should only be called once per '
            f'process.',
            stacklevel=2,
        )

    agent_id = registration.agent_id

    spec = _RunSpec(
        agent=agent,
        config=RuntimeConfig() if config is None else config,
        exchange_factory=self.exchange_factory,
        registration=registration,
        agent_args=() if args is None else args,
        agent_kwargs={} if kwargs is None else kwargs,
        submit_kwargs={} if submit_kwargs is None else submit_kwargs,
        init_logging=init_logging,
        loglevel=loglevel,
        logfile=logfile,
    )

    task = asyncio.create_task(
        self._run_agent(executor_instance, spec),
        name=f'manager-{executor}-run-{agent_id}',
    )

    acb = _ACB(agent_id=agent_id, executor=executor, task=task)
    self._acbs[agent_id] = acb
    handle = self.get_handle(agent_id)
    logger.info(
        'Launched agent (%s; %s)',
        agent_id,
        agent,
        extra={'academy.agent_id': agent_id, 'academy.agent': agent},
    )
    if executor_instance is not None:
        self._warn_executor_overloaded(executor_instance, executor)
    return handle

get_handle ¶

get_handle(
    agent: AgentId[AgentT] | AgentRegistration[AgentT],
) -> Handle[AgentT]

Create a new handle to an agent.

A handle acts like a reference to a remote agent, enabling a user to manage the agent or asynchronously invoke actions.

Parameters:

agent (AgentId[AgentT] | AgentRegistration[AgentT]) –

Agent ID or registration indicating the agent to create a handle to. The agent must be registered with the same exchange that this manager is a client of.

Returns:

Handle[AgentT] –

Handle to the agent.

Source code in academy/manager.py

def get_handle(
    self,
    agent: AgentId[AgentT] | AgentRegistration[AgentT],
) -> Handle[AgentT]:
    """Create a new handle to an agent.

    A handle acts like a reference to a remote agent, enabling a user
    to manage the agent or asynchronously invoke actions.

    Args:
        agent: Agent ID or registration indicating the agent to create
            a handle to. The agent must be registered with the same
            exchange that this manager is a client of.

    Returns:
        Handle to the agent.
    """
    agent_id = agent if isinstance(agent, AgentId) else agent.agent_id
    handle = self._handles.get(agent_id, None)
    if handle is not None:
        return handle
    handle = Handle(agent_id)
    self._handles[agent_id] = handle
    return handle

register_agent `async` ¶

register_agent(
    agent: type[AgentT], *, name: str | None = None
) -> AgentRegistration[AgentT]

Register a new agent with the exchange.

Parameters:

agent (type[AgentT]) –

Agent type of the agent.
name (str | None, default: None ) –

Optional display name for the agent.

Returns:

AgentRegistration[AgentT] –

Agent registration info that can be passed to
AgentRegistration[AgentT] –

launch().

Source code in academy/manager.py

async def register_agent(
    self,
    agent: type[AgentT],
    *,
    name: str | None = None,
) -> AgentRegistration[AgentT]:
    """Register a new agent with the exchange.

    Args:
        agent: Agent type of the agent.
        name: Optional display name for the agent.

    Returns:
        Agent registration info that can be passed to
        [`launch()`][academy.manager.Manager.launch].
    """
    return await self.exchange_client.register_agent(agent, name=name)

running ¶

running() -> set[AgentId[Any]]

Get a set of IDs of all running agents.

Returns:

set[AgentId[Any]] –

Set of agent IDs corresponding to all agents launched by this manager that have not completed yet.

Source code in academy/manager.py

def running(self) -> set[AgentId[Any]]:
    """Get a set of IDs of all running agents.

    Returns:
        Set of agent IDs corresponding to all agents launched by this \
        manager that have not completed yet.
    """
    running: set[AgentId[Any]] = set()
    for acb in self._acbs.values():
        if not acb.task.done():
            running.add(acb.agent_id)
    return running

shutdown `async` ¶

shutdown(
    agent: AgentId[Any] | Handle[Any],
    *,
    blocking: bool = True,
    raise_error: bool = True,
    terminate: bool | None = None,
    timeout: float | None = None
) -> None

Shutdown a launched agent.

Parameters:

agent (AgentId[Any] | Handle[Any]) –

ID or handle to the launched agent.
blocking (bool, default: True ) –

Wait for the agent to exit before returning.
raise_error (bool, default: True ) –

Raise the error returned by the agent if blocking=True.
terminate (bool | None, default: None ) –

Override the termination behavior of the agent defined in the RuntimeConfig.
timeout (float | None, default: None ) –

Optional timeout is seconds when blocking=True.

Raises:

BadEntityIdError –

If an agent with agent_id was not launched by this manager.
TimeoutError –

If timeout was exceeded while blocking for agent.

Source code in academy/manager.py

async def shutdown(
    self,
    agent: AgentId[Any] | Handle[Any],
    *,
    blocking: bool = True,
    raise_error: bool = True,
    terminate: bool | None = None,
    timeout: float | None = None,
) -> None:
    """Shutdown a launched agent.

    Args:
        agent: ID or handle to the launched agent.
        blocking: Wait for the agent to exit before returning.
        raise_error: Raise the error returned by the agent if
            `blocking=True`.
        terminate: Override the termination behavior of the agent defined
            in the [`RuntimeConfig`][academy.runtime.RuntimeConfig].
        timeout: Optional timeout is seconds when `blocking=True`.

    Raises:
        BadEntityIdError: If an agent with `agent_id` was not
            launched by this manager.
        TimeoutError: If `timeout` was exceeded while blocking for agent.
    """
    agent_id = agent.agent_id if isinstance(agent, Handle) else agent

    if agent_id not in self._acbs:
        raise BadEntityIdError(agent_id) from None
    if self._acbs[agent_id].task.done():
        return

    handle = self.get_handle(agent_id)
    with contextlib.suppress(AgentTerminatedError):
        await handle.shutdown(terminate=terminate)

    if blocking:
        await self.wait(
            {agent_id},
            raise_error=raise_error,
            timeout=timeout,
        )

wait `async` ¶

wait(
    agents: Iterable[AgentId[Any] | Handle[Any]],
    *,
    raise_error: bool = False,
    return_when: str = ALL_COMPLETED,
    timeout: float | None = None
) -> None

Wait for launched agents to complete.

Note

Calling wait() is only valid after launch() has succeeded.

Parameters:

agents (Iterable[AgentId[Any] | Handle[Any]]) –

An iterable of agent IDs or handles to wait on.
raise_error (bool, default: False ) –

Raise errors returned by completed agents.
return_when (str, default: ALL_COMPLETED ) –

Indicate when this function should return. The same as asyncio.wait().
timeout (float | None, default: None ) –

Optional timeout in seconds to wait for agents.

Raises:

BadEntityIdError –

If an agent was not launched by this manager.
TimeoutError –

If timeout was exceeded while waiting for agents.
Exception –

Any exception raised by an agent that completed due to a failure and raise_error=True is set.

Source code in academy/manager.py

async def wait(
    self,
    agents: Iterable[AgentId[Any] | Handle[Any]],
    *,
    raise_error: bool = False,
    return_when: str = asyncio.ALL_COMPLETED,
    timeout: float | None = None,
) -> None:
    """Wait for launched agents to complete.

    Note:
        Calling `wait()` is only valid after `launch()` has succeeded.

    Args:
        agents: An iterable of agent IDs or handles to wait on.
        raise_error: Raise errors returned by completed agents.
        return_when: Indicate when this function should return. The
            same as [`asyncio.wait()`][asyncio.wait].
        timeout: Optional timeout in seconds to wait for agents.

    Raises:
        BadEntityIdError: If an agent was not launched by this manager.
        TimeoutError: If `timeout` was exceeded while waiting for agents.
        Exception: Any exception raised by an agent that completed due
            to a failure and `raise_error=True` is set.
    """
    agent_ids = {
        agent if isinstance(agent, AgentId) else agent.agent_id
        for agent in agents
    }

    if len(agent_ids) == 0:
        return

    agent_tasks: list[asyncio.Task[None]] = []
    for agent_id in agent_ids:
        try:
            agent_tasks.append(self._acbs[agent_id].task)
        except KeyError:
            raise BadEntityIdError(agent_id) from None

    done, pending = await asyncio.wait(
        agent_tasks,
        return_when=return_when,
        timeout=timeout,
    )

    if len(done) == 0:
        raise TimeoutError(
            f'No agents completed within {timeout} seconds: '
            f'{len(pending)} pending agent(s).',
        )
    elif return_when == asyncio.ALL_COMPLETED and len(pending) > 0:
        raise TimeoutError(
            f'Not all agents completed within {timeout} seconds: '
            f'{len(pending)} pending agent(s).',
        )

    if raise_error:
        exceptions = (task.exception() for task in agent_tasks)
        exceptions_only = tuple(e for e in exceptions if e is not None)
        raise_exceptions(
            exceptions_only,
            message='Waited agents raised the following exceptions.',
        )

academy.manager¶

Manager ¶

exchange_client property ¶

exchange_factory property ¶

user_id property ¶

from_exchange_factory async classmethod ¶

close async ¶

add_executor ¶

set_default_executor ¶

launch async ¶

get_handle ¶

register_agent async ¶

running ¶

shutdown async ¶

wait async ¶

exchange_client `property` ¶

exchange_factory `property` ¶

user_id `property` ¶

from_exchange_factory `async` `classmethod` ¶

close `async` ¶

launch `async` ¶

register_agent `async` ¶

shutdown `async` ¶

wait `async` ¶