MPIEvaluator feedback & future development

[EwoutH]

A discussion to gather feedback on the MPIEvaluator and discuss future development.

This discussion will be linked in the official documentation as long as the MPIEvaluator is experimental.

Previous development and discussion

• Discussion Multi-node (HPC) evaluator discussion #266: Multi-node (HPC) evaluator discussion
  Includes most of the requirements and design discussions
• Draft PR Prototype of MPIEvaluator for multi-node workloads #292: Prototype of MPIEvaluator for multi-node workloads
  Includes the prototyping and iteration of the first draft implementation
• PR Introducing MPIEvaluator: Run on multi-node HPC systems using mpi4py #299: Introducing MPIEvaluator: Run on multi-node HPC systems using mpi4py
  Includes the final implementation delivered, including extensive documentation
• PR Docs: Add MPIEvaluator tutorial for multi-node HPC systems, including DelftBlue #308: Docs: Add MPIEvaluator tutorial for multi-node HPC systems, including DelftBlue
  Includes a tutorial for using the MPIEvaluator, including an example for running on the DelftBlue SLURM-based system.

An extensive technical report on the MPIEvaluator is available: MPIEvaluatorTechnicalReport - 2023-11.pdf

[quaquel]

I fixed the initializer problem. A minimum working example is given below. My main change is to let __enter__ return the MPIEvaluator instance instead of the underlying MPIPoolExecutor instance. This then means that __exit__ is also called on the MPIEvaluator instance. On my Mac, this runs fine, and the BrokenExecutor error is gone.

import time
from mpi4py.futures import MPIPoolExecutor


def simple_task(x):
    time.sleep(0.2)
    return x * x


def common_initializer():
    # A basic initializer, doing almost nothing for now.
    pass


def common_initializer2():
    # A basic initializer, doing almost nothing for now.
    pass


class MPIEvaluator:
    def __init__(self, n_workers=None, initializer=common_initializer):
        self._pool = None
        self.n_workers = n_workers
        self.initializer = initializer

    def __enter__(self):
        self._pool = MPIPoolExecutor(max_workers=self.n_workers, initializer=self.initializer)
        print(f"MPI pool started with {self._pool._max_workers} workers")
        return self

    def map(self, fn, *iterables, timeout=None, chunksize=1, unordered=False):
        return self._pool.map(fn, *iterables, timeout=timeout, chunksize=chunksize, unordered=unordered)

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._pool:
            self._pool.shutdown(wait=True)
            print("MPI pool has been shut down")
            self._pool = None


def main():
    # First use of the MPIEvaluator
    with MPIEvaluator(n_workers=4) as pool:
        results1 = list(pool.map(simple_task, range(10)))
        print("First pool completed")

    # Second use of the MPIEvaluator
    with MPIEvaluator(n_workers=4, initializer=None) as pool:
        results2 = list(pool.map(simple_task, range(20)))
        print("Second pool completed")


if __name__ == "__main__":
    main()

[EwoutH]

Interesting approach!

[EwoutH]

Ah, you return self instead of self._pool. So self._pool didn't exist in the __exit__ method, because self was already the pool, not the evaluator.

Seems so obvious now.

You can really develop a blindspot on such things. Nicely spotted!

[quaquel]

So, I made a mistake in the code. It does, in fact, not work. Note how the initializer is None in the second call to MPIEvaluator.

However, this has given me a hint of what might be going on. As described in the warning for command line usage of mpi4py.futures, If we use mpi4py.futures to run the script, the total set of processes is automagically separated into 1 main process and a set of worker processes. In these worker processes. However, we can only run an initializer once. Since the workers are spawned by MPI, rather than by the pool as in multiprocessing, creating another MPIEvaluator simply uses the already available and fully initialized workers. So, if we pass an initializer to the second instantiation of the MPIEvaluator, all worker processes have already been initialized, and hence, we get an error.

Since we need initializers to deal with working directories and avoid having to send the entire model over using pickle every time, my current inclination is to bring back the initializer but carefully document the caveats with which it comes.

[quaquel]

I have started a new branch to do some testing with the initializer. A quick test on the lake model with 4 processes and 4000 experiments shows that the runtime goes from 14 seconds with the current version to just 4 seconds with the initializer. So, from a performance point of view, sending over the model only once and creating the experiment_runner only once really make a big difference.

I am not yet sure about the drawbacks of being able to call an initializer only once.

[quaquel]

Another twist in investigating this. We have been using -m mpi4py.futures when running our code. However, with MPI-2 this is no longer needed. We can just use mpiexec -n 1 -usize 5 python3 <python file>. In this example, we have an MPI universe of size 5, so 4 worker processes. If we switch to this, the executor will dynamically spawn the worker processes, and we can in fact run multiple pools one after another just fine.

This implies we need to update the tutorial.