Task scheduling customization with MPIPoolExecutor

[FabioLuporini]

Hi again,

I guess the answer to the following question will be "No, you're just trying to use the wrong tool for your problem" -- fundamentally because I'm after data-parallelism, not task-parallelism -- but I'll try anyway :)

Long story short: I would need executor.map(args, ...) to work such that args[0] is always assigned to rank0, args[1] to rank1, etc.

Context: In my library, I will perform several executor.map(...), interleaved by user code; there's a certain amount of state that I need to preserve in the ranks, which is essentially args-specific (basically, rank0 will construct some stuff based on what it sees in args[0]).
Also, when I execute a map, I'm sure by construction the previous map has terminated.

I looked into mpi4py and I see the task scheduling is handled by the executor's underlying pool. So perhaps through some form of subclassing and specialization I may achieve what I need?

Any thoughts or is there something I've missed in mpi4py that I could rather use ?

Thanks (again!)

[dalcinl]

In the master branch, mpi4py has gained some support for the execution of parallel talks, and that has some resemblance to what you are looking after.
https://mpi4py.readthedocs.io/en/latest/mpi4py.futures.html#parallel-tasks
https://mpi4py.readthedocs.io/en/latest/mpi4py.futures.html#computing-pi-parallel-task

However, I think you will not be able to control where exactly each task runs. I mean, mapping args[0] exactly to rank 0 for every successive execution of map. Without that mapping guaranteed, the thing may be useless to you. Or maybe that's not really important, or you may be able to refactor things on your side.

Unfortunately, I don't think subclassing mpi4py.futures stuff will help you with anything. The scheduler is hidden deep in the implementation and uses a rather functional style of programming that is not amenable to the sort of hacks that you would like to inject.

I will think a bit whether I can improve the parallel task implementation to support your use case, but I'm a bit skeptical I'll be able to do so without disrupting or penalizing the primary focus of the library.

fundamentally because I'm after data-parallelism

Perhaps you should move away of mpi4py.futures, it is indeed the wrong abstraction for data parallelism. Plain MPI is much better suited for data-parallel algorithms. In fact, I would argue that most MPI codes out there deal with data parallelism. mpi4py.futures was just my humble attempt to provide a simple solution of those in need of task parallelism. Looks like my thing is so good at hiding the MPI mess from users that people want to use if for everything, but at some point these users will hit a problem.

What you are trying to do should be quite easy to implement using MPI spawn. 1) You spawn children processes, 2) you broadcast the callable function from one side to the other, 3) you send and recv input data with arg = intercomm.scatter(args), 4) workers process the input to produce a result, 5) you collect the results with results = intercomm.gather(result).

Do you have some previous experience with MPI such that you can implement the thing I just described?

[FabioLuporini]

I'll start backwards...

Do you have some previous experience with MPI such that you can implement the thing I just described?

Yes I do, thanks

What you say makes a lot of sense. Just before reading your reply, I managed to achieve what I needed, but admittedly it's with some horrible hacks and unsure how robust it is... though it seems to be doing the trick for now.

The gist of it (pseudocode) is the following:

class Fleet(MPIPoolExecutor):
    _make_pool = staticmethod(my_own_pool)

def my_own_pool(executor):
    # Same Pool as in mpi4py.futures._lib, but now passing in my own manager)
    return Pool(executor, my_own_manager)  

class OrderedDeque(Stack):

    def put(self, item):
        # Horrible and hacky
        # basically keeping the structure sorted each time there's an insertion
        # so that the next time we execute `map` we'll get ranks 0,1,2,... in order
        elems = list(self) + [item]
        self.clear()
        self.extend(sorted(elems, reverse=True))

def my_own_manager(pool, options):
     # essentially a concatenation of mpi4py's _manager_spawn and 
     # _manager_comm, where the *only* difference is the use of OrderedDeque
     # in place of Stack

Yup -- this is ugly. But it was 2-3 hours of work to dive deep inside mpi4py and find a work around. I agree, as I said above, that a solution based on MPI_Spawn is immensely cleaner. I'll think about it as I refine the implementation.

I will think a bit whether I can improve the parallel task implementation to support your use case, but I'm a bit skeptical I'll be able to do so without disrupting or penalizing the primary focus of the library.

Thanks for even just thinking about this.

Looks like my thing is so good at hiding the MPI mess from users that people want to use if for everything, but at some point these users will hit a problem.

Basically yes, I think we ended up trying MPIPoolExecutor because, despite the initial feeling of not being the perfect fit as an abstraction to our needs, it would have allowed us to forget about a lot of complexity.

[dalcinl]

The gist of it (pseudocode) is the following:

You did quite a bit of hacking, much deeper that what I thought you would go. And looking at your sketch, looks like it should definitely work.

Yup -- this is ugly.

python -m this | grep practicality

Basically yes, I think we ended up trying MPIPoolExecutor

Maybe I should consider growing the executor offering in mpi4py.futures and provide something with a much simpler scheduler to support the usual SPMD task style used for data-parallel processing, all implemented around bcast/scatter/gather. The usual trouble is to define a proper API, that is how to define/redefine the submit/map(/reduce?) functionality within the SPMD style.

[FabioLuporini]

Excellent, thanks for the feedback!

I probably would be just a waste of time for you, but if you fancy a high-level discussion about the "why" and the "what", I'd be happy to. I see you're based in Kaust and I'll be there for a talk in September, I could pop by in case.

But basically this:

Maybe I should consider growing the executor offering in mpi4py.futures and provide something with a much simpler scheduler to support the usual SPMD task style used for data-parallel processing, all implemented around bcast/scatter/gather. The usual trouble is to define a proper API, that is how to define/redefine the submit/map(/reduce?) functionality within the SPMD style.

yes 💯

[dalcinl]

I see you're based in Kaust and I'll be there for a talk in September, I could pop by in case.

Sure thing... just drop me an email before your arrival.