UMAP memory parallelisation issues

Nobody has claimed this yet.

Hi UMAP team,

First off, thanks for the great work! I am trying to use UMAP to reduce the dimensionality of a dataset, where initial tests have shown very promising results. I’ve run into several issues (parallelization, memory management..) where some help would be great.

Context:
I’m working on a node in a cluster with ca 270GB available memory and 16 cores. The dataset I am currently working on is (3069899, 5), see below for example snippet. I would like to scale that up 12 times if possible.
Running UMAP in serial works for the (3069899, 5) dataset using only two dims and setting the random seed, but takes hours. Ideally I would like to run parameter sweeps to ensure meaningful results and reproducibility.

Issues:
- Removing the random seed doesn’t make it use more cores.
- Following https://github.com/lmcinnes/umap/issues/379, we tried installing pynndescent. The script seemed to use multiple cores on one of the nodes, after setting n_jobs=-1 in umap. But, due to memory issues, it did not run to completion.
-Example errors while invoking UMAP from within a DASK cluster are as follows:
- distributed.worker - WARNING - Memory use is high but worker has no data to store to disk. Perhaps some other process is leaking memory? Process memory: 264.16 GB -- Worker memory limit: 270.51 GB
- distributed.core - INFO - Event loop was unresponsive in Worker for 14.13s. This is often caused by long-running GIL-holding functions or moving large chunks of data. This can cause timeouts and instability.
- When the pandas dataframe was much smaller than the original, the calls to UMAP seem to run to completion (till saving results to a numpy array).
-low_memory=False in a couple of recent tests continued to throw memory exceptions.
- Dealing with a pending ‘small memory’, it has not worked testing on smaller machines.
- Calling UMAP defined inside a function to use Sciris to parallelise serial runs the following crop up:
- It doesn’t save the results if called inside a function (see below). I tried np.save, pickle, turning it into an array again etc.. no luck.
-It eats all available memory.

I have occasionally also (trying out e.g. different numba versions and whatnot) run into segfaults and ‘BUS error’ I believe are memory related.

My initial results were hugely promising, so I would very much appreciate any help you can offer!

Example code snippet that works:

clusterable_embedding = umap.UMAP(
n_neighbors=100,
min_dist=0.0,
n_components=2,
random_state=48).fit_transform(pd_BV)

np.save('umap_data.npy', clusterable_embedding)

Defined in a function/save issue:

def runUMAP(neighbours, run):
print(neighbours)
clusterable_embedding = umap.UMAP(
n_neighbors=neighbours,
min_dist=0.0,
n_components=2
).fit_transform(pd_BV)
return clusterable_embedding

if __name__ == "__main__":

neighbours = [60,80,120,150,180, 200,250]
run = np.arange(1,5)
[aa, bb] = np.meshgrid(neighbours,run)

iterargs = list(zip(aa.flatten(), bb.flatten()))
results=sc.parallelize(runUMAP, iterargs, ncpus=8)
print('done')

for ix, result in enumerate(results):
np.save(f'save_{ix}.npy', result)

Start by reproducing the example UMAP call on the reported (3069899, 5) dataset, then compare serial execution with n_jobs=-1 and the DASK or Sciris parallel cases. Investigate the reported memory exhaustion, unresponsive worker, save failures, and occasional segfaults; the issue does not define a single expected fix or completion test.

Written by the indexing model from the issue text.