transports

mpi

class anemoi.inference.transports.mpi.MPITransport(couplings: Configuration, tasks: dict[str, Task], *args: Any, **kwargs: Any)

Bases: Transport

Transport implementation using MPI.

Initialize the MPITransport.

Parameters:

  • couplings (Configuration) – The couplings for the transport.

  • tasks (Dict[str, Any]) – The tasks to be executed.

start() → None

Start the transport by initializing MPI tasks.

wait() → None

Wait for all MPI tasks to complete.

send(sender: Task, target: Task, state: dict[str, Any], tag: int) → None

Send a state from the sender to the target.

Parameters:

  • sender (Task) – The task sending the state.

  • target (Task) – The task receiving the state.

  • state (State) – The state to be sent.

  • tag (int) – The tag associated with the state.

receive(receiver: Task, source: Task, tag: int) → dict[str, Any]

Receive a state from the source to the receiver.

Parameters:

  • receiver (Any) – The task receiving the state.

  • source (Any) – The task sending the state.

  • tag (int) – The tag associated with the state.

Returns:

The received state.

Return type:

Any

processes

class anemoi.inference.transports.processes.ProcessesTransport(couplings: Any, tasks: dict[str, Any], *args: Any, **kwargs: Any)

Bases: Transport

Transport implementation using processes.

Initialize the ProcessesTransport.

Parameters:

  • couplings (Any) – The couplings for the transport.

  • tasks (Dict[str, Any]) – The tasks to be executed.

child_process(task: Any) → int

Run the task in a child process.

Parameters:

task (Any) – The task to be run.

Returns:

The exit status of the child process.

Return type:

int

start() → None

Start the transport by forking processes for each task.

wait() → None

Wait for all child processes to complete and handle any errors.

send(sender: Task, target: Task, state: dict[str, Any], tag: int) → None

Send a state from the sender to the target.

Parameters:

  • sender (Any) – The task sending the state.

  • target (Any) – The task receiving the state.

  • state (Any) – The state to be sent.

  • tag (int) – The tag associated with the state.

receive(receiver: Task, source: Task, tag: int) → dict[str, Any]

Receive a state from the source to the receiver.

Parameters:

  • receiver (Any) – The task receiving the state.

  • source (Any) – The task sending the state.

  • tag (int) – The tag associated with the state.

Returns:

The received state.

Return type:

Any

threads

class anemoi.inference.transports.threads.TaskWrapper(task: Task)

Bases: object

Wraps a task to be executed in a thread.

Initialize the TaskWrapper.

Parameters:

task (Task) – The task to be wrapped.

run(transport: ThreadsTransport) → None

Run the task within the given transport.

Parameters:

transport (ThreadsTransport) – The transport in which the task is run.

class anemoi.inference.transports.threads.ThreadsTransport(couplings: Any, tasks: dict[str, Task], *args: Any, **kwargs: Any)

Bases: Transport

Transport implementation using threads.

Initialize the ThreadsTransport.

Parameters:

  • couplings (Any) – The couplings for the transport.

  • tasks (Dict[str, Any]) – The tasks to be executed.

start() → None

Start the transport by initializing and starting threads for each task.

wait() → None

Wait for all threads to complete and handle any errors.

send(sender: Task, target: Task, state: dict[str, Any], tag: int) → None

Send a state from the sender to the target.

Parameters:

  • sender (Any) – The task sending the state.

  • target (Any) – The task receiving the state.

  • state (Any) – The state to be sent.

  • tag (int) – The tag associated with the state.

receive(receiver: Task, source: Task, tag: int) → dict[str, Any]

Receive a state from the source to the receiver.

Parameters:

  • receiver (Any) – The task receiving the state.

  • source (Any) – The task sending the state.

  • tag (int) – The tag associated with the state.

Returns:

The received state.

Return type:

Any