Rich with multiprocessing.Pool.starmap() ?

[gatto]

I have used the following code but it doesn't work as expected; the progress bar only moves at all after all the processes are done with the computation.

with Progress() as progress:
    task_id = progress.add_task("[cyan]Completed...", total=len(comb))
    with multiprocessing.Pool() as pool:
        miaos = pool.starmap(do_kmeans, kmeans_params)
        for miao in miaos:
            progress.advance(task_id)

I tried to use some code similar to this, however the difference is I need to use starmap() because I need multiple parameters.

Does anyone know how to make this work? Thanks.

[willmcgugan]

I think the problem is that starmap returns a list which mean it has to wait for all the results before returning.

Try starmap_async with a callback. The callback is called after each individual task is complete. Something like the following:

pool.starmap(do_kmeans, kmeans_params, callback=lambda: progress.advance(task_id))

[gatto]

Thanks for the answer. I tried the following, and it's a bit better because the progress bar gets updated, but only once. (for one task).

with Progress() as progress:
    task_id = progress.add_task("[cyan]Working...", total=len(comb))
    with multiprocessing.Pool() as pool:
        miaos = pool.starmap_async(do_kmeans, algo_params, callback=lambda x: progress.advance(task_id))
        results = results.append(miaos.get(), ignore_index=True)

[willmcgugan]

It's unlikely to be an issue with Rich. Could your do_kmeans be throwing an error? Try setting error_callback=print. Are you setting the correct total?

If after debugging you still think the issue may be the Progress bars, try to produce a working piece of code you can paste here. Substitute your do_kmeans for a function which sleeps 10 seconds.

[gatto]

I did not even think the issue would be in rich, I just wrote because I'm a student programmer trying to learn. But here is code that only updates the progress bar once.

import time
import multiprocessing
from rich.progress import Progress


def do_kmeans(param1):
    time.sleep(5)
    return 2


if __name__ == "__main__":
    algo_params = []
    results = []
    for x in range(100):
        algo_params.append([x])

    with Progress() as progress:
        task_id = progress.add_task("[cyan]Working...", total=len(algo_params))
        with multiprocessing.Pool() as pool:
            miaos = pool.starmap_async(do_kmeans, algo_params, callback=lambda x: progress.advance(task_id))
            results.extend(miaos.get())

    print(results)

[willmcgugan]

I was mistaken about that callback. I thought it would be called for each result, but it seems to be called at the end with all results.

You can use imap which yields each result as they are calculated. The following works for me:

import time
import multiprocessing
from rich.progress import Progress


def do_kmeans(param1):
    time.sleep(5)    
    return param1

def call_do_kmeans(params):
    return do_kmeans(*params)

if __name__ == "__main__":
    algo_params = []
    results = []
    for x in range(100):
        algo_params.append([x])

    with Progress() as progress:
        task_id = progress.add_task("[cyan]Working...", total=len(algo_params))
        with multiprocessing.Pool(processes=5) as pool:
            for result in pool.imap(call_do_kmeans, algo_params):
                print(result)
                results.append(result)            
                progress.advance(task_id)

    print(results)

[gatto]

Thank you very much. I appreciate the help and the very beautiful library.