imap in SerialPool and MPIPool

[edwardreed81]

It'd be sweet if imap was implemented in MPIPool and SerialPool like in MultiPool (inherited) so that I could get a progress bar for big jobs using tqdm.

Right now I'm using:

with schwimmbad.choose_pool(mpi=args.mpi, processes=args.np) as pool:
    # make it so we can use imap in serial and mpi mode
    if not isinstance(pool, schwimmbad.MultiPool):
        pool.imap = pool.map
    main(pool)

and

with tqdm(total=nmc) as pbar:
    for i, d in enumerate(pool.imap(worker, arg_gen)):
        dat.append(d)
        pbar.update()

This works fine for SerialPool since map doesn't block (I guess?), but when using MPI it doesn't update the progress bar until it is entirely finished.

Previously I had MPI implemented just using mpi4py directly and was able to get the behavior I wanted, but I like the idea of using a single code for all possible run scenarios.

[vsmagalhaes]

I am also interested in this enhancement, is there any progress on this?

Best

[AlecThomson]

+1 on this feature. Would really love to have it for tqdm usage. Thanks!