Implementing Parallelization within the MUSES framework

Hi everyone!

We are working to upload our quarkyonic module to MUSES, but I have a quick question. I am currently parallelizing the python script using joblib, like this:

from joblib import Parallel, delayed

def compute_point(i, N, ann, bnn, apn, bpn, para, steps, res, delta):

    binding_en, quark_fract, charge_fract, press = total_en_slope(N, ann, bnn, apn, bpn, para, steps, res, delta)

    return i, binding_en, quark_fract, charge_fract, press



results = Parallel(n_jobs = -2)(

    delayed(compute_point)(i, denmin + i\*delt, x0\[0\], x0\[1\], x0\[2\], x0\[3\], x0\[4\], res, fr_acc, der_step)

    for i in range(incr))

Here the n_jobs = -2 tells the system to use all but 1 CPU. However, I am suspicious that this won’t mesh with you celery worker implementation. Can I get a bit of guidance here on how to assign jobs to specific workers?

Your MUSES module container will be a child process of a Celery worker “worker process”. Each worker process can execute one Celery task at a time, and the number of worker processes is set by the Celery worker’s concurrency value. Our cluster worker nodes (I know, a lot of overloading of the word “worker”, but not my fault :grin:) have 16 cpu each, so I typically set the concurrency to 12 (constraining one Celery worker per node). Since MUSES workflows run a module container as a Celery task, there can be up to 12 concurrent MUSES modules running per node. While there is no core limit per task, there is a core limit per Celery worker, which I set to 12. Thus, if each MUSES module only ran single-core, then all 12 modules would have an entire core to themselves. However, if some of the concurrently running modules try to use more than 1 core, they will either get those cores if available, or the Celery worker processes will begin multi-threading to stay within the 12-core cap imposed by the Kubernetes process ultimately ruling all these subprocesses.

Your module has no control over the Celery worker that will execute it, nor is that machinery even visible to it. It’s not clear to me how many cores your process will attempt to use if you set n_jobs at -2, but it is probably either num_node_cores-1 (15) or num_concurrency-1 (11). Either way, you may as well try to be greedy and gobble up the cores if they are available rather than sacrifice them by a lower core request.

The trickier part might be the memory if it consumes a lot. Unlike multi-threading when too many cores are requested, there is no graceful fallback for memory. If the Celery worker exceeds the limit enforced by Kubernetes, that Celery worker is simply killed and restarted, taking out all the MUSES modules running with it.