Use ray for parallelization of nearest neighbor search - python-3.x

Question

Suppose I have a big sparse matrix. I want to take each row vector of that matrix and compute the cosine distance to its nearest neighbor among the previous window_size rows of the matrix.
Since sklearn.neighbors does not improve run time when using parallelization (see this issue on github ), I tried to parallelize the process using ray. My code does better than sklearn with multiprocessing, but it's still slower than just serial distance computation.
My code is below. Is there something I have done wrong and should improve?
import scipy.sparse
from sklearn.metrics.pairwise import cosine_distances
from sklearn.neighbors import NearestNeighbors
import numpy as np
import timeit
import ray
import math
n = 4000
m = 100
big_matrix = scipy.sparse.random(n, m).tocsr()
window_size = 1000
retry = 1
n_jobs = 4
ray.init(num_cpus=n_jobs)
def simple_cosine_distance():
distances = np.zeros(n)
for i in range(1, n):
past = big_matrix[max(0, i - window_size):i]
query = big_matrix[i]
distances[i] = cosine_distances(query, past).min(axis=1)
def sklearn_nneighbor():
distances = np.zeros(n)
for i in range(1, n):
past = big_matrix[max(0, i - window_size):i]
nn = NearestNeighbors(metric="cosine", n_neighbors=1, n_jobs=1)
nn.fit(X=past)
query = big_matrix[i]
distance, _ = nn.kneighbors(query)
distances[i] = distance[0]
def sklearn_nneighbor_parallel():
distances = np.zeros(n)
for i in range(1, n):
past = big_matrix[max(0, i - window_size):i]
nn = NearestNeighbors(metric="cosine", n_neighbors=1, n_jobs=n_jobs)
nn.fit(X=past)
query = big_matrix[i]
distance, _ = nn.kneighbors(query)
distances[i] = distance[0]
#ray.remote
def get_chunk_min(data, indices, indptr, shape, slice_i, slice_j, query):
matrix = scipy.sparse.csr_matrix((data, indices, indptr), shape=shape)
past = matrix[slice_i:slice_j]
query = matrix[query]
chunk_min = cosine_distances(query, past).min(axis=1)
return chunk_min
def ray_parallel():
distances = np.zeros(n)
data = ray.put(big_matrix.data)
indices = ray.put(big_matrix.indices)
indptr = ray.put(big_matrix.indptr)
shape = ray.put(big_matrix.shape)
for i in range(1, n):
chunk_size = math.ceil((i - max(0, i - window_size)) / n_jobs)
chunk_mins = ray.get([
get_chunk_min.remote(
data,
indices,
indptr,
shape,
enum,
enum + chunk_size,
i
) for enum in range(max(0, i - window_size), i, chunk_size)])
distances[i] = min(chunk_mins)
for method in ["simple_cosine_distance", "sklearn_nneighbor", "sklearn_nneighbor_parallel", "ray_parallel"]:
print(method)
print(timeit.timeit(method + "()", setup="from __main__ import " + method, number=retry))
print("*"*50)
Output:
simple_cosine_distance
3.978868665999471
sklearn_nneighbor
4.265772191996803
sklearn_nneighbor_parallel
28.664759318002325
ray_parallel
17.89882287799992

Answer 1

https://docs.python.org/3/library/concurrent.futures.html
This documentation is a good resource to understand some mistakes you can do while are parallelizing workload using Ray.
Also note that, if your remote function is simple enough, the overhead of distributed computing (e.g. scheduling time, serialization time, etc.) can be bigger than function itself's computing time.