Performance improvements for Regress Out

[dsm-72]

What kind of feature would you like to request?

Additional function parameters / changed functionality / changed defaults?

Please describe your wishes

I am using regress_out and it is painfully slow. Even on a system where I set n_jobs=36 and sc.settings.n_jobs = 36, each core of which has 36Gib of memory, I find that

sc.pp.regress_out(adata, ['total_counts', ], n_jobs=n_jobs)

is practically unusable. At the moment that calculation is at 985 minutes. Looking at htop while the memory is certainly allocated (191 Gib / 1.48Tb), it feels like setting n_jobs at all actually hinders performance....

I've checked the source code so I know that n_jobs
should be set correctly.

n_jobs = sett.n_jobs if n_jobs is None else n_jobs # NOTE: sett.n_jobs defaults to 1

# ...

res = Parallel(n_jobs=n_jobs)(delayed(_regress_out_chunk)(task) for task in tasks)

Looking at _regress_out_chunk, there really doesn't seem to be anything necessarily bottlenecking the performance except for either setting the chunk length in regress_out, or just the size of the dataset...

len_chunk = np.ceil(min(1000, X.shape[1]) / n_jobs).astype(int)

Am I missing something?

[flying-sheep]

@fidelram implemented this, and he’s sadly no longer active in open source software. I personally don’t have any experience with joblib, so I don’t know what could be the reason. 🤔

Generally, Python’s GIL means that Python code can’t run in parallel. I assume @fidelram used it for a good reason and there’s enough parallelizable code in there that it matters, but that’s the only lead I can think of.

[dawe]

There are two spawning approaches in joblib, processes and threads. Default is processes but sometimes it doesn't work properly (at least it happened to me), so you may want to change. Also, I see there's now dask support, so it may be worth a try.
latest docs here

[dsm-72]

@dawe any idea how to just use Dask for this without having to reimplement this function? is there something like numba's @jit and @njit decorators?

[dawe]

Not sure, according to this page function code should be changed. It would be handy to add the backend as option to regress_out

[flying-sheep]

Done in #3110