Umap: parallelization not working for nearest neighbors computation in umap 0.4

[dylkot]

Hi There,

I'm analyzing single-cell rna-seq data using scanpy on a Ubuntu virtual machine with 16 cpus and the following package versions:

• umap-learn 0.4.0 (installed from the 0.4dev branch earlier today)
• pynndescent 0.3.3
  
  
  
  • scanpy 1.4.5.dev175+g64f04d8 (installed from the master branch earlier today)
  
  

I'm using the development 0.4 version of umap because it is supposed to have support for parallelized computation of the nearest neighbors. Specifically, scanpy calls:

from umap.umap_ import nearest_neighbors

random_state = check_random_state(random_state)

knn_indices, knn_dists, forest = nearest_neighbors(
        X, n_neighbors, random_state=random_state,
        metric=metric, metric_kwds=metric_kwds,
        angular=angular, verbose=verbose,
)

the code is working but, empirically is just using a single CPU. It also gives the following warning message:

/opt/miniconda3/envs/py37_2/lib/python3.7/site-packages/numba/compiler.py:602: NumbaPerformanceWarning: 
The keyword argument 'parallel=True' was specified but no transformation for parallel execution was possible.

To find out why, try turning on parallel diagnostics, see http://numba.pydata.org/numba-doc/latest/user/parallel.html#diagnostics for help.

File "../../../../../opt/miniconda3/envs/py37_2/lib/python3.7/site-packages/umap/nndescent.py", line 47:
    @numba.njit(parallel=True)
    def nn_descent(
    ^

  self.func_ir.loc))

and when I turn on numba parallel diagnostics, that gives the report below. Any idea what is going on or if parallelized approximate nearest neighbors computations is supposed to be supported? Thanks!

================================================================================
 Parallel Accelerator Optimizing:  Function make_nn_descent.<locals>.nn_descent,
 /opt/miniconda3/envs/py37_2/lib/python3.7/site-packages/umap/nndescent.py (46)

================================================================================


Parallel loop listing for  Function make_nn_descent.<locals>.nn_descent, /opt/miniconda3/envs/py37_2/lib/python3.7/site-packages/umap/nndescent.py (46) 
------------------------------------------------------------------------------------------|loop #ID
    @numba.njit(parallel=True)                                                            | 
    def nn_descent(                                                                       | 
        data,                                                                             | 
        n_neighbors,                                                                      | 
        rng_state,                                                                        | 
        max_candidates=50,                                                                | 
        n_iters=10,                                                                       | 
        delta=0.001,                                                                      | 
        rho=0.5,                                                                          | 
        rp_tree_init=True,                                                                | 
        leaf_array=None,                                                                  | 
        verbose=False,                                                                    | 
    ):                                                                                    | 
        n_vertices = data.shape[0]                                                        | 
                                                                                          | 
        current_graph = make_heap(data.shape[0], n_neighbors)                             | 
        for i in range(data.shape[0]):                                                    | 
            indices = rejection_sample(n_neighbors, data.shape[0], rng_state)             | 
            for j in range(indices.shape[0]):                                             | 
                d = dist(data[i], data[indices[j]], *dist_args)                           | 
                heap_push(current_graph, i, d, indices[j], 1)                             | 
                heap_push(current_graph, indices[j], d, i, 1)                             | 
                                                                                          | 
        if rp_tree_init:                                                                  | 
            for n in range(leaf_array.shape[0]):                                          | 
                for i in range(leaf_array.shape[1]):                                      | 
                    if leaf_array[n, i] < 0:                                              | 
                        break                                                             | 
                    for j in range(i + 1, leaf_array.shape[1]):                           | 
                        if leaf_array[n, j] < 0:                                          | 
                            break                                                         | 
                        d = dist(                                                         | 
                            data[leaf_array[n, i]], data[leaf_array[n, j]], *dist_args    | 
                        )                                                                 | 
                        heap_push(                                                        | 
                            current_graph, leaf_array[n, i], d, leaf_array[n, j], 1       | 
                        )                                                                 | 
                        heap_push(                                                        | 
                            current_graph, leaf_array[n, j], d, leaf_array[n, i], 1       | 
                        )                                                                 | 
                                                                                          | 
        for n in range(n_iters):                                                          | 
            if verbose:                                                                   | 
                print("\t", n, " / ", n_iters)                                            | 
                                                                                          | 
            candidate_neighbors = build_candidates(                                       | 
                current_graph, n_vertices, n_neighbors, max_candidates, rng_state         | 
            )                                                                             | 
                                                                                          | 
            c = 0                                                                         | 
            for i in range(n_vertices):                                                   | 
                for j in range(max_candidates):                                           | 
                    p = int(candidate_neighbors[0, i, j])                                 | 
                    if p < 0 or tau_rand(rng_state) < rho:                                | 
                        continue                                                          | 
                    for k in range(max_candidates):                                       | 
                        q = int(candidate_neighbors[0, i, k])                             | 
                        if (                                                              | 
                            q < 0                                                         | 
                            or not candidate_neighbors[2, i, j]                           | 
                            and not candidate_neighbors[2, i, k]                          | 
                        ):                                                                | 
                            continue                                                      | 
                                                                                          | 
                        d = dist(data[p], data[q], *dist_args)                            | 
                        c += heap_push(current_graph, p, d, q, 1)                         | 
                        c += heap_push(current_graph, q, d, p, 1)                         | 
                                                                                          | 
            if c <= delta * n_neighbors * data.shape[0]:                                  | 
                break                                                                     | 
                                                                                          | 
        return deheap_sort(current_graph)                                                 | 
--------------------------------- Fusing loops ---------------------------------
Attempting fusion of parallel loops (combines loops with similar properties)...
----------------------------- Before Optimisation ------------------------------
--------------------------------------------------------------------------------
------------------------------ After Optimisation ------------------------------
Parallel structure is already optimal.
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

---------------------------Loop invariant code motion---------------------------

Instruction hoisting:
No instruction hoisting found
--------------------------------------------------------------------------------

[lmcinnes]

Hi -- there is good and bad news here. The parallelisation of pynndescent fitting with UMAP is not entirely done yet. Or, at least, it works on my machine with the my current versions of everything, but it probably isn't ready for prime time yet. If you want to try it out I would suggest grabbing the new_search branch of pynndescent from github and installing that (it should install as version 0.4.0). I would definitely do that in a separate environment (conda or virtualenv) as it is a "work in progress" version of pynndescent. On the machines I've tested on the nearest neighbor computation goes back and forth between many cores, and just one (there are some necessary serial portions).

If you just want to make use of 16 cores then the other option is in umap/umap_.py where it calls pynndescent add an extra option n_jobs=-1, which will turn on the (slightly memory intensive) threaded implementation of nndescent that exists in pynndescent v0.3.

[dylkot]

Ah I see, thank you so much for the quick response. And apologies about opening the issue. From the conversation on the scanpy repo, I was under the impression that this update was already completed. Thank you for the temporary workaround and for all of your amazing work.

[lmcinnes]

Opening an issues is the right thing to do. It helps remind me that I need to get this updated/fixed. Hopefully you can get one of the two options working.

[dylkot]

OK great. Well unfortunately option 2 doesn't work with a dataset of ~300,000 cells on a VM with 104 GB of RAM as it tries to allocate an array with shape (16, 771240000, 4) and data type float32. I will probably hold off on option 1 for now and will continue to work with downsampled data while iterating on my pipeline.

[flying-sheep]

How will the API look like? Will we be able to simply specify umap(..., n_jobs=4)? We’d like to use scanpy’s n_jobs setting here when calling umap from scanpy.

A thing I won’t get tired to say: Setting the default of any package to n_jobs=number_of_cores is a bad idea, setting that without the ability to change it when calling is an even worse one. Multi-user servers have many many cores and selfishly using all of them is very bad behavior, so libraries should make it hard for people to accidentally do that.

[lmcinnes]

I am using Numba for the parallelisation, and that (currently; see https://github.com/numba/numba/issues/2713) does not allow dynamic handling on how many threads to use. You can, however, set the environment variable NUMBA_NUM_THREADS to restrict the thread pool size similar to OpenMP. This is one of the things that I would like to see improved before the parallel stuff becomes standard.

[lmcinnes]

I see that @tomwhite gave the correct answer over in the ScanPy issues: https://github.com/theislab/scanpy/issues/913#issuecomment-553420798

Sorry -- I've got a little bit of tunnel vision from spending all my time recently working on a numba parallel version of pynndescent that scales better with lower core counts, and forgot how to properly use Tom's code. Sorry for any confusion.

[millerh1]

Hello @lmcinnes -- I was hoping you might be willing to provide an update on this issue. I am coming from R and I was hoping that there might be a stable umap-learn parallelization of the NN search similar to the one from RcppParallel. From following these conversations, it seems that this will be implemented in the 0.5 version pynndescent. Is that correct? Thank you for your time!

[lmcinnes]

If you just want multicore NN search then the current pynndescent will get you that. If you want distributed/cluster support for NN search then you may have to wait a while longer.

[millerh1]

Thank you for the quick reply! It appears made the mistake of not including pynndescent initially...

That being said, even with pynndescent running, the multicore doesn't utilize more than 2-3 cores at once. I have a machine with 192 cores available. The embedding is perfectly capable of using all of them, but it seems the NN search doesn't for some reason. Any idea what might cause that? I should mention I'm using HDDs with somewhat limited read/write speed -- which sometimes poses a problem for programs that continually read/write from disk. I have 500GB of RAM so it would be nice (if the disk I/O is the problem) if I could do more in memory for the NN search.

As someone who isn't very knowledgeable about these algorithms or optimization in general, it would be nice if you could offer some more insight into this issue. Thank you for your time!

[lmcinnes]

No -- I don't really know what the issue could be. I have used a 128 core SMP successfully using all the cores. Check the NUMBA_NUM_THREADS environment variable and related options perhaps?

[millerh1]

numba.config.NUMBA_NUM_THREADS is set to 192. Not sure what could cause this then.

[lmcinnes]

Sorry I can't be more help. It doesn't max out all the cores all the time, but you should see significant core utilization...